Introduces independent lock storage and vqueue scheduler resource manager

Author: AhmedSoliman

crates/partition-store/src/tests/vqueue_table_test/mod.rs

@@ -280,11 +280,7 @@ fn verify_empty_queue_returns_none(db: &crate::PartitionDb) {
 async fn vqueue_isolation<W: WriteVQueueTable>(txn: &mut W) {
     let pkey = PartitionKey::from(5000u64);
 
-    let qid1 = VQueueId::new(
-        VQueueParent::from_raw(1),
-        pkey,
-        VQueueInstance::from_raw(1),
-    );
+    let qid1 = VQueueId::new(VQueueParent::from_raw(1), pkey, VQueueInstance::from_raw(1));
     let qid2 = VQueueId::new(
         VQueueParent::from_raw(1),
         pkey,
@@ -304,21 +300,9 @@ async fn vqueue_isolation<W: WriteVQueueTable>(txn: &mut W) {
 fn verify_vqueue_isolation(db: &crate::PartitionDb) {
     let pkey = PartitionKey::from(5000u64);
 
-    let qid1 = VQueueId::new(
-        VQueueParent::from_raw(1),
-        pkey,
-        VQueueInstance::from_raw(1),
-    );
-    let qid2 = VQueueId::new(
-        VQueueParent::from_raw(1),
-        pkey,
-        VQueueInstance::from_raw(2),
-    );
-    let qid3 = VQueueId::new(
-        VQueueParent::from_raw(2),
-        pkey,
-        VQueueInstance::from_raw(1),
-    );
+    let qid1 = VQueueId::new(VQueueParent::from_raw(1), pkey, VQueueInstance::from_raw(1));
+    let qid2 = VQueueId::new(VQueueParent::from_raw(1), pkey, VQueueInstance::from_raw(2));
+    let qid3 = VQueueId::new(VQueueParent::from_raw(2), pkey, VQueueInstance::from_raw(1));
 
     // Each queue should only see its own entry
     let mut reader1 = db.new_inbox_reader(&qid1);

crates/partition-store/src/vqueue_table/mod.rs

@@ -332,6 +332,23 @@ impl WriteVQueueTable for PartitionStoreTransaction<'_> {
         }
     }
 
+    fn delete_inbox_entry(&mut self, qid: &VQueueId, stage: Stage, card: &EntryCard) {
+        let key_buffer = InboxKey {
+            partition_key: qid.partition_key(),
+            parent: qid.parent,
+            instance: qid.instance,
+            stage,
+            visible_at: card.visible_at,
+            priority: card.priority,
+            created_at: card.created_at,
+            kind: card.kind,
+            id: card.id,
+        }
+        .to_bytes();
+
+        self.raw_delete_cf(KeyKind::VQueueInbox, key_buffer);
+    }
+
     fn mark_vqueue_as_active(&mut self, qid: &restate_types::vqueue::VQueueId) {
         let mut key_buffer = [0u8; ActiveKey::serialized_length_fixed()];
         ActiveKey {

crates/storage-api/src/vqueue_table/entry.rs

@@ -96,6 +96,7 @@ impl From<Lsn> for EntryId {
 
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
 pub struct EntryCard {
+    // pub has_lock: bool,
     pub priority: EffectivePriority,
     pub visible_at: VisibleAt,
     /// The unique timestamp of the initial creation of the entry.
@@ -120,6 +121,11 @@ impl EntryCard {
         + std::mem::size_of::<EntryId>()
     }
 
+    pub fn has_lock(&self) -> bool {
+        // todo: change
+        self.priority.token_held()
+    }
+
     /// A unique hash of the entry card.
     ///
     /// Do not use this for any stored data as it changes across version/restarts.

crates/storage-api/src/vqueue_table/metadata.rs

@@ -13,10 +13,10 @@ use smallvec::SmallVec;
 use restate_clock::WallClock;
 use restate_clock::time::MillisSinceEpoch;
 use restate_limiter::LimitKey;
-use restate_types::Scope;
 use restate_types::clock::UniqueTimestamp;
 use restate_types::vqueue::EffectivePriority;
-use restate_util_string::{InternedReString, ReString};
+use restate_types::{LockName, Scope};
+use restate_util_string::ReString;
 
 use super::{Stage, VisibleAt};
 
@@ -88,10 +88,16 @@ pub struct VQueueMeta {
     pub(crate) scope: Option<Scope>,
     #[bilrost(tag(8))]
     pub(crate) limit_key: LimitKey<ReString>,
+    #[bilrost(tag(9))]
+    lock_name: Option<LockName>,
 }
 
 impl VQueueMeta {
-    pub fn new(scope: Option<Scope>, limit_key: LimitKey<ReString>) -> Self {
+    pub fn new(
+        scope: Option<Scope>,
+        limit_key: LimitKey<ReString>,
+        lock_name: Option<LockName>,
+    ) -> Self {
         Self {
             is_paused: false,
             length: 0,
@@ -101,13 +107,26 @@ impl VQueueMeta {
             stats: VQueueStatistics::new(WallClock::recent_ms()),
             scope,
             limit_key,
+            lock_name,
         }
     }
 
-    pub fn scope(&self) -> &Option<InternedReString> {
+    pub fn scope(&self) -> &Option<Scope> {
         &self.scope
     }
 
+    pub fn scope_ref(&self) -> &Option<Scope> {
+        &self.scope
+    }
+
+    pub fn requires_locking(&self) -> bool {
+        self.lock_name.is_some()
+    }
+
+    pub fn lock_name(&self) -> &Option<LockName> {
+        &self.lock_name
+    }
+
     pub fn limit_key(&self) -> &LimitKey<ReString> {
         &self.limit_key
     }

crates/storage-api/src/vqueue_table/mod.rs

@@ -101,7 +101,7 @@ pub struct WaitStats {
     pub vqueue_start_throttling_ms: u32,
     /// Total milliseconds the item was throttled on global "run" token bucket
     #[bilrost(tag(3))]
-    pub global_throttling_ms: u32,
+    pub global_invoker_throttling_ms: u32,
     /// Total milliseconds the item spent waiting on invoker memory pool
     #[bilrost(tag(4))]
     pub blocked_on_invoker_memory_ms: u32,

crates/storage-api/src/vqueue_table/store.rs

@@ -30,7 +30,7 @@ pub trait VQueueEntry: PartialOrd + PartialEq + Eq + Clone + std::fmt::Debug {
     fn priority(&self) -> EffectivePriority;
     fn visible_at(&self) -> VisibleAt;
     fn kind(&self) -> EntryKind;
-    fn is_token_held(&self) -> bool;
+    fn has_lock(&self) -> bool;
     // Weight is as a proxy for _costing_ dequeueing an item. Lower weight means
     // lower cost and higher chances to be dequeued.
     fn weight(&self) -> NonZeroU16;
@@ -88,7 +88,8 @@ impl VQueueEntry for EntryCard {
     }
 
     #[inline(always)]
-    fn is_token_held(&self) -> bool {
+    fn has_lock(&self) -> bool {
+        // todo: change to use has_lock directly
         self.priority.token_held()
     }
 }

crates/storage-api/src/vqueue_table/tables.rs

@@ -59,6 +59,9 @@ pub trait WriteVQueueTable {
     /// successful or not.
     fn pop_inbox_entry(&mut self, qid: &VQueueId, stage: Stage, card: &EntryCard) -> Result<bool>;
 
+    /// Deletes entry from inbox unconditionally
+    fn delete_inbox_entry(&mut self, qid: &VQueueId, stage: Stage, card: &EntryCard);
+
     /// Adds a vqueue to the list of active vqueues
     ///
     /// A vqueue is considered active when it's of interest to the scheduler.

crates/vqueues/src/cache.rs

@@ -10,15 +10,14 @@
 
 use hashbrown::{HashMap, hash_map};
 use slotmap::SlotMap;
+use tokio::task::JoinHandle;
 use tracing::{debug, trace};
 
 use restate_storage_api::StorageError;
 use restate_storage_api::vqueue_table::metadata::{self, VQueueMeta};
 use restate_storage_api::vqueue_table::{ReadVQueueTable, ScanVQueueTable};
 use restate_types::vqueue::VQueueId;
 
-use crate::vqueue_config::ConfigPool;
-
 type Result<T> = std::result::Result<T, StorageError>;
 
 slotmap::new_key_type! { pub struct VQueueCacheKey; }
@@ -35,10 +34,6 @@ impl<'a> VQueuesMeta<'a> {
         Self { inner: cache }
     }
 
-    pub(crate) fn config_pool(&'a self) -> &'a ConfigPool {
-        &self.inner.config
-    }
-
     pub fn get(&self, key: VQueueCacheKey) -> Option<&Slot> {
         self.inner.slab.get(key)
     }
@@ -102,7 +97,6 @@ impl Slot {
 // Needs rewriting after the workload pattern becomes more clear.
 #[derive(Clone)]
 pub struct VQueuesMetaCache {
-    config: ConfigPool,
     queues: HashMap<VQueueId, VQueueCacheKey>,
     slab: SlotMap<VQueueCacheKey, Slot>,
 }
@@ -121,31 +115,36 @@ impl VQueuesMetaCache {
         Self {
             slab: Default::default(),
             queues: Default::default(),
-            config: ConfigPool::default(),
         }
     }
 
     /// Initializes the vqueue cache by loading all active vqueues into the cache.
     ///
     /// From this point on, the cache remains in-sync with the storage state by
     /// using the "apply_updates" method.
-    pub async fn create<S: ScanVQueueTable>(storage: &S) -> Result<Self> {
-        let mut cache = Self {
-            slab: Default::default(),
-            queues: Default::default(),
-            config: ConfigPool::default(),
-        };
-        // find and load all active vqueues.
-        storage.scan_active_vqueues(|qid, meta| {
-            let key = cache.slab.insert(Slot {
-                qid: qid.clone(),
-                meta,
-            });
-            // SAFETY: at batch load time we are guaranteed to observe every vqueue id only once.
-            unsafe { cache.queues.insert_unique_unchecked(qid, key) };
-        })?;
-
-        Ok(cache)
+    pub async fn create<S: ScanVQueueTable + Send + Sync + 'static>(storage: S) -> Result<Self> {
+        let handle: JoinHandle<Result<_>> = tokio::task::spawn_blocking({
+            move || {
+                let mut slab = SlotMap::default();
+                let mut queues = HashMap::default();
+                // find and load all active vqueues.
+                storage.scan_active_vqueues(|qid, meta| {
+                    let key = slab.insert(Slot {
+                        qid: qid.clone(),
+                        meta,
+                    });
+                    // SAFETY: at batch load time we are guaranteed to observe every vqueue id only once.
+                    unsafe { queues.insert_unique_unchecked(qid, key) };
+                })?;
+                Ok((slab, queues))
+            }
+        });
+
+        let (slab, queues) = handle
+            .await
+            .map_err(|e| StorageError::Generic(e.into()))??;
+
+        Ok(Self { slab, queues })
     }
 
     pub async fn load<S: ReadVQueueTable>(