fix(rivetkit): fire abort signal on shutdown grace

Author: NathanFlurry

docs-internal/engine/rivetkit-core-internals.md

@@ -85,7 +85,7 @@ Two-phase:
 - `SleepGrace` fires `onSleep` immediately and keeps dispatch/save timers live.
 - `SleepFinalize` gates dispatch, suspends alarms, and runs teardown.
 
-Sleep grace must fire the actor abort signal on entry and wait for the run handler to exit before finalize. Destroy abort firing remains unchanged.
+Sleep grace fires the actor abort signal on entry and waits for the run handler to exit before finalize.
 
 Finalize:
 

docs-internal/engine/sleep-sequence.md

@@ -39,7 +39,7 @@ 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.
-- For destroy, the abort signal may fire earlier than grace entry because `ctx.destroy()` cancels the abort token immediately via `mark_destroy_requested(...)`.
+- Destroy requests also use the normal grace path. The actor abort signal fires when destroy grace starts.
 
 ## Grace deadline enforcement
 

rivetkit-rust/packages/rivetkit-core/src/actor/context.rs

@@ -478,19 +478,17 @@ impl ActorContext {
 		self.flush_on_shutdown();
 		self.0.destroy_requested.store(true, Ordering::SeqCst);
 		self.0.destroy_completed.store(false, Ordering::SeqCst);
-		self.0.abort_signal.lock().cancel();
 	}
 
 	#[cfg(feature = "wasm-runtime")]
 	fn mark_destroy_requested_without_spawn(&self) {
 		self.cancel_sleep_timer();
 		self.0.destroy_requested.store(true, Ordering::SeqCst);
 		self.0.destroy_completed.store(false, Ordering::SeqCst);
-		self.0.abort_signal.lock().cancel();
 	}
 
 	#[doc(hidden)]
-	pub fn cancel_abort_signal_for_sleep(&self) {
+	pub fn cancel_actor_abort_signal(&self) {
 		self.0.abort_signal.lock().cancel();
 	}
 
@@ -500,9 +498,9 @@ impl ActorContext {
 			return;
 		}
 
-		// Sleep cancels the generation abort signal to break queue waits and the
-		// run loop out of blocking calls. A restarted actor needs a fresh signal
-		// so the next generation can wait normally.
+		// Sleep or destroy cancels the generation abort signal to break actor
+		// scoped waits. A restarted actor needs a fresh signal so the next
+		// generation can wait normally.
 		let next_signal = CancellationToken::new();
 		*abort_signal = next_signal.clone();
 		*self.0.queue_abort_signal.lock() = next_signal;

rivetkit-rust/packages/rivetkit-core/src/actor/task.rs

@@ -1528,9 +1528,7 @@ impl ActorTask {
 		let grace_period = self.factory.config().effective_sleep_grace_period();
 		self.sleep_deadline = None;
 		self.ctx.cancel_sleep_timer();
-		// Entering grace cancels the actor abort signal so user code blocked on
-		// queues or other actor scoped waits can unwind and let sleep finalize.
-		self.ctx.cancel_abort_signal_for_sleep();
+		self.ctx.cancel_actor_abort_signal();
 		self.sleep_grace = Some(SleepGraceState {
 			deadline: Instant::now() + grace_period,
 			reason,
@@ -1745,15 +1743,18 @@ impl ActorTask {
 			ShutdownKind::Sleep => LifecycleState::SleepFinalize,
 			ShutdownKind::Destroy => LifecycleState::Destroying,
 		});
-		self.save_final_state().await?;
-		self.close_actor_event_channel();
-		self.join_aborted_run_handle().await;
-		Self::finish_shutdown_cleanup_with_ctx(self.ctx.clone(), reason).await?;
-		if matches!(reason, ShutdownKind::Destroy) {
+		let result: Result<()> = async {
+			self.save_final_state().await?;
+			self.close_actor_event_channel();
+			self.join_aborted_run_handle().await;
+			Self::finish_shutdown_cleanup_with_ctx(self.ctx.clone(), reason).await
+		}
+		.await;
+		if result.is_ok() && matches!(reason, ShutdownKind::Destroy) {
 			self.ctx.mark_destroy_completed();
 		}
 		self.ctx.record_shutdown_wait(reason, started_at.elapsed());
-		Ok(())
+		result
 	}
 
 	async fn save_final_state(&mut self) -> Result<()> {

rivetkit-rust/packages/rivetkit-core/tests/queue.rs

@@ -162,7 +162,7 @@ mod moved_tests {
 		});
 
 		yield_now().await;
-		queue.mark_destroy_requested();
+		queue.cancel_actor_abort_signal();
 
 		let error = wait
 			.await

rivetkit-typescript/packages/rivetkit/src/actor/config.ts

@@ -1227,7 +1227,7 @@ interface BaseActorConfig<
 	 * shutdown window (`sleepGracePeriod`) cover deferred work.
 	 *
 	 * The handler receives an abort signal via `c.abortSignal` and a
-	 * `c.aborted` alias for loop checks. Use these to gracefully exit.
+	 * `c.aborted` alias. Use these to gracefully exit when shutdown starts.
 	 *
 	 * If this handler exits, the actor will follow the normal idle sleep timeout
 	 * once it becomes idle.
@@ -1770,7 +1770,7 @@ export const DocActorOptionsSchema = z
 			.number()
 			.optional()
 			.describe(
-				`Max time in ms for the graceful shutdown window. Covers lifecycle hooks (onSleep, onDestroy), the run handler abort wait, async raw WebSocket handlers, disconnect callbacks, and final state serialization. Default: ${DEFAULT_SLEEP_GRACE_PERIOD}.`,
+				`Max time in ms for the graceful shutdown window. Covers lifecycle hooks (onSleep, onDestroy), the run handler wait, async raw WebSocket handlers, disconnect callbacks, and final state serialization. Default: ${DEFAULT_SLEEP_GRACE_PERIOD}.`,
 			),
 		onDestroyTimeout: z
 			.number()

rivetkit-typescript/packages/rivetkit/src/actor/instance/mod.ts

@@ -957,14 +957,7 @@ export class ActorInstance<
 			// Scheduled events are persisted and will be re-initialized
 			// on wake via initializeAlarms().
 			this.driver.cancelAlarm?.(this.#actorId);
-
-			// Abort listeners in the canonical stop path.
-			// This must run for all stop modes, including sleep and remote stop.
-			// Destroy may have already triggered an early abort, but repeating abort
-			// is intentional and safe.
-			try {
-				this.#abortController.abort();
-			} catch {}
+			this.#abortActorSignal();
 
 			// The run-handler join, lifecycle hooks, and remaining shutdown
 			// tasks all share the single sleepGracePeriod budget.
@@ -1030,11 +1023,8 @@ export class ActorInstance<
 			}
 
 			this.driver.cancelAlarm?.(this.#actorId);
+			this.#abortActorSignal();
 			this.stateManager.clearPendingSaveTimeout();
-
-			try {
-				this.#abortController.abort();
-			} catch {}
 		} finally {
 			this.#shutdownComplete = true;
 			await this.#cleanupDatabase();
@@ -1086,14 +1076,7 @@ export class ActorInstance<
 			return;
 		}
 		this.#destroyCalled = true;
-
-		// Abort immediately so in flight waits can exit before the driver stop
-		// handshake completes.
-		// The onStop path will call abort again as a safety net for all stop
-		// modes.
-		try {
-			this.#abortController.abort();
-		} catch {}
+		this.#abortActorSignal();
 
 		const destroy = this.driver.startDestroy.bind(
 			this.driver,
@@ -1108,6 +1091,15 @@ export class ActorInstance<
 		});
 	}
 
+	#abortActorSignal() {
+		if (this.#abortController.signal.aborted) {
+			return;
+		}
+		try {
+			this.#abortController.abort();
+		} catch {}
+	}
+
 	// MARK: - HTTP Request Tracking
 	beginHonoHttpRequest() {
 		this.#activeHonoHttpRequests++;
@@ -2070,7 +2062,7 @@ export class ActorInstance<
 
 		if (timeoutMs <= 0) {
 			this.#rLog.warn({
-				msg: "run handler did not complete in time, it may have leaked - ensure you use c.aborted (or the abort signal c.abortSignal) to exit gracefully",
+				msg: "run handler did not complete in time, it may have leaked - ensure long-running work settles before shutdown finalizes",
 				timeoutMs,
 			});
 			return;
