diff --git a/crates/core/src/host/host_controller.rs b/crates/core/src/host/host_controller.rs index 9fc3e012fa0..cc706fa01da 100644 --- a/crates/core/src/host/host_controller.rs +++ b/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); } diff --git a/crates/core/src/host/scheduler.rs b/crates/core/src/host/scheduler.rs index 29ec9c75834..fb5dec90ac6 100644 --- a/crates/core/src/host/scheduler.rs +++ b/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>, queue: DelayQueue, key_map: FxHashMap, + active_calls: FuturesUnordered, + 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>; + +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) { + fn handle_queued(&mut self, id: Expired) { 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)] diff --git a/crates/core/src/worker_metrics/mod.rs b/crates/core/src/worker_metrics/mod.rs index 5cecd3c40ab..c8cd7bc48d7 100644 --- a/crates/core/src/worker_metrics/mod.rs +++ b/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)] diff --git a/modules/sdk-test-procedure-concurrency/src/lib.rs b/modules/sdk-test-procedure-concurrency/src/lib.rs index e8874309f8e..8c2c15a6e01 100644 --- a/modules/sdk-test-procedure-concurrency/src/lib.rs +++ b/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(), }); } diff --git a/sdks/rust/tests/procedure-concurrency-client/src/module_bindings/scheduled_reducer_row_type.rs b/sdks/rust/tests/procedure-concurrency-client/src/module_bindings/scheduled_reducer_row_type.rs index 92af01cb5ba..b91b4caa816 100644 --- a/sdks/rust/tests/procedure-concurrency-client/src/module_bindings/scheduled_reducer_row_type.rs +++ b/sdks/rust/tests/procedure-concurrency-client/src/module_bindings/scheduled_reducer_row_type.rs @@ -9,6 +9,7 @@ use spacetimedb_sdk::__codegen::{self as __sdk, __lib, __sats, __ws}; pub struct ScheduledReducerRow { pub scheduled_id: u64, pub scheduled_at: __sdk::ScheduleAt, + pub insertion_context: String, } impl __sdk::InModule for ScheduledReducerRow { @@ -21,6 +22,7 @@ impl __sdk::InModule for ScheduledReducerRow { pub struct ScheduledReducerRowCols { pub scheduled_id: __sdk::__query_builder::Col, pub scheduled_at: __sdk::__query_builder::Col, + pub insertion_context: __sdk::__query_builder::Col, } impl __sdk::__query_builder::HasCols for ScheduledReducerRow { @@ -29,6 +31,7 @@ impl __sdk::__query_builder::HasCols for ScheduledReducerRow { ScheduledReducerRowCols { scheduled_id: __sdk::__query_builder::Col::new(table_name, "scheduled_id"), scheduled_at: __sdk::__query_builder::Col::new(table_name, "scheduled_at"), + insertion_context: __sdk::__query_builder::Col::new(table_name, "insertion_context"), } } } diff --git a/sdks/rust/tests/procedure-concurrency-client/src/test_handlers.rs b/sdks/rust/tests/procedure-concurrency-client/src/test_handlers.rs index a7ab7df2e86..7ecfd4fddf5 100644 --- a/sdks/rust/tests/procedure-concurrency-client/src/test_handlers.rs +++ b/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, reducer: Option, scheduled_reducer: Option, + scheduled_reducer_1: Option, + scheduled_reducer_2: Option, procedure_after: Option, scheduled_procedure_before: Option, scheduled_procedure_after: Option, @@ -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,9 +497,15 @@ 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()); } @@ -506,26 +513,27 @@ async fn exec_scheduled_procedure_scheduled_reducer_not_interleaved(db_name: &st } 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(()) })(), diff --git a/sdks/rust/tests/test.rs b/sdks/rust/tests/test.rs index af7c535824b..7203c8b1e4a 100644 --- a/sdks/rust/tests/test.rs +++ b/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() } }