[DataFusion] Unify partition leader status and leader query protocol

Author: AhmedSoliman

Cargo.lock

@@ -26,6 +26,18 @@ use crate::PartitionKey;
 #[repr(transparent)]
 pub struct KeyRange(std::range::RangeInclusive<PartitionKey>);
 
+impl PartialOrd for KeyRange {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl Ord for KeyRange {
+    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+        (self.start(), self.end()).cmp(&(other.start(), other.end()))
+    }
+}
+
 impl NetSerde for KeyRange {}
 
 impl KeyRange {

crates/sharding/src/key_range.rs

@@ -23,6 +23,7 @@ restate-service-protocol-v4 = { workspace = true, features = ["entry-codec"]  }
 restate-sharding = { workspace = true }
 restate-storage-api = { workspace = true }
 restate-types = { workspace = true }
+restate-worker-api = { workspace = true }
 
 ahash = { workspace = true }                                                    # Required to due a yanked version used by datafusion
 anyhow = { workspace = true }

crates/storage-query-datafusion/Cargo.toml

@@ -30,6 +30,7 @@ use codederror::CodedError;
 use restate_core::{Metadata, TaskCenter};
 use restate_invoker_api::StatusHandle;
 use restate_partition_store::PartitionStoreManager;
+use restate_sharding::KeyRange;
 use restate_types::cluster::cluster_state::LegacyClusterState;
 use restate_types::config::QueryEngineOptions;
 use restate_types::errors::GenericError;
@@ -39,6 +40,7 @@ use restate_types::partition_table::Partition;
 use restate_types::partitions::state::PartitionReplicaSetStates;
 use restate_types::schema::deployment::DeploymentResolver;
 use restate_types::schema::service::ServiceMetadataResolver;
+use restate_worker_api::SchedulerStatusEntry;
 
 use crate::node_fan_out::NodeWarnings;
 use crate::remote_query_scanner_manager::RemoteScannerManager;
@@ -130,6 +132,28 @@ pub trait RegisterTable: Send + Sync + 'static {
     fn register(&self, ctx: &QueryContext) -> impl Future<Output = Result<(), BuildError>>;
 }
 
+/// A leader-state introspection handle that extends invoker status queries with
+/// additional query methods for future leader-owned components.
+pub trait PartitionLeaderStatusHandle:
+    StatusHandle + Send + Sync + Debug + Clone + 'static
+{
+    type SchedulerStatus;
+    type SchedulerStatusIterator: Iterator<Item = Self::SchedulerStatus> + Send;
+
+    type UserLimitCounter;
+    type UserLimitCounterIterator: Iterator<Item = Self::UserLimitCounter> + Send;
+
+    fn read_scheduler_status(
+        &self,
+        keys: KeyRange,
+    ) -> impl Future<Output = Self::SchedulerStatusIterator> + Send;
+
+    fn read_user_limit_counters(
+        &self,
+        keys: KeyRange,
+    ) -> impl Future<Output = Self::UserLimitCounterIterator> + Send;
+}
+
 /// A no-op registerer that creates a minimal query context with no tables.
 /// Useful for nodes that only need to serve remote scanner RPCs (e.g.,
 /// log-server-only nodes), where only `task_ctx()` is needed.
