[VQueues] Introduce sys_vqueue_stats datafusion table

Author: AhmedSoliman

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

@@ -18,20 +18,23 @@ mod running_reader;
 mod waiting_reader;
 
 use std::io::Cursor;
+use std::ops::RangeInclusive;
 
 pub use entry::{EntryStatusKey, StatusHeaderRaw};
 pub use inbox::InboxKey;
 pub use items::InputPayloadKey;
 pub use metadata::*;
 
 use anyhow::Context;
-use bilrost::{Message, OwnedMessage};
+use bilrost::{BorrowedMessage, Message, OwnedMessage};
 use bytes::BytesMut;
 use rocksdb::{DBRawIteratorWithThreadMode, ReadOptions};
 use tracing::error;
 
+use restate_rocksdb::Priority;
 use restate_storage_api::StorageError;
-use restate_storage_api::vqueue_table::metadata::{VQueueMeta, VQueueMetaUpdates};
+use restate_storage_api::vqueue_table::ScanVQueueMetaTable;
+use restate_storage_api::vqueue_table::metadata::{VQueueMeta, VQueueMetaRef, VQueueMetaUpdates};
 use restate_storage_api::vqueue_table::{
     EntryKey, EntryMetadata, EntryStatusHeader, EntryValue, LazyEntryStatus, ReadVQueueTable,
     ScanVQueueTable, Stage, Status, WriteVQueueTable, stats::EntryStatistics,
@@ -41,7 +44,11 @@ use restate_types::vqueues::{EntryId, Seq, VQueueId};
 
 use self::entry::{LazyEntryStatusHolder, OwnedEntryStatusHeader, StatusHeaderRawRef};
 use crate::keys::{DecodeTableKey, EncodeTableKey, EncodeTableKeyPrefix, KeyKind};
-use crate::{PartitionDb, PartitionStoreTransaction, Result, StorageAccess, TableKind};
+use crate::scan::TableScan;
+use crate::{
+    PartitionDb, PartitionStore, PartitionStoreTransaction, Result, StorageAccess, TableKind,
+    break_on_err,
+};
 
 impl ScanVQueueTable for PartitionDb {
     fn scan_active_vqueues(
@@ -438,6 +445,35 @@ impl ReadVQueueTable for PartitionStoreTransaction<'_> {
     }
 }
 
+impl ScanVQueueMetaTable for PartitionStore {
+    fn for_each_vqueue_meta<
+        F: for<'a> FnMut((&'a VQueueId, &'a VQueueMetaRef<'a>)) -> std::ops::ControlFlow<()>
+            + Send
+            + Sync
+            + 'static,
+    >(
+        &self,
+        range: RangeInclusive<PartitionKey>,
+        mut f: F,
+    ) -> Result<impl Future<Output = Result<()>> + Send> {
+        self.iterator_for_each(
+            "df-vqueue-meta",
+            Priority::Low,
+            TableScan::FullScanPartitionKeyRange::<MetaKey>(range),
+            move |(mut key, value)| {
+                let meta_key = break_on_err(MetaKey::deserialize_from(&mut key))?;
+                let meta = break_on_err(
+                    VQueueMetaRef::decode_borrowed(value).map_err(StorageError::BilrostDecode),
+                )?;
+
+                let (vqueue_id,) = meta_key.split();
+                f((&vqueue_id, &meta)).map_break(Ok)
+            },
+        )
+        .map_err(|_| StorageError::OperationalError)
+    }
+}
+
 // ## Safety
 // The iterator is guaranteed to be dropped before the database is dropped, we hold to the
 // PartitionDb in this struct for as long as the iterator is alive.

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

@@ -8,11 +8,13 @@
 // the Business Source License, use of this software will be governed
 // by the Apache License, Version 2.0.
 
+use std::ops::RangeInclusive;
+
 use restate_types::identifiers::PartitionKey;
 use restate_types::vqueues::{Seq, VQueueId};
 
 use super::Status;
-use super::metadata::VQueueMeta;
+use super::metadata::{VQueueMeta, VQueueMetaRef};
 use super::{
     EntryId, EntryKey, EntryMetadata, EntryStatusHeader, EntryValue, LazyEntryStatus,
     stats::EntryStatistics,
@@ -207,3 +209,17 @@ pub trait ScanVQueueTable {
         on_item: impl FnMut(VQueueId, super::metadata::VQueueMeta),
     ) -> Result<()>;
 }
