Scheduled functions: Concurrent execution by not awaiting on current run.

crates/core/src/host/host_controller.rs

@@ -1614,5 +1614,8 @@ where
     V8HeapMetrics::remove_all_metric_label_values_for_database(db);
 
     let _ = WORKER_METRICS.v8_request_queue_length.remove_label_values(db);
+    let _ = WORKER_METRICS
+        .scheduler_active_scheduled_functions
+        .remove_label_values(db);
     let _ = DB_METRICS.http_response_size_bytes.remove_label_values(db);
 }

crates/core/src/host/scheduler.rs

@@ -6,8 +6,11 @@ use crate::db::relational_db::RelationalDB;
 use crate::host::module_host::{CallProcedureParams, ModuleInfo};
 use crate::host::wasm_common::module_host_actor::{InstanceCommon, WasmInstance};
 use crate::host::{InvalidProcedureArguments, InvalidReducerArguments, NoSuchModule};
+use crate::worker_metrics::WORKER_METRICS;
 use anyhow::anyhow;
 use core::time::Duration;
+use futures::future::BoxFuture;
+use futures::stream::FuturesUnordered;
 use futures::{FutureExt, StreamExt};
 use rustc_hash::FxHashMap;
 use spacetimedb_client_api_messages::energy::FunctionBudget;
@@ -16,7 +19,7 @@ use spacetimedb_datastore::locking_tx_datastore::MutTxId;
 use spacetimedb_datastore::system_tables::{StScheduledFields, ST_SCHEDULED_ID};
 use spacetimedb_datastore::traits::IsolationLevel;
 use spacetimedb_lib::scheduler::ScheduleAt;
-use spacetimedb_lib::{TimeDuration, Timestamp};
+use spacetimedb_lib::{Identity, TimeDuration, Timestamp};
 use spacetimedb_primitives::{ColId, TableId};
 use spacetimedb_sats::bsatn::ToBsatn as _;
 use spacetimedb_sats::AlgebraicValue;
@@ -151,6 +154,8 @@ impl SchedulerStarter {
                 rx: self.rx,
                 queue,
                 key_map,
+                active_calls: FuturesUnordered::new(),
+                database_identity: module_host.info().database_identity,
                 module_host: module_host.downgrade(),
             }
             .run(),
@@ -276,6 +281,8 @@ struct SchedulerActor {
     rx: mpsc::UnboundedReceiver<MsgOrExit<SchedulerMessage>>,
     queue: DelayQueue<QueueItem>,
     key_map: FxHashMap<ScheduledFunctionId, delay_queue::Key>,
+    active_calls: FuturesUnordered<ScheduledFunctionFuture>,
+    database_identity: Identity,
     module_host: WeakModuleHost,
 }
 
@@ -300,6 +307,17 @@ enum ScheduledFunctionKind {
     Procedure,
 }
 
