[3/N][VQueues] Partition Store Implementation

Author: AhmedSoliman

Cargo.lock

@@ -23,6 +23,7 @@ restate-types = { workspace = true }
 
 ahash = { workspace = true }
 anyhow = { workspace = true }
+bilrost = { workspace = true }
 bytes = { workspace = true }
 bytestring = { workspace = true }
 codederror = { workspace = true }

crates/partition-store/Cargo.toml

@@ -14,7 +14,11 @@ use anyhow::anyhow;
 use bytes::{Buf, BufMut, Bytes, BytesMut};
 use bytestring::ByteString;
 use prost::encoding::encoded_len_varint;
+use rocksdb::MergeOperands;
 use strum::EnumIter;
+use tracing::{error, trace};
+
+use restate_types::clock::UniqueTimestamp;
 
 /// Every table key needs to have a key kind. This allows to multiplex different keys in the same
 /// column family and to evolve a key if necessary.
@@ -42,6 +46,17 @@ pub enum KeyKind {
     State,
     Timers,
     Promise,
+    // VQueues --> owned by restate-vqueues
+    //
+    // todo: split this into empty and non-empty, or add the status in the key prefix
+    // for instance, make this VQueueStatus (S | VQueueId)
+    // or have empty_vqueues that carry the empty_since in their key prefix. Note that
+    // doing so would require that we know the empty_since when we attempt to delete it
+    VQueueActive,
+    VQueueInbox,
+    VQueueMeta,
+    // Resources' canonical key(s)
+    VQueueEntryState,
 }
 
 impl KeyKind {
@@ -86,6 +101,13 @@ impl KeyKind {
             KeyKind::State => b"st",
             KeyKind::Timers => b"ti",
             KeyKind::Promise => b"pr",
+            // ** VQueues ** //
+            // VQueues own all keys that start with b"q".
+            KeyKind::VQueueActive => b"qa",
+            KeyKind::VQueueInbox => b"qi",
+            KeyKind::VQueueMeta => b"qm",
+            // Queue Entry State (canonical state of vqueue entries)
+            KeyKind::VQueueEntryState => b"qe",
         }
     }
 
@@ -115,6 +137,11 @@ impl KeyKind {
             b"st" => Some(KeyKind::State),
             b"ti" => Some(KeyKind::Timers),
             b"pr" => Some(KeyKind::Promise),
+            // VQueues own all keys that start with b"q"
+            b"qa" => Some(KeyKind::VQueueActive),
+            b"qi" => Some(KeyKind::VQueueInbox),
+            b"qm" => Some(KeyKind::VQueueMeta),
+            b"qe" => Some(KeyKind::VQueueEntryState),
             _ => None,
         }
     }
@@ -134,6 +161,50 @@ impl KeyKind {
         Self::from_bytes(&bytes)
             .ok_or_else(|| StorageError::Generic(anyhow::anyhow!("unknown key kind: {:x?}", bytes)))
     }
+
+    // Rocksdb merge operator function (full merge)
+    pub fn full_merge(
+        key: &[u8],
+        existing_val: Option<&[u8]>,
+        operands: &MergeOperands,
+    ) -> Option<Vec<u8>> {
+        let mut kind_buf = key;
+        let kind = match KeyKind::deserialize(&mut kind_buf) {
+            Ok(kind) => kind,
+            Err(e) => {
+                error!("Cannot apply merge operator; {e}");
+                return None;
+            }
+        };
+        trace!(?kind, "full merge");
+
+        match kind {
+            KeyKind::VQueueMeta => vqueue_meta_merge::full_merge(key, existing_val, operands),
+            _ => None,
+        }
+    }
+
+    // Rocksdb merge operator function (partial merge)
+    pub fn partial_merge(
+        key: &[u8],
+        _unused: Option<&[u8]>,
+        operands: &MergeOperands,
+    ) -> Option<Vec<u8>> {
+        let mut kind_buf = key;
+        let kind = match KeyKind::deserialize(&mut kind_buf) {
+            Ok(kind) => kind,
+            Err(e) => {
+                error!("Cannot apply merge operator; {e}");
+                return None;
+            }
+        };
+        trace!(?kind, "partial merge");
+
+        match kind {
+            KeyKind::VQueueMeta => vqueue_meta_merge::partial_merge(key, operands),
+            _ => None,
+        }
+    }
 }
 
 pub trait TableKey: Sized + std::fmt::Debug + Send + 'static {
@@ -350,6 +421,7 @@ macro_rules! define_table_key {
 
 use crate::PaddedPartitionId;
 use crate::TableKind;
+use crate::vqueue_table::vqueue_meta_merge;
 pub(crate) use define_table_key;
 use restate_storage_api::StorageError;
 use restate_storage_api::deduplication_table::ProducerId;
@@ -411,6 +483,41 @@ impl KeyCodec for PaddedPartitionId {
     }
 }
 