+
+pub trait ScanVQueueMetaTable {
+    /// Used for data-fusion queries
+    fn for_each_vqueue_meta<
+        F: for<'a> FnMut((&'a VQueueId, &'a VQueueMetaRef<'a>)) -> std::ops::ControlFlow<()>
+            + Send
+            + Sync
+            + 'static,
+    >(
+        &self,
+        range: RangeInclusive<PartitionKey>,
+        f: F,
+    ) -> Result<impl Future<Output = Result<()>> + Send>;
+}

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

@@ -238,6 +238,13 @@ where
             self.partition_store_manager.clone(),
             &self.remote_scanner_manager,
         )?;
+        // VQueues Tables
+        crate::vqueue_stats::register_self(
+            ctx,
+            self.partition_selector.clone(),
+            self.partition_store_manager.clone(),
+            &self.remote_scanner_manager,
+        )?;
 
         ctx.datafusion_context.sql(SYS_INVOCATION_VIEW).await?;
 

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

@@ -39,6 +39,7 @@ mod statistics;
 pub mod table_docs;
 mod table_macro;
 mod table_providers;
+mod vqueue_stats;
 pub use table_providers::Scan;
 pub mod table_util;
 

crates/storage-query-datafusion/src/vqueue_stats/mod.rs

@@ -0,0 +1,15 @@
+// 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.
+
+mod row;
+pub(crate) mod schema;
+mod table;
+
+pub(crate) use table::register_self;

crates/storage-query-datafusion/src/vqueue_stats/row.rs

@@ -0,0 +1,124 @@
+// 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.
+
+use restate_storage_api::vqueue_table::metadata::VQueueMetaRef;
+use restate_types::vqueues::VQueueId;
+
+use super::schema::SysVqueueStatsBuilder;
+
+#[inline]
+pub(crate) fn append_vqueues_meta_row<'a>(
+    builder: &mut SysVqueueStatsBuilder,
+    qid: &'a VQueueId,
+    meta: &'a VQueueMetaRef<'a>,
+) {
+    let mut row = builder.row();
+
+    row.partition_key(qid.partition_key());
+    if row.is_id_defined() {
+        row.fmt_id(qid);
+    }
+    if row.is_scope_defined()
+        && let Some(scope) = meta.scope
+    {
+        row.scope(scope);
+    }
+
+    if row.is_service_name_defined()
+        && let Some(service_name) = meta.service_name()
+    {
+        row.service_name(service_name);
+    }
+
+    if row.is_is_active_defined() {
+        row.is_active(meta.is_active());
+    }
+    if row.is_queue_is_paused_defined() {
+        row.queue_is_paused(meta.queue_is_paused);
+    }
+
+    if row.is_limit_key_defined() {
+        row.fmt_limit_key(&meta.limit_key);
+    }
+
+    if row.is_lock_name_defined()
+        && let Some(lock_name) = meta.lock_name()
+    {
+        row.fmt_lock_name(lock_name);
+    }
+
+    if row.is_created_at_defined() {
+        row.created_at(meta.stats.created_at().as_u64() as i64);
+    }
+
+    if row.is_last_enqueued_at_defined()
+        && let Some(last_enqueued_at) = meta.stats.last_enqueued_at()
+    {
+        row.last_enqueued_at(last_enqueued_at.as_u64() as i64);
+    }
+
+    if row.is_last_start_at_defined()
+        && let Some(last_start_at) = meta.stats.last_start_at()
+    {
+        row.last_start_at(last_start_at.as_u64() as i64);
+    }
+
+    if row.is_last_attempt_at_defined()
+        && let Some(last_attempt_at) = meta.stats.last_attempt_at()
+    {
+        row.last_attempt_at(last_attempt_at.as_u64() as i64);
+    }
+
+    if row.is_last_finish_at_defined()
+        && let Some(last_finish_at) = meta.stats.last_finish_at()
+    {
+        row.last_finish_at(last_finish_at.as_u64() as i64);
+    }
+
+    if row.is_avg_queue_duration_defined() {
+        row.avg_queue_duration(meta.stats.avg_queue_duration_ms() as i64);
+    }
+
+    if row.is_avg_inbox_duration_defined() {
+        row.avg_inbox_duration(meta.stats.avg_inbox_duration_ms() as i64);
+    }
+
+    if row.is_avg_run_duration_defined() {
+        row.avg_run_duration(meta.stats.avg_run_duration_ms() as i64);
+    }
+
+    if row.is_avg_suspension_duration_defined() {
+        row.avg_suspension_duration(meta.stats.avg_suspension_duration_ms() as i64);
+    }
+
+    if row.is_avg_end_to_end_duration_defined() {
+        row.avg_end_to_end_duration(meta.stats.avg_end_to_end_duration_ms() as i64);
+    }
+
+    if row.is_num_inbox_defined() {
+        row.num_inbox(meta.stats.num_inbox());
+    }
+
+    if row.is_num_running_defined() {
+        row.num_running(meta.stats.num_running());
+    }
+
+    if row.is_num_suspended_defined() {
+        row.num_suspended(meta.stats.num_suspended());
+    }
+
+    if row.is_num_paused_defined() {
+        row.num_paused(meta.stats.num_paused());
+    }
+
+    if row.is_num_finished_defined() {
+        row.num_finished(meta.stats.num_finished());
+    }
+}