@@ -145,7 +169,7 @@ impl RegisterTable for NoTables {
 pub struct UserTables<P, S, D> {
     partition_selector: P,
     partition_store_manager: Arc<PartitionStoreManager>,
-    status: Option<S>,
+    partition_leader_status: Option<S>,
     schemas: Live<D>,
     remote_scanner_manager: RemoteScannerManager,
 }
@@ -154,14 +178,14 @@ impl<P, S, D> UserTables<P, S, D> {
     pub fn new(
         partition_selector: P,
         partition_store_manager: Arc<PartitionStoreManager>,
-        status: Option<S>,
+        partition_leader_status: Option<S>,
         schemas: Live<D>,
         remote_scanner_manager: RemoteScannerManager,
     ) -> Self {
         Self {
             partition_selector,
             partition_store_manager,
-            status,
+            partition_leader_status,
             schemas,
             remote_scanner_manager,
         }
@@ -171,7 +195,7 @@ impl<P, S, D> UserTables<P, S, D> {
 impl<P, S, D> RegisterTable for UserTables<P, S, D>
 where
     P: SelectPartitions + Clone,
-    S: StatusHandle + Send + Sync + Debug + Clone + 'static,
+    S: PartitionLeaderStatusHandle<SchedulerStatus = SchedulerStatusEntry>,
     D: DeploymentResolver + ServiceMetadataResolver + Send + Sync + Debug + Clone + 'static,
 {
     async fn register(&self, ctx: &QueryContext) -> Result<(), BuildError> {
@@ -182,7 +206,7 @@ where
         crate::invocation_state::register_self(
             ctx,
             self.partition_selector.clone(),
-            self.status.clone(),
+            self.partition_leader_status.clone(),
             self.partition_store_manager.clone(),
             &self.remote_scanner_manager,
         )?;
@@ -350,7 +374,9 @@ impl QueryContext {
         options: &QueryEngineOptions,
         partition_selector: impl SelectPartitions + Clone,
         partition_store_manager: Arc<PartitionStoreManager>,
-        status: Option<impl StatusHandle + Send + Sync + Debug + Clone + 'static>,
+        partition_leader_status: Option<
+            impl PartitionLeaderStatusHandle<SchedulerStatus = SchedulerStatusEntry>,
+        >,
         schemas: Live<
             impl DeploymentResolver + ServiceMetadataResolver + Send + Sync + Debug + Clone + 'static,
         >,
@@ -359,7 +385,7 @@ impl QueryContext {
         let tables = UserTables::new(
             partition_selector,
             partition_store_manager,
-            status,
+            partition_leader_status,
             schemas,
             remote_scanner_manager,
         );

crates/storage-query-datafusion/src/context.rs

@@ -8,10 +8,14 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
+use std::future::Future;
 use std::{future, iter};
 
 use restate_invoker_api::{InvocationStatusReport, StatusHandle};
 use restate_types::sharding::KeyRange;
+use restate_worker_api::SchedulerStatusEntry;
+
+use crate::context::PartitionLeaderStatusHandle;
 
 #[derive(Clone, Debug)]
 pub struct EmptyInvokerStatusHandle;
@@ -23,3 +27,25 @@ impl StatusHandle for EmptyInvokerStatusHandle {
         future::ready(iter::empty())
     }
 }
+
+impl PartitionLeaderStatusHandle for EmptyInvokerStatusHandle {
+    type SchedulerStatus = SchedulerStatusEntry;
+    type SchedulerStatusIterator = std::iter::Empty<Self::SchedulerStatus>;
+
+    type UserLimitCounter = ();
+    type UserLimitCounterIterator = std::iter::Empty<Self::UserLimitCounter>;
+
+    fn read_scheduler_status(
+        &self,
+        _keys: KeyRange,
+    ) -> impl Future<Output = Self::SchedulerStatusIterator> + Send {
+        future::ready(iter::empty())
+    }
+
+    fn read_user_limit_counters(
+        &self,
+        _keys: KeyRange,
+    ) -> impl Future<Output = Self::UserLimitCounterIterator> + Send {
+        future::ready(iter::empty())
+    }
+}

crates/storage-query-datafusion/src/empty_invoker_status_handle.rs

@@ -9,6 +9,7 @@
 // by the Apache License, Version 2.0.
 
 use std::fmt::Debug;
+use std::future::Future;
 use std::marker::PhantomData;
 use std::sync::Arc;
 
@@ -20,7 +21,6 @@ use datafusion::common::DataFusionError;
 use googletest::matcher::{Matcher, MatcherResult};
 use serde_json::Value;
 
-use restate_invoker_api::StatusHandle;
 use restate_invoker_api::status_handle::test_util::MockStatusHandle;
 use restate_partition_store::{PartitionStore, PartitionStoreManager};
 use restate_rocksdb::RocksDbManager;
@@ -38,9 +38,10 @@ use restate_types::schema::deployment::{Deployment, DeploymentResolver};
 use restate_types::schema::service::test_util::MockServiceMetadataResolver;
 use restate_types::schema::service::{ServiceMetadata, ServiceMetadataResolver};
 use restate_types::sharding::KeyRange;
+use restate_worker_api::SchedulerStatusEntry;
 
 use super::context::QueryContext;
-use crate::context::SelectPartitions;
+use crate::context::{PartitionLeaderStatusHandle, SelectPartitions};
 use crate::remote_query_scanner_client::{RemoteScanner, RemoteScannerService};
 use crate::remote_query_scanner_manager::{
     PartitionLocation, PartitionLocator, RemoteScannerManager,
@@ -74,6 +75,28 @@ impl ServiceMetadataResolver for MockSchemas {
     }
 }
 
+impl PartitionLeaderStatusHandle for MockStatusHandle {
+    type SchedulerStatus = SchedulerStatusEntry;
+    type SchedulerStatusIterator = std::iter::Empty<Self::SchedulerStatus>;
+
+    type UserLimitCounter = ();
+    type UserLimitCounterIterator = std::iter::Empty<Self::UserLimitCounter>;
+
+    fn read_scheduler_status(
+        &self,
+        _keys: KeyRange,
+    ) -> impl Future<Output = Self::SchedulerStatusIterator> + Send {
+        std::future::ready(std::iter::empty())
+    }
+
+    fn read_user_limit_counters(
+        &self,
+        _keys: KeyRange,
+    ) -> impl Future<Output = Self::UserLimitCounterIterator> + Send {
+        std::future::ready(std::iter::empty())
+    }
+}
+
 impl DeploymentResolver for MockSchemas {
     fn resolve_latest_deployment_for_service(
         &self,
@@ -149,7 +172,7 @@ impl PartitionLocator for AlwaysLocalPartitionLocator {
 
 impl MockQueryEngine {
     pub async fn create_with(
-        status: impl StatusHandle + Send + Sync + Debug + Clone + 'static,
+        status: impl PartitionLeaderStatusHandle<SchedulerStatus = SchedulerStatusEntry>,
         schemas: impl DeploymentResolver
         + ServiceMetadataResolver
         + Send

crates/storage-query-datafusion/src/mocks.rs

@@ -12,6 +12,8 @@ restate-workspace-hack = { workspace = true }
 
 restate-clock = { workspace = true }
 restate-core = { workspace = true }
+restate-futures-util = { workspace = true }
+restate-invoker-api = { workspace = true }
 restate-limiter = { workspace = true }
 restate-storage-api = { workspace = true }
 restate-types = { workspace = true }

crates/worker-api/Cargo.toml

@@ -0,0 +1,61 @@
+// Copyright (c) 2023 - 2026 Restate Software, Inc., Restate GmbH.
+// All rights reserved.
+//
+// Use of this software is governed by the Business Source License
+// included in the LICENSE file.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0.
+
+#![allow(dead_code)]
+
+use tokio::sync::mpsc;
+
+use restate_futures_util::command::{Command, UnboundedCommandReceiver, UnboundedCommandSender};
+use restate_types::sharding::KeyRange;
+use restate_types::vqueues::VQueueId;
+
+use crate::VQueueSchedulerStatus;
+
+/// Queries that route through the partition processor's main `select!` loop.
+///
+/// Invoker status queries do NOT travel through this channel: DataFusion calls
+/// the invoker's `ChannelStatusReader` directly via the partition-leader handles
+/// registry. See `PartitionLeaderHandlesRegistry` in the worker crate.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum LeaderQueryKind {
+    SchedulerStatus,
+    UserLimitCounters,
+}
+
+#[derive(Debug, Clone)]
+pub enum LeaderQueryRequest {
+    SchedulerStatus { keys: KeyRange },
+    UserLimitCounters { keys: KeyRange },
+}
+
+impl LeaderQueryRequest {
+    pub fn kind(&self) -> LeaderQueryKind {
+        match self {
+            Self::SchedulerStatus { .. } => LeaderQueryKind::SchedulerStatus,
+            Self::UserLimitCounters { .. } => LeaderQueryKind::UserLimitCounters,
+        }
+    }
+}
+
+pub type SchedulerStatusEntry = (VQueueId, VQueueSchedulerStatus);
+
+#[derive(Debug, Clone)]
+pub enum LeaderQueryResponse {
+    SchedulerStatus(Vec<SchedulerStatusEntry>),
+    NotLeader(LeaderQueryKind),
+}
+
+pub type LeaderQueryCommand = Command<LeaderQueryRequest, LeaderQueryResponse>;
+pub type LeaderQuerySender = UnboundedCommandSender<LeaderQueryRequest, LeaderQueryResponse>;
+pub type LeaderQueryReceiver = UnboundedCommandReceiver<LeaderQueryRequest, LeaderQueryResponse>;
+
+pub fn channel() -> (LeaderQuerySender, LeaderQueryReceiver) {
+    mpsc::unbounded_channel()
+}

crates/worker-api/src/leader_query.rs

@@ -8,10 +8,12 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
+mod leader_query;
 mod partition_processor_manager;
 mod partition_processor_rpc_client;
 mod scheduler_status;
 
+pub use leader_query::*;
 pub use partition_processor_manager::*;
 pub use partition_processor_rpc_client::*;
 pub use scheduler_status::*;

crates/worker-api/src/lib.rs

@@ -37,6 +37,7 @@ restate-memory = { workspace = true }
 restate-metadata-server = { workspace = true }
 restate-metadata-store = { workspace = true }
 restate-partition-store = { workspace = true }
+restate-platform = { workspace = true }
 restate-rocksdb = { workspace = true }
 restate-service-protocol = { workspace = true, features = ["codec", "message"] }
 restate-service-protocol-v4 = { workspace = true, features = ["entry-codec"] }