@@ -2085,7 +2077,7 @@ export class ActorInstance<
 
 		if (timedOut) {
 			this.#rLog.warn({
-				msg: "run handler did not complete in time, it may have leaked - ensure you use c.aborted (or the abort signal c.abortSignal) to exit gracefully",
+				msg: "run handler did not complete in time, it may have leaked - ensure long-running work settles before shutdown finalizes",
 				timeoutMs,
 			});
 		} else {

website/src/content/docs/actors/actions.mdx

@@ -325,7 +325,7 @@ Actions have a single return value. To stream realtime data in response to an ac
 
 ## Canceling Long-Running Actions
 
-For operations that should be cancelable on-demand, create your own `AbortController` and chain it with `c.abortSignal` for automatic cleanup on actor shutdown.
+For operations that should be cancelable on-demand, create your own `AbortController`. Chain it with `c.abortSignal` so actor shutdown also cancels the operation.
 
 ```typescript
 import { actor } from "rivetkit";

website/src/content/docs/actors/index.mdx

@@ -466,7 +466,7 @@ const userAccount = actor({
 
 ### Lifecycle Hooks
 
-Actors support hooks for initialization, background processing, connections, networking, and state changes. Use `run` for long-lived background loops, and exit cleanly on shutdown with `c.aborted` or `c.abortSignal`.
+Actors support hooks for initialization, background processing, connections, networking, and state changes. Use `run` for long-lived background loops, and use `c.aborted` or `c.abortSignal` for graceful shutdown.
 
 ```ts
 import { actor, event, queue } from "rivetkit";

website/src/content/docs/actors/lifecycle.mdx

@@ -304,7 +304,7 @@ const tickActor = actor({
       c.state.tickCount++;
       c.log.info({ msg: "tick", count: c.state.tickCount });
 
-      // Wait 1 second, but exit early if aborted
+      // Wait 1 second. Final shutdown also resolves this wait.
       await new Promise<void>((resolve) => {
         const timeout = setTimeout(resolve, 1000);
         c.abortSignal.addEventListener("abort", () => {
@@ -876,7 +876,7 @@ When an actor sleeps or is destroyed, it enters the graceful shutdown window:
 
 1. `c.abortSignal` fires and `c.aborted` becomes `true`. New connections and dispatch are rejected. Alarm timeouts are cancelled. On sleep, scheduled events are persisted and will be re-armed when the actor wakes.
 2. `onSleep` or `onDestroy` and `onDisconnect` for each closing connection run during the same window. User `waitUntil` promises and async raw WebSocket handlers are drained. Hibernatable WebSocket connections are preserved on sleep and closed on destroy.
-3. Once graceful work has completed, state is saved and the database is cleaned up.
+3. Once graceful work has completed, state is saved and final cleanup runs.
 
 The entire window is bounded by `sleepGracePeriod` on both sleep and destroy. Defaults to 15 seconds. If the window is exceeded, the actor proceeds to state save anyway.