From bc8a66d90b4ef7c3fe45d1036014f93a0e08b822 Mon Sep 17 00:00:00 2001 From: Sree Narayanan Date: Sat, 10 Oct 2026 01:53:08 +0400 Subject: [PATCH] fix(rivetkit): run onSleep after in-flight work and refuse dispatch once it starts --- docs-internal/engine/sleep-sequence.md | 4 +- .../packages/rivetkit-core/src/actor/task.rs | 77 ++++- .../packages/rivetkit-core/tests/task.rs | 304 +++++++++++++----- .../fixtures/driver-test-suite/sleep-db.ts | 11 + .../packages/rivetkit/src/registry/native.ts | 9 +- .../tests/driver/actor-sleep-db.test.ts | 31 ++ 6 files changed, 344 insertions(+), 92 deletions(-) diff --git a/docs-internal/engine/sleep-sequence.md b/docs-internal/engine/sleep-sequence.md index 8c3d2d7291..08bde1c78a 100644 --- a/docs-internal/engine/sleep-sequence.md +++ b/docs-internal/engine/sleep-sequence.md @@ -40,7 +40,9 @@ Removing `preventSleep` deleted both predicate branches. Any future sleep-affect - `start_grace(reason)` fires at the start of `SleepGrace` / `DestroyGrace`. It cancels the sleep idle timer, cancels the actor abort signal (`actor_abort_signal`), installs a `SleepGraceState` with the effective grace deadline, and resets the sleep timer to arm the grace tick. - The actor abort signal is a soft signal: "shutdown has started, please wrap up." User code observes it via `c.abortSignal`. It does not force-stop work. - Destroy requests also use the normal grace path. The actor abort signal fires when destroy grace starts. -- Sleep grace admits new work only while its hooks still run: the app's `onSleep` and connection disconnects, counted in `core_dispatched_hooks`. After they finish, the app has released its resources (the TypeScript sleep hook also closes SQLite and clears runtime state), so dispatch refuses with `actor/restarting` and the client retries against the next generation. `SleepFinalize` refuses the same way. +- Sleep grace runs `onSleep` and connection disconnects only after tracked work that drains shutdown grace (actions, `keepAwake`, `waitUntil`, registered tasks, WebSocket callbacks) finishes, or at the hooks deadline, which keeps the last quarter of the grace period for the hooks and the final save. Raw `onRequest` handlers and `run` are not waited for. Destroy grace runs its hooks at once. +- Sleep grace admits new work until the hooks run. After that, dispatch refuses with `actor/restarting` and the client retries against the next generation. `SleepFinalize` refuses the same way. +- Core owns SQLite teardown: it syncs the alarm and then closes SQLite in final cleanup. The TypeScript sleep cleanup closes SQLite itself only when tracked work outlives the grace deadline, so that work cannot write after shutdown. ## Grace deadline enforcement diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs index f2afa24b2f..46e100d233 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs @@ -68,6 +68,10 @@ pub type ActionDispatchResult = std::result::Result, ActionDispatchError pub type HttpDispatchResult = Result; const SERIALIZE_STATE_SHUTDOWN_SANITY_CAP: Duration = Duration::from_secs(15); +/// Sleep grace keeps the last quarter of the grace period for onSleep, connection +/// disconnects, and the final save, so in-flight work cannot use up the whole window. +/// It scales with the configured grace period: the 15 second default leaves 3.75 seconds. +const SLEEP_HOOK_RESERVE_DIVISOR: u32 = 4; #[cfg(test)] const LONG_SHUTDOWN_DRAIN_WARNING_THRESHOLD: Duration = Duration::from_secs(1); const INSPECTOR_SERIALIZE_STATE_INTERVAL: Duration = Duration::from_millis(50); @@ -315,6 +319,18 @@ enum LiveExit { struct SleepGraceState { deadline: Instant, reason: ShutdownKind, + hooks: GraceHooks, +} + +/// Whether onSleep and the connection disconnect hooks have run in this grace period. +#[derive(Clone, Copy)] +enum GraceHooks { + /// Sleep waits for work that started before grace, and runs the hooks at + /// `deadline` if that work is still running. + Pending { + deadline: Instant, + }, + Emitted, } struct PersistedStartup { @@ -486,6 +502,7 @@ impl ActorTask { } // TODO: Sample inbox depths periodically instead of on every loop iteration. self.record_inbox_depths(); + let hooks_deadline = self.pending_hooks_deadline(); tokio::select! { biased; lifecycle_command = self.lifecycle_inbox.recv() => { @@ -522,6 +539,12 @@ impl ActorTask { return exit; } } + _ = Self::sleep_grace_tick(hooks_deadline), if hooks_deadline.is_some() => { + self.emit_grace_hooks(); + if let Some(exit) = self.try_finish_grace() { + return exit; + } + } _ = Self::sleep_grace_tick(self.sleep_grace.as_ref().map(|grace| grace.deadline)), if self.sleep_grace.is_some() => { if let Some(exit) = self.on_sleep_grace_deadline().await { return exit; @@ -672,8 +695,12 @@ impl ActorTask { let activity = activity_notify.notified(); tokio::pin!(activity); + let hooks_deadline = self.pending_hooks_deadline(); tokio::select! { _ = &mut activity => {} + _ = Self::sleep_grace_tick(hooks_deadline), if hooks_deadline.is_some() => { + self.emit_grace_hooks(); + } _ = Self::sleep_grace_tick(Some(deadline)) => { if let Some(exit) = self.on_sleep_grace_deadline().await { let LiveExit::Shutdown { reason } = exit else { @@ -778,9 +805,39 @@ impl ActorTask { ShutdownKind::Destroy => LifecycleState::DestroyGrace, }); self.start_grace(reason); + match reason { + ShutdownKind::Sleep => self.try_emit_grace_hooks(), + ShutdownKind::Destroy => self.emit_grace_hooks(), + } + } + + /// Runs the grace hooks once work that started before grace has finished. + /// The hooks deadline arm runs them otherwise. + fn try_emit_grace_hooks(&mut self) { + if self.ctx.shutdown_task_count() == 0 && self.ctx.websocket_callback_count() == 0 { + self.emit_grace_hooks(); + } + } + + fn emit_grace_hooks(&mut self) { + let Some(grace) = self.sleep_grace.as_mut() else { + return; + }; + if matches!(grace.hooks, GraceHooks::Emitted) { + return; + } + grace.hooks = GraceHooks::Emitted; + let reason = grace.reason; self.emit_grace_events(reason); } + fn pending_hooks_deadline(&self) -> Option { + match self.sleep_grace.as_ref()?.hooks { + GraceHooks::Pending { deadline } => Some(deadline), + GraceHooks::Emitted => None, + } + } + fn emit_grace_events(&mut self, reason: ShutdownKind) { let conns: Vec<_> = self.ctx.conns().collect(); for conn in conns { @@ -1159,10 +1216,10 @@ impl ActorTask { match self.lifecycle { LifecycleState::Started => None, - // The app's onSleep and connection cleanup still run, so the actor still serves work. - LifecycleState::SleepGrace if self.ctx.core_dispatched_hook_count() > 0 => None, - // The app has released its resources. The work never ran here, so the client - // retries it and the next generation serves it. + // Work admitted before onSleep runs is waited for before onSleep starts. + LifecycleState::SleepGrace if self.pending_hooks_deadline().is_some() => None, + // onSleep has started releasing the app's resources. The work never ran here, + // so the client retries it and the next generation serves it. LifecycleState::SleepGrace | LifecycleState::SleepFinalize => { self.ctx.warn_work_sent_to_stopping_instance("dispatch"); Some(ActorLifecycleError::Restarting.build()) @@ -1604,9 +1661,13 @@ impl ActorTask { self.sleep_deadline = None; self.ctx.cancel_sleep_timer(); self.ctx.cancel_actor_abort_signal(); + let deadline = Instant::now() + grace_period; self.sleep_grace = Some(SleepGraceState { - deadline: Instant::now() + grace_period, + deadline, reason, + hooks: GraceHooks::Pending { + deadline: deadline - grace_period / SLEEP_HOOK_RESERVE_DIVISOR, + }, }); self.ctx.reset_sleep_timer(); } @@ -1638,10 +1699,14 @@ impl ActorTask { } fn try_finish_grace(&mut self) -> Option { + self.try_emit_grace_hooks(); let Some(grace) = self.sleep_grace.as_ref() else { return None; }; - if self.ctx.can_finalize_shutdown(grace.reason) { + // Before the hooks run, the hook count is zero, so finalizing would skip onSleep. + if matches!(grace.hooks, GraceHooks::Emitted) + && self.ctx.can_finalize_shutdown(grace.reason) + { let reason = grace.reason; self.sleep_grace = None; return Some(LiveExit::Shutdown { reason }); diff --git a/rivetkit-rust/packages/rivetkit-core/tests/task.rs b/rivetkit-rust/packages/rivetkit-core/tests/task.rs index 1000365ffa..f7ffafb900 100644 --- a/rivetkit-rust/packages/rivetkit-core/tests/task.rs +++ b/rivetkit-rust/packages/rivetkit-core/tests/task.rs @@ -70,17 +70,6 @@ pub(crate) mod moved_tests { assert_eq!(counter.load(Ordering::SeqCst), expected); } - async fn wait_for_grace_hooks_done(ctx: &ActorContext) { - for _ in 0..50 { - if ctx.core_dispatched_hook_count() == 0 { - return; - } - sleep(Duration::from_millis(10)).await; - } - - assert_eq!(ctx.core_dispatched_hook_count(), 0); - } - async fn wait_for_state(ctx: &ActorContext, expected: &[u8]) { for _ in 0..50 { if ctx.state() == expected { @@ -848,31 +837,60 @@ pub(crate) mod moved_tests { })) } - /// Like `sleep_grace_factory`, but the graceful cleanup replies only after - /// `release` is notified, so a test can act while the app's onSleep still runs. - fn held_cleanup_factory( + #[derive(Default)] + struct SleepOrderSignals { + log: parking_lot::Mutex>, + slow_started: Notify, + release_slow: Notify, + on_sleep_started: Notify, + release_on_sleep: Notify, + } + + impl SleepOrderSignals { + fn record(&self, event: &'static str) { + self.log.lock().push(event); + } + + fn log(&self) -> Vec<&'static str> { + self.log.lock().clone() + } + } + + /// The `slow` action replies after `release_slow`. When `hold_on_sleep` is set, + /// onSleep replies after `release_on_sleep`. Every step is recorded in order. + fn sleep_order_factory( config: ActorConfig, - cleanup_started: Arc, - release: Arc, - action_count: Arc, + signals: Arc, + hold_on_sleep: bool, ) -> Arc { Arc::new(ActorFactory::new(config, move |start| { - let cleanup_started = cleanup_started.clone(); - let release = release.clone(); - let action_count = action_count.clone(); + let signals = signals.clone(); Box::pin(async move { let mut events = start.events; while let Some(event) = events.recv().await { match event { + ActorEvent::Action { name, reply, .. } if name == "slow" => { + signals.record("slow_started"); + signals.slow_started.notify_one(); + let signals = signals.clone(); + tokio::spawn(async move { + signals.release_slow.notified().await; + signals.record("slow_finished"); + reply.send(Ok(Vec::new())); + }); + } ActorEvent::Action { reply, .. } => { - action_count.fetch_add(1, Ordering::SeqCst); - reply.send(Ok(vec![7, 7, 7])); + signals.record("ping"); + reply.send(Ok(Vec::new())); } ActorEvent::RunGracefulCleanup { reply, .. } => { - cleanup_started.fetch_add(1, Ordering::SeqCst); - let release = release.clone(); + signals.record("on_sleep"); + signals.on_sleep_started.notify_one(); + let signals = signals.clone(); tokio::spawn(async move { - release.notified().await; + if hold_on_sleep { + signals.release_on_sleep.notified().await; + } reply.send(Ok(())); }); } @@ -884,6 +902,23 @@ pub(crate) mod moved_tests { })) } + fn dispatch_action( + dispatch_tx: &mpsc::UnboundedSender, + name: &str, + ) -> oneshot::Receiver>> { + let (reply_tx, reply_rx) = oneshot::channel(); + dispatch_tx + .send(DispatchCommand::Action { + name: name.to_owned(), + args: Vec::new(), + incoming: IncomingInvocationContext::default(), + conn: ConnHandle::new("conn-sleep-order", Vec::new(), Vec::new(), false), + reply: reply_tx, + }) + .expect("action should send"); + reply_rx + } + #[derive(Default)] struct MessageVisitor { message: Option, @@ -3755,8 +3790,8 @@ pub(crate) mod moved_tests { .expect("task run should succeed"); } - #[tokio::test] - async fn sleep_grace_refuses_dispatch_after_cleanup_and_second_sleep_is_idempotent() { + #[tokio::test(start_paused = true)] + async fn sleep_grace_runs_on_sleep_after_keep_awake_work_and_second_sleep_is_idempotent() { let ctx = new_with_kv( "actor-sleep-grace-dispatch", "task-sleep-grace-dispatch", @@ -3808,9 +3843,6 @@ pub(crate) mod moved_tests { reply: sleep_tx, }) .expect("sleep stop should send"); - wait_for_count(&begin_sleep_count, 1).await; - assert!(ctx.actor_aborted()); - let (sleep_again_tx, sleep_again_rx) = oneshot::channel(); lifecycle_tx .send(LifecycleCommand::Stop { @@ -3822,28 +3854,21 @@ pub(crate) mod moved_tests { .await .expect("second sleep reply should send") .expect("second sleep should no-op"); - assert_eq!(begin_sleep_count.load(Ordering::SeqCst), 1); - // The cleanup replied, so the app's onSleep is done; keep-awake work still holds grace. - wait_for_grace_hooks_done(&ctx).await; + assert!(ctx.actor_aborted()); - let (action_tx, action_rx) = oneshot::channel(); - dispatch_tx - .send(DispatchCommand::Action { - name: "ping".to_owned(), - args: Vec::new(), - incoming: IncomingInvocationContext::default(), - conn: ConnHandle::new("conn-grace", Vec::new(), Vec::new(), false), - reply: action_tx, - }) - .expect("action should send during sleep grace"); - let error = action_rx - .await - .expect("action reply should send") - .expect_err("sleep grace should refuse dispatch once cleanup finished"); - let error = rivet_error::RivetError::extract(&error); - assert_eq!((error.group(), error.code()), ("actor", "restarting")); - assert_eq!(action_count.load(Ordering::SeqCst), 0); - assert_eq!(destroy_count.load(Ordering::SeqCst), 0); + // Paused time runs every ready task before the clock moves, so onSleep + // would have run by now if it did not wait for the keep-awake work. + sleep(Duration::from_millis(100)).await; + assert_eq!(begin_sleep_count.load(Ordering::SeqCst), 0); + + assert_eq!( + dispatch_action(&dispatch_tx, "ping") + .await + .expect("action reply should send") + .expect("sleep grace should accept dispatch before onSleep runs"), + vec![7, 7, 7] + ); + assert_eq!(action_count.load(Ordering::SeqCst), 1); release_tx.send(()).expect("keep-awake release should send"); sleep_rx @@ -3853,34 +3878,32 @@ pub(crate) mod moved_tests { keep_awake .await .expect("keep-awake task should finish after release"); + assert_eq!(begin_sleep_count.load(Ordering::SeqCst), 1); assert_eq!(destroy_count.load(Ordering::SeqCst), 0); run.await .expect("task run should finish") .expect("task run should succeed"); } - #[tokio::test] - async fn sleep_grace_accepts_dispatch_while_cleanup_runs() { + #[tokio::test(start_paused = true)] + async fn sleep_runs_on_sleep_after_in_flight_action_replies() { let ctx = new_with_kv( - "actor-sleep-grace-cleanup-dispatch", - "task-sleep-grace-cleanup-dispatch", + "actor-sleep-order-action", + "task-sleep-order-action", Vec::new(), "local", new_in_memory(), ); let (mut task, lifecycle_tx, dispatch_tx, _events_tx) = new_task_with_senders(ctx.clone()); - let cleanup_started = Arc::new(AtomicUsize::new(0)); - let release = Arc::new(Notify::new()); - let action_count = Arc::new(AtomicUsize::new(0)); - task.factory = held_cleanup_factory( + let signals = Arc::new(SleepOrderSignals::default()); + task.factory = sleep_order_factory( ActorConfig { sleep_grace_period: Duration::from_secs(5), sleep_grace_period_overridden: true, ..ActorConfig::default() }, - cleanup_started.clone(), - release.clone(), - action_count.clone(), + signals.clone(), + false, ); let run = tokio::spawn(task.run()); @@ -3893,6 +3916,9 @@ pub(crate) mod moved_tests { .expect("start reply should send") .expect("start should succeed"); + let slow = dispatch_action(&dispatch_tx, "slow"); + signals.slow_started.notified().await; + let (sleep_tx, sleep_rx) = oneshot::channel(); lifecycle_tx .send(LifecycleCommand::Stop { @@ -3900,28 +3926,136 @@ pub(crate) mod moved_tests { reply: sleep_tx, }) .expect("sleep stop should send"); - wait_for_count(&cleanup_started, 1).await; + // Paused time runs every ready task before the clock moves, so onSleep + // would have run by now if it did not wait for the action. + sleep(Duration::from_millis(100)).await; + assert_eq!(signals.log(), vec!["slow_started"]); + + signals.release_slow.notify_one(); + slow.await + .expect("slow action reply should send") + .expect("slow action should finish"); + sleep_rx + .await + .expect("sleep reply should send") + .expect("sleep should succeed"); + assert_eq!( + signals.log(), + vec!["slow_started", "slow_finished", "on_sleep"] + ); + run.await + .expect("task run should finish") + .expect("task run should succeed"); + } - let (action_tx, action_rx) = oneshot::channel(); - dispatch_tx - .send(DispatchCommand::Action { - name: "ping".to_owned(), - args: Vec::new(), - incoming: IncomingInvocationContext::default(), - conn: ConnHandle::new("conn-cleanup", Vec::new(), Vec::new(), false), - reply: action_tx, + #[tokio::test(start_paused = true)] + async fn sleep_runs_on_sleep_at_the_hook_reserve_when_work_outlives_it() { + let ctx = new_with_kv( + "actor-sleep-order-reserve", + "task-sleep-order-reserve", + Vec::new(), + "local", + new_in_memory(), + ); + let (mut task, lifecycle_tx, dispatch_tx, _events_tx) = new_task_with_senders(ctx.clone()); + let signals = Arc::new(SleepOrderSignals::default()); + task.factory = sleep_order_factory( + ActorConfig { + sleep_grace_period: Duration::from_secs(4), + sleep_grace_period_overridden: true, + ..ActorConfig::default() + }, + signals.clone(), + false, + ); + let run = tokio::spawn(task.run()); + + let (start_tx, start_rx) = oneshot::channel(); + lifecycle_tx + .send(LifecycleCommand::Start { reply: start_tx }) + .expect("start command should send"); + start_rx + .await + .expect("start reply should send") + .expect("start should succeed"); + + let _slow = dispatch_action(&dispatch_tx, "slow"); + signals.slow_started.notified().await; + + let stopped_at = tokio::time::Instant::now(); + let (sleep_tx, sleep_rx) = oneshot::channel(); + lifecycle_tx + .send(LifecycleCommand::Stop { + reason: ShutdownKind::Sleep, + reply: sleep_tx, }) - .expect("action should send during sleep grace"); - assert_eq!( - action_rx - .await - .expect("action reply should send") - .expect("sleep grace should accept dispatch while cleanup runs"), - vec![7, 7, 7] + .expect("sleep stop should send"); + signals.on_sleep_started.notified().await; + let waited = stopped_at.elapsed(); + // The last quarter of the 4 second grace period is kept for onSleep. + assert!( + waited >= Duration::from_secs(3) && waited < Duration::from_secs(4), + "onSleep should run at the hook reserve, ran after {waited:?}" ); - assert_eq!(action_count.load(Ordering::SeqCst), 1); - release.notify_one(); + sleep_rx + .await + .expect("sleep reply should send") + .expect("sleep should succeed"); + run.await + .expect("task run should finish") + .expect("task run should succeed"); + } + + #[tokio::test] + async fn sleep_refuses_dispatch_once_on_sleep_runs() { + let ctx = new_with_kv( + "actor-sleep-order-refuse", + "task-sleep-order-refuse", + Vec::new(), + "local", + new_in_memory(), + ); + let (mut task, lifecycle_tx, dispatch_tx, _events_tx) = new_task_with_senders(ctx.clone()); + let signals = Arc::new(SleepOrderSignals::default()); + task.factory = sleep_order_factory( + ActorConfig { + sleep_grace_period: Duration::from_secs(5), + sleep_grace_period_overridden: true, + ..ActorConfig::default() + }, + signals.clone(), + true, + ); + let run = tokio::spawn(task.run()); + + let (start_tx, start_rx) = oneshot::channel(); + lifecycle_tx + .send(LifecycleCommand::Start { reply: start_tx }) + .expect("start command should send"); + start_rx + .await + .expect("start reply should send") + .expect("start should succeed"); + + let (sleep_tx, sleep_rx) = oneshot::channel(); + lifecycle_tx + .send(LifecycleCommand::Stop { + reason: ShutdownKind::Sleep, + reply: sleep_tx, + }) + .expect("sleep stop should send"); + signals.on_sleep_started.notified().await; + + let error = dispatch_action(&dispatch_tx, "ping") + .await + .expect("action reply should send") + .expect_err("dispatch should be refused once onSleep runs"); + let error = rivet_error::RivetError::extract(&error); + assert_eq!((error.group(), error.code()), ("actor", "restarting")); + assert_eq!(signals.log(), vec!["on_sleep"]); + + signals.release_on_sleep.notify_one(); sleep_rx .await .expect("sleep reply should send") @@ -4311,6 +4445,16 @@ pub(crate) mod moved_tests { "a crash must not request a destroy" ); + // The errored stop request runs as tracked shutdown work, and sleep hooks + // wait for that work. This test drives the task by hand, so let it finish. + timeout(Duration::from_secs(2), async { + while ctx.shutdown_task_count() != 0 { + yield_now().await; + } + }) + .await + .expect("errored stop request should finish"); + let (stop_tx, stop_rx) = oneshot::channel(); task.handle_lifecycle(LifecycleCommand::Stop { reason: ShutdownKind::Sleep, diff --git a/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/sleep-db.ts b/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/sleep-db.ts index fbfb56f2de..6cbd1ac8a7 100644 --- a/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/sleep-db.ts +++ b/rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/sleep-db.ts @@ -348,6 +348,7 @@ export const sleepWithDbAction = actor({ }), events: { sleeping: event(), + slowStarted: event(), }, onWake: async (c) => { c.state.startCount += 1; @@ -400,6 +401,16 @@ export const sleepWithDbAction = actor({ `INSERT INTO sleep_log (event, created_at) VALUES ('${event}', ${Date.now()})`, ); }, + slowInsert: async (c, event: string, ms: number) => { + await c.db.execute( + `INSERT INTO sleep_log (event, created_at) VALUES ('${event}-start', ${Date.now()})`, + ); + c.broadcast("slowStarted", undefined); + await new Promise((resolve) => setTimeout(resolve, ms)); + await c.db.execute( + `INSERT INTO sleep_log (event, created_at) VALUES ('${event}-end', ${Date.now()})`, + ); + }, }, options: { sleepTimeout: SLEEP_DB_TIMEOUT, diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/native.ts b/rivetkit-typescript/packages/rivetkit/src/registry/native.ts index 0c5c55648e..731544355e 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/native.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/native.ts @@ -531,9 +531,10 @@ async function cleanupNativeSleepRuntimeState( ): Promise { // The bounded wait gives shutdown work one grace-period chance to finish. // Drained means all tracked shutdown work completed before the deadline, so - // we can save final state and clear runtime state immediately. If it did not - // drain, close database handles now, then defer the final save and clear until - // the tracked work finishes without a deadline. + // we can save final state and clear runtime state immediately. Core then + // syncs the alarm and closes SQLite. If it did not drain, close database + // handles now so late work cannot write after shutdown, then defer the + // final save and clear until the tracked work finishes without a deadline. const drained = await runtime.actorWaitForTrackedShutdownWork(ctx); if (!drained) { await closeNativeDatabaseClient(runtime, ctx); @@ -554,8 +555,6 @@ async function cleanupNativeSleepRuntimeState( } await afterTrackedWorkDrained?.(); - await closeNativeDatabaseClient(runtime, ctx); - await closeNativeSqlDatabase(runtime, ctx); clearNativeRuntimeState(runtime, ctx); } diff --git a/rivetkit-typescript/packages/rivetkit/tests/driver/actor-sleep-db.test.ts b/rivetkit-typescript/packages/rivetkit/tests/driver/actor-sleep-db.test.ts index 7e84779e7d..8e220740d6 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/driver/actor-sleep-db.test.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/driver/actor-sleep-db.test.ts @@ -434,6 +434,37 @@ describeDriverMatrix( expect(events).toContain("sleep-end"); }); + test("an action in flight when sleep starts finishes before onSleep runs", async (c) => { + const { client } = await setupDriverTest(c, driverTestConfig); + + const handle = client.sleepWithDbAction.getOrCreate([ + "in-flight-action-before-on-sleep", + ]); + const connection = handle.connect(); + await waitForConnectionReady(connection); + + const slowStarted = new Promise((resolve) => { + connection.once("slowStarted", () => resolve()); + }); + const slow = connection.slowInsert("slow", 300); + await slowStarted; + await connection.triggerSleep(); + await slow; + await connection.dispose(); + + await waitForAction(handle.getCounts, (counts) => { + expect(counts.sleepCount).toBe(1); + expect(counts.startCount).toBe(2); + }); + const events = (await handle.getLogEntries()).map( + (entry: LogEntry) => entry.event, + ); + expect(events).toContain("slow-end"); + expect(events.indexOf("slow-end")).toBeLessThan( + events.indexOf("sleep-start"), + ); + }); + test("action via handle during sleep shutdown is not queued", async (c) => { const { client } = await setupDriverTest(c, driverTestConfig);