+impl KeyCodec for UniqueTimestamp {
+    fn encode<B: BufMut>(&self, target: &mut B) {
+        // store u64 in big-endian order to support byte-wise increment operation. See `crate::scan::try_increment`.
+        target.put_u64(self.as_u64());
+    }
+
+    fn decode<B: Buf>(source: &mut B) -> crate::Result<Self> {
+        UniqueTimestamp::try_from(source.get_u64()).map_err(|e| StorageError::Conversion(e.into()))
+    }
+
+    fn serialized_length(&self) -> usize {
+        std::mem::size_of::<Self>()
+    }
+}
+
+impl<const L: usize> KeyCodec for [u8; L] {
+    fn encode<B: BufMut>(&self, target: &mut B) {
+        // stores the array as is.
+        target.put_slice(self.as_ref());
+    }
+
+    fn decode<B: Buf>(source: &mut B) -> crate::Result<Self> {
+        if source.remaining() < L {
+            return Err(StorageError::DataIntegrityError);
+        }
+        let mut buf = [0u8; L];
+        source.copy_to_slice(&mut buf);
+        Ok(buf)
+    }
+
+    fn serialized_length(&self) -> usize {
+        L
+    }
+}
+
 impl KeyCodec for u64 {
     fn encode<B: BufMut>(&self, target: &mut B) {
         // store u64 in big-endian order to support byte-wise increment operation. See `crate::scan::try_increment`.

crates/partition-store/src/keys.rs

@@ -32,6 +32,7 @@ pub mod service_status_table;
 pub mod snapshots;
 pub mod state_table;
 pub mod timer_table;
+pub mod vqueue_table;
 
 #[cfg(test)]
 mod tests;

crates/partition-store/src/lib.rs

@@ -28,7 +28,9 @@ use restate_types::config::Configuration;
 use restate_types::logs::Lsn;
 use restate_types::partitions::{CfName, Partition};
 
+use crate::TableKind;
 use crate::durable_lsn_tracking::{AppliedLsnCollectorFactory, DurableLsnEventListener};
+use crate::keys::KeyKind;
 use crate::memory::MemoryBudget;
 use crate::snapshots::LocalPartitionSnapshot;
 
@@ -80,6 +82,10 @@ impl PartitionDb {
         &self.cf.0
     }
 
+    pub(crate) fn table_cf_handle(&self, _table_kind: TableKind) -> &Arc<BoundColumnFamily<'_>> {
+        self.cf_handle()
+    }
+
     pub fn cf_names(&self) -> Vec<SmartString> {
         vec![self.meta.cf_name().into_inner()]
     }
@@ -515,6 +521,12 @@ impl CfConfigurator for RocksConfigurator<AllDataCf> {
             Some(global_cache),
         );
         cf_options.set_block_based_table_factory(&block_options);
+        cf_options.set_merge_operator(
+            "PartitionMerge",
+            KeyKind::full_merge,
+            KeyKind::partial_merge,
+        );
+        cf_options.set_max_successive_merges(100);
 
         // Actually, we would love to use CappedPrefixExtractor but unfortunately it's neither exposed
         // in the C API nor the rust binding. That's okay and we can change it later.

crates/partition-store/src/partition_db.rs

@@ -139,6 +139,7 @@ pub enum TableKind {
     Journal,
     JournalEvent,
     Promise,
+    VQueue,
 }
 
 impl TableKind {
@@ -162,6 +163,12 @@ impl TableKind {
             ],
             Self::JournalEvent => &[KeyKind::JournalEvent],
             Self::Promise => &[KeyKind::Promise],
+            Self::VQueue => &[
+                KeyKind::VQueueMeta,
+                KeyKind::VQueueInbox,
+                KeyKind::VQueueActive,
+                KeyKind::VQueueEntryState,
+            ],
         }
     }
 
@@ -250,8 +257,8 @@ impl PartitionStore {
         self.db.partition().key_range.contains(&key)
     }
 
-    fn table_handle(&self, _table_kind: TableKind) -> &Arc<BoundColumnFamily<'_>> {
-        self.db.cf_handle()
+    pub(crate) fn table_handle(&self, table_kind: TableKind) -> &Arc<BoundColumnFamily<'_>> {
+        self.db.table_cf_handle(table_kind)
     }
 
     fn new_prefix_iterator_opts(&self, _key_kind: KeyKind, prefix: Bytes) -> ReadOptions {
@@ -822,6 +829,34 @@ impl PartitionStoreTransaction<'_> {
         opts
     }
 
+    #[inline]
+    pub fn raw_put_cf(
+        &mut self,
+        _key_kind: KeyKind,
+        key: impl AsRef<[u8]>,
+        value: impl AsRef<[u8]>,
+    ) {
+        self.write_batch_with_index
+            .put_cf(self.data_cf_handle, key, value);
+    }
+
+    #[inline]
+    pub fn raw_merge_cf(
+        &mut self,
+        _key_kind: KeyKind,
+        key: impl AsRef<[u8]>,
+        value: impl AsRef<[u8]>,
+    ) {
+        self.write_batch_with_index
+            .merge_cf(self.data_cf_handle, key, value);
+    }
+
+    #[inline]
+    pub fn raw_delete_cf(&mut self, _key_kind: KeyKind, key: impl AsRef<[u8]>) {
+        self.write_batch_with_index
+            .delete_cf(self.data_cf_handle, key);
+    }
+
     pub(crate) fn prefix_iterator(
         &self,
         table: TableKind,