[VQueues] Provisions to add user limits to datafusion

Author: AhmedSoliman

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

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

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

@@ -0,0 +1,23 @@
+// 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)]
+mod row;
+#[allow(dead_code)]
+pub(crate) mod schema;
+
+// TODO: table.rs with register_self() — requires plumbing a handle from the
+// scheduler's UserLimiter to the DataFusion context, similar to the StatusHandle
+// pattern used by invocation_state/table.rs.
+//
+// The scan should:
+// 1. Call UserLimiter::scan_counters() to get CounterView entries
+// 2. For each entry, resolve the rule handle via UserLimiter::resolve_rule()
+// 3. Build UserLimitRow and append to the table builder

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

@@ -0,0 +1,75 @@
+// 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_types::identifiers::PartitionKey;
+
+use super::schema::SysUserLimitsBuilder;
+
+/// A single row of enriched counter data ready to be appended to the table.
+///
+/// This is produced by the scanner from a `CounterView` + rule resolution.
+#[allow(dead_code)]
+pub(crate) struct UserLimitRow {
+    pub partition_key: PartitionKey,
+    pub scope: String,
+    pub l1: Option<String>,
+    pub l2: Option<String>,
+    pub level: &'static str,
+    pub usage: u32,
+    pub concurrency_limit: Option<u64>,
+    pub rule_pattern: Option<String>,
+    pub num_waiters: u64,
+}
+
+#[inline]
+#[allow(dead_code)]
+pub(crate) fn append_user_limit_row(builder: &mut SysUserLimitsBuilder, row_data: &UserLimitRow) {
+    let mut row = builder.row();
+
+    row.partition_key(row_data.partition_key);
+
+    if row.is_scope_defined() {
+        row.scope(&row_data.scope);
+    }
+    if row.is_l1_defined()
+        && let Some(l1) = &row_data.l1
+    {
+        row.l1(l1);
+    }
+    if row.is_l2_defined()
+        && let Some(l2) = &row_data.l2
+    {
+        row.l2(l2);
+    }
+    if row.is_level_defined() {
+        row.level(row_data.level);
+    }
+    if row.is_usage_defined() {
+        row.usage(row_data.usage);
+    }
+    if row.is_concurrency_limit_defined()
+        && let Some(limit) = row_data.concurrency_limit
+    {
+        row.concurrency_limit(limit);
+    }
+    if row.is_rule_pattern_defined()
+        && let Some(pattern) = &row_data.rule_pattern
+    {
+        row.rule_pattern(pattern);
+    }
+    if row.is_available_defined()
+        && let Some(limit) = row_data.concurrency_limit
+    {
+        row.available(limit.saturating_sub(row_data.usage as u64));
+    }
+    if row.is_num_waiters_defined() {
+        row.num_waiters(row_data.num_waiters);
+    }
+}

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

@@ -0,0 +1,49 @@
+// 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_user_limits(partition_key, scope, l1, l2));
+
+define_table!(sys_user_limits(
+    /// Internal column that is used for partitioning. Can be ignored.
+    partition_key: DataType::UInt64,
+
+    /// The scope this counter belongs to.
+    scope: DataType::Utf8,
+
+    /// The level-1 key component (null for scope-level counters).
+    l1: DataType::Utf8,
+
+    /// The level-2 key component (null for scope-level and L1-level counters).
+    l2: DataType::Utf8,
+
+    /// The hierarchy level: "Scope", "Level1", or "Level2".
+    level: DataType::Utf8,
+
+    /// Current concurrency usage at this counter.
+    usage: DataType::UInt32,
+
+    /// The configured concurrency limit (null if unlimited).
+    concurrency_limit: DataType::UInt64,
+
+    /// The rule pattern that defines the limit (null if unlimited).
+    /// Resolved from the rule handle; shows "[removed]" if the rule was
+    /// deleted since the counter was created.
+    rule_pattern: DataType::Utf8,
+
+    /// Available capacity (limit - usage). Null if unlimited.
+    available: DataType::UInt64,
+
+    /// Number of vqueues currently waiting behind this counter.
+    num_waiters: DataType::UInt64,
+));