Add user-limits and lock stats to WaitStats

Author: AhmedSoliman

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

@@ -27,6 +27,12 @@ pub struct WaitStats {
     /// Total milliseconds the item spent waiting on invoker memory pool
     #[bilrost(tag(4))]
     pub blocked_on_invoker_memory_ms: u32,
+    /// Total milliseconds the item spent waiting on user-defined concurrency limits
+    #[bilrost(tag(5))]
+    pub blocked_on_user_limit_ms: u32,
+    /// Total milliseconds the item spent waiting to acquire a virtual object lock
+    #[bilrost(tag(6))]
+    pub blocked_on_lock_ms: u32,
 }
 
 #[derive(Debug, Clone, bilrost::Message)]

crates/vqueues/src/metric_definitions.rs

@@ -23,6 +23,8 @@ pub const VQUEUE_GLOBAL_THROTTLE_WAIT_MS: &str =
     "restate.vqueue.scheduler.global_throttle_ms.total";
 // Per vqueue start throttling (affects starts only)
 pub const VQUEUE_LOCAL_THROTTLE_WAIT_MS: &str = "restate.vqueue.scheduler.vqueue_throttle_ms.total";
+pub const VQUEUE_USER_LIMIT_WAIT_MS: &str = "restate.vqueue.scheduler.user_limit_wait_ms.total";
+pub const VQUEUE_LOCK_WAIT_MS: &str = "restate.vqueue.scheduler.lock_wait_ms.total";
 
 pub const ACTION_YIELD: &str = "yield";
 pub const ACTION_RUN: &str = "run";
@@ -69,6 +71,18 @@ pub fn describe_metrics() {
         Unit::Count,
         "Cumulative number of seconds vqueues waited because of their self-imposed start throttling"
     );
+
+    describe_counter!(
+        VQUEUE_USER_LIMIT_WAIT_MS,
+        Unit::Count,
+        "Cumulative number of milliseconds spent waiting for user-defined concurrency limits"
+    );
+
+    describe_counter!(
+        VQUEUE_LOCK_WAIT_MS,
+        Unit::Count,
+        "Cumulative number of milliseconds spent waiting to acquire virtual object locks"
+    );
 }
 
 pub fn publish_scheduler_decision_metrics(num_run: u32, num_yield: u32) {

crates/vqueues/src/scheduler/vqueue_state.rs

@@ -20,7 +20,8 @@ use restate_types::vqueues::VQueueId;
 
 use crate::metric_definitions::{
     VQUEUE_GLOBAL_THROTTLE_WAIT_MS, VQUEUE_INVOKER_CONCURRENCY_WAIT_MS,
-    VQUEUE_INVOKER_MEMORY_WAIT_MS, VQUEUE_LOCAL_THROTTLE_WAIT_MS,
+    VQUEUE_INVOKER_MEMORY_WAIT_MS, VQUEUE_LOCAL_THROTTLE_WAIT_MS, VQUEUE_LOCK_WAIT_MS,
+    VQUEUE_USER_LIMIT_WAIT_MS,
 };
 use crate::scheduler::queue::QueueItem;
 
@@ -92,6 +93,10 @@ struct Stats {
     blocked_on_invoker_memory_micros: u32,
     local_start_throttling_micros: u32,
     global_throttling_micros: u32,
+    last_blocked_on_user_limit: Option<tokio::time::Instant>,
+    blocked_on_user_limit_micros: u32,
+    last_blocked_on_lock: Option<tokio::time::Instant>,
+    blocked_on_lock_micros: u32,
 }
 impl Stats {
     fn reset(&mut self) {
@@ -101,19 +106,26 @@ impl Stats {
         self.blocked_on_invoker_memory_micros = 0;
         self.local_start_throttling_micros = 0;
         self.global_throttling_micros = 0;
+        self.last_blocked_on_user_limit = None;
+        self.blocked_on_user_limit_micros = 0;
+        self.last_blocked_on_lock = None;
+        self.blocked_on_lock_micros = 0;
     }
 
     fn finalize(&mut self) -> WaitStats {
-        // ensures that the last capacity/memory-blocked segment is accounted for
+        // ensures that the last blocked segments are accounted for
         self.record_invoker_memory_delay(false);
-        // ensures that the last capacity-blocked segment is accounted for
         self.record_invoker_concurrency_delay(false);
+        self.record_user_limit_delay(false);
+        self.record_lock_delay(false);
 
         let stats = WaitStats {
             blocked_on_global_capacity_ms: self.blocked_on_invoker_concurrency_micros / 1000,
             vqueue_start_throttling_ms: self.local_start_throttling_micros / 1000,
             blocked_on_invoker_memory_ms: self.blocked_on_invoker_memory_micros / 1000,
             global_invoker_throttling_ms: self.global_throttling_micros / 1000,
+            blocked_on_user_limit_ms: self.blocked_on_user_limit_micros / 1000,
+            blocked_on_lock_ms: self.blocked_on_lock_micros / 1000,
         };
         self.reset();
 
@@ -139,11 +151,29 @@ impl Stats {
                 self.blocked_on_invoker_memory_micros
             };
 
+        let blocked_on_user_limit_micros = if let Some(last) = self.last_blocked_on_user_limit {
+            let delay = last.elapsed();
+            self.blocked_on_user_limit_micros
+                .saturating_add(delay.as_micros().try_into().unwrap_or(u32::MAX))
+        } else {
+            self.blocked_on_user_limit_micros
+        };
+
+        let blocked_on_lock_micros = if let Some(last) = self.last_blocked_on_lock {
+            let delay = last.elapsed();
+            self.blocked_on_lock_micros
+                .saturating_add(delay.as_micros().try_into().unwrap_or(u32::MAX))
+        } else {
+            self.blocked_on_lock_micros
+        };
+
         WaitStats {
             blocked_on_global_capacity_ms: blocked_on_global_capacity_micros / 1000,
             vqueue_start_throttling_ms: self.local_start_throttling_micros / 1000,
             global_invoker_throttling_ms: self.global_throttling_micros / 1000,
             blocked_on_invoker_memory_ms: blocked_on_invoker_memory_micros / 1000,
+            blocked_on_user_limit_ms: blocked_on_user_limit_micros / 1000,
+            blocked_on_lock_ms: blocked_on_lock_micros / 1000,
         }
     }
 }
@@ -198,6 +228,38 @@ impl Stats {
                 .saturating_add(delay.as_micros().try_into().unwrap_or(u32::MAX))
         }
     }
+
+    fn record_user_limit_delay(&mut self, is_now_blocked: bool) {
+        let last = self.last_blocked_on_user_limit.take();
+        self.last_blocked_on_user_limit = if is_now_blocked {
+            Some(Instant::now())
+        } else {
+            None
+        };
+        if let Some(last) = last {
+            let delay = last.elapsed();
+            counter!(VQUEUE_USER_LIMIT_WAIT_MS).increment(delay.as_millis() as u64);
+            self.blocked_on_user_limit_micros = self
+                .blocked_on_user_limit_micros
+                .saturating_add(delay.as_micros().try_into().unwrap_or(u32::MAX))
+        }
+    }
+
+    fn record_lock_delay(&mut self, is_now_blocked: bool) {
+        let last = self.last_blocked_on_lock.take();
+        self.last_blocked_on_lock = if is_now_blocked {
+            Some(Instant::now())
+        } else {
+            None
+        };
+        if let Some(last) = last {
+            let delay = last.elapsed();
+            counter!(VQUEUE_LOCK_WAIT_MS).increment(delay.as_millis() as u64);
+            self.blocked_on_lock_micros = self
+                .blocked_on_lock_micros
+                .saturating_add(delay.as_micros().try_into().unwrap_or(u32::MAX))
+        }
+    }
 }
 
 #[derive(derive_more::Debug)]
@@ -318,6 +380,12 @@ impl<S: VQueueStore> VQueueState<S> {
                     ResourceKind::InvokerThrottling => {
                         // self.head_stats.record_start_throttling_delay(delay);
                     }
+                    ResourceKind::LimitKeyConcurrency { .. } => {
+                        self.head_stats.record_user_limit_delay(true);
+                    }
+                    ResourceKind::Lock { .. } => {
+                        self.head_stats.record_lock_delay(true);
+                    }
                     _ => {}
                 }
                 Ok(Pop::Blocked(resource))