Initial commit

Signed-off-by: Filippo Costa <[email protected]>

.github/workflows/ci.yml

@@ -0,0 +1,74 @@
+name: CI
+
+on:
+  push:
+    branches: ["main"]
+  pull_request:
+    branches: ["main"]
+
+env:
+  CARGO_TERM_COLOR: always
+  RUSTFLAGS: -D warnings
+
+jobs:
+  check-aggregate:
+    # Follows the guide at <https://github.com/re-actors/alls-green>.
+    runs-on: ubuntu-latest
+    if: always()
+    needs:
+      # If you're modifying this workflow file and you're adding/removing a job
+      # which should be required to pass before merging a PR, don't forget to
+      # update this list!
+      - check
+      - features
+      - test
+    steps:
+      - name: Compute whether the needed jobs succeeded or failed
+        uses: re-actors/alls-green@release/v1
+        with:
+          allowed-skips: deploy-github-pages
+          jobs: ${{ toJSON(needs) }}
+
+  check:
+    name: check
+    runs-on: buildjet-4vcpu-ubuntu-2204
+    timeout-minutes: 60
+    steps:
+      - uses: actions/checkout@v3
+      - name: Install Rust
+        run: rustup show
+      - uses: Swatinem/rust-cache@v2
+      - name: Check formatting
+        run: cargo fmt --all --check
+      - run: cargo check --all-targets --all-features
+      - run: cargo clippy --all-targets --all-features
+
+  # Check that every combination of features is working properly.
+  features:
+    name: features
+    runs-on: buildjet-8vcpu-ubuntu-2204
+    timeout-minutes: 120
+    steps:
+      - uses: actions/checkout@v3
+      - name: Install Rust
+        run: rustup show
+      - name: cargo install cargo-hack
+        uses: taiki-e/install-action@cargo-hack
+      - uses: Swatinem/rust-cache@v2
+      - name: cargo hack
+        run: cargo hack check --workspace --feature-powerset --all-targets
+  test:
+    name: test
+    runs-on: buildjet-4vcpu-ubuntu-2204
+    timeout-minutes: 60
+    steps:
+      - uses: actions/checkout@v3
+      - uses: taiki-e/install-action@nextest
+      - name: Install Rust
+        run: rustup show
+      - uses: Swatinem/rust-cache@v2
+      - run: cargo nextest run --workspace --all-features
+      # `cargo-nextest` does not support doctests (yet?), so we have to run them
+      # separately.
+      # TODO: https://github.com/nextest-rs/nextest/issues/16
+      - run: cargo test --workspace --doc --all-features

.gitignore