+type ScheduledFunctionFuture = BoxFuture<'static, ScheduledFunctionCompletion>;
+
+// The outer result is from `catch_unwind`, so `Err` means the scheduled call panicked.
+// The inner result is the normal module-host call result, such as `NoSuchModule`.
+type ScheduledFunctionCallResult = std::thread::Result<Result<CallScheduledFunctionResult, CallScheduledFunctionError>>;
+
+struct ScheduledFunctionCompletion {
+    item: QueueItem,
+    result: ScheduledFunctionCallResult,
+}
+
 impl ScheduledFunctionParams {
     fn function_name(&self) -> &str {
         match &self.0 {
@@ -328,16 +346,26 @@ spacetimedb_table::static_assert_size!(QueueItem, 64);
 
 impl SchedulerActor {
     async fn run(mut self) {
+        let mut closing = false;
+        self.update_active_calls_metric();
         loop {
+            if closing && self.active_calls.is_empty() {
+                self.update_active_calls_metric();
+                break;
+            }
+
             tokio::select! {
-                msg = self.rx.recv() => match msg {
+                msg = self.rx.recv(), if !closing => match msg {
                     Some(MsgOrExit::Msg(msg)) => self.handle_message(msg),
                     // it's fine to just drop any messages in the queue because they've
                     // already been stored in the database
-                    Some(MsgOrExit::Exit) | None => break,
+                    Some(MsgOrExit::Exit) | None => closing = true,
+                },
+                Some(scheduled) = self.queue.next(), if !closing => {
+                    self.handle_queued(scheduled);
                 },
-                Some(scheduled) = self.queue.next() => {
-                    self.handle_queued(scheduled).await;
+                Some(completion) = self.active_calls.next(), if !self.active_calls.is_empty() => {
+                    self.handle_completion(completion);
                 }
             }
         }
@@ -374,7 +402,7 @@ impl SchedulerActor {
         }
     }
 
-    async fn handle_queued(&mut self, id: Expired<QueueItem>) {
+    fn handle_queued(&mut self, id: Expired<QueueItem>) {
         let item = id.into_inner();
         let id = match &item {
             QueueItem::Id { id, .. } => Some(*id),
@@ -388,19 +416,14 @@ impl SchedulerActor {
             return;
         };
 
-        let params = ScheduledFunctionParams(item.clone());
-        let result = match params.kind(module_host.info()) {
-            ScheduledFunctionKind::Procedure => {
-                panic::AssertUnwindSafe(module_host.call_scheduled_procedure(params))
-                    .catch_unwind()
-                    .await
-            }
-            ScheduledFunctionKind::Reducer => {
-                panic::AssertUnwindSafe(module_host.call_scheduled_reducer(params))
-                    .catch_unwind()
-                    .await
-            }
-        };
+        self.active_calls.push(call_scheduled_function(module_host, item));
+        self.update_active_calls_metric();
+    }
+
+    fn handle_completion(&mut self, completion: ScheduledFunctionCompletion) {
+        self.update_active_calls_metric();
+        let ScheduledFunctionCompletion { item, result } = completion;
+
         let result = match result {
             Ok(result) => result,
             Err(_) => {
@@ -420,20 +443,50 @@ impl SchedulerActor {
                 reschedule: Some(Reschedule { at_ts, at_real }),
             }) => {
                 if let QueueItem::Id { id, function_name, .. } = item {
-                    // If this was repeated, we need to add it back to the queue.
-                    let key = self.queue.insert_at(
-                        QueueItem::Id {
-                            id,
-                            function_name,
-                            at: at_ts,
-                        },
-                        at_real,
-                    );
-                    self.key_map.insert(id, key);
+                    // A schedule-table update may have queued a newer entry while
+                    // this call was running. Keep that newer entry authoritative.
+                    if !self.key_map.contains_key(&id) {
+                        let key = self.queue.insert_at(
+                            QueueItem::Id {
+                                id,
+                                function_name,
+                                at: at_ts,
+                            },
+                            at_real,
+                        );
+                        self.key_map.insert(id, key);
+                    }
                 }
             }
         }
     }
+
+    fn update_active_calls_metric(&self) {
+        WORKER_METRICS
+            .scheduler_active_scheduled_functions
+            .with_label_values(&self.database_identity)
+            .set(self.active_calls.len() as i64);
+    }
+}
+
+fn call_scheduled_function(module_host: ModuleHost, item: QueueItem) -> ScheduledFunctionFuture {
+    async move {
+        let params = ScheduledFunctionParams(item.clone());
+        let result = match params.kind(module_host.info()) {
+            ScheduledFunctionKind::Procedure => {
+                panic::AssertUnwindSafe(module_host.call_scheduled_procedure(params))
+                    .catch_unwind()
+                    .await
+            }
+            ScheduledFunctionKind::Reducer => {
+                panic::AssertUnwindSafe(module_host.call_scheduled_reducer(params))
+                    .catch_unwind()
+                    .await
+            }
+        };
+        ScheduledFunctionCompletion { item, result }
+    }
+    .boxed()
 }
 
 #[derive(Debug)]

crates/core/src/worker_metrics/mod.rs

@@ -540,6 +540,11 @@ metrics_group!(
         #[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_scheduler_active_scheduled_functions]
+        #[help = "The number of scheduled functions dispatched by the scheduler and not yet completed"]
+        #[labels(db: Identity)]
+        pub scheduler_active_scheduled_functions: IntGaugeVec,
+
         #[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)]

modules/sdk-test-procedure-concurrency/src/lib.rs

@@ -76,13 +76,14 @@ struct ScheduledReducerRow {
     #[auto_inc]
     scheduled_id: u64,
     scheduled_at: ScheduleAt,
+    insertion_context: String,
 }
 
 #[reducer]
-fn insert_scheduled_reducer(ctx: &ReducerContext, _schedule: ScheduledReducerRow) {
+fn insert_scheduled_reducer(ctx: &ReducerContext, schedule: ScheduledReducerRow) {
     ctx.db.procedure_concurrency_row().insert(ProcedureConcurrencyRow {
         insertion_order: 0,
-        insertion_context: "scheduled_reducer".into(),
+        insertion_context: schedule.insertion_context,
     });
 }
 
@@ -93,6 +94,7 @@ fn procedure_schedule_reducer_between_inserts(ctx: &mut ProcedureContext) {
         ctx.db.scheduled_reducer_row().insert(ScheduledReducerRow {
             scheduled_id: 0,
             scheduled_at: ctx.timestamp.into(),
+            insertion_context: "scheduled_reducer".into(),
         });
     });
     poll_until_tx_true(
@@ -119,9 +121,8 @@ struct ScheduledProcedureRow {
 #[procedure]
 fn scheduled_procedure_sleep_between_inserts(ctx: &mut ProcedureContext, _schedule: ScheduledProcedureRow) {
     ctx.with_tx(|ctx| insert_procedure_concurrency_row(ctx, "scheduled_procedure_before"));
-    // Unfortunately, we can't poll and wake on event here,
-    // as (until https://github.com/clockworklabs/SpacetimeDB/pull/5224 is fixed)
-    // the scheduled reducer actually won't run until after this procedure fully completes.
+    // Sleep long enough for the later scheduled reducer to run while this
+    // procedure is still suspended.
     ctx.sleep_until(ctx.timestamp + Duration::from_secs(10));
     ctx.with_tx(|ctx| insert_procedure_concurrency_row(ctx, "scheduled_procedure_after"));
 }
@@ -135,5 +136,11 @@ fn schedule_procedure_then_reducer(ctx: &ReducerContext) {
     ctx.db.scheduled_reducer_row().insert(ScheduledReducerRow {
         scheduled_id: 0,
         scheduled_at: (ctx.timestamp + Duration::from_secs(2)).into(),
+        insertion_context: "scheduled_reducer_1".into(),
+    });
+    ctx.db.scheduled_reducer_row().insert(ScheduledReducerRow {
+        scheduled_id: 0,
+        scheduled_at: (ctx.timestamp + Duration::from_secs(3)).into(),
+        insertion_context: "scheduled_reducer_2".into(),
     });
 }

sdks/rust/tests/procedure-concurrency-client/src/test_handlers.rs

@@ -11,8 +11,8 @@ pub async fn dispatch(test: &str, db_name: &str) {
         "procedure-concurrent-with-scheduled-reducer" => {
             exec_procedure_concurrent_with_scheduled_reducer(db_name).await
         }
-        "scheduled-procedure-scheduled-reducer-not-interleaved" => {
-            exec_scheduled_procedure_scheduled_reducer_not_interleaved(db_name).await
+        "scheduled-procedure-scheduled-reducer-interleaved" => {
+            exec_scheduled_procedure_scheduled_reducer_interleaved(db_name).await
         }
         _ => panic!("Unknown test: {test}"),
     }
@@ -91,6 +91,8 @@ struct ConnectionRowObservation {
     procedure_before: Option<u32>,
     reducer: Option<u32>,
     scheduled_reducer: Option<u32>,
+    scheduled_reducer_1: Option<u32>,
+    scheduled_reducer_2: Option<u32>,
     procedure_after: Option<u32>,
     scheduled_procedure_before: Option<u32>,
     scheduled_procedure_after: Option<u32>,
@@ -467,8 +469,7 @@ async fn exec_procedure_concurrent_with_scheduled_reducer(db_name: &str) {
 }
 
 /// Like [`exec_procedure_reducer_same_client_not_interleaved`], but with the scheduler instead of a client.
-/// Tracks a behavior that we'd like to change.
-async fn exec_scheduled_procedure_scheduled_reducer_not_interleaved(db_name: &str) {
+async fn exec_scheduled_procedure_scheduled_reducer_interleaved(db_name: &str) {
     let test_counter = TestCounter::new();
     let sub_applied_nothing_result = test_counter.add_test("on_subscription_applied_nothing");
     let mut reducer_callback_result = Some(test_counter.add_test("schedule_procedure_then_reducer_callback"));
@@ -496,36 +497,43 @@ async fn exec_scheduled_procedure_scheduled_reducer_not_interleaved(db_name: &st
                                     .replace(row.insertion_order)
                                     .is_none());
                             }
-                            "scheduled_reducer" => {
+                            "scheduled_reducer_1" => {
                                 assert!(observation
-                                    .scheduled_reducer
+                                    .scheduled_reducer_1
+                                    .replace(row.insertion_order)
+                                    .is_none());
+                            }
+                            "scheduled_reducer_2" => {
+                                assert!(observation
+                                    .scheduled_reducer_2
                                     .replace(row.insertion_order)
                                     .is_none());
                             }
                             unexpected => panic!("Unexpected insertion context: {unexpected}"),
                         }
                         match (
                             observation.scheduled_procedure_before,
-                            observation.scheduled_reducer,
+                            observation.scheduled_reducer_1,
+                            observation.scheduled_reducer_2,
                             observation.scheduled_procedure_after,
                         ) {
-                            (Some(before), Some(scheduled_reducer), Some(after))
+                            (Some(before), Some(reducer_1), Some(reducer_2), Some(after))
                                 if !observation.ordering_checked =>
                             {
                                 observation.ordering_checked = true;
-                                Some((before, scheduled_reducer, after))
+                                Some((before, reducer_1, reducer_2, after))
                             }
                             _ => None,
                         }
                     };
 
-                    if let Some((before, scheduled_reducer, after)) = maybe_ordering {
+                    if let Some((before, reducer_1, reducer_2, after)) = maybe_ordering {
                         (ordering_result.take().expect("Ordering result should only be reported once"))(
                             #[allow(clippy::redundant_closure_call)]
                             (|| {
                                 anyhow::ensure!(
-                                    before < after && after < scheduled_reducer,
-                                    "Expected scheduled procedure insertion order scheduled_procedure_before < scheduled_procedure_after < scheduled_reducer, got {before} < {after} < {scheduled_reducer}"
+                                    before < reducer_1 && reducer_1 < reducer_2 && reducer_2 < after,
+                                    "Expected scheduled procedure/reducer insertion order scheduled_procedure_before < scheduled_reducer_1 < scheduled_reducer_2 < scheduled_procedure_after, got {before} < {reducer_1} < {reducer_2} < {after}"
                                 );
                                 Ok(())
                             })(),

sdks/rust/tests/test.rs

@@ -545,14 +545,11 @@ mod rust_procedure_concurrency {
         make_test("procedure-concurrent-with-scheduled-reducer").run()
     }
 
-    /// Test that the scheduler has only a single active execution slot,
-    /// which can be occupied by a long-running or suspended procedure.
-    ///
-    /// We're not attached to this behavior, and in fact it should be changed.
-    /// At that time, this test should be altered to demonstrate that the execution is interleaved.
+    /// Test that the scheduler can dispatch a scheduled reducer while a previous
+    /// scheduled procedure is still running.
     #[test]
-    fn scheduled_procedure_scheduled_reducer_not_interleaved() {
-        make_test("scheduled-procedure-scheduled-reducer-not-interleaved").run()
+    fn scheduled_procedure_scheduled_reducer_interleaved() {
+        make_test("scheduled-procedure-scheduled-reducer-interleaved").run()
     }
 }