diff --git a/crates/core/src/host/scheduler.rs b/crates/core/src/host/scheduler.rs index 996b9d90ab3..da42f404d63 100644 --- a/crates/core/src/host/scheduler.rs +++ b/crates/core/src/host/scheduler.rs @@ -24,11 +24,12 @@ use spacetimedb_primitives::{ColId, TableId}; use spacetimedb_sats::bsatn::ToBsatn as _; use spacetimedb_sats::AlgebraicValue; use spacetimedb_table::table::RowRef; +use std::collections::HashMap; use std::panic; use std::sync::Arc; use tokio::sync::mpsc; use tokio::time::Instant; -use tokio_util::time::delay_queue::{DelayQueue, Expired}; +use tokio_util::time::delay_queue::{DelayQueue, Expired, Key}; #[derive(Debug, Copy, Clone, Eq, PartialEq, Hash)] pub struct ScheduledFunctionId { @@ -91,6 +92,7 @@ impl SchedulerStarter { // time to make it better right now. pub fn start(mut self, module_host: &ModuleHost) -> anyhow::Result<()> { let mut queue: DelayQueue = DelayQueue::new(); + let mut queued_ids = HashMap::new(); let tx = self.db.begin_tx(Workload::Internal); // Draining rx before processing schedules from the DB to ensure there are no in-flight messages, @@ -128,7 +130,9 @@ impl SchedulerStarter { id_column, at_column, }; - queue.insert_at( + insert_scheduled_queue_item( + &mut queue, + &mut queued_ids, QueueItem::Id { id, function_name: function_name.clone(), @@ -144,6 +148,7 @@ impl SchedulerStarter { SchedulerActor { rx: self.rx, queue, + queued_ids, inflight_calls: FuturesUnordered::new(), active_calls_metric: WORKER_METRICS .scheduler_active_scheduled_functions @@ -274,6 +279,7 @@ impl Scheduler { struct SchedulerActor { rx: mpsc::UnboundedReceiver>, queue: DelayQueue, + queued_ids: HashMap, inflight_calls: FuturesUnordered, active_calls_metric: IntGauge, module_host: WeakModuleHost, @@ -371,7 +377,7 @@ impl SchedulerActor { effective_at, real_at, } => { - self.queue.insert_at( + self.insert_scheduled( QueueItem::Id { id, function_name, @@ -392,6 +398,9 @@ impl SchedulerActor { fn handle_queued(&mut self, expired: Expired) { let item = expired.into_inner(); + if let QueueItem::Id { id, .. } = item { + self.queued_ids.remove(&id); + } let Some(module_host) = self.module_host.upgrade() else { return; @@ -430,7 +439,7 @@ impl SchedulerActor { .. } = item { - self.queue.insert_at( + self.insert_scheduled( QueueItem::Id { id, function_name, @@ -444,11 +453,30 @@ impl SchedulerActor { } } + fn insert_scheduled(&mut self, item: QueueItem, real_at: Instant) { + insert_scheduled_queue_item(&mut self.queue, &mut self.queued_ids, item, real_at); + } + fn update_active_calls_metric(&self) { self.active_calls_metric.set(self.inflight_calls.len() as i64); } } +fn insert_scheduled_queue_item( + queue: &mut DelayQueue, + queued_ids: &mut HashMap, + item: QueueItem, + real_at: Instant, +) { + if let QueueItem::Id { id, .. } = item { + if let Some(key) = queued_ids.remove(&id) { + queue.remove(&key); + } + let key = queue.insert_at(item, real_at); + queued_ids.insert(id, key); + } +} + fn call_scheduled_function(module_host: ModuleHost, item: QueueItem) -> ScheduledFunctionFuture { async move { let params = ScheduledFunctionParams(item.clone()); @@ -1082,4 +1110,33 @@ mod tests { let next = next_interval_tick_after(ts(1_000), TimeDuration::from_micros(100), ts(1_300)); assert_eq!(next, ts(1_400)); } + + #[tokio::test] + async fn scheduling_same_row_replaces_pending_queue_entry() { + let mut queue = DelayQueue::new(); + let mut queued_ids = HashMap::new(); + let id = ScheduledFunctionId { + table_id: TableId(1), + schedule_id: 1, + id_column: ColId(0), + at_column: ColId(1), + }; + + for at in [ts(1_000), ts(2_000)] { + insert_scheduled_queue_item( + &mut queue, + &mut queued_ids, + QueueItem::Id { + id, + function_name: Arc::from("tick"), + at, + row_hash: Hash::ZERO, + }, + Instant::now() + Duration::from_secs(60), + ); + } + + assert_eq!(queue.len(), 1); + assert_eq!(queued_ids.len(), 1); + } } diff --git a/crates/smoketests/modules/schedule-subscribe/src/lib.rs b/crates/smoketests/modules/schedule-subscribe/src/lib.rs index 6a7f63a5a74..785ffeecf18 100644 --- a/crates/smoketests/modules/schedule-subscribe/src/lib.rs +++ b/crates/smoketests/modules/schedule-subscribe/src/lib.rs @@ -18,6 +18,15 @@ pub struct FailingScheduledTable { prev: Timestamp, } +#[spacetimedb::table(accessor = capped_scheduled_table, public, scheduled(capped_self_update, at = sched_at))] +pub struct CappedScheduledTable { + #[primary_key] + #[auto_inc] + scheduled_id: u64, + sched_at: spacetimedb::ScheduleAt, + n: u16, +} + #[spacetimedb::table(accessor = player_entity, public)] pub struct PlayerEntity { #[primary_key] @@ -77,6 +86,15 @@ fn schedule_repeated_reducer(ctx: &ReducerContext) { }); } +#[spacetimedb::reducer] +fn schedule_capped_self_update(ctx: &ReducerContext) { + ctx.db.capped_scheduled_table().insert(CappedScheduledTable { + scheduled_id: 1, + sched_at: duration!(50ms).into(), + n: 0, + }); +} + #[spacetimedb::reducer] fn seed_player_entity(ctx: &ReducerContext, entity_id: u64) { ctx.db.player_entity().entity_id().delete(&entity_id); @@ -99,3 +117,12 @@ pub fn my_reducer(ctx: &ReducerContext, arg: ScheduledTable) { pub fn failing_reducer(_ctx: &ReducerContext, _arg: FailingScheduledTable) -> Result<(), String> { Err("scheduled reducer failed".into()) } + +#[spacetimedb::reducer] +pub fn capped_self_update(ctx: &ReducerContext, arg: CappedScheduledTable) { + log::info!("CappedSelfUpdate"); + ctx.db.capped_scheduled_table().scheduled_id().update(CappedScheduledTable { + n: arg.n.saturating_add(1).min(10), + ..arg + }); +} diff --git a/crates/smoketests/tests/cluster/schedule_reducer.rs b/crates/smoketests/tests/cluster/schedule_reducer.rs index f78d6eda0b7..4c94b45aab8 100644 --- a/crates/smoketests/tests/cluster/schedule_reducer.rs +++ b/crates/smoketests/tests/cluster/schedule_reducer.rs @@ -132,6 +132,21 @@ fn test_scheduled_table_subscription_repeated_reducer() { ); } +#[test] +fn test_capped_self_update_repeated_reducer_does_not_run_away() { + let test = Smoketest::builder().precompiled_module("schedule-subscribe").build(); + + test.call("schedule_capped_self_update", &[]).unwrap(); + thread::sleep(Duration::from_secs(2)); + + let logs = test.logs(200).unwrap(); + let invoked_count = logs.iter().filter(|line| line.contains("CappedSelfUpdate")).count(); + assert!( + (10..80).contains(&invoked_count), + "Expected capped self-updating reducer to run linearly, got {invoked_count}. Logs: {logs:?}" + ); +} + /// Scheduled *procedure* subscription: expect insert + delete for both table and view. #[test] fn test_scheduled_procedure_table_and_view_subscription() {