[VQueues] Move scheduler status types to worker-api

This is to allow datafusion storage engine to access these types without direct dependency on restate-vqueues.

Author: AhmedSoliman

Cargo.lock

@@ -23,6 +23,7 @@ restate-serde-util = { workspace = true }
 restate-storage-api = { workspace = true }
 restate-types = { workspace = true }
 restate-util-string = { workspace = true }
+restate-worker-api = { workspace = true }
 
 arrayvec = { workspace = true }
 bilrost = { workspace = true, features = ["smallvec"] }

crates/vqueues/Cargo.toml

@@ -16,9 +16,10 @@ mod util;
 // Re-exports
 pub use cache::{VQueuesMeta, VQueuesMetaCache};
 pub use metric_definitions::describe_metrics;
-pub use scheduler::{
-    ResourceManager, SchedulerService, SchedulingStatus, ThrottleScope, VQueueSchedulerStatus,
+pub use restate_worker_api::{
+    ResourceKind, SchedulingStatus, ThrottleScope, VQueueSchedulerStatus,
 };
+pub use scheduler::{ResourceManager, SchedulerService};
 use smallvec::SmallVec;
 use tracing::debug;
 pub use util::*;

crates/vqueues/src/lib.rs

@@ -15,19 +15,18 @@ use std::pin::Pin;
 use std::future::poll_fn;
 use std::task::Poll;
 
-use restate_clock::RoughTimestamp;
 use restate_storage_api::StorageError;
 use restate_storage_api::vqueue_table::scheduler::{RunAction, SchedulerAction, YieldAction};
-use restate_storage_api::vqueue_table::{EntryKey, ScanVQueueTable, VQueueStore, stats::WaitStats};
-use restate_types::time::MillisSinceEpoch;
+use restate_storage_api::vqueue_table::{EntryKey, ScanVQueueTable, VQueueStore};
 use restate_types::vqueues::VQueueId;
+use restate_worker_api::{SchedulingStatus, VQueueSchedulerStatus};
 
 use crate::VQueueEvent;
 use crate::VQueuesMetaCache;
 use crate::metric_definitions::publish_scheduler_decision_metrics;
 
 use self::drr::DRRScheduler;
-use self::resource_manager::{PermitBuilder, ReservedResources, ResourceKind};
+use self::resource_manager::{PermitBuilder, ReservedResources};
 use self::vqueue_state::DetailedEligibility;
 
 mod clock;
@@ -46,79 +45,16 @@ type UnconfirmedAssignments = hashbrown::HashMap<EntryKey, PermitBuilder>;
 
 slotmap::new_key_type! { pub(crate) struct VQueueHandle; }
 
-/// A public view of the scheduler's status of a single vqueue.
-///
-/// This struct provides introspection into the current scheduling state, and
-/// wait statistics for a vqueue.
-#[derive(Debug, Clone, Default)]
-pub struct VQueueSchedulerStatus {
-    /// Statistics about the wait time experienced by the head item in the vqueue.
-    pub wait_stats: WaitStats,
-    /// Number of items remaining in the running stage.
-    pub remaining_running: u32,
-    /// Number of items waiting in the inbox stage.
-    pub waiting_inbox: u64,
-    /// The current scheduling status of this vqueue.
-    pub status: SchedulingStatus,
-}
-
-/// The current scheduling status of a vqueue.
-///
-/// This enum represents the various states a vqueue can be in from the
-/// scheduler's perspective.
-#[derive(Debug, Clone, Default, PartialEq, Eq)]
-pub enum SchedulingStatus {
-    #[default]
-    /// The vqueue is not tracked by the scheduler (e.g., it has no items).
-    Dormant,
-    /// The vqueue is empty.
-    Empty,
-    /// The vqueue head is ready to be scheduled and it's in the inbox/running stage.
-    Ready,
-    /// The vqueue is throttled until the specified time.
-    Throttled {
-        /// When the throttle expires.
-        until: MillisSinceEpoch,
-        /// The scope of throttling (global or per-vqueue).
-        scope: ThrottleScope,
-    },
-    /// The vqueue is scheduled to be woken up at the given time because the head
-    /// item is scheduled to run at that time.
-    Scheduled {
-        /// When the head item becomes visible.
-        at: RoughTimestamp,
-    },
-    /// The vqueue is blocked on invoker global capacity.
-    BlockedOn(ResourceKind),
-    /// The vqueue is waiting to acquire a lock of a VO.
-    BlockedOnLock,
-    /// The vqueue is waiting for concurrency tokens. Concurrency tokens are released
-    /// when currently running items are completed or (in some cases) when running items
-    /// are parked.
-    WaitingConcurrencyTokens,
-}
-
-impl From<DetailedEligibility> for SchedulingStatus {
-    fn from(value: DetailedEligibility) -> Self {
-        match value {
-            DetailedEligibility::EligibleRunning | DetailedEligibility::EligibleInbox => {
-                SchedulingStatus::Ready
-            }
-            DetailedEligibility::Scheduled(ts) => SchedulingStatus::Scheduled { at: ts },
-            DetailedEligibility::Empty => SchedulingStatus::Empty,
+fn status_from_detailed_eligibility(value: DetailedEligibility) -> SchedulingStatus {
+    match value {
+        DetailedEligibility::EligibleRunning | DetailedEligibility::EligibleInbox => {
+            SchedulingStatus::Ready
         }
+        DetailedEligibility::Scheduled(ts) => SchedulingStatus::Scheduled { at: ts },
+        DetailedEligibility::Empty => SchedulingStatus::Empty,
     }
 }
 
-/// The scope at which throttling is applied.
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
-pub enum ThrottleScope {
-    /// Throttling is applied globally across all vqueues.
-    Global,
-    /// Throttling is applied to a specific vqueue.
-    VQueue,
-}
-
 #[derive(Default, Debug)]
 pub struct Decisions {
     // Vqueue ids are ordered by partition key, this enables us to group

crates/vqueues/src/scheduler.rs

@@ -458,9 +458,9 @@ mod tests {
     use restate_types::sharding::KeyRange;
     use restate_types::vqueues::VQueueId;
     use restate_types::vqueues::{EntryId, EntryKind};
+    use restate_worker_api::ResourceKind;
 
     use crate::cache::VQueuesMetaCache;
-    use crate::scheduler::resource_manager::ResourceKind;
     use crate::{SchedulingStatus, VQueue, VQueueEvent};
 
     use super::*;

crates/vqueues/src/scheduler/drr.rs

@@ -20,11 +20,11 @@ use tracing::trace;
 use restate_storage_api::StorageError;
 use restate_storage_api::vqueue_table::VQueueStore;
 use restate_types::time::MillisSinceEpoch;
+use restate_worker_api::{ResourceKind, SchedulingStatus, ThrottleScope};
 
+use super::VQueueHandle;
 use super::clock::SchedulerClock;
-use super::resource_manager::ResourceKind;
 use super::vqueue_state::VQueueState;
-use super::{SchedulingStatus, ThrottleScope, VQueueHandle};
 use crate::scheduler::vqueue_state::Eligibility;
 
 #[derive(Debug, Copy, Clone)]
@@ -80,7 +80,7 @@ impl EligibilityTracker {
     pub fn get_status<S: VQueueStore>(&self, qstate: &VQueueState<S>) -> SchedulingStatus {
         match self.states.get(qstate.handle) {
             None | Some(State::NeedsPoll) | Some(State::Ready) => {
-                SchedulingStatus::from(qstate.check_eligibility())
+                super::status_from_detailed_eligibility(qstate.check_eligibility())
             }
             Some(State::Throttled { wake_up, scope }) => SchedulingStatus::Throttled {
                 until: wake_up.ts,

crates/vqueues/src/scheduler/eligible.rs

@@ -25,14 +25,13 @@ use tokio::sync::mpsc;
 use tracing::trace;
 
 use restate_futures_util::concurrency::Concurrency;
-use restate_limiter::{Level, LimitKey, RuleHandle};
 use restate_memory::{MemoryPool, NonZeroByteCount};
 use restate_storage_api::StorageError;
 use restate_storage_api::lock_table::LoadLocks;
 use restate_storage_api::vqueue_table::{EntryKey, EntryMetadata};
 use restate_types::vqueues::EntryKind;
 use restate_types::{LockName, Scope};
-use restate_util_string::ReString;
+use restate_worker_api::ResourceKind;
 
 use self::invoker::InvokerConcurrencyLimiter;
 use self::invoker_memory::InvokerMemoryLimiter;
@@ -62,33 +61,6 @@ pub struct ResourceManager {
     tx: mpsc::UnboundedSender<ResourceManagerUpdate>,
 }
 
-#[derive(Debug, Clone, PartialEq, Eq)]
-pub enum ResourceKind {
-    /// Waiting to acquire a lock of a VO.
-    Lock {
-        scope: Option<Scope>,
-        lock_name: LockName,
-    },
-    /// Waiting to acquire invoker concurrency capacity
-    InvokerConcurrency,
-    /// Waiting to acquire invoker being throttled
-    InvokerThrottling,
-    /// Invoker needs to allocate memory for an invocation
-    InvokerMemory,
-    /// Waiting for deployment-level concurrency tokens to be available
-    DeploymentConcurrency,
-    /// Waiting for user-defined concurrency to be acquired.
-    /// Carries routing info so the eligibility tracker can return it for waiter removal.
-    LimitKeyConcurrency {
-        scope: Scope,
-        limit_key: LimitKey<ReString>,
-        blocked_level: Level,
-        /// Handle to the blocking rule. Resolve via the rules store for display.
-        /// May be stale if the rule was removed since blocking.
-        blocked_rule: Option<RuleHandle>,
-    },
-}
-
 #[allow(dead_code)]
 enum ResourceManagerUpdate {
     PermitReleased(SmallVec<[UserPermitKind; 1]>),

crates/vqueues/src/scheduler/resource_manager.rs

@@ -18,6 +18,7 @@ use restate_clock::RoughTimestamp;
 use restate_storage_api::StorageError;
 use restate_storage_api::vqueue_table::{EntryKey, EntryValue, VQueueStore, stats::WaitStats};
 use restate_types::vqueues::VQueueId;
+use restate_worker_api::{ResourceKind, ThrottleScope};
 
 use crate::metric_definitions::{
     VQUEUE_CONCURRENCY_RULES_WAIT_MS, VQUEUE_DEPLOYMENT_CONCURRENCY_WAIT_MS,
@@ -30,10 +31,8 @@ use super::clock::SchedulerClock;
 use super::queue::Queue;
 use super::queue_meta::MetaLiteUpdate;
 pub use super::queue_meta::VQueueMetaLite;
-use super::resource_manager::{AcquireOutcome, PermitBuilder, ResourceKind};
-use super::{
-    ResourceManager, RunAction, ThrottleScope, UnconfirmedAssignments, VQueueHandle, YieldAction,
-};
+use super::resource_manager::{AcquireOutcome, PermitBuilder};
+use super::{ResourceManager, RunAction, UnconfirmedAssignments, VQueueHandle, YieldAction};
 
 const QUANTUM: i32 = 1;
 

crates/vqueues/src/scheduler/vqueue_state.rs

@@ -10,8 +10,12 @@ publish = false
 [dependencies]
 restate-workspace-hack = { workspace = true }
 
+restate-clock = { workspace = true }
 restate-core = { workspace = true }
+restate-limiter = { workspace = true }
+restate-storage-api = { workspace = true }
 restate-types = { workspace = true }
+restate-util-string = { workspace = true }
 
 thiserror = { workspace = true }
 tokio = { workspace = true }

crates/worker-api/Cargo.toml

@@ -10,6 +10,8 @@
 
 mod partition_processor_manager;
 mod partition_processor_rpc_client;
+mod scheduler_status;
 
 pub use partition_processor_manager::*;
 pub use partition_processor_rpc_client::*;
+pub use scheduler_status::*;