@@ -0,0 +1,16 @@
+# From <https://github.com/github/gitignore/blob/main/Rust.gitignore>.
+
+# Generated by Cargo
+# will have compiled files and executables
+debug/
+target/
+
+# Remove Cargo.lock from gitignore if creating an executable, leave it for libraries
+# More information here https://doc.rust-lang.org/cargo/guide/cargo-toml-vs-cargo-lock.html
+Cargo.lock
+
+# These are backup files generated by rustfmt
+**/*.rs.bk
+
+# MSVC Windows builds of rustc generate these, which store debugging information
+*.pdb

Cargo.toml

@@ -0,0 +1,36 @@
+[package]
+name = "rockbound"
+version = "1.0.0"
+edition = "2021"
+authors = ["Sovereign Labs"]
+homepage = "https://github.com/sovereign-labs/rockbound"
+repository = "https://github.com/sovereign-labs/rockbound"
+description = "A low level interface transforming RocksDB into a type-oriented data store"
+
+# Some of the code is derived from Aptos <https://github.com/aptos-labs/aptos-core>
+# which is licensed under Apache-2.0.
+license = "Apache-2.0"
+
+readme = "README.md"
+
+[dependencies]
+anyhow = "1"
+byteorder = { version = "1", default-features = true, optional = true }
+once_cell = "1"
+prometheus = { version = "0.13", optional = true }
+proptest = { version = "1", optional = true }
+proptest-derive = { version = "0.4", optional = true }
+rocksdb = { version = "0.21" }
+thiserror = "1"
+tracing = "0.1"
+
+[dev-dependencies]
+rockbound = { path = ".", features = ["test-utils"] }
+tempfile = "3"
+
+[features]
+default = []
+
+arbitrary = ["dep:proptest", "dep:proptest-derive"]
+prometheus = ["dep:prometheus"]
+test-utils = ["dep:byteorder"]

README.md

@@ -0,0 +1,60 @@
+# Rockbound
+
+This package is a low-level wrapper transforming [RocksDB](https://rocksdb.org/) from a byte-oriented key value store into a
+type-oriented store. It's adapted from a similar package in Aptos-Core.
+
+The most important concept exposed by Rockbound is a `Schema`, which maps a column-family name
+codec's for the key and value.
+
+```rust
+pub trait Schema  {
+    /// The column family name associated with this struct.
+    /// Note: all schemas within the same DB must have distinct column family names.
+    const COLUMN_FAMILY_NAME: &'static str;
+
+    /// Type of the key.
+    type Key: KeyCodec<Self>;
+
+    /// Type of the value.
+    type Value: ValueCodec<Self>;
+}
+```
+
+Using this schema, we can write generic functions for storing and fetching typed data, and ensure that
+it's always encoded/decoded in the way we expect.
+
+```rust
+impl SchemaDB {
+	pub fn put<S: Schema>(&self, key: &impl KeyCodec<S>, value: &impl ValueCodec<S>) -> Result<()> {
+		let key_bytes = key.encode_key()?;
+        let value_bytes = value.encode_value()?;
+		self.rocks_db_handle.put(S::COLUMN_FAMILY_NAME, key_bytes, value_bytes)
+	}
+}
+```
+
+To actually store and retrieve data, all we need to do is to implement a Schema:
+
+```rust
+pub struct AccountBalanceSchema;
+
+impl Schema for AccountBalanceSchema {
+	const COLUMN_FAMILY_NAME: &str = "account_balances";
+	type Key = Account;
+	type Value = u64;
+}
+
+impl KeyCodec<AccountBalanceSchema> for Account {
+	fn encode_key(&self) -> Vec<u8> {
+		bincode::to_vec(self)
+	}
+
+	fn decode_key(key: Vec<u8>) -> Self {
+		// elided
+	}
+}
+
+impl ValueCode<AccountBlanceSchema> for u64 {
+	// elided
+}
+```

src/iterator.rs

@@ -0,0 +1,192 @@
+use std::iter::FusedIterator;
+use std::marker::PhantomData;
+
+use anyhow::Result;
+
+use crate::metrics::{ROCKBOUND_ITER_BYTES, ROCKBOUND_ITER_LATENCY_SECONDS};
+use crate::schema::{KeyDecoder, Schema, ValueCodec};
+use crate::{SchemaKey, SchemaValue};
+
+/// This defines a type that can be used to seek a [`SchemaIterator`], via
+/// interfaces like [`SchemaIterator::seek`]. Mind you, not all
+/// [`KeyEncoder`](crate::schema::KeyEncoder)s shall be [`SeekKeyEncoder`]s, and
+/// vice versa. E.g.:
+///
+/// - Some key types don't use an encoding that results in sensible
+/// seeking behavior under lexicographic ordering (what RocksDB uses by
+/// default), which means you shouldn't implement [`SeekKeyEncoder`] at all.
+/// - Other key types might maintain full lexicographic order, which means the
+/// original key type can also be [`SeekKeyEncoder`].
+/// - Other key types may be composite, and the first field alone may be
+/// a good candidate for [`SeekKeyEncoder`].
+pub trait SeekKeyEncoder<S: Schema + ?Sized>: Sized {
+    /// Converts `self` to bytes which is used to seek the underlying raw
+    /// iterator.
+    ///
+    /// If `self` is also a [`KeyEncoder`](crate::schema::KeyEncoder), then
+    /// [`SeekKeyEncoder::encode_seek_key`] MUST return the same bytes as
+    /// [`KeyEncoder::encode_key`](crate::schema::KeyEncoder::encode_key).
+    fn encode_seek_key(&self) -> crate::schema::Result<Vec<u8>>;
+}
+
+pub(crate) enum ScanDirection {
+    Forward,
+    Backward,
+}
+
+/// DB Iterator parameterized on [`Schema`] that seeks with [`Schema::Key`] and yields
+/// [`Schema::Key`] and [`Schema::Value`] pairs.
+pub struct SchemaIterator<'a, S> {
+    db_iter: rocksdb::DBRawIterator<'a>,
+    direction: ScanDirection,
+    phantom: PhantomData<S>,
+}
+
+impl<'a, S> SchemaIterator<'a, S>
+where
+    S: Schema,
+{
+    pub(crate) fn new(db_iter: rocksdb::DBRawIterator<'a>, direction: ScanDirection) -> Self {
+        SchemaIterator {
+            db_iter,
+            direction,
+            phantom: PhantomData,
+        }
+    }
+
+    /// Seeks to the first key.
+    pub fn seek_to_first(&mut self) {
+        self.db_iter.seek_to_first();
+    }
+
+    /// Seeks to the last key.
+    pub fn seek_to_last(&mut self) {
+        self.db_iter.seek_to_last();
+    }
+
+    /// Seeks to the first key whose binary representation is equal to or greater than that of the
+    /// `seek_key`.
+    pub fn seek(&mut self, seek_key: &impl SeekKeyEncoder<S>) -> Result<()> {
+        let key = seek_key.encode_seek_key()?;
+        self.db_iter.seek(&key);
+        Ok(())
+    }
+
+    /// Seeks to the last key whose binary representation is less than or equal to that of the
+    /// `seek_key`.
+    ///
+    /// See example in [`RocksDB doc`](https://github.com/facebook/rocksdb/wiki/SeekForPrev).
+    pub fn seek_for_prev(&mut self, seek_key: &impl SeekKeyEncoder<S>) -> Result<()> {
+        let key = seek_key.encode_seek_key()?;
+        self.db_iter.seek_for_prev(&key);
+        Ok(())
+    }
+
+    /// Reverses iterator direction.
+    pub fn rev(self) -> Self {
+        let new_direction = match self.direction {
+            ScanDirection::Forward => ScanDirection::Backward,
+            ScanDirection::Backward => ScanDirection::Forward,
+        };
+        SchemaIterator {
+            db_iter: self.db_iter,
+            direction: new_direction,
+            phantom: Default::default(),
+        }
+    }
+
+    fn next_impl(&mut self) -> Result<Option<IteratorOutput<S::Key, S::Value>>> {
+        let _timer = ROCKBOUND_ITER_LATENCY_SECONDS
+            .with_label_values(&[S::COLUMN_FAMILY_NAME])
+            .start_timer();
+
+        if !self.db_iter.valid() {
+            self.db_iter.status()?;
+            return Ok(None);
+        }
+
+        let raw_key = self.db_iter.key().expect("db_iter.key() failed.");
+        let raw_value = self.db_iter.value().expect("db_iter.value() failed.");
+        let value_size_bytes = raw_value.len();
+        ROCKBOUND_ITER_BYTES
+            .with_label_values(&[S::COLUMN_FAMILY_NAME])
+            .observe((raw_key.len() + raw_value.len()) as f64);
+
+        let key = <S::Key as KeyDecoder<S>>::decode_key(raw_key)?;
+        let value = <S::Value as ValueCodec<S>>::decode_value(raw_value)?;
+
+        match self.direction {
+            ScanDirection::Forward => self.db_iter.next(),
+            ScanDirection::Backward => self.db_iter.prev(),
+        }
+
+        Ok(Some(IteratorOutput {
+            key,
+            value,
+            value_size_bytes,
+        }))
+    }
+}
+
+/// The output of [`SchemaIterator`]'s next_impl
+pub struct IteratorOutput<K, V> {
+    pub key: K,
+    pub value: V,
+    pub value_size_bytes: usize,
+}
+
+impl<K, V> IteratorOutput<K, V> {
+    pub fn into_tuple(self) -> (K, V) {
+        (self.key, self.value)
+    }
+}
+
+impl<'a, S> Iterator for SchemaIterator<'a, S>
+where
+    S: Schema,
+{
+    type Item = Result<IteratorOutput<S::Key, S::Value>>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        self.next_impl().transpose()
+    }
+}
+
+impl<'a, S> FusedIterator for SchemaIterator<'a, S> where S: Schema {}
+
+/// Iterates over given column backwards
+pub struct RawDbReverseIterator<'a> {
+    db_iter: rocksdb::DBRawIterator<'a>,
+}
+
+impl<'a> RawDbReverseIterator<'a> {
+    pub(crate) fn new(mut db_iter: rocksdb::DBRawIterator<'a>) -> Self {
+        db_iter.seek_to_last();
+        RawDbReverseIterator { db_iter }
+    }
+
+    /// Navigate iterator go given key
+    pub fn seek(&mut self, seek_key: SchemaKey) -> Result<()> {
+        self.db_iter.seek_for_prev(&seek_key);
+        Ok(())
+    }
+}
+
+impl<'a> Iterator for RawDbReverseIterator<'a> {
+    type Item = (SchemaKey, SchemaValue);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if !self.db_iter.valid() {
+            self.db_iter.status().ok()?;
+            return None;
+        }
+
+        let next_item = self.db_iter.item().expect("db_iter.key() failed.");
+        // Have to allocate to fix lifetime issue
+        let next_item = (next_item.0.to_vec(), next_item.1.to_vec());
+
+        self.db_iter.prev();
+
+        Some(next_item)
+    }
+}