crates/storage-query-datafusion/src/vqueue_stats/schema.rs

@@ -0,0 +1,96 @@
+// 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.
+
+use crate::table_macro::*;
+
+use datafusion::arrow::datatypes::DataType;
+
+define_sort_order!(sys_vqueue_stats(partition_key, scope, lock_name));
+
+define_table!(sys_vqueue_stats(
+    /// Internal column that is used for partitioning. Can be ignored.
+    partition_key: DataType::UInt64,
+
+    /// The VQueue Identifier (vq_...)
+    id: DataType::Utf8,
+
+    /// Whether this vqueue is active or not. An active vqueue
+    /// is a vqueue that is not paused, and has non-finished
+    /// items.
+    is_active: DataType::Boolean,
+
+    /// Whether this vqueue is paused.
+    queue_is_paused: DataType::Boolean,
+
+    /// Service name linked to this vqueue
+    service_name: DataType::Utf8,
+
+    /// The scope of this vqueue.
+    scope: DataType::Utf8,
+
+    /// The name of the limit-key assigned to this vqueue
+    limit_key: DataType::Utf8,
+
+    /// The name of the lock (in the format of `service/key`)
+    ///
+    /// This is only set if this is a vqueue for a virtual object.
+    lock_name: DataType::Utf8,
+
+    /// When was this vqueue first created
+    created_at: TimestampMillisecond,
+
+    /// Last timestamp an entry moved into Inbox.
+    ///
+    /// This covers items enqueued for the first time only.
+    last_enqueued_at: TimestampMillisecond,
+
+    /// Last timestamp an entry first transitioned to Run.
+    last_start_at: TimestampMillisecond,
+
+    /// Last timestamp an entry transitioned to Run.
+    last_attempt_at: TimestampMillisecond,
+
+    /// Last timestamp an entry transitioned to Finished.
+    last_finish_at: TimestampMillisecond,
+
+    /// Exponential moving average (EMA) of first-attempt queue wait time.
+    avg_queue_duration: DataType::Duration,
+
+    /// Exponential moving average (EMA) of how long entries stayed in inbox.
+    avg_inbox_duration: DataType::Duration,
+
+    /// Exponential moving average (EMA) of how long entries stayed running.
+    avg_run_duration: DataType::Duration,
+
+    /// Exponential moving average (EMA) of how long entries stayed suspended.
+    avg_suspension_duration: DataType::Duration,
+
+    /// Exponential moving average (EMA) of end-to-end entry lifetime from first-runnable time to completion.
+    /// Note that this only tracks entries that were not killed/cancelled or failed/paused.
+    avg_end_to_end_duration: DataType::Duration,
+
+    /// The number of entries that are in the inbox. The inbox is the priority
+    /// queue that the scheduler uses to choose which entries to run next.
+    num_inbox: DataType::UInt64,
+
+    /// The number of entries that are currently running.
+    num_running: DataType::UInt64,
+
+    /// The number of entries that are suspended.
+    num_suspended: DataType::UInt64,
+
+    /// The number of entries that are paused.
+    num_paused: DataType::UInt64,
+
+    /// The number of entries that have finished processing and are pending
+    /// deletion or archival.
+    num_finished: DataType::UInt64,
+
+));