Track scheduled function delay (#5592)

## Summary
fixes: clockworklabs/SpacetimeDBPrivate#3441.

It chooses an arbitrary threshold of `50ms`, if a scheduled function
starts more than that threshold after its expected time, we emit a
warning and record a metric.


Note that, metric does not necessarily tells If something is wrong with
Host as drift can also happen due to module's previous schedule function
taking long time to finish.

Author: Shubham8287

crates/core/src/host/module_host.rs

@@ -36,6 +36,7 @@ use indexmap::IndexSet;
 use itertools::Itertools;
 use parking_lot::Mutex;
 use prometheus::{Histogram, HistogramTimer, IntGauge};
+use rustc_hash::FxHashMap;
 use scopeguard::ScopeGuard;
 use smallvec::SmallVec;
 use spacetimedb_auth::identity::ConnectionAuthCtx;
@@ -254,6 +255,45 @@ pub struct ModuleMetrics {
     pub request_round_trip_subscribe: Histogram,
     pub request_round_trip_unsubscribe: Histogram,
     pub request_round_trip_sql: Histogram,
+    scheduled_function_delay: ScheduledFunctionDelayMetrics,
+}
+
+#[derive(Debug)]
+struct ScheduledFunctionDelayMetrics {
+    database_identity: Identity,
+    metrics_by_function: Mutex<FxHashMap<Box<str>, Histogram>>,
+}
+
+impl ScheduledFunctionDelayMetrics {
+    fn new(database_identity: &Identity) -> Self {
+        Self {
+            database_identity: *database_identity,
+            metrics_by_function: Mutex::new(FxHashMap::default()),
+        }
+    }
+
+    fn observe(&self, function_name: &str, delay: Duration) {
+        let mut metrics_by_function = self.metrics_by_function.lock();
+        let metric = metrics_by_function
+            .entry(Box::<str>::from(function_name))
+            .or_insert_with(|| {
+                WORKER_METRICS
+                    .scheduled_function_delay
+                    .with_label_values(&self.database_identity, function_name)
+            });
+        metric.observe(delay.as_secs_f64());
+    }
+}
+
+impl Drop for ScheduledFunctionDelayMetrics {
+    fn drop(&mut self) {
+        let metrics_by_function = std::mem::take(self.metrics_by_function.get_mut());
+        for function_name in metrics_by_function.keys() {
+            let _ = WORKER_METRICS
+                .scheduled_function_delay
+                .remove_label_values(&self.database_identity, function_name);
+        }
+    }
 }
 
 impl ModuleMetrics {
@@ -272,15 +312,21 @@ impl ModuleMetrics {
         let request_round_trip_sql = WORKER_METRICS
             .request_round_trip
             .with_label_values(&WorkloadType::Sql, db, "");
+        let scheduled_function_delay = ScheduledFunctionDelayMetrics::new(db);
         Self {
             connected_clients,
             ws_clients_spawned,
             ws_clients_aborted,
             request_round_trip_subscribe,
             request_round_trip_unsubscribe,
             request_round_trip_sql,
+            scheduled_function_delay,
         }
     }
+
+    pub(in crate::host) fn observe_scheduled_function_delay(&self, function_name: &str, delay: Duration) {
+        self.scheduled_function_delay.observe(function_name, delay);
+    }
 }
 
 impl ModuleInfo {

crates/core/src/host/scheduler.rs

@@ -190,6 +190,9 @@ const MAX_SCHEDULE_DELAY: Duration = Duration::from_millis(
     (1 << (6 * 6)) - 1,
 );
 
+/// Warn when a scheduled function starts more than this long after it was due.
+const SCHEDULED_FUNCTION_DELAY_WARNING_THRESHOLD: Duration = Duration::from_millis(30);
+
 #[derive(thiserror::Error, Debug)]
 pub enum ScheduleError {
     #[error("Unable to schedule with long delay at {0:?}")]
@@ -447,8 +450,9 @@ struct Reschedule {
 enum ScheduledProcedureStep {
     Done(CallScheduledFunctionResult, bool),
     Procedure {
-        params: CallProcedureParams,
+        params: Box<CallProcedureParams>,
         reschedule: Option<Reschedule>,
+        delay: Option<(Arc<str>, Duration)>,
     },
 }
 
@@ -465,9 +469,17 @@ pub(super) async fn call_scheduled_procedure(
     // even though it has been already moved during `delete_scheduled_function_row` call.
     match next_step {
         ScheduledProcedureStep::Done(result, trapped) => (result, trapped),
-        ScheduledProcedureStep::Procedure { params, reschedule } => {
+        ScheduledProcedureStep::Procedure {
+            params,
+            reschedule,
+            delay,
+        } => {
+            if let Some((function_name, delay)) = delay.as_ref() {
+                record_scheduled_function_delay(module_info, function_name, *delay);
+            }
+
             // Execute the procedure. See above for commentary on `catch_unwind()`.
-            let result = panic::AssertUnwindSafe(inst_common.call_procedure(params, inst))
+            let result = panic::AssertUnwindSafe(inst_common.call_procedure(*params, inst))
                 .catch_unwind()
                 .await;
 
@@ -504,6 +516,7 @@ fn prepare_scheduled_procedure_call(
     inst: &mut impl WasmInstance,
 ) -> ScheduledProcedureStep {
     let ScheduledFunctionParams(item) = params;
+    let delay = scheduled_function_delay_context_for_item(&item);
     let id = scheduled_item_id(&item);
     let db = &**module_info.relational_db();
     let tx = db.begin_mut_tx(IsolationLevel::Serializable, Workload::Internal);
@@ -530,7 +543,13 @@ fn prepare_scheduled_procedure_call(
     let reschedule = id.and_then(|id| {
         delete_scheduled_function_row(module_info, db, id, Some(tx), (timestamp, instant), inst_common, inst)
     });
-    ScheduledProcedureStep::Procedure { params, reschedule }
+    let delay =
+        delay.map(|(function_name, requested_at)| (function_name, scheduled_function_delay(timestamp, requested_at)));
+    ScheduledProcedureStep::Procedure {
+        params: Box::new(params),
+        reschedule,
+        delay,
+    }
 }
 
 fn call_scheduled_reducer_until_done(
@@ -540,6 +559,7 @@ fn call_scheduled_reducer_until_done(
     inst: &mut impl WasmInstance,
 ) -> (CallScheduledFunctionResult, bool) {
     let ScheduledFunctionParams(item) = params;
+    let delay = scheduled_function_delay_context_for_item(&item);
     let id = scheduled_item_id(&item);
     let db = &**module_info.relational_db();
     let tx = db.begin_mut_tx(IsolationLevel::Serializable, Workload::Internal);
@@ -561,6 +581,10 @@ fn call_scheduled_reducer_until_done(
         }
     };
 
+    if let Some((function_name, requested_at)) = delay {
+        let delay = scheduled_function_delay(timestamp, requested_at);
+        record_scheduled_function_delay(module_info, &function_name, delay);
+    }
     call_scheduled_reducer_with_tx(module_info, db, id, tx, (timestamp, instant), params, inst_common, inst)
 }
 
@@ -571,6 +595,35 @@ fn scheduled_item_id(item: &QueueItem) -> Option<ScheduledFunctionId> {
     }
 }
 
+fn scheduled_function_delay_context_for_item(item: &QueueItem) -> Option<(Arc<str>, Timestamp)> {
+    match item {
+        QueueItem::Id { function_name, at, .. } => Some((function_name.clone(), *at)),
+        QueueItem::VolatileNonatomicImmediate { .. } => None,
+    }
+}
+
+fn record_scheduled_function_delay(module_info: &ModuleInfo, function_name: &str, delay: Duration) {
+    module_info
+        .metrics
+        .observe_scheduled_function_delay(function_name, delay);
+
+    if delay <= SCHEDULED_FUNCTION_DELAY_WARNING_THRESHOLD {
+        return;
+    }
+
+    log::warn!(
+        "scheduled function `{}` for database {} is delayed by {:.3}s, exceeding the {:.3}s threshold",
+        function_name,
+        module_info.database_identity,
+        delay.as_secs_f64(),
+        SCHEDULED_FUNCTION_DELAY_WARNING_THRESHOLD.as_secs_f64(),
+    );
+}
+
+fn scheduled_function_delay(actual: Timestamp, requested: Timestamp) -> Duration {
+    actual.duration_since(requested).unwrap_or(Duration::ZERO)
+}
+
 #[allow(clippy::too_many_arguments)]
 fn call_scheduled_reducer_with_tx(
     module_info: &ModuleInfo,
@@ -755,7 +808,9 @@ fn call_params_for_queued_item<T>(
     ) -> anyhow::Result<(Timestamp, Instant, T)>,
 ) -> anyhow::Result<Option<(Timestamp, Instant, T)>> {
     Ok(Some(match item {
-        QueueItem::Id { id, function_name, at } => {
+        QueueItem::Id {
+            id, function_name, at, ..
+        } => {
             let Some(schedule_row) = get_schedule_row_mut(tx, db, id)? else {
                 // If the row is not found, it means the schedule is cancelled by the user.
                 return Ok(None);

crates/core/src/worker_metrics/mod.rs

@@ -534,6 +534,12 @@ metrics_group!(
         #[buckets(100e-6, 500e-6, 0.001, 0.005, 0.01, 0.05, 0.1, 0.5, 1, 5, 10)]
         pub reducer_wait_time: HistogramVec,
 
+        #[name = spacetime_scheduled_function_delay_seconds]
+        #[help = "The amount of time (in seconds) between when a scheduled function was due and when the scheduler began invoking it"]
+        #[labels(db: Identity, function: str)]
+        #[buckets(0.001, 0.005, 0.01, 0.05, 0.1, 0.5, 1, 5, 10, 30, 60, 300)]
+        pub scheduled_function_delay: HistogramVec,
+
         #[name = spacetime_worker_wasm_instance_errors_total]
         #[help = "The number of fatal WASM instance errors, such as reducer panics."]
         #[labels(database_identity: Identity, module_hash: Hash, reducer_symbol: str)]