diff --git a/CLAUDE.md b/CLAUDE.md index 36aaf5e3e3..4d4faed996 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -262,6 +262,7 @@ When the user asks to track something in a note, store it in `~/.agents/notes/` - In `rivetkit-core` `ActorTask::run`, bind inbox `recv()` calls as raw `Option`s and log the closed channel before terminating. `Some(...) = recv()` plus `else => break` hides which inbox died. - In `rivetkit-typescript/packages/rivetkit/src/common/utils.ts::deconstructError`, only passthrough canonical structured errors (`instanceof RivetError` or tagged `__type: "RivetError"` with full fields). Plain-object lookalikes must still be classified and sanitized. - Actor-owned lifecycle / dispatch / lifecycle-event inbox producers use `try_reserve` helpers and return `actor.overloaded`. Do not await bounded `mpsc::Sender::send`. +- A Lost actor generation must stop immediately: no user sleep hooks, no final save, and no storage writes after the lost signal fires. Never treat Lost as a graceful sleep. ## Performance diff --git a/docs-internal/engine/rivetkit-core-internals.md b/docs-internal/engine/rivetkit-core-internals.md index 7d4312b25e..b780791cb2 100644 --- a/docs-internal/engine/rivetkit-core-internals.md +++ b/docs-internal/engine/rivetkit-core-internals.md @@ -66,7 +66,11 @@ Legacy traces under `[7, 1, ...]` are intentionally skipped and are not copied t ## Registry + dispatch - Registry startup builds configured `ActorContext`s with `ActorContext::build(...)` so state, queue, and connection managers inherit the actor config before lifecycle startup runs. `ActorContext::build(...)` must seed owned queue, connection, and sleep config storage from its `ActorConfig`; do not initialize those fields with `ActorConfig::default()`. -- Registry actor task handles live in one `actor_instances: SccHashMap`. Use `entry_async` for Active/Stopping transitions. +- Registry records (`actor_instances`, `starting_instances`, `pending_stops`) are keyed by `(actor_id, generation)`, so a lost generation that is still tearing down stays reachable after the next generation starts. `current_generations` holds the newest registered generation per actor id, and dispatch routes only to it. Use `entry_async` for Active/Stopping transitions. +- A stop for a starting generation parks under that generation's starting entry and is consumed by its own startup. A stop for an older generation with no record completes immediately, and registering a newer generation completes stops parked for older ones. +- A starting generation marks older generations of the same actor on the runner, starting or registered, lost and waits up to 10s for them to finish (`ActorContext::wait_for_generation_finished`). A generation finishes once every `GenerationHold` is dropped: the actor task holds one, and a lost generation's storage cleanup holds another until a strict SQLite close returns. Runtime wrappers that spawn per-generation tasks outside the actor task take a hold for each and stop them on the lost signal. If one does not finish, the new start fails instead of overlapping it. The engine only starts a generation after giving up on the older ones. +- A generation lost during startup stays registered as starting until its task finishes aborting, so envoy-client does not report it stopped early. +- A stopped generation's record is removed only once the generation has finished, and a generation refuses to start if a newer generation of the same actor is already starting or registered, since start callbacks can arrive out of order. - `RegistryDispatcher::handle_fetch` owns framework HTTP routes `/metadata`, `/health`, `/inspector/*`, `/action/*`, and `/queue/*`. Prometheus metrics are process-wide and are never served by core itself; `metrics::render_prometheus_metrics()` only backs the NAPI `metrics()` binding and `Registry::prometheus_metrics()` so hosts can mount the rendered response on their own router. TypeScript NAPI callbacks keep action/queue schema validation and queue `canPublish`. - Raw `onRequest` HTTP fetches bypass `maxIncomingMessageSize` / `maxOutgoingMessageSize`. Those message-size guards apply only to `/action/*` and `/queue/*` framework routes, not unmatched user `onRequest` paths. - Framework HTTP error payloads omit absent `metadata` for JSON/CBOR responses so missing metadata stays `undefined`. Only explicit metadata `null` serializes as `null`. diff --git a/docs-internal/engine/sleep-sequence.md b/docs-internal/engine/sleep-sequence.md index 9f9d5c7637..d562937ddf 100644 --- a/docs-internal/engine/sleep-sequence.md +++ b/docs-internal/engine/sleep-sequence.md @@ -49,6 +49,15 @@ When the grace deadline elapses before `can_finalize_sleep()` returns true: - The NAPI `RunGracefulCleanup` task observes `shutdown_deadline_token()` via `tokio::select!` and aborts its in-flight `onSleep` / `onDestroy` call so SQLite and KV cleanup in `teardown_sleep_state` do not race against mid-commit user work. - Foreign-runtime adapters that run user cleanup callbacks must observe the shutdown deadline token the same way. +## Lost generations + +- `StopActor { reason: Lost }` or an envoy-side lost timeout means the engine has given up on the generation and may already be running the next one. envoy-client cancels a per-generation lost `CancellationToken` synchronously, before the stop is queued, and passes it to core through `on_actor_start_with_lost_signal`. +- Core adopts the token as the context lost signal (`ActorContext::is_lost`). Once it fires, `SqliteDb` rejects every statement except `ROLLBACK` and `LegacyActorKv` rejects writes, so state, KV, queue, and workflow writes from the lost generation are revoked without waiting on the task. +- `ActorTask` observes the token in `run_live`, in graceful shutdown, and in serialization waits, and jumps to `abort_lost()`. That path runs no user sleep hooks, skips the final save and alarm sync, dispatches runtime cleanup once, and bounds every wait by `LOST_SHUTDOWN_GRACE_PERIOD`. Cleanup that has not finished stays owned by its retained task. +- The final `ActorStateStopped` for a lost generation is always `StopCode::Error` with "envoy connection lost", including when Lost escalates a pending graceful stop. `StopCode::Ok` would make the engine destroy the actor. +- Already-submitted SQLite commits cannot be recalled. Depot generation checks are the backstop for those. +- A lost generation starts no new user callback. Every runtime callback passes through `refuseCallbacksAfterLost` in `registry/native.ts` when JS enters it, except the `onSleep` and `onDestroy` cleanup wrappers, which release runtime state and skip the user hook themselves. Callbacks the framework starts on its own (workflow steps, run, inspector workflow calls, deferred `onStateChange`, WebSocket open and message listeners, database provider hooks) check the lost state first. Continuations inside a user call that is already running, and abort or close listeners, are not blocked. + ## Guarding lifecycle requests - `ctx.sleep()` and `ctx.destroy()` return `Result<()>`. They fail with `actor/starting` if called before startup completes and `actor/stopping` if the request flag has already been swapped to true for this generation. An atomic `swap(true, ...)` on `sleep_requested` / `destroy_requested` enforces single-shot request semantics per generation. diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs index e22daaba7b..7e41bf28e4 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/context.rs @@ -169,6 +169,13 @@ pub(crate) struct ActorContextInner { destroy_requested: AtomicBool, destroy_completed: AtomicBool, destroy_completion_notify: Notify, + /// Outstanding holds on this generation. The actor task holds one, and a lost generation's + /// storage cleanup holds another until SQLite is released. The registry waits for zero + /// before starting a newer generation. + generation_holds: AtomicUsize, + generation_finished_notify: Notify, + // Forced-sync: read by synchronous accessors and bound before the actor task starts. + lost: Mutex, abort_signal: Mutex, shutdown_deadline: CancellationToken, // Forced-sync: runtime wiring slots are configured through synchronous @@ -429,6 +436,9 @@ impl ActorContext { let max_schedules = config.max_schedules; let abort_signal = CancellationToken::new(); let shutdown_deadline = CancellationToken::new(); + let lost = CancellationToken::new(); + sql.set_lost_signal(lost.clone()); + legacy_kv.set_lost_signal(lost.clone()); let sleep = SleepState::new(config.clone()); let user_kv = ActorKv { sql: sql.clone() }; let inner = Arc::new(ActorContextInner { @@ -508,6 +518,9 @@ impl ActorContext { destroy_requested: AtomicBool::new(false), destroy_completed: AtomicBool::new(false), destroy_completion_notify: Notify::new(), + generation_holds: AtomicUsize::new(0), + generation_finished_notify: Notify::new(), + lost: Mutex::new(lost), abort_signal: Mutex::new(abort_signal), shutdown_deadline, inspector: RwLock::new(None), @@ -774,6 +787,31 @@ impl ActorContext { Ok(()) } + /// Adopts the lost signal envoy-client created for this generation. envoy-client cancels it + /// synchronously when the engine or the envoy declares the generation lost, which revokes + /// SQLite writes immediately and wakes the actor task into its abort path. + pub(crate) fn configure_lost_signal(&self, lost: CancellationToken) { + self.sql().set_lost_signal(lost.clone()); + self.0.legacy_kv.set_lost_signal(lost.clone()); + *self.0.lost.lock() = lost; + } + + /// Returns this generation's lost signal. + pub fn lost_signal(&self) -> CancellationToken { + self.0.lost.lock().clone() + } + + /// Declares this generation lost. Idempotent. + pub(crate) fn mark_lost(&self) { + self.0.lost.lock().cancel(); + } + + /// Returns whether this generation was declared lost. Foreign-runtime shutdown hooks must + /// skip user callbacks and state saves when this is set, and only release runtime resources. + pub fn is_lost(&self) -> bool { + self.0.lost.lock().is_cancelled() + } + pub fn mark_destroy_requested(&self) { self.cancel_sleep_timer(); self.flush_on_shutdown(); @@ -1531,6 +1569,31 @@ impl ActorContext { self.destroy_requested() } + /// Keeps this generation from counting as finished until the returned hold is dropped. + /// Runtime wrappers hold it for tasks they spawn outside the actor task. + #[doc(hidden)] + pub fn hold_generation(&self) -> GenerationHold { + self.0.generation_holds.fetch_add(1, Ordering::SeqCst); + GenerationHold(self.clone()) + } + + pub(crate) fn is_generation_finished(&self) -> bool { + self.0.generation_holds.load(Ordering::SeqCst) == 0 + } + + pub(crate) async fn wait_for_generation_finished(&self) { + loop { + let notified = self.0.generation_finished_notify.notified(); + tokio::pin!(notified); + // Register before checking so a finish between the check and the await is not missed. + notified.as_mut().enable(); + if self.is_generation_finished() { + return; + } + notified.await; + } + } + pub(crate) async fn wait_for_destroy_completion(&self) { if self.0.destroy_completed.load(Ordering::SeqCst) { return; @@ -2172,6 +2235,17 @@ impl std::fmt::Debug for ActorContext { } // Test shim keeps moved tests in crate-root tests/ with private-module access. +/// A hold on an actor generation. See [`ActorContext::hold_generation`]. +pub struct GenerationHold(ActorContext); + +impl Drop for GenerationHold { + fn drop(&mut self) { + if self.0.0.generation_holds.fetch_sub(1, Ordering::SeqCst) == 1 { + self.0.0.generation_finished_notify.notify_waiters(); + } + } +} + #[cfg(test)] #[path = "../../tests/context.rs"] pub(crate) mod tests; diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/kv.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/kv.rs index 066c127723..d4fd9ce7c0 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/kv.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/kv.rs @@ -1,13 +1,15 @@ #[cfg(test)] use std::collections::BTreeMap; -#[cfg(test)] use std::sync::Arc; -use anyhow::Result; +use anyhow::{Context, Result}; +use parking_lot::Mutex; #[cfg(test)] -use parking_lot::{Mutex, RwLock}; +use parking_lot::RwLock; use rivet_envoy_client::handle::EnvoyHandle; +use tokio_util::sync::CancellationToken; +use crate::error::SqliteRuntimeError; use crate::types::ListOpts; /// Narrow access to the actor's pre-SQLite KV namespace. @@ -18,6 +20,8 @@ use crate::types::ListOpts; #[derive(Clone)] pub(crate) struct LegacyActorKv { backend: LegacyActorKvBackend, + /// The owning generation's lost signal. Writes are revoked once it fires. + lost: Arc>>, } #[derive(Clone)] @@ -45,6 +49,7 @@ impl LegacyActorKv { handle, actor_id: actor_id.into(), }, + lost: Default::default(), } } @@ -57,9 +62,27 @@ impl LegacyActorKv { list_limit_cap: Mutex::new(None), range_start_inclusive: Mutex::new(true), })), + lost: Default::default(), } } + pub(crate) fn set_lost_signal(&self, lost: CancellationToken) { + *self.lost.lock() = Some(lost); + } + + fn ensure_not_lost(&self) -> Result<()> { + let lost = self + .lost + .lock() + .as_ref() + .is_some_and(CancellationToken::is_cancelled); + if lost { + return Err(SqliteRuntimeError::Closed.build()) + .context("actor generation was declared lost; storage writes are revoked"); + } + Ok(()) + } + #[cfg(test)] pub(crate) async fn get(&self, key: &[u8]) -> Result>> { Ok(self.batch_get(&[key]).await?.pop().flatten()) @@ -88,17 +111,20 @@ impl LegacyActorKv { } pub(crate) async fn batch_put(&self, entries: &[(&[u8], &[u8])]) -> Result<()> { + self.ensure_not_lost()?; match &self.backend { LegacyActorKvBackend::Envoy { handle, actor_id } => { - handle - .kv_put( - actor_id.clone(), - entries - .iter() - .map(|(key, value)| (key.to_vec(), value.to_vec())) - .collect(), - ) - .await + let entries = entries + .iter() + .map(|(key, value)| (key.to_vec(), value.to_vec())) + .collect(); + // The write can sit in envoy's queue, so envoy drops it if the generation is + // declared lost before it is sent. + let lost = self.lost.lock().clone(); + match lost { + Some(lost) => handle.kv_put_fenced(actor_id.clone(), entries, lost).await, + None => handle.kv_put(actor_id.clone(), entries).await, + } } #[cfg(test)] LegacyActorKvBackend::InMemory(store) => { diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/sleep.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/sleep.rs index e224e5127d..54b5452683 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/sleep.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/sleep.rs @@ -294,6 +294,18 @@ impl ActorContext { && self.pending_disconnect_count() == 0 } + /// Waits until every core-dispatched hook has replied, or `deadline` passes. Waits on the hook + /// counter itself. The sleep activity signal only fires when its dirty flag flips, which the + /// lost abort path never acknowledges, so it can miss a hook completing. + pub(crate) async fn wait_for_core_dispatched_hooks(&self, deadline: Instant) -> bool { + self.0 + .sleep + .work + .core_dispatched_hooks + .wait_zero(deadline) + .await + } + /// Spawn the fallback sleep timer used by `ActorContext`s that are not /// bound to an `ActorTask`. /// @@ -530,6 +542,13 @@ impl ActorContext { } let policy = kind.policy(); + if policy.aborts_at_shutdown_deadline && self.is_lost() { + tracing::warn!( + kind = kind.label(), + "actor work spawned after the generation was lost; dropping it" + ); + return false; + } if policy.aborts_at_shutdown_deadline { let mut shutdown_tasks = self.0.sleep.work.shutdown_tasks.lock(); if self.0.sleep.work.teardown_started.load(Ordering::Acquire) { @@ -582,7 +601,19 @@ impl ActorContext { let _region = region; if policy.aborts_at_shutdown_deadline { let shutdown_deadline = ctx.shutdown_deadline_token(); + // A lost generation runs no more user code, so its work is dropped as soon as the + // engine gives up on it. Loss wins over a future that is also ready. + let lost = ctx.lost_signal(); tokio::select! { + biased; + _ = lost.cancelled() => { + tracing::debug!( + actor_id = %ctx.actor_id(), + kind = kind.label(), + reason = "generation_lost", + "actor work cancelled because the generation was lost" + ); + } _ = fut => {} _ = shutdown_deadline.cancelled() => { tracing::warn!( @@ -636,6 +667,9 @@ impl ActorContext { let task = async move { let _region = region; if policy.aborts_at_shutdown_deadline { + // Wasm work wraps JS promises that loss cannot stop, so it keeps only the + // shutdown-deadline policy. Lost wasm generations are fenced by storage + // revocation and the generation hold instead. let shutdown_deadline = ctx_for_task.shutdown_deadline_token(); tokio::select! { _ = fut => {} diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/sqlite/mod.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/sqlite/mod.rs index 70d7f11d9c..3357045e68 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/sqlite/mod.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/sqlite/mod.rs @@ -450,13 +450,23 @@ impl SqliteDb { #[cfg(feature = "sqlite-local")] async fn local_exec(&self, sql: String) -> Result { - self.open().await?; + self.ensure_not_lost(&sql)?; + let opened = self.open().await; + // Opening can wait on the open lock and the depot, so the generation may have been lost + // meanwhile. Recheck right before submitting to the native worker. + self.ensure_not_lost(&sql)?; + opened?; self.map_local_worker_result(self.native_db_handle()?.exec(sql).await) } #[cfg(feature = "sqlite-local")] async fn local_exec_profiled(&self, sql: String) -> Result> { - self.open().await?; + self.ensure_not_lost(&sql)?; + let opened = self.open().await; + // Opening can wait on the open lock and the depot, so the generation may have been lost + // meanwhile. Recheck right before submitting to the native worker. + self.ensure_not_lost(&sql)?; + opened?; self.map_local_worker_result(self.native_db_handle()?.exec_profiled(sql).await) } @@ -471,7 +481,12 @@ impl SqliteDb { sql: String, params: Option>, ) -> Result { - self.open().await?; + self.ensure_not_lost(&sql)?; + let opened = self.open().await; + // Opening can wait on the open lock and the depot, so the generation may have been lost + // meanwhile. Recheck right before submitting to the native worker. + self.ensure_not_lost(&sql)?; + opened?; self.map_local_worker_result(self.native_db_handle()?.execute(sql, params).await) } @@ -481,7 +496,12 @@ impl SqliteDb { sql: String, params: Option>, ) -> Result> { - self.open().await?; + self.ensure_not_lost(&sql)?; + let opened = self.open().await; + // Opening can wait on the open lock and the depot, so the generation may have been lost + // meanwhile. Recheck right before submitting to the native worker. + self.ensure_not_lost(&sql)?; + opened?; self.map_local_worker_result(self.native_db_handle()?.execute_profiled(sql, params).await) } @@ -1130,6 +1150,7 @@ impl SqliteDb { } async fn remote_exec(&self, sql: String) -> Result { + self.ensure_not_lost(&sql)?; let config = self.remote_config()?; let response = config .handle @@ -1157,6 +1178,7 @@ impl SqliteDb { sql: String, expected_session: Option, ) -> Result<(QueryResult, u64)> { + self.ensure_not_lost(&sql)?; let config = self.remote_config()?; let (response, session) = config .handle @@ -1186,6 +1208,7 @@ impl SqliteDb { sql: String, params: Option>, ) -> Result { + self.ensure_not_lost(&sql)?; let config = self.remote_config()?; let response = config .handle @@ -1213,6 +1236,9 @@ impl SqliteDb { &self, statements: Vec, ) -> Result> { + for statement in &statements { + self.ensure_not_lost(&statement.sql)?; + } let config = self.remote_config()?; let response = config .handle @@ -1249,6 +1275,7 @@ impl SqliteDb { params: Option>, expected_session: Option, ) -> Result<(ExecuteResult, u64)> { + self.ensure_not_lost(&sql)?; let config = self.remote_config()?; let (response, session) = config .handle diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/sqlite/tx.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/sqlite/tx.rs index 1102395883..2e95feddf2 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/sqlite/tx.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/sqlite/tx.rs @@ -20,6 +20,7 @@ use tokio::sync::{ }; use tokio_util::sync::CancellationToken; +use crate::error::SqliteRuntimeError; #[cfg(not(target_arch = "wasm32"))] use crate::runtime::RuntimeSpawner; use crate::telemetry::SqliteOperation; @@ -80,6 +81,8 @@ pub(super) struct TransactionCoordinator { pub(super) state: AsyncMutex, epoch: AtomicU64, waiters: AtomicU64, + // Forced-sync: read on every statement submission, including from synchronous paths. + lost: parking_lot::Mutex>, } pub(super) struct TransactionCoordinatorState { @@ -173,11 +176,44 @@ impl Default for TransactionCoordinator { }), epoch: AtomicU64::new(0), waiters: AtomicU64::new(0), + lost: parking_lot::Mutex::new(None), } } } impl SqliteDb { + /// Binds the actor generation's lost signal. Once it fires, every statement except `ROLLBACK` + /// is rejected at submission, including statements of transactions that were admitted + /// earlier, so a lost generation can no longer write actor storage. + /// Whether the owning generation was declared lost. + pub fn is_lost(&self) -> bool { + self.transaction_coordinator + .lost + .lock() + .as_ref() + .is_some_and(CancellationToken::is_cancelled) + } + + pub(crate) fn set_lost_signal(&self, lost: CancellationToken) { + *self.transaction_coordinator.lost.lock() = Some(lost); + } + + /// Rejects a statement from a lost generation. `ROLLBACK` stays allowed because it only + /// discards uncommitted work. + pub(super) fn ensure_not_lost(&self, sql: &str) -> Result<()> { + let lost = self + .transaction_coordinator + .lost + .lock() + .as_ref() + .is_some_and(CancellationToken::is_cancelled); + if lost && !is_rollback_statement(sql) { + return Err(SqliteRuntimeError::Closed.build()) + .context("actor generation was declared lost; storage writes are revoked"); + } + Ok(()) + } + pub(super) fn try_transaction_admission(&self) -> Result { match Arc::clone(&self.transaction_coordinator.admission).try_acquire_owned() { Ok(permit) => Ok(permit), @@ -1406,3 +1442,10 @@ fn transaction_terminal_state_error(key: &str, state: &'static str) -> anyhow::E state, }) } + +fn is_rollback_statement(sql: &str) -> bool { + sql.trim() + .trim_end_matches(';') + .trim() + .eq_ignore_ascii_case("ROLLBACK") +} diff --git a/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs b/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs index d00e8be57c..4d79506ca8 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/actor/task.rs @@ -40,6 +40,7 @@ use futures::FutureExt; use parking_lot::Mutex; use tokio::sync::{broadcast, mpsc, oneshot}; use tokio::task::{JoinError, JoinHandle}; +use tokio_util::sync::CancellationToken; use tracing::{Instrument, instrument::WithSubscriber}; use crate::actor::action::ActionDispatchError; @@ -70,6 +71,10 @@ pub type HttpDispatchResult = Result; const SERIALIZE_STATE_SHUTDOWN_SANITY_CAP: Duration = Duration::from_secs(15); #[cfg(test)] const LONG_SHUTDOWN_DRAIN_WARNING_THRESHOLD: Duration = Duration::from_secs(1); +// Bounds a lost generation's whole teardown: runtime cleanup, joining its event loop, and +// resource cleanup. The engine may already be running the next generation, so user work is never +// awaited past this. +const LOST_SHUTDOWN_GRACE_PERIOD: Duration = Duration::from_secs(2); const INSPECTOR_SERIALIZE_STATE_INTERVAL: Duration = Duration::from_millis(50); const INSPECTOR_OVERLAY_CHANNEL_CAPACITY: usize = 32; @@ -308,7 +313,11 @@ impl PartialEq for LifecycleEvent { impl Eq for LifecycleEvent {} enum LiveExit { - Shutdown { reason: ShutdownKind }, + Shutdown { + reason: ShutdownKind, + }, + /// The generation's lost signal fired; tear down through `abort_lost`. + Lost, Terminated, } @@ -393,6 +402,17 @@ pub struct ActorTask { /// Active sleep-grace idle wait. Polled by the main loop so grace keeps the /// same inbox/timer handling as the started actor. sleep_grace: Option, + + // === LOST === + /// Fires when this generation is declared lost. Observed by the live loop and by graceful + /// finalization, both of which hand off to `abort_lost`. + lost: CancellationToken, + /// Whether the runtime cleanup event (`RunGracefulCleanup`) was already dispatched, so the + /// abort path does not dispatch it twice. + runtime_cleanup_dispatched: bool, + /// Final resource cleanup, spawned once and retained so an interrupted graceful shutdown and + /// the abort path share it instead of starting a second close. + shutdown_cleanup: Option>>, } impl ActorTask { @@ -406,6 +426,7 @@ impl ActorTask { ctx: ActorContext, start_input: Option>, ) -> Self { + let lost = ctx.lost_signal(); let (actor_event_tx, actor_event_rx) = mpsc::unbounded_channel(); let (inspector_overlay_tx, _) = broadcast::channel(INSPECTOR_OVERLAY_CHANNEL_CAPACITY); let inspector_attach_count = Arc::new(AtomicU32::new(0)); @@ -442,6 +463,9 @@ impl ActorTask { sleep_deadline: None, shutdown_reply: None, sleep_grace: None, + lost, + runtime_cleanup_dispatched: false, + shutdown_cleanup: None, } } @@ -455,18 +479,33 @@ impl ActorTask { )] pub async fn run(mut self) -> Result<()> { let exit = self.run_live().await; - let LiveExit::Shutdown { reason } = exit else { - self.record_inbox_depths(); - self.ctx.metrics().record_actor_stopped(); - return Ok(()); - }; - - let result = match AssertUnwindSafe(self.run_shutdown(reason)) - .catch_unwind() - .await - { - Ok(result) => result, - Err(_) => Err(anyhow!("shutdown panicked during {reason:?}")), + let (reason, result) = match exit { + LiveExit::Terminated => { + self.record_inbox_depths(); + self.ctx.metrics().record_actor_stopped(); + return Ok(()); + } + LiveExit::Lost => (ShutdownKind::Sleep, self.abort_lost().await), + LiveExit::Shutdown { reason } => { + // Graceful finalization is abandoned for the abort path if the generation is + // declared lost part way through. Every step of `run_shutdown` leaves its + // resources owned by `self`, so dropping it here loses nothing. + let lost = self.lost.clone(); + let graceful = { + let shutdown = AssertUnwindSafe(self.run_shutdown(reason)).catch_unwind(); + tokio::select! { + biased; + _ = lost.cancelled() => None, + result = shutdown => Some(result), + } + }; + let result = match graceful { + Some(Ok(result)) => result, + Some(Err(_)) => Err(anyhow!("shutdown panicked during {reason:?}")), + None => self.abort_lost().await, + }; + (reason, result) + } }; self.deliver_shutdown_reply(reason, &result); self.transition_to(LifecycleState::Terminated); @@ -477,22 +516,39 @@ impl ActorTask { async fn run_live(&mut self) -> LiveExit { let activity_notify = self.ctx.sleep_activity_notify(); + let lost = self.lost.clone(); loop { if self.ctx.acknowledge_activity_dirty() { - if let Some(exit) = self.on_activity_signal().await { - return exit; + match unless_lost(&lost, self.on_activity_signal()).await { + Some(Some(exit)) => return exit, + Some(None) => {} + None => return LiveExit::Lost, } } // TODO: Sample inbox depths periodically instead of on every loop iteration. self.record_inbox_depths(); tokio::select! { biased; + _ = lost.cancelled() => { + return LiveExit::Lost; + } lifecycle_command = self.lifecycle_inbox.recv() => { match lifecycle_command { Some(command) => { - if let Some(exit) = self.handle_lifecycle(command).await { - return exit; + // Startup owns resources that nothing else retains until it returns, + // such as a SQLite database that is still opening, so it is not + // dropped on loss. It observes loss itself while waiting on the + // runtime preamble. + let exit = if matches!(command, LifecycleCommand::Start { .. }) { + Some(self.handle_lifecycle(command).await) + } else { + unless_lost(&lost, self.handle_lifecycle(command)).await + }; + match exit { + Some(Some(exit)) => return exit, + Some(None) => {} + None => return LiveExit::Lost, } } None => { @@ -506,7 +562,11 @@ impl ActorTask { } lifecycle_event = self.lifecycle_events.recv() => { match lifecycle_event { - Some(event) => self.handle_event(event).await, + Some(event) => { + if unless_lost(&lost, self.handle_event(event)).await.is_none() { + return LiveExit::Lost; + } + } None => { self.log_closed_channel( "lifecycle_events", @@ -518,18 +578,26 @@ impl ActorTask { } _ = activity_notify.notified() => { self.ctx.acknowledge_activity_dirty(); - if let Some(exit) = self.on_activity_signal().await { - return exit; + match unless_lost(&lost, self.on_activity_signal()).await { + Some(Some(exit)) => return exit, + Some(None) => {} + None => return LiveExit::Lost, } } _ = 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; + match unless_lost(&lost, self.on_sleep_grace_deadline()).await { + Some(Some(exit)) => return exit, + Some(None) => {} + None => return LiveExit::Lost, } } dispatch_command = self.dispatch_inbox.recv(), if self.accepting_dispatch() => { match dispatch_command { - Some(command) => self.handle_dispatch(command).await, + Some(command) => { + if unless_lost(&lost, self.handle_dispatch(command)).await.is_none() { + return LiveExit::Lost; + } + } None => { self.log_closed_channel( "dispatch_inbox", @@ -545,13 +613,19 @@ impl ActorTask { } } _ = Self::state_save_tick(self.state_save_deadline), if self.state_save_timer_active() => { - self.on_state_save_tick().await; + if unless_lost(&lost, self.on_state_save_tick()).await.is_none() { + return LiveExit::Lost; + } } _ = Self::inspector_serialize_state_tick(self.inspector_serialize_state_deadline), if self.inspector_serialize_timer_active() => { - self.on_inspector_serialize_state_tick().await; + if unless_lost(&lost, self.on_inspector_serialize_state_tick()).await.is_none() { + return LiveExit::Lost; + } } _ = Self::sleep_tick(self.sleep_deadline), if self.sleep_timer_active() => { - self.on_sleep_tick().await; + if unless_lost(&lost, self.on_sleep_tick()).await.is_none() { + return LiveExit::Lost; + } } } @@ -577,6 +651,12 @@ impl ActorTask { LifecycleCommand::Start { reply } => { let result = self.start_actor().await; let failed = result.is_err(); + if failed && self.ctx.is_lost() { + // A lost generation tears down through the bounded abort path rather than + // waiting on a run handler that may never finish its preamble. + self.reply_lifecycle_command(command_kind, reason, reply, result); + return Some(LiveExit::Lost); + } if failed { self.finish_failed_startup().await; } @@ -801,6 +881,7 @@ impl ActorTask { } } + self.runtime_cleanup_dispatched = true; self.ctx.begin_core_dispatched_hook(); let reply = self.core_dispatched_hook_reply("run_graceful_cleanup"); if let Err(error) = self.send_actor_event( @@ -857,6 +938,9 @@ impl ActorTask { } async fn flush_workflow_state(&mut self, writes: Vec) -> Result<()> { + if self.ctx.is_lost() { + return Err(ActorLifecycleError::Stopping.build()); + } if !matches!( self.lifecycle, LifecycleState::Started | LifecycleState::SleepGrace @@ -874,9 +958,10 @@ impl ActorTask { reply: Reply::from(reply_tx), }, )?; - let deltas = reply_rx - .await - .context("receive workflow flush serialize-state reply")??; + let Some(reply) = unless_lost(&self.lost, reply_rx).await else { + return Err(ActorLifecycleError::Stopping.build()); + }; + let deltas = reply.context("receive workflow flush serialize-state reply")??; if self .ctx .save_state_and_workflow_batch_at_transaction_epoch( @@ -1152,6 +1237,10 @@ impl ActorTask { fn dispatch_lifecycle_error(&self) -> Option { // TODO: Share admission policy with RegistryDispatcher::active_actor. + if self.ctx.is_lost() { + self.ctx.warn_work_sent_to_stopping_instance("dispatch"); + return Some(ActorLifecycleError::Stopping.build()); + } if self.ctx.destroy_requested() { self.ctx.warn_work_sent_to_stopping_instance("dispatch"); return Some(ActorLifecycleError::Destroying.build()); @@ -1387,8 +1476,15 @@ impl ActorTask { .with_subscriber(run_dispatch), )); if let Some(startup_ready_rx) = startup_ready_rx { - startup_ready_rx - .await + // The run handle is retained in `self.run_handle`, so giving up on the preamble only + // cancels this wait. + let lost = self.lost.clone(); + let Some(ready) = unless_lost(&lost, startup_ready_rx).await else { + return Err(ActorLifecycleError::Stopping.build()).context( + "actor generation was declared lost during the runtime startup preamble", + ); + }; + ready .context("receive runtime startup ready reply")? .context("runtime startup preamble")?; } @@ -1675,21 +1771,167 @@ impl ActorTask { }) } - async fn join_aborted_run_handle(&mut self) { - let Some(mut run_handle) = self.run_handle.take() else { + /// Joins the actor event loop without taking the handle out first, so an interrupted join + /// leaves it owned by `self` for the abort path. + async fn join_run_handle_in_place(&mut self) { + let Some(run_handle) = self.run_handle.as_mut() else { return; }; - match (&mut run_handle).await { - Ok(Ok(())) => {} - Ok(Err(error)) => { - log_actor_error(&error, "actor run handler failed during shutdown"); + let result = run_handle.await; + self.run_handle = None; + log_run_handle_shutdown_result(result); + } + + /// Awaits final resource cleanup. The cleanup runs in a spawned task retained on `self`, so + /// an interrupted wait neither cancels it nor lets a later caller start a second one. + async fn await_shutdown_cleanup(&mut self, reason: ShutdownKind) -> Result<()> { + let ctx = self.ctx.clone(); + let cleanup = self + .shutdown_cleanup + .get_or_insert_with(|| spawn_shutdown_cleanup(ctx, reason)); + let result = (&mut *cleanup).await; + self.shutdown_cleanup = None; + shutdown_cleanup_result(reason, result) + } + + /// Tears down a generation the engine declared lost. The engine may already be running the + /// next generation, so no user hook runs, no state is written, and every wait is bounded by + /// one deadline. Storage writes are already revoked by the lost signal. Work that does not + /// finish in time keeps running in retained tasks; it is not treated as finished. + async fn abort_lost(&mut self) -> Result<()> { + let started_at = Instant::now(); + let deadline = started_at + LOST_SHUTDOWN_GRACE_PERIOD; + tracing::warn!( + actor_id = %self.ctx.actor_id(), + lifecycle = ?self.lifecycle, + "actor generation declared lost, aborting it without user hooks" + ); + + self.sleep_grace = None; + self.state_save_deadline = None; + self.inspector_serialize_state_deadline = None; + self.sleep_deadline = None; + self.ctx.suspend_alarm_dispatch(); + self.ctx.cancel_local_alarm_timeouts(); + self.ctx.set_local_alarm_callback(None); + self.ctx.cancel_sleep_timer(); + self.ctx.cancel_actor_abort_signal(); + // Teardown of tracked work switches to cancellation instead of draining. + self.ctx.mark_shutdown_deadline_reached(); + if !matches!( + self.lifecycle, + LifecycleState::SleepFinalize | LifecycleState::Destroying | LifecycleState::Terminated + ) { + self.transition_to(LifecycleState::SleepFinalize); + } + + // Runtime adapters release their per-generation resources in the cleanup hook. They + // check `ActorContext::is_lost` and skip user callbacks. + // A generation lost during startup may have no run handler yet to receive the event. + if !self.runtime_cleanup_dispatched + && self.actor_event_tx.is_some() + && self.run_handle.is_some() + { + self.runtime_cleanup_dispatched = true; + self.ctx.begin_core_dispatched_hook(); + let reply = self.core_dispatched_hook_reply("run_lost_cleanup"); + if let Err(error) = self.send_actor_event( + "lost_run_cleanup", + ActorEvent::RunGracefulCleanup { + reason: ShutdownKind::Sleep, + reply, + }, + ) { + tracing::error!(?error, "failed to enqueue lost runtime cleanup event"); } - Err(error) => { - if !error.is_cancelled() { - tracing::error!(?error, "actor run handler join failed during shutdown"); + } + if !self.wait_for_core_dispatched_hooks(deadline).await { + tracing::warn!( + actor_id = %self.ctx.actor_id(), + core_dispatched_hook_count = self.ctx.core_dispatched_hook_count(), + "lost actor runtime cleanup did not finish before the lost deadline" + ); + } + // Adapters still waiting on a user hook stop waiting. + self.ctx.cancel_shutdown_deadline(); + + self.close_actor_event_channel(); + self.join_lost_run_handle(deadline).await; + + // The storage fence holds the generation until SQLite is really released, which can + // outlive this task when a statement is stuck. A bounded close can return with the native + // worker still running, so finish with a strict close. + let ctx = self.ctx.clone(); + let cleanup = self + .shutdown_cleanup + .take() + .unwrap_or_else(|| spawn_shutdown_cleanup(ctx.clone(), ShutdownKind::Sleep)); + let generation_hold = ctx.hold_generation(); + let mut storage_fence = RuntimeSpawner::spawn( + async move { + let _generation_hold = generation_hold; + let result = shutdown_cleanup_result(ShutdownKind::Sleep, cleanup.await); + if let Err(error) = ctx.sql().close_and_wait().await { + tracing::warn!( + actor_id = %ctx.actor_id(), + ?error, + "lost actor sqlite close failed" + ); } + result + } + .in_current_span() + .with_current_subscriber(), + ); + let result = match timeout( + deadline.saturating_duration_since(Instant::now()), + &mut storage_fence, + ) + .await + { + Ok(result) => shutdown_cleanup_result(ShutdownKind::Sleep, result), + Err(_) => { + tracing::warn!( + actor_id = %self.ctx.actor_id(), + "lost actor storage cleanup is still running after the lost deadline, likely a stuck sqlite statement; newer generations wait for it" + ); + Ok(()) } }; + self.ctx + .record_shutdown_wait(ShutdownKind::Sleep, started_at.elapsed()); + result + } + + async fn wait_for_core_dispatched_hooks(&self, deadline: Instant) -> bool { + self.ctx.wait_for_core_dispatched_hooks(deadline).await + } + + /// Joins the event loop of a lost generation, aborting it at `deadline`. Runtime adapters + /// drain registered tasks after their event channel closes, which can wait on user promises + /// indefinitely. + async fn join_lost_run_handle(&mut self, deadline: Instant) { + let Some(run_handle) = self.run_handle.as_mut() else { + return; + }; + let result = match timeout( + deadline.saturating_duration_since(Instant::now()), + &mut *run_handle, + ) + .await + { + Ok(result) => result, + Err(_) => { + tracing::warn!( + actor_id = %self.ctx.actor_id(), + "lost actor event loop did not exit before the lost deadline, aborting it" + ); + run_handle.abort(); + (&mut *run_handle).await + } + }; + self.run_handle = None; + log_run_handle_shutdown_result(result); } #[cfg(test)] @@ -1814,8 +2056,8 @@ impl ActorTask { 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 + self.join_run_handle_in_place().await; + self.await_shutdown_cleanup(reason).await } .await; if result.is_ok() && matches!(reason, ShutdownKind::Destroy) { @@ -1836,7 +2078,7 @@ impl ActorTask { if let Some(run_handle) = self.run_handle.as_mut() { run_handle.abort(); } - self.join_aborted_run_handle().await; + self.join_run_handle_in_place().await; // Failed startup has no grace period. Treat it as an elapsed shutdown // deadline so abortable waitUntil/registered work cannot hold teardown. self.ctx.mark_shutdown_deadline_reached(); @@ -1920,13 +2162,16 @@ impl ActorTask { step = "wait_for_pending_state_writes", "actor shutdown cleanup step completed" ); - ctx.sync_alarm_logged().await; - tracing::debug!( - actor_id = %actor_id, - reason = reason_label, - step = "sync_alarm", - "actor shutdown cleanup step completed" - ); + // A lost generation leaves the persisted alarm to the next generation. + if !ctx.is_lost() { + ctx.sync_alarm_logged().await; + tracing::debug!( + actor_id = %actor_id, + reason = reason_label, + step = "sync_alarm", + "actor shutdown cleanup step completed" + ); + } // Destroy cancels the engine alarm here so the persist it spawns is awaited by // `wait_for_pending_alarm_writes` below and cannot race the SQLite teardown. match reason { @@ -2057,10 +2302,11 @@ impl ActorTask { async fn on_state_save_tick(&mut self) { self.state_save_deadline = None; self.inspector_serialize_state_deadline = None; - if !matches!( - self.lifecycle, - LifecycleState::Started | LifecycleState::SleepGrace - ) || !self.ctx.save_requested() + if self.ctx.is_lost() + || !matches!( + self.lifecycle, + LifecycleState::Started | LifecycleState::SleepGrace + ) || !self.ctx.save_requested() { return; } @@ -2083,7 +2329,10 @@ impl ActorTask { } } - match reply_rx.await { + let Some(reply) = unless_lost(&self.lost, reply_rx).await else { + return; + }; + match reply { Ok(Ok(deltas)) => { let serialized_bytes = state_delta_payload_bytes(&deltas); tracing::debug!( @@ -2168,7 +2417,10 @@ impl ActorTask { } } - match reply_rx.await { + let Some(reply) = unless_lost(&self.lost, reply_rx).await else { + return; + }; + match reply { Ok(Ok(deltas)) => { tracing::debug!( actor_id = %self.ctx.actor_id(), @@ -2335,6 +2587,54 @@ impl ActorTask { } } +/// Spawns final resource cleanup so its owner survives an interrupted wait. The task keeps the +/// caller's span and tracing subscriber so its logs stay attributed to the actor. +fn spawn_shutdown_cleanup(ctx: ActorContext, reason: ShutdownKind) -> JoinHandle> { + RuntimeSpawner::spawn( + ActorTask::finish_shutdown_cleanup_with_ctx(ctx, reason) + .in_current_span() + .with_current_subscriber(), + ) +} + +/// Maps the cleanup task's join result to the shutdown result, keeping the historical error for +/// a panic during shutdown. +fn shutdown_cleanup_result( + reason: ShutdownKind, + result: std::result::Result, JoinError>, +) -> Result<()> { + match result { + Ok(result) => result, + Err(error) if error.is_panic() => Err(anyhow!("shutdown panicked during {reason:?}")), + Err(error) => Err(error).context("join actor shutdown cleanup"), + } +} + +/// Runs `future` unless the generation is declared lost first. Live-loop handlers run inside a +/// `select!` arm, where the loop's own lost arm cannot interrupt them. Only use this where every +/// resource the future owns is also retained elsewhere, so dropping it cancels only the wait. +async fn unless_lost(lost: &CancellationToken, future: F) -> Option { + tokio::select! { + biased; + _ = lost.cancelled() => None, + output = future => Some(output), + } +} + +fn log_run_handle_shutdown_result(result: std::result::Result, JoinError>) { + match result { + Ok(Ok(())) => {} + Ok(Err(error)) => { + log_actor_error(&error, "actor run handler failed during shutdown"); + } + Err(error) => { + if !error.is_cancelled() { + tracing::error!(?error, "actor run handler join failed during shutdown"); + } + } + } +} + fn shutdown_reason_label(reason: ShutdownKind) -> &'static str { match reason { ShutdownKind::Sleep => "sleep", diff --git a/rivetkit-rust/packages/rivetkit-core/src/registry/envoy_callbacks.rs b/rivetkit-rust/packages/rivetkit-core/src/registry/envoy_callbacks.rs index dd7a3ca215..abf6d3f31a 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/registry/envoy_callbacks.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/registry/envoy_callbacks.rs @@ -1,3 +1,4 @@ +use tokio_util::sync::CancellationToken; use tracing::Instrument; use super::*; @@ -6,12 +7,31 @@ use crate::runtime::RuntimeSpawner; impl EnvoyCallbacks for RegistryCallbacks { fn on_actor_start( + &self, + handle: EnvoyHandle, + actor_id: String, + generation: u32, + config: protocol::ActorConfig, + preloaded_kv: Option, + ) -> EnvoyBoxFuture> { + self.on_actor_start_with_lost_signal( + handle, + actor_id, + generation, + config, + preloaded_kv, + CancellationToken::new(), + ) + } + + fn on_actor_start_with_lost_signal( &self, handle: EnvoyHandle, actor_id: String, generation: u32, config: protocol::ActorConfig, _preloaded_kv: Option, + lost: CancellationToken, ) -> EnvoyBoxFuture> { let dispatcher = self.dispatcher.clone(); let actor_name = config.name.clone(); @@ -34,6 +54,7 @@ impl EnvoyCallbacks for RegistryCallbacks { key, factory.as_ref(), )?; + ctx.configure_lost_signal(lost); dispatcher .start_actor(StartActorRequest { diff --git a/rivetkit-rust/packages/rivetkit-core/src/registry/mod.rs b/rivetkit-rust/packages/rivetkit-core/src/registry/mod.rs index f63cc207e4..c83818adf2 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/registry/mod.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/registry/mod.rs @@ -27,7 +27,7 @@ use scc::{HashMap as SccHashMap, hash_map::Entry as SccEntry}; use serde::{Deserialize, Serialize}; use serde_bytes::ByteBuf; use serde_json::{Value as JsonValue, json}; -use tokio::sync::{Mutex as TokioMutex, Notify, broadcast, mpsc, oneshot}; +use tokio::sync::{Mutex as TokioMutex, broadcast, mpsc, oneshot}; use tokio::task::JoinHandle; use tokio_util::sync::CancellationToken; use url::Url; @@ -165,23 +165,29 @@ struct PendingStop { stop_handle: ActorStopHandle, } -/// Outcome of attempting to transition the actor instance under an id to stopping -/// for a specific generation. +/// Identifies one generation of an actor. Registry records are per generation so a lost +/// generation that is still shutting down stays reachable after the next generation starts. +type InstanceKey = (String, u32); + +fn instance_key(actor_id: &str, generation: u32) -> InstanceKey { + (actor_id.to_owned(), generation) +} + +/// Outcome of attempting to transition one actor generation to stopping. enum TransitionResult { - /// The current instance matches the requested generation and was moved to stopping. + /// The generation was registered and is now stopping. Transitioned(ActiveActorInstance), - /// An instance exists but for a different generation than the stop targets, so it - /// was left untouched. The stop must not be applied to it. - Stale, - /// No instance is registered for the actor id. + /// No instance is registered for the generation. Vacant, } pub(crate) struct RegistryDispatcher { pub(crate) factories: HashMap>, - actor_instances: SccHashMap, - starting_instances: SccHashMap>, - pending_stops: SccHashMap, + actor_instances: SccHashMap, + /// Newest registered generation per actor id. Dispatch routes only to this generation. + current_generations: SccHashMap, + starting_instances: SccHashMap, + pending_stops: SccHashMap, region: String, handle_inspector_http_in_runtime: bool, } @@ -478,6 +484,11 @@ struct InspectorSummaryJson { workflow_history: Option, } +/// How long a new generation waits for older generations of the same actor on this runner to +/// finish. Lost generations abort within two seconds, so this only expires if an older task is +/// stuck, and the new generation then fails to start rather than overlap it. +const OLDER_GENERATION_STOP_TIMEOUT: Duration = Duration::from_secs(10); + const WS_PROTOCOL_ENCODING: &str = "rivet_encoding."; const WS_PROTOCOL_CONN_PARAMS: &str = "rivet_conn_params."; @@ -737,6 +748,7 @@ impl RegistryDispatcher { Self { factories, actor_instances: SccHashMap::new(), + current_generations: SccHashMap::new(), starting_instances: SccHashMap::new(), pending_stops: SccHashMap::new(), region: env::var("RIVET_REGION").unwrap_or_default(), @@ -803,26 +815,43 @@ fn build_actor_metadata_map_from_factories( impl RegistryDispatcher { async fn start_actor(self: &Arc, request: StartActorRequest) -> Result<()> { - let startup_notify = Arc::new(Notify::new()); + let key = instance_key(&request.actor_id, request.generation); + // The actor task owns this hold. It is released when the task ends, or when startup is + // abandoned before a task exists. + let generation_hold = request.ctx.hold_generation(); let _ = self .starting_instances - .insert_async(request.actor_id.clone(), startup_notify.clone()) + .insert_async(key.clone(), request.ctx.clone()) .await; - // Test-only seam: lets a test hold a generation in the "starting" window so it - // can deterministically deliver a stop for a previous generation that parks - // under the actor id and gets consumed by this startup. - #[cfg(test)] - test_hooks::wait_for_startup_gate(&request.actor_id).await; - let factory = self - .factories - .get(&request.actor_name) - .cloned() - .ok_or_else(|| { - ActorRuntime::NotRegistered { - actor_name: request.actor_name.clone(), - } - .build() - })?; + let prepared = async { + self.wait_for_older_generations(&request.actor_id, request.generation) + .await?; + // Test-only seam: lets a test hold a generation in the "starting" window so it + // can deterministically deliver stops while the generation is starting. + #[cfg(test)] + test_hooks::wait_for_startup_gate(&request.actor_id).await; + if request.ctx.is_lost() { + return Err(ActorLifecycleError::Stopping.build()) + .context("actor generation was declared lost before it started"); + } + self.factories + .get(&request.actor_name) + .cloned() + .ok_or_else(|| { + ActorRuntime::NotRegistered { + actor_name: request.actor_name.clone(), + } + .build() + }) + } + .await; + let factory = match prepared { + Ok(factory) => factory, + Err(error) => { + self.abandon_startup(&key).await; + return Err(error); + } + }; let (lifecycle_tx, lifecycle_rx) = mpsc::unbounded_channel(); let (dispatch_tx, dispatch_rx) = mpsc::unbounded_channel(); let (lifecycle_events_tx, lifecycle_events_rx) = mpsc::unbounded_channel(); @@ -857,7 +886,10 @@ impl RegistryDispatcher { request.ctx.clone(), request.input, ); - let join = RuntimeSpawner::spawn(task.run()); + let join = RuntimeSpawner::spawn(async move { + let _generation_hold = generation_hold; + task.run().await + }); let (start_tx, start_rx) = oneshot::channel(); let result: Result> = async { @@ -886,104 +918,253 @@ impl RegistryDispatcher { match result { Ok(instance) => { + // Hold the starting entry while consuming the parked stop and registering the + // instance. `stop_actor` parks under the same entry, so a stop either parks + // before this point and is consumed here, or finds the registered instance. + let starting = self.starting_instances.entry_async(key.clone()).await; let pending_stop = self .pending_stops - .remove_async(&request.actor_id.clone()) + .remove_async(&key) .await .map(|(_, pending_stop)| pending_stop); - // Only apply a parked stop if it targets the generation we just started. - // A stop parked for a previous generation is stale: complete its handle - // so teardown finalizes, but leave the new generation running. - let pending_stop = match pending_stop { - Some(pending_stop) if pending_stop.generation == request.generation => { - Some(pending_stop) - } - Some(stale_stop) => { - let _ = stale_stop.stop_handle.complete(); - None + match pending_stop { + Some(pending_stop) => { + let actor_id = request.actor_id.clone(); + let stop_reason = map_envoy_stop_reason(&pending_stop.reason); + if matches!(stop_reason, ShutdownKind::Destroy) { + instance.ctx.mark_destroy_requested(); + } + self.register_actor_instance( + key.clone(), + ActorInstanceState::Stopping { + instance: instance.clone(), + reason: stop_reason, + }, + ) + .await; + remove_starting_entry(starting); + + let dispatcher = self.clone(); + RuntimeSpawner::spawn(async move { + if let Err(error) = dispatcher + .shutdown_started_instance( + &actor_id, + instance.clone(), + pending_stop.reason, + pending_stop.stop_handle, + ) + .await + { + tracing::error!( + actor_id, + ?error, + "failed to stop actor queued during startup" + ); + } + dispatcher + .remove_stopping_actor_instance_when_finished(&actor_id, &instance) + .await; + }); } - None => None, - }; - if let Some(pending_stop) = pending_stop { - let actor_id = request.actor_id.clone(); - let stop_reason = map_envoy_stop_reason(&pending_stop.reason); - if matches!(stop_reason, ShutdownKind::Destroy) { - instance.ctx.mark_destroy_requested(); + None => { + self.register_actor_instance( + key.clone(), + ActorInstanceState::Active(instance), + ) + .await; + remove_starting_entry(starting); } - self.set_actor_instance_state( - actor_id.clone(), - ActorInstanceState::Stopping { - instance: instance.clone(), - reason: stop_reason, - }, - ) + } + self.complete_stops_for_older_generations(&request.actor_id, request.generation) .await; - let _ = self - .starting_instances - .remove_async(&request.actor_id.clone()) - .await; - + Ok(()) + } + Err(error) => { + // A generation lost during startup drops its start reply while its task is still + // aborting. Keep it registered as starting until the task finishes so envoy-client + // does not report it stopped, and a newer generation does not start, too early. + if request.ctx.is_lost() + && timeout( + OLDER_GENERATION_STOP_TIMEOUT, + request.ctx.wait_for_generation_finished(), + ) + .await + .is_err() + { + tracing::warn!( + actor_id = %request.actor_id, + generation = request.generation, + "lost actor generation did not finish aborting its startup in time, likely stuck storage; keeping it registered until it finishes" + ); + // Keep the starting record so newer generations still see this one and + // refuse to start until it really finishes. let dispatcher = self.clone(); + let ctx = request.ctx.clone(); RuntimeSpawner::spawn(async move { - if let Err(error) = dispatcher - .shutdown_started_instance( - &actor_id, - instance.clone(), - pending_stop.reason, - pending_stop.stop_handle, - ) - .await - { - tracing::error!( - actor_id, - ?error, - "failed to stop actor queued during startup" - ); - } - dispatcher - .remove_stopping_actor_instance(&actor_id, &instance) - .await; + ctx.wait_for_generation_finished().await; + dispatcher.abandon_startup(&key).await; }); - startup_notify.notify_waiters(); + return Err(error); + } + self.abandon_startup(&key).await; + Err(error) + } + } + } - Ok(()) - } else { - self.set_actor_instance_state( - request.actor_id.clone(), - ActorInstanceState::Active(instance), - ) - .await; - let _ = self - .starting_instances - .remove_async(&request.actor_id.clone()) - .await; - startup_notify.notify_waiters(); - Ok(()) + /// Removes a generation whose startup failed. A stop parked while the start was in flight + /// would otherwise leak its ActorStopHandle in the map and hang the caller, so complete it + /// since there is no instance to stop. + async fn abandon_startup(&self, key: &InstanceKey) { + let starting = self.starting_instances.entry_async(key.clone()).await; + if let Some((_, pending_stop)) = self.pending_stops.remove_async(key).await { + let _ = pending_stop.stop_handle.complete(); + } + remove_starting_entry(starting); + } + + /// Registers a started generation and makes it the dispatch target unless a newer + /// generation is already registered. + async fn register_actor_instance(&self, key: InstanceKey, state: ActorInstanceState) { + let (actor_id, generation) = key.clone(); + self.set_actor_instance_state(key, state).await; + match self.current_generations.entry_async(actor_id).await { + SccEntry::Occupied(mut entry) => { + if *entry.get() < generation { + entry.insert(generation); } } - Err(error) => { - let _ = self - .starting_instances - .remove_async(&request.actor_id.clone()) - .await; - // A stop parked while this start was in flight would otherwise leak its - // ActorStopHandle in the map (the map holds the sender alive, hanging the - // caller). Drain and complete it since there is no instance to stop. - if let Some((_, pending_stop)) = self - .pending_stops - .remove_async(&request.actor_id.clone()) - .await - { - let _ = pending_stop.stop_handle.complete(); + SccEntry::Vacant(entry) => { + entry.insert_entry(generation); + } + } + } + + /// Completes stops parked for generations older than `generation`. Generations only + /// increase, so such a generation will never start and its stop has nothing left to do. + async fn complete_stops_for_older_generations(&self, actor_id: &str, generation: u32) { + let mut stale_stops = Vec::new(); + self.pending_stops + .retain_async(|(pending_actor_id, pending_generation), pending_stop| { + if pending_actor_id == actor_id && *pending_generation < generation { + stale_stops.push(pending_stop.stop_handle.clone()); + false + } else { + true } - startup_notify.notify_waiters(); - Err(error) + }) + .await; + for stop_handle in stale_stops { + let _ = stop_handle.complete(); + } + } + + async fn newer_generation_exists(&self, actor_id: &str, generation: u32) -> bool { + let is_newer = |(other_actor_id, other_generation): &InstanceKey| { + other_actor_id == actor_id && *other_generation > generation + }; + // Each check releases its entry before the next one runs. + let newer_starting = self + .starting_instances + .any_async(|key, _| is_newer(key)) + .await + .is_some(); + if newer_starting { + return true; + } + self.actor_instances + .any_async(|key, _| is_newer(key)) + .await + .is_some() + } + + /// Contexts of older generations of `actor_id` that are starting or registered on this + /// runner. + async fn older_generation_contexts( + &self, + actor_id: &str, + generation: u32, + ) -> Vec { + let mut older = Vec::new(); + self.starting_instances + .iter_async(|(starting_actor_id, starting_generation), ctx| { + if starting_actor_id == actor_id && *starting_generation < generation { + older.push(ctx.clone()); + } + true + }) + .await; + self.actor_instances + .iter_async(|(instance_actor_id, instance_generation), state| { + if instance_actor_id == actor_id && *instance_generation < generation { + older.push(state.instance().ctx.clone()); + } + true + }) + .await; + older + } + + /// Keeps two generations of one actor from running on this runner at once. The engine only + /// starts a generation after giving up on every older one, so older generations still here + /// are lost: mark them lost so they abort, then wait for their tasks to finish. If one does + /// not finish in time, fail this start rather than overlap it. + async fn wait_for_older_generations(&self, actor_id: &str, generation: u32) -> Result<()> { + // Start callbacks run in independent tasks, so a newer generation can register first. + // Generations only increase, so this one has already been superseded and must not run. + // Both starts insert their starting record before scanning, so at least one sees the + // other. + if self.newer_generation_exists(actor_id, generation).await { + tracing::warn!( + actor_id, + generation, + "refusing to start an actor generation that a newer generation already superseded" + ); + return Err(ActorLifecycleError::Stopping.build()) + .context("a newer generation of this actor is already running on this runner"); + } + let older = self.older_generation_contexts(actor_id, generation).await; + if older.is_empty() { + return Ok(()); + } + for ctx in &older { + if !ctx.is_lost() { + tracing::warn!( + actor_id, + generation, + "newer actor generation is starting while an older one still runs; marking the older generation lost" + ); + ctx.mark_lost(); } } + let all_finished = async { + for ctx in &older { + ctx.wait_for_generation_finished().await; + } + }; + if timeout(OLDER_GENERATION_STOP_TIMEOUT, all_finished) + .await + .is_err() + { + tracing::error!( + actor_id, + generation, + "an older actor generation is still running on this runner, likely a stuck task; refusing to start the new generation" + ); + return Err(ActorLifecycleError::Stopping.build()) + .context("an older generation of this actor is still running on this runner"); + } + Ok(()) + } + + async fn current_generation(&self, actor_id: &str) -> Option { + self.current_generations + .read_async(actor_id, |_, generation| *generation) + .await } - async fn set_actor_instance_state(&self, actor_id: String, state: ActorInstanceState) { - match self.actor_instances.entry_async(actor_id).await { + async fn set_actor_instance_state(&self, key: InstanceKey, state: ActorInstanceState) { + match self.actor_instances.entry_async(key).await { SccEntry::Occupied(mut entry) => { entry.insert(state); } @@ -999,16 +1180,13 @@ impl RegistryDispatcher { generation: u32, reason: ShutdownKind, ) -> TransitionResult { - match self.actor_instances.entry_async(actor_id.to_owned()).await { + match self + .actor_instances + .entry_async(instance_key(actor_id, generation)) + .await + { SccEntry::Occupied(mut entry) => { let instance = entry.get().instance(); - // A stop is scoped to the generation it was issued for. If the currently - // registered instance is a different generation (e.g. a lost previous - // generation whose replacement is already running), the stop is stale and - // must not tear down the newer generation. - if instance.generation != generation { - return TransitionResult::Stale; - } if matches!(entry.get(), ActorInstanceState::Active(_)) { entry.insert(ActorInstanceState::Stopping { instance: instance.clone(), @@ -1028,8 +1206,36 @@ impl RegistryDispatcher { } } + /// Removes a stopped generation's record once the generation has finished. A lost + /// generation's storage cleanup can outlive its task, and the record must stay visible to + /// newer generations until it does. + async fn remove_stopping_actor_instance_when_finished( + self: &Arc, + actor_id: &str, + expected: &ActiveActorInstance, + ) { + if expected.ctx.is_generation_finished() { + self.remove_stopping_actor_instance(actor_id, expected) + .await; + return; + } + let dispatcher = self.clone(); + let actor_id = actor_id.to_owned(); + let expected = expected.clone(); + RuntimeSpawner::spawn(async move { + expected.ctx.wait_for_generation_finished().await; + dispatcher + .remove_stopping_actor_instance(&actor_id, &expected) + .await; + }); + } + async fn remove_stopping_actor_instance(&self, actor_id: &str, expected: &ActiveActorInstance) { - match self.actor_instances.entry_async(actor_id.to_owned()).await { + let removed = match self + .actor_instances + .entry_async(instance_key(actor_id, expected.generation)) + .await + { SccEntry::Occupied(entry) => { let should_remove = match entry.get() { ActorInstanceState::Stopping { instance, .. } => { @@ -1040,20 +1246,37 @@ impl RegistryDispatcher { if should_remove { let _ = entry.remove_entry(); } + should_remove } SccEntry::Vacant(entry) => { drop(entry); + false } + }; + if removed { + // Only clear the dispatch target if no newer generation has replaced it. + let _ = self + .current_generations + .remove_if_async(actor_id, |generation| *generation == expected.generation) + .await; } } async fn active_actor(&self, actor_id: &str) -> Result> { - if let Some(instance) = self.actor_instances.get_async(&actor_id.to_owned()).await { + let instance = match self.current_generation(actor_id).await { + Some(generation) => { + self.actor_instances + .get_async(&instance_key(actor_id, generation)) + .await + } + None => None, + }; + if let Some(instance) = instance { match instance.get() { ActorInstanceState::Active(instance) => { let instance = instance.clone(); // TODO: Share admission policy with ActorTask::dispatch_lifecycle_error. - if instance.ctx.started() { + if instance.ctx.started() && !instance.ctx.is_lost() { if instance.ctx.destroy_requested() { instance .ctx @@ -1068,7 +1291,7 @@ impl RegistryDispatcher { .warn_work_sent_to_stopping_instance("active_actor"); return Err(if instance.ctx.destroy_requested() { ActorLifecycleError::Destroying.build() - } else if instance.ctx.sleep_requested() { + } else if instance.ctx.sleep_requested() || instance.ctx.is_lost() { ActorLifecycleError::Stopping.build() } else { ActorLifecycleError::Starting.build() @@ -1077,7 +1300,13 @@ impl RegistryDispatcher { ActorInstanceState::Stopping { instance, reason } => { let instance = instance.clone(); match reason { - ShutdownKind::Sleep if instance.ctx.started() => return Ok(instance), + // A lost generation takes no new work; the engine routes it to the next + // generation. + ShutdownKind::Sleep + if instance.ctx.started() && !instance.ctx.is_lost() => + { + return Ok(instance); + } ShutdownKind::Sleep => { instance .ctx @@ -1104,33 +1333,36 @@ impl RegistryDispatcher { } async fn stop_actor( - &self, + self: &Arc, actor_id: &str, generation: u32, reason: protocol::StopActorReason, stop_handle: ActorStopHandle, ) -> Result<()> { - if self - .starting_instances - .get_async(&actor_id.to_owned()) - .await - .is_some() - { - // The target generation is still starting. Park the stop with its generation - // so startup can decide whether it belongs to the generation being started. - let _ = self - .pending_stops - .insert_async( - actor_id.to_owned(), - PendingStop { - generation, - reason, - stop_handle, - }, - ) - .await; - return Ok(()); - } + let key = instance_key(actor_id, generation); + let pending_stop = PendingStop { + generation, + reason, + stop_handle, + }; + let pending_stop = match self.starting_instances.entry_async(key.clone()).await { + SccEntry::Occupied(starting) => { + // The generation is still starting. Park the stop under the starting entry so + // its startup consumes it. + self.park_stop(key, pending_stop).await; + drop(starting); + return Ok(()); + } + SccEntry::Vacant(starting) => { + drop(starting); + pending_stop + } + }; + let PendingStop { + reason, + stop_handle, + .. + } = pending_stop; let task_stop_reason = map_envoy_stop_reason(&reason); match self @@ -1141,34 +1373,48 @@ impl RegistryDispatcher { let result = self .shutdown_started_instance(actor_id, instance.clone(), reason, stop_handle) .await; - self.remove_stopping_actor_instance(actor_id, &instance) + self.remove_stopping_actor_instance_when_finished(actor_id, &instance) .await; result } - TransitionResult::Stale => { - // The running instance is a different generation; this stop targets a - // generation that is already gone. Complete the handle so envoy-client - // finalizes teardown cleanly instead of warning about a dropped handle. - let _ = stop_handle.complete(); - Ok(()) - } TransitionResult::Vacant => { - let _ = self - .pending_stops - .insert_async( - actor_id.to_owned(), - PendingStop { - generation, - reason, - stop_handle, - }, - ) - .await; + match self.current_generation(actor_id).await { + Some(current) if current > generation => { + // A newer generation is registered, so this generation already finished + // tearing down. Complete the handle so envoy-client finalizes the stop. + let _ = stop_handle.complete(); + } + Some(_) | None => { + // The stop can arrive before its generation's start reaches the + // registry. Park it for that startup. + self.park_stop( + key, + PendingStop { + generation, + reason, + stop_handle, + }, + ) + .await; + } + } Ok(()) } } } + async fn park_stop(&self, key: InstanceKey, pending_stop: PendingStop) { + if let Err((_, duplicate)) = self.pending_stops.insert_async(key, pending_stop).await { + // envoy-client sends one stop per generation. A duplicate adds nothing, so complete + // it rather than dropping its handle. + tracing::warn!( + generation = duplicate.generation, + "duplicate stop parked for an actor generation" + ); + let _ = duplicate.stop_handle.complete(); + } + } + async fn shutdown_started_instance( &self, actor_id: &str, @@ -1181,6 +1427,11 @@ impl RegistryDispatcher { if matches!(task_stop_reason, ShutdownKind::Destroy) { instance.ctx.mark_destroy_requested(); } + if matches!(reason, protocol::StopActorReason::Lost) { + // envoy-client normally fires the lost signal before this stop arrives. Marking it + // here as well covers stops that reach core by other paths. + instance.ctx.mark_lost(); + } tracing::debug!( actor_id, @@ -1206,6 +1457,19 @@ impl RegistryDispatcher { .and_then(|result| result), Err(error) => Err(error), }; + // A lost generation aborts on its lost signal without waiting for this stop, so the task + // may already have exited and dropped the command. Its outcome is its join result below. + let shutdown_result = match shutdown_result { + Err(error) if instance.ctx.is_lost() => { + tracing::debug!( + actor_id, + %error, + "lost actor task finished before its stop command was handled" + ); + Ok(()) + } + result => result, + }; if matches!(task_stop_reason, ShutdownKind::Destroy) { let shutdown_deadline = @@ -1274,15 +1538,34 @@ impl RegistryDispatcher { } } +fn remove_starting_entry(starting: SccEntry<'_, InstanceKey, ActorContext>) { + match starting { + SccEntry::Occupied(entry) => { + let _ = entry.remove_entry(); + } + SccEntry::Vacant(entry) => { + drop(entry); + } + } +} + impl RegistryDispatcher { fn can_hibernate(&self, actor_id: &str, request: &HttpRequest) -> bool { if matches!(is_actor_connect_path(&request.path), Ok(true)) { return true; } + let Some(generation) = self + .current_generations + .read_sync(actor_id, |_, generation| *generation) + else { + return false; + }; let Some(instance) = self .actor_instances - .read_sync(actor_id, |_, state| state.active_instance()) + .read_sync(&instance_key(actor_id, generation), |_, state| { + state.active_instance() + }) .flatten() else { return false; diff --git a/rivetkit-rust/packages/rivetkit-core/tests/registry.rs b/rivetkit-rust/packages/rivetkit-core/tests/registry.rs index 05564197f3..952dd38f80 100644 --- a/rivetkit-rust/packages/rivetkit-core/tests/registry.rs +++ b/rivetkit-rust/packages/rivetkit-core/tests/registry.rs @@ -60,10 +60,10 @@ async fn stop_for_previous_generation_does_not_kill_freshly_started_generation() assert!( dispatcher .pending_stops - .get_async(&actor_id.to_owned()) + .get_async(&instance_key(actor_id, 49)) .await .is_some(), - "gen-49 stop should be parked under the actor id", + "gen-49 stop should be parked for gen 49", ); // gen 50 is started to take over the actor on the same runner. @@ -83,7 +83,7 @@ async fn stop_for_previous_generation_does_not_kill_freshly_started_generation() assert!( dispatcher .pending_stops - .get_async(&actor_id.to_owned()) + .get_async(&instance_key(actor_id, 49)) .await .is_none(), "stale gen-49 stop should be cleared during gen 50 startup", @@ -93,7 +93,7 @@ async fn stop_for_previous_generation_does_not_kill_freshly_started_generation() let is_active = is_actor_active( dispatcher .actor_instances - .get_async(&actor_id.to_owned()) + .get_async(&instance_key(actor_id, 50)) .await .as_ref() .map(|entry| entry.get()), @@ -150,7 +150,7 @@ async fn parked_previous_generation_stop_does_not_kill_starting_generation() { loop { if dispatcher .starting_instances - .get_async(&actor_id.to_owned()) + .get_async(&instance_key(actor_id, 50)) .await .is_some() { @@ -162,7 +162,7 @@ async fn parked_previous_generation_stop_does_not_kill_starting_generation() { .await .expect("gen 50 should register as starting"); - // gen 49's `Lost` stop arrives while gen 50 is starting: it parks under the actor id. + // gen 49's `Lost` stop arrives while gen 50 is starting and gen 49 has no record: it parks. dispatcher .stop_actor( actor_id, @@ -175,7 +175,7 @@ async fn parked_previous_generation_stop_does_not_kill_starting_generation() { assert!( dispatcher .pending_stops - .get_async(&actor_id.to_owned()) + .get_async(&instance_key(actor_id, 49)) .await .is_some(), "gen-49 stop should be parked while gen 50 is starting", @@ -191,7 +191,7 @@ async fn parked_previous_generation_stop_does_not_kill_starting_generation() { assert!( dispatcher .pending_stops - .get_async(&actor_id.to_owned()) + .get_async(&instance_key(actor_id, 49)) .await .is_none(), "stale gen-49 stop should be cleared during gen 50 startup", @@ -199,7 +199,7 @@ async fn parked_previous_generation_stop_does_not_kill_starting_generation() { let is_active = is_actor_active( dispatcher .actor_instances - .get_async(&actor_id.to_owned()) + .get_async(&instance_key(actor_id, 50)) .await .as_ref() .map(|entry| entry.get()), @@ -245,7 +245,7 @@ async fn stop_for_current_generation_stops_it() { is_actor_active( dispatcher .actor_instances - .get_async(&actor_id.to_owned()) + .get_async(&instance_key(actor_id, 50)) .await .as_ref() .map(|entry| entry.get()), @@ -268,7 +268,7 @@ async fn stop_for_current_generation_stops_it() { !is_actor_active( dispatcher .actor_instances - .get_async(&actor_id.to_owned()) + .get_async(&instance_key(actor_id, 50)) .await .as_ref() .map(|entry| entry.get()), @@ -276,3 +276,467 @@ async fn stop_for_current_generation_stops_it() { "a stop for the current generation must stop it", ); } + +fn counter_dispatcher() -> Arc { + let mut factories = HashMap::new(); + factories.insert( + "counter".to_owned(), + Arc::new(ActorFactory::new(ActorConfig::default(), |_start| { + Box::pin(async { Ok(()) }) + })), + ); + Arc::new(RegistryDispatcher::new(factories, false)) +} + +async fn start_counter( + dispatcher: &Arc, + actor_id: &str, + generation: u32, +) -> crate::actor::context::ActorContext { + let ctx = crate::actor::context::ActorContext::new(actor_id, "counter", Vec::new(), "local"); + dispatcher + .start_actor(StartActorRequest { + actor_id: actor_id.to_owned(), + generation, + actor_name: "counter".to_owned(), + input: None, + ctx: ctx.clone(), + }) + .await + .expect("generation should start"); + ctx +} + +async fn instance_generation( + dispatcher: &RegistryDispatcher, + actor_id: &str, + generation: u32, +) -> Option { + dispatcher + .actor_instances + .get_async(&instance_key(actor_id, generation)) + .await + .map(|entry| entry.get().instance().generation) +} + +/// A Lost stop for generation N can reach the registry after N+1 has registered, because N+1 +/// only waits for N's task to finish, not for N's stop. The stop must still reach N instead of +/// being dropped as stale, and it must leave N+1 routed. +#[tokio::test(start_paused = true)] +async fn stop_for_older_generation_after_newer_registered_stops_the_older_generation() { + use rivet_envoy_client::config::ActorStopHandle; + use rivet_envoy_client::protocol::StopActorReason; + + let dispatcher = counter_dispatcher(); + let actor_id = "actor-overlap"; + let old_ctx = start_counter(&dispatcher, actor_id, 1).await; + let new_ctx = start_counter(&dispatcher, actor_id, 2).await; + assert_eq!(instance_generation(&dispatcher, actor_id, 1).await, Some(1)); + assert_eq!(dispatcher.current_generation(actor_id).await, Some(2)); + + dispatcher + .stop_actor( + actor_id, + 1, + StopActorReason::Lost, + ActorStopHandle::detached(), + ) + .await + .expect("stopping the older generation should succeed"); + + assert!( + old_ctx.is_lost(), + "the older generation must receive the Lost stop" + ); + assert_eq!( + instance_generation(&dispatcher, actor_id, 1).await, + None, + "the older generation's record is removed after its teardown", + ); + assert!( + !new_ctx.is_lost(), + "the newer generation must not be touched" + ); + assert_eq!(instance_generation(&dispatcher, actor_id, 2).await, Some(2)); + assert_eq!(dispatcher.current_generation(actor_id).await, Some(2)); + let routed = dispatcher + .active_actor(actor_id) + .await + .expect("dispatch should route to the newer generation"); + assert_eq!(routed.generation, 2); +} + +/// A stop for a registered generation must stop it immediately even while the next +/// generation is starting. Only a stop for the starting generation parks. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn parked_stops_are_scoped_to_the_starting_generation() { + use std::time::Duration; + + use rivet_envoy_client::config::ActorStopHandle; + use rivet_envoy_client::protocol::StopActorReason; + + let dispatcher = counter_dispatcher(); + let actor_id = "actor-scoped-park"; + let old_ctx = start_counter(&dispatcher, actor_id, 1).await; + + test_hooks::arm_startup_gate(actor_id); + let start_dispatcher = dispatcher.clone(); + let new_ctx = + crate::actor::context::ActorContext::new(actor_id, "counter", Vec::new(), "local"); + let start_ctx = new_ctx.clone(); + let start_task = tokio::spawn(async move { + start_dispatcher + .start_actor(StartActorRequest { + actor_id: actor_id.to_owned(), + generation: 2, + actor_name: "counter".to_owned(), + input: None, + ctx: start_ctx, + }) + .await + }); + tokio::time::timeout(Duration::from_secs(5), async { + loop { + if dispatcher + .starting_instances + .get_async(&instance_key(actor_id, 2)) + .await + .is_some() + { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("gen 2 should register as starting"); + + dispatcher + .stop_actor( + actor_id, + 1, + StopActorReason::Lost, + ActorStopHandle::detached(), + ) + .await + .expect("stopping gen 1 should succeed"); + assert!(old_ctx.is_lost(), "gen 1 must be stopped, not parked"); + assert!( + dispatcher + .pending_stops + .get_async(&instance_key(actor_id, 1)) + .await + .is_none(), + "the gen-1 stop must not park behind gen 2's startup", + ); + assert_eq!(instance_generation(&dispatcher, actor_id, 1).await, None); + + dispatcher + .stop_actor( + actor_id, + 2, + StopActorReason::SleepIntent, + ActorStopHandle::detached(), + ) + .await + .expect("a stop for the starting generation parks"); + assert!( + dispatcher + .pending_stops + .get_async(&instance_key(actor_id, 2)) + .await + .is_some(), + "the gen-2 stop parks for gen 2's startup", + ); + + test_hooks::release_startup_gate(actor_id); + start_task + .await + .expect("start task joins") + .expect("gen 2 should start"); + assert!( + dispatcher + .pending_stops + .get_async(&instance_key(actor_id, 2)) + .await + .is_none(), + "gen 2's startup consumes its parked stop", + ); + assert!(!new_ctx.is_lost()); +} + +/// A stop for a generation that already finished while a newer generation is registered +/// completes immediately instead of parking forever. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn stop_for_finished_older_generation_completes_immediately() { + use rivet_envoy_client::config::ActorStopHandle; + use rivet_envoy_client::protocol::StopActorReason; + + let dispatcher = counter_dispatcher(); + let actor_id = "actor-finished"; + start_counter(&dispatcher, actor_id, 2).await; + + dispatcher + .stop_actor( + actor_id, + 1, + StopActorReason::Lost, + ActorStopHandle::detached(), + ) + .await + .expect("a stale stop returns Ok"); + assert!( + dispatcher + .pending_stops + .get_async(&instance_key(actor_id, 1)) + .await + .is_none(), + "a stop for a finished older generation must not park", + ); + assert_eq!(instance_generation(&dispatcher, actor_id, 2).await, Some(2)); +} + +/// A generation starting while an older generation of the same actor still runs on this runner +/// marks the older one lost and waits for its task to finish, so the two never overlap. +#[tokio::test(start_paused = true)] +async fn newer_generation_waits_for_older_generation_to_finish() { + use rivet_envoy_client::config::ActorStopHandle; + use rivet_envoy_client::protocol::StopActorReason; + + let dispatcher = counter_dispatcher(); + let actor_id = "actor-gated-by-older"; + let old_ctx = start_counter(&dispatcher, actor_id, 1).await; + assert!(!old_ctx.is_generation_finished()); + + start_counter(&dispatcher, actor_id, 2).await; + + assert!( + old_ctx.is_lost(), + "the older generation must be marked lost" + ); + assert!( + old_ctx.is_generation_finished(), + "the newer generation must not start before the older task finished", + ); + assert_eq!(instance_generation(&dispatcher, actor_id, 2).await, Some(2)); + assert_eq!(dispatcher.current_generation(actor_id).await, Some(2)); + + // The engine's Lost stop for the older generation still removes its record. + dispatcher + .stop_actor( + actor_id, + 1, + StopActorReason::Lost, + ActorStopHandle::detached(), + ) + .await + .expect("stopping the older generation should succeed"); + assert_eq!(instance_generation(&dispatcher, actor_id, 1).await, None); + assert_eq!(instance_generation(&dispatcher, actor_id, 2).await, Some(2)); +} + +/// If an older generation never finishes, the newer one fails to start instead of overlapping +/// it. A generation held in startup models a stuck older task. +#[tokio::test(start_paused = true)] +async fn newer_generation_fails_closed_when_older_generation_never_finishes() { + use std::time::Duration; + + let dispatcher = counter_dispatcher(); + let actor_id = "actor-stuck-older"; + + test_hooks::arm_startup_gate(actor_id); + let old_ctx = + crate::actor::context::ActorContext::new(actor_id, "counter", Vec::new(), "local"); + let start_dispatcher = dispatcher.clone(); + let start_old_ctx = old_ctx.clone(); + let old_start = tokio::spawn(async move { + start_dispatcher + .start_actor(StartActorRequest { + actor_id: actor_id.to_owned(), + generation: 1, + actor_name: "counter".to_owned(), + input: None, + ctx: start_old_ctx, + }) + .await + }); + tokio::time::timeout(Duration::from_secs(1), async { + while dispatcher + .starting_instances + .get_async(&instance_key(actor_id, 1)) + .await + .is_none() + { + tokio::task::yield_now().await; + } + }) + .await + .expect("gen 1 should register as starting"); + + let new_ctx = + crate::actor::context::ActorContext::new(actor_id, "counter", Vec::new(), "local"); + let error = dispatcher + .start_actor(StartActorRequest { + actor_id: actor_id.to_owned(), + generation: 2, + actor_name: "counter".to_owned(), + input: None, + ctx: new_ctx.clone(), + }) + .await + .expect_err("gen 2 must not start while gen 1 is still running"); + assert!( + format!("{error:#}").contains("older generation of this actor is still running"), + "unexpected error: {error:#}" + ); + assert!( + old_ctx.is_lost(), + "the stuck older generation is marked lost" + ); + assert!(new_ctx.is_generation_finished()); + assert!( + dispatcher + .starting_instances + .get_async(&instance_key(actor_id, 2)) + .await + .is_none(), + "the refused generation must not stay registered as starting", + ); + + // Released, the lost older generation refuses to start. + test_hooks::release_startup_gate(actor_id); + old_start + .await + .expect("gen 1 start joins") + .expect_err("a generation lost before starting must not start"); + assert!(old_ctx.is_generation_finished()); +} + +/// A generation lost while its runtime preamble hangs drops its start reply. The registry must +/// keep it until its task finishes aborting, so envoy-client does not report it stopped early. +#[tokio::test(start_paused = true)] +async fn lost_startup_is_released_only_after_its_task_finishes() { + use std::time::Duration; + + let (entered_tx, entered_rx) = oneshot::channel::<()>(); + let entered_tx = Arc::new(std::sync::Mutex::new(Some(entered_tx))); + let mut factories = HashMap::new(); + factories.insert( + "stuck".to_owned(), + Arc::new(ActorFactory::new_with_manual_startup_ready( + ActorConfig::default(), + move |mut start| { + let entered_tx = entered_tx.clone(); + Box::pin(async move { + let _startup_ready = start.startup_ready.take(); + if let Some(entered_tx) = entered_tx.lock().expect("entered lock").take() { + let _ = entered_tx.send(()); + } + std::future::pending::<()>().await; + Ok(()) + }) + }, + )), + ); + let dispatcher = Arc::new(RegistryDispatcher::new(factories, false)); + let actor_id = "actor-lost-startup"; + let ctx = crate::actor::context::ActorContext::new(actor_id, "stuck", Vec::new(), "local"); + let start_dispatcher = dispatcher.clone(); + let start_ctx = ctx.clone(); + let start = tokio::spawn(async move { + start_dispatcher + .start_actor(StartActorRequest { + actor_id: actor_id.to_owned(), + generation: 1, + actor_name: "stuck".to_owned(), + input: None, + ctx: start_ctx, + }) + .await + }); + entered_rx.await.expect("runtime preamble should start"); + + ctx.mark_lost(); + tokio::time::timeout(Duration::from_secs(10), start) + .await + .expect("a lost startup must finish within the lost bound") + .expect("start task joins") + .expect_err("a generation lost during startup must not start"); + assert!( + ctx.is_generation_finished(), + "startup must not be released before the generation's task finished" + ); + assert!( + dispatcher + .starting_instances + .get_async(&instance_key(actor_id, 1)) + .await + .is_none() + ); +} + +/// Start callbacks run in independent tasks, so an older generation's start can arrive after a +/// newer generation registered. It must refuse to run instead of overlapping the newer one. +#[tokio::test(start_paused = true)] +async fn superseded_generation_refuses_to_start() { + let dispatcher = counter_dispatcher(); + let actor_id = "actor-superseded"; + let new_ctx = start_counter(&dispatcher, actor_id, 2).await; + + let old_ctx = + crate::actor::context::ActorContext::new(actor_id, "counter", Vec::new(), "local"); + let error = dispatcher + .start_actor(StartActorRequest { + actor_id: actor_id.to_owned(), + generation: 1, + actor_name: "counter".to_owned(), + input: None, + ctx: old_ctx.clone(), + }) + .await + .expect_err("a superseded generation must not start"); + assert!( + format!("{error:#}").contains("newer generation of this actor is already running"), + "unexpected error: {error:#}" + ); + assert!(old_ctx.is_generation_finished()); + assert!(!new_ctx.is_lost(), "the newer generation must be untouched"); + assert_eq!(instance_generation(&dispatcher, actor_id, 1).await, None); + assert_eq!(dispatcher.current_generation(actor_id).await, Some(2)); +} + +/// A lost generation's storage cleanup can outlive its task. Its record must stay visible until +/// that cleanup releases the generation, so a newer generation keeps waiting on it. +#[tokio::test(start_paused = true)] +async fn stopped_generation_stays_registered_until_its_storage_is_released() { + use rivet_envoy_client::config::ActorStopHandle; + use rivet_envoy_client::protocol::StopActorReason; + + let dispatcher = counter_dispatcher(); + let actor_id = "actor-storage-fence"; + let ctx = start_counter(&dispatcher, actor_id, 1).await; + // Stands in for a storage cleanup that is still closing SQLite. + let storage_hold = ctx.hold_generation(); + + dispatcher + .stop_actor( + actor_id, + 1, + StopActorReason::Lost, + ActorStopHandle::detached(), + ) + .await + .expect("stopping the generation should succeed"); + assert_eq!( + instance_generation(&dispatcher, actor_id, 1).await, + Some(1), + "the record must stay while storage is still held", + ); + + drop(storage_hold); + ctx.wait_for_generation_finished().await; + for _ in 0..10 { + tokio::task::yield_now().await; + } + assert_eq!(instance_generation(&dispatcher, actor_id, 1).await, None); +} diff --git a/rivetkit-rust/packages/rivetkit-core/tests/sqlite.rs b/rivetkit-rust/packages/rivetkit-core/tests/sqlite.rs index 018e7da300..eb7c99a9b8 100644 --- a/rivetkit-rust/packages/rivetkit-core/tests/sqlite.rs +++ b/rivetkit-rust/packages/rivetkit-core/tests/sqlite.rs @@ -1570,3 +1570,37 @@ fn remote_head_fence_mismatch_stops_actor_once() { }); assert!(envoy_rx.try_recv().is_err()); } + +/// A statement can pass the lost check and then wait on the open lock. If the generation is +/// lost while it waits, it must not be submitted to the native worker afterward. +#[cfg(feature = "sqlite-local")] +#[tokio::test] +async fn lost_while_waiting_to_open_rejects_local_submission() { + let db = Arc::new(SqliteDb::default()); + let lost = tokio_util::sync::CancellationToken::new(); + db.set_lost_signal(lost.clone()); + + let open_guard = db.open_lock.clone().lock_owned().await; + let execute = tokio::spawn({ + let db = db.clone(); + async move { + db.local_execute("INSERT INTO lost_probe VALUES (1)".to_owned(), None) + .await + } + }); + // Let the statement pass the first lost check and park on the open lock. + for _ in 0..10 { + tokio::task::yield_now().await; + } + lost.cancel(); + drop(open_guard); + + let error = execute + .await + .expect("execute task joins") + .expect_err("a statement from a lost generation must not be submitted"); + assert!( + format!("{error:#}").contains("storage writes are revoked"), + "unexpected error: {error:#}" + ); +} diff --git a/rivetkit-rust/packages/rivetkit-core/tests/task.rs b/rivetkit-rust/packages/rivetkit-core/tests/task.rs index cd3515877b..a7b9c8c199 100644 --- a/rivetkit-rust/packages/rivetkit-core/tests/task.rs +++ b/rivetkit-rust/packages/rivetkit-core/tests/task.rs @@ -45,8 +45,8 @@ pub(crate) mod moved_tests { encode_persisted_actor, }; use crate::actor::task::{ - ActorTask, DispatchCommand, LONG_SHUTDOWN_DRAIN_WARNING_THRESHOLD, LifecycleCommand, - LifecycleEvent, LifecycleState, LiveExit, + ActorTask, DispatchCommand, LONG_SHUTDOWN_DRAIN_WARNING_THRESHOLD, + LOST_SHUTDOWN_GRACE_PERIOD, LifecycleCommand, LifecycleEvent, LifecycleState, LiveExit, }; use crate::actor::task_types::ShutdownKind; use crate::kv::tests::new_in_memory; @@ -2372,6 +2372,413 @@ pub(crate) mod moved_tests { .expect("destroy stop should succeed"); } + #[derive(Clone, Default)] + struct LostTestCounters { + serialize: Arc, + action: Arc, + cleanup: Arc, + } + + #[derive(Clone, Copy, PartialEq, Eq)] + enum LostTestHold { + None, + /// Holds the runtime cleanup hook open, like a slow `onSleep`. + Cleanup, + /// Never answers state serialization, like a stuck state serializer. + Serialize, + /// Behaves normally and exits its event loop once core closes the channel. + ExitOnClose, + } + + /// Factory for lost-generation tests. Its event loop counts serialization, actions, and + /// cleanup, can hold one of them open, and never exits after its event channel closes, like a + /// runtime adapter stuck draining user promises. + fn lost_test_factory(counters: LostTestCounters, hold: LostTestHold) -> Arc { + let config = ActorConfig { + sleep_grace_period: Duration::from_secs(60), + sleep_grace_period_overridden: true, + ..Default::default() + }; + Arc::new(ActorFactory::new(config, move |start| { + let counters = counters.clone(); + Box::pin(async move { + let mut events = start.events; + // Replies held open for the lifetime of the loop. + let mut held_serialize_replies = Vec::new(); + while let Some(event) = events.recv().await { + if let ActorEvent::SerializeState { reply, .. } = event { + counters.serialize.fetch_add(1, Ordering::SeqCst); + if hold == LostTestHold::Serialize { + held_serialize_replies.push(reply); + } else { + reply.send(Ok(Vec::new())); + } + } else if let ActorEvent::Action { reply, .. } = event { + counters.action.fetch_add(1, Ordering::SeqCst); + reply.send(Ok(Vec::new())); + } else if let ActorEvent::RunGracefulCleanup { reply, .. } = event { + counters.cleanup.fetch_add(1, Ordering::SeqCst); + if hold == LostTestHold::Cleanup { + futures::future::pending::<()>().await; + } + reply.send(Ok(())); + } + } + if hold == LostTestHold::ExitOnClose { + return Ok(()); + } + futures::future::pending::<()>().await; + Ok(()) + }) + })) + } + + async fn start_task_for_lost_test( + ctx: &ActorContext, + counters: LostTestCounters, + hold: LostTestHold, + ) -> ( + tokio::task::JoinHandle>, + mpsc::UnboundedSender, + mpsc::UnboundedSender, + mpsc::UnboundedSender, + ) { + let (mut task, lifecycle_tx, dispatch_tx, events_tx) = new_task_with_senders(ctx.clone()); + task.factory = lost_test_factory(counters, hold); + 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"); + // The caller keeps every sender alive; dropping one closes the task inbox and ends it. + (run, lifecycle_tx, dispatch_tx, events_tx) + } + + fn lost_test_ctx(name: &str) -> ActorContext { + new_with_kv(name, name, Vec::new(), "local", new_in_memory()) + } + + #[tokio::test] + async fn lost_signal_cancels_tracked_user_work_and_rejects_new_work() { + struct DropFlag(Arc); + impl Drop for DropFlag { + fn drop(&mut self) { + self.0.store(true, Ordering::SeqCst); + } + } + + let ctx = lost_test_ctx("actor-lost-tracked-work"); + let dropped = Arc::new(AtomicBool::new(false)); + let resumed = Arc::new(AtomicBool::new(false)); + let release = Arc::new(Notify::new()); + ctx.wait_until({ + let dropped = DropFlag(dropped.clone()); + let resumed = resumed.clone(); + let release = release.clone(); + async move { + let _dropped = dropped; + release.notified().await; + resumed.store(true, Ordering::SeqCst); + } + }); + tokio::task::yield_now().await; + + ctx.mark_lost(); + // The work becomes ready at the same time, but cancellation must win. + release.notify_waiters(); + for _ in 0..100 { + if dropped.load(Ordering::SeqCst) { + break; + } + tokio::task::yield_now().await; + } + assert!(dropped.load(Ordering::SeqCst), "lost work must be dropped"); + assert!( + !resumed.load(Ordering::SeqCst), + "lost work must not resume user code" + ); + + let ran = Arc::new(AtomicBool::new(false)); + ctx.register_task({ + let ran = ran.clone(); + async move { + ran.store(true, Ordering::SeqCst); + } + }); + for _ in 0..100 { + tokio::task::yield_now().await; + } + assert!( + !ran.load(Ordering::SeqCst), + "work registered after Lost must not run" + ); + } + + #[tokio::test] + async fn lost_signal_revokes_legacy_kv_writes() { + let ctx = lost_test_ctx("actor-lost-legacy-kv"); + ctx.legacy_kv() + .put(b"probe", b"before") + .await + .expect("writes succeed before the generation is lost"); + + ctx.mark_lost(); + + let error = ctx + .legacy_kv() + .put(b"probe", b"after") + .await + .expect_err("a lost generation must not write legacy kv"); + assert!( + format!("{error:#}").contains("storage writes are revoked"), + "unexpected error: {error:#}" + ); + } + + #[tokio::test(start_paused = true)] + async fn lost_signal_interrupts_a_startup_stuck_in_the_runtime_preamble() { + let ctx = lost_test_ctx("actor-lost-startup"); + let (entered_tx, entered_rx) = oneshot::channel(); + let entered_tx = Arc::new(Mutex::new(Some(entered_tx))); + let factory = Arc::new(ActorFactory::new_with_manual_startup_ready( + Default::default(), + move |mut start| { + let entered_tx = entered_tx.clone(); + Box::pin(async move { + // Models a `createVars` or `onWake` that never returns. The preamble never + // reports startup ready. + let _startup_ready = start.startup_ready.take(); + entered_tx + .lock() + .expect("entered lock poisoned") + .take() + .expect("entered sender should exist") + .send(()) + .expect("entered receiver should exist"); + std::future::pending::<()>().await; + Ok(()) + }) + }, + )); + let (mut task, lifecycle_tx, _dispatch_tx, _events_tx) = new_task_with_senders(ctx.clone()); + task.factory = factory; + 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"); + entered_rx.await.expect("runtime preamble should start"); + + let lost_at = Instant::now(); + ctx.mark_lost(); + + timeout(Duration::from_secs(10), run) + .await + .expect("a generation lost during startup must tear down on its own") + .expect("task run should join") + .expect("task run should succeed"); + assert!( + lost_at.elapsed() <= LOST_SHUTDOWN_GRACE_PERIOD + Duration::from_millis(100), + "teardown took {:?}", + lost_at.elapsed() + ); + assert!( + !matches!(start_rx.await, Ok(Ok(()))), + "a generation lost during startup must not report a successful start" + ); + } + + #[tokio::test(start_paused = true)] + async fn lost_teardown_finishes_as_soon_as_runtime_cleanup_replies() { + let ctx = lost_test_ctx("actor-lost-prompt"); + let counters = LostTestCounters::default(); + let (run, _lifecycle_tx, _dispatch_tx, _events_tx) = + start_task_for_lost_test(&ctx, counters.clone(), LostTestHold::ExitOnClose).await; + + // Paused time only advances when every task is idle, so this measures how long the + // teardown waited, not wall time. + let lost_at = tokio::time::Instant::now(); + ctx.mark_lost(); + timeout(Duration::from_secs(10), run) + .await + .expect("a lost generation must tear down on its own") + .expect("task run should join") + .expect("task run should succeed"); + + assert_eq!( + counters.cleanup.load(Ordering::SeqCst), + 1, + "runtime cleanup should run once" + ); + assert!( + lost_at.elapsed() < LOST_SHUTDOWN_GRACE_PERIOD / 2, + "lost teardown waited {:?} although runtime cleanup replied immediately", + lost_at.elapsed() + ); + } + + #[tokio::test(start_paused = true)] + async fn lost_signal_aborts_a_running_generation_without_a_stop_command() { + let ctx = lost_test_ctx("actor-lost-running"); + let counters = LostTestCounters::default(); + let (run, _lifecycle_tx, dispatch_tx, _events_tx) = + start_task_for_lost_test(&ctx, counters.clone(), LostTestHold::None).await; + + // A pending save tick and in-flight user work would normally be honored. + ctx.request_save(RequestSaveOpts::default()); + let _user_work = ctx.keep_awake_region(); + + let lost_at = Instant::now(); + ctx.mark_lost(); + let (action_tx, action_rx) = oneshot::channel(); + dispatch_tx + .send(DispatchCommand::Action { + name: "after-lost".to_owned(), + args: Vec::new(), + incoming: IncomingInvocationContext::default(), + conn: ConnHandle::new("conn-lost", Vec::new(), Vec::new(), false), + reply: action_tx, + }) + .expect("action should send"); + + timeout(Duration::from_secs(10), run) + .await + .expect("a lost generation must tear down on its own, even with a stuck event loop") + .expect("task run should join") + .expect("task run should succeed"); + + assert!( + lost_at.elapsed() <= 2 * LOST_SHUTDOWN_GRACE_PERIOD, + "lost teardown took {:?}", + lost_at.elapsed(), + ); + assert!( + action_rx.await.map_or(true, |reply| reply.is_err()), + "a lost generation must not run new actions" + ); + assert_eq!(counters.action.load(Ordering::SeqCst), 0); + assert_eq!( + counters.serialize.load(Ordering::SeqCst), + 0, + "a lost generation must not serialize state" + ); + assert_eq!( + counters.cleanup.load(Ordering::SeqCst), + 1, + "a lost generation still runs runtime cleanup exactly once" + ); + assert!(ctx.actor_aborted()); + } + + #[tokio::test(start_paused = true)] + async fn lost_signal_interrupts_a_graceful_sleep_stuck_in_cleanup() { + let ctx = lost_test_ctx("actor-lost-cleanup"); + let counters = LostTestCounters::default(); + let (run, lifecycle_tx, _dispatch_tx, _events_tx) = + start_task_for_lost_test(&ctx, counters.clone(), LostTestHold::Cleanup).await; + + let (sleep_tx, sleep_rx) = oneshot::channel(); + lifecycle_tx + .send(LifecycleCommand::Stop { + reason: ShutdownKind::Sleep, + reply: sleep_tx, + }) + .expect("sleep stop should send"); + sleep(Duration::from_millis(100)).await; + assert_eq!(counters.cleanup.load(Ordering::SeqCst), 1); + + let lost_at = Instant::now(); + ctx.mark_lost(); + timeout(Duration::from_secs(10), sleep_rx) + .await + .expect("the graceful sleep must finish within the lost bound") + .expect("sleep reply should send") + .expect("sleep stop should succeed"); + timeout(Duration::from_secs(10), run) + .await + .expect("task run should finish") + .expect("task run should join") + .expect("task run should succeed"); + + assert!(lost_at.elapsed() <= 2 * LOST_SHUTDOWN_GRACE_PERIOD); + assert_eq!( + counters.cleanup.load(Ordering::SeqCst), + 1, + "the abort path must not dispatch runtime cleanup a second time" + ); + assert_eq!( + counters.serialize.load(Ordering::SeqCst), + 0, + "an interrupted sleep must not serialize final state" + ); + } + + #[tokio::test(start_paused = true)] + async fn lost_signal_interrupts_a_save_waiting_on_serialization() { + let ctx = lost_test_ctx("actor-lost-serialize"); + let counters = LostTestCounters::default(); + let (run, _lifecycle_tx, _dispatch_tx, events_tx) = + start_task_for_lost_test(&ctx, counters.clone(), LostTestHold::Serialize).await; + + // A save tick asks the runtime to serialize state and waits for the reply, which the + // test factory never sends. + ctx.request_save(RequestSaveOpts { + immediate: true, + max_wait_ms: None, + }); + events_tx + .send(LifecycleEvent::SaveRequested { immediate: true }) + .expect("save request should send"); + timeout(Duration::from_secs(30), async { + while counters.serialize.load(Ordering::SeqCst) == 0 { + sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("the save tick should request state serialization"); + + let lost_at = Instant::now(); + ctx.mark_lost(); + timeout(Duration::from_secs(10), run) + .await + .expect("a save stuck on serialization must not hold a lost generation") + .expect("task run should join") + .expect("task run should succeed"); + assert!(lost_at.elapsed() <= 2 * LOST_SHUTDOWN_GRACE_PERIOD); + } + + #[tokio::test] + async fn lost_signal_revokes_sqlite_writes() { + let ctx = lost_test_ctx("actor-lost-sqlite"); + ctx.sql() + .execute("CREATE TABLE lost_probe (id INTEGER)", None) + .await + .expect("writes succeed before the generation is lost"); + + ctx.mark_lost(); + + let error = ctx + .sql() + .execute("INSERT INTO lost_probe VALUES (1)", None) + .await + .expect_err("a lost generation must not write storage"); + assert!( + format!("{error:#}").contains("storage writes are revoked"), + "unexpected error: {error:#}" + ); + let rollback = ctx.sql().execute("ROLLBACK", None).await; + if let Err(error) = rollback { + assert!( + !format!("{error:#}").contains("storage writes are revoked"), + "rollback must stay allowed after loss: {error:#}" + ); + } + } + #[tokio::test(start_paused = true)] async fn destroy_waits_for_tracked_schedule_persistence() { let kv = new_in_memory(); diff --git a/rivetkit-rust/packages/rivetkit/src/registry.rs b/rivetkit-rust/packages/rivetkit/src/registry.rs index 41196c3606..7b3ff6d96e 100644 --- a/rivetkit-rust/packages/rivetkit/src/registry.rs +++ b/rivetkit-rust/packages/rivetkit/src/registry.rs @@ -238,7 +238,18 @@ where match wrap_start::(core_start) { Ok(mut start) => { start.startup_ready = startup_ready; - entry(start).await + // A lost generation runs no more user code, so the custom entry is dropped + // as soon as the engine gives up on it. + let lost = start.ctx.inner().lost_signal(); + // Calling the entry can run user code before its future is polled. + if lost.is_cancelled() { + return Ok(()); + } + tokio::select! { + biased; + _ = lost.cancelled() => Ok(()), + result = entry(start) => result, + } } Err(error) => { if let Some(reply) = startup_ready { diff --git a/rivetkit-rust/packages/rivetkit/src/sqlite.rs b/rivetkit-rust/packages/rivetkit/src/sqlite.rs index 8d0526d9df..4a866275f3 100644 --- a/rivetkit-rust/packages/rivetkit/src/sqlite.rs +++ b/rivetkit-rust/packages/rivetkit/src/sqlite.rs @@ -1,6 +1,7 @@ use std::{future::Future, time::Duration}; use anyhow::Result; +use rivetkit_core::error::ActorLifecycle; use rivetkit_core::{SqliteDb, SqliteTransaction}; #[derive(Clone, Copy, Debug, Default)] @@ -37,6 +38,13 @@ impl SqliteDbExt for SqliteDb { let transaction = self .begin_named_transaction(options.name, options.timeout) .await?; + // Beginning awaits, and the generation can be lost meanwhile. Do not start the callback. + if self.is_lost() { + let _ = transaction.rollback().await; + return Err(ActorLifecycle::Stopping + .build() + .context("actor generation was declared lost")); + } match callback(transaction.clone()).await { Ok(value) => { if let Err(error) = transaction.commit().await { diff --git a/rivetkit-rust/packages/rivetkit/src/start.rs b/rivetkit-rust/packages/rivetkit/src/start.rs index b6502f25b3..2ac2609706 100644 --- a/rivetkit-rust/packages/rivetkit/src/start.rs +++ b/rivetkit-rust/packages/rivetkit/src/start.rs @@ -216,14 +216,26 @@ pub async fn run_actor(start: Start) -> Result<()> { }; ctx.set_initial_state(state); + // Each hook can finish after the generation was lost, so check before starting the next. + ensure_not_lost(&ctx)?; let actor = Arc::new(A::create(&ctx).await?); if is_new { + ensure_not_lost(&ctx)?; actor.clone().on_create(ctx.clone()).await?; } + ensure_not_lost(&ctx)?; actor.clone().on_start(ctx.clone()).await?; Ok::<_, anyhow::Error>(actor) - } - .await; + }; + // A lost generation runs no more user code, so in-flight callbacks are dropped as soon as the + // engine gives up on it instead of when core aborts this adapter. + let lost = ctx.inner().lost_signal(); + let startup = tokio::select! { + biased; + _ = lost.cancelled() => Err(ActorLifecycle::Stopping.build()) + .context("actor generation was declared lost during startup"), + startup = startup => startup, + }; let actor = match startup { Ok(actor) => { @@ -245,8 +257,11 @@ pub async fn run_actor(start: Start) -> Result<()> { let http_pools = HttpCallbackPools::for_actor::(); while let Some(event) = events.recv_raw().await { - let should_stop = - handle_actor_event(actor.clone(), ctx.clone(), http_pools.as_ref(), event).await?; + let should_stop = tokio::select! { + biased; + _ = lost.cancelled() => break, + should_stop = handle_actor_event(actor.clone(), ctx.clone(), http_pools.as_ref(), event) => should_stop?, + }; if should_stop { break; } @@ -301,8 +316,16 @@ fn spawn_run_task( ctx: Ctx, cancel: CancellationToken, ) -> JoinHandle> { + // Core can abort this adapter's event loop when the generation is lost, which detaches this + // task. It stops on the lost signal itself, and holds the generation so a newer generation + // does not start until it has ended. + let generation_hold = ctx.inner().hold_generation(); + let lost = ctx.inner().lost_signal(); tokio::spawn(async move { + let _generation_hold = generation_hold; let result = tokio::select! { + biased; + _ = lost.cancelled() => return Ok(()), _ = cancel.cancelled() => return Ok(()), result = AssertUnwindSafe(actor.run(ctx.clone())).catch_unwind() => result, }; @@ -401,6 +424,10 @@ async fn handle_actor_event( return Ok(false); }; + if let Err(error) = ensure_not_lost(&ctx) { + reply.send(Err(error)); + return Ok(false); + } let admission = actor.admit_http_request(&ctx, &request); let release = match admission { HttpCallbackAdmission::Untracked => None, @@ -469,9 +496,11 @@ async fn handle_actor_event( .clone() .on_before_connect(ctx.clone(), ¶ms) .await?; + ensure_not_lost(&ctx)?; let conn_state = actor.clone().create_conn_state(ctx.clone(), params).await?; let conn = ConnCtx::from(conn); conn.set_state(&conn_state)?; + ensure_not_lost(&ctx)?; actor.on_connect(ctx, conn).await } .await; @@ -508,6 +537,9 @@ async fn handle_actor_event( } ActorEvent::RunGracefulCleanup { reason, reply } => { let result = match reason { + // A lost generation runs no user hooks because the engine may already be + // running the next generation. + ShutdownKind::Sleep if ctx.inner().is_lost() => Ok(()), ShutdownKind::Sleep => actor.on_sleep(ctx).await, ShutdownKind::Destroy => actor.on_destroy(ctx).await, }; @@ -550,9 +582,17 @@ fn spawn_action_reply( reply: Reply>, future: crate::action::BoxActionFuture, ) { + // Like the run task, this detached task stops on Lost and holds the generation until it ends. + let generation_hold = ctx.inner().hold_generation(); + let lost = ctx.inner().lost_signal(); tokio::spawn(async move { + let _generation_hold = generation_hold; let abort = ctx.abort_signal(); tokio::select! { + biased; + _ = lost.cancelled() => { + reply.send(Err(ActorLifecycle::Stopping.build())); + } _ = abort.cancelled() => { reply.send(Err(ActorLifecycle::Stopping.build())); } @@ -568,9 +608,15 @@ fn spawn_queue_reply( reply: Reply, future: crate::queue::BoxQueueFuture, ) { + // Like the run task, this detached task stops on Lost and holds the generation until it ends. + let generation_hold = ctx.inner().hold_generation(); + let lost = ctx.inner().lost_signal(); tokio::spawn(async move { + let _generation_hold = generation_hold; let abort = ctx.abort_signal(); let result = tokio::select! { + biased; + _ = lost.cancelled() => Err(ActorLifecycle::Stopping.build()), _ = abort.cancelled() => Err(ActorLifecycle::Stopping.build()), result = future => result.map(|response| QueueSendResult { status: QueueSendStatus::Completed, @@ -581,6 +627,16 @@ fn spawn_queue_reply( }); } +/// Refuses to start another user callback for a generation that was declared lost. +fn ensure_not_lost(ctx: &Ctx) -> Result<()> { + if ctx.inner().is_lost() { + return Err(ActorLifecycle::Stopping + .build() + .context("actor generation was declared lost")); + } + Ok(()) +} + fn not_configured(component: impl Into) -> anyhow::Error { ActorRuntime::NotConfigured { component: component.into(), @@ -654,7 +710,7 @@ mod tests { use std::pin::Pin; use std::sync::{ OnceLock, - atomic::{AtomicUsize, Ordering}, + atomic::{AtomicBool, AtomicUsize, Ordering}, }; use async_trait::async_trait; @@ -2242,6 +2298,183 @@ mod tests { } } + static HUNG_SLEEP_STARTED: AtomicBool = AtomicBool::new(false); + static HUNG_SLEEP_DROPPED: AtomicBool = AtomicBool::new(false); + + struct HungSleepDropFlag; + + impl Drop for HungSleepDropFlag { + fn drop(&mut self) { + HUNG_SLEEP_DROPPED.store(true, Ordering::SeqCst); + } + } + + struct HungSleepActor; + + #[async_trait] + impl Actor for HungSleepActor { + type State = (); + type Input = (); + type Actions = (); + type Events = (); + type Queue = (); + type ConnParams = (); + type ConnState = (); + type Action = action::Raw; + + async fn create_state(_ctx: &Ctx, (): Self::Input) -> Result { + Ok(()) + } + + async fn create(_ctx: &Ctx) -> Result { + Ok(Self) + } + + async fn on_sleep(self: Arc, _ctx: Ctx) -> Result<()> { + let _dropped = HungSleepDropFlag; + HUNG_SLEEP_STARTED.store(true, Ordering::SeqCst); + std::future::pending::<()>().await; + Ok(()) + } + } + + #[tokio::test] + async fn lost_generation_drops_an_in_flight_sleep_hook() { + let (tx, rx) = unbounded_channel(); + let (start, ctx) = unit_start_with_ctx::(rx.into()); + let actor = tokio::spawn(run_actor::(start)); + let (reply_tx, _reply_rx) = oneshot::channel(); + tx.send(ActorEvent::RunGracefulCleanup { + reason: ShutdownKind::Sleep, + reply: reply_tx.into(), + }) + .expect("send sleep cleanup"); + for _ in 0..1000 { + if HUNG_SLEEP_STARTED.load(Ordering::SeqCst) { + break; + } + tokio::task::yield_now().await; + } + assert!( + HUNG_SLEEP_STARTED.load(Ordering::SeqCst), + "on_sleep should start" + ); + + ctx.inner().lost_signal().cancel(); + + tokio::time::timeout(Duration::from_secs(1), actor) + .await + .expect("the adapter must stop as soon as the generation is lost") + .expect("run_actor joins") + .expect("run_actor returns"); + assert!( + HUNG_SLEEP_DROPPED.load(Ordering::SeqCst), + "the in-flight sleep hook must be dropped, not resumed" + ); + drop(tx); + } + + static LOST_RUN_STARTED: AtomicBool = AtomicBool::new(false); + static LOST_RUN_DROPPED: AtomicBool = AtomicBool::new(false); + + struct LostRunDropFlag; + + impl Drop for LostRunDropFlag { + fn drop(&mut self) { + LOST_RUN_DROPPED.store(true, Ordering::SeqCst); + } + } + + struct LostRunActor; + + #[async_trait] + impl Actor for LostRunActor { + type State = (); + type Input = (); + type Actions = (); + type Events = (); + type Queue = (); + type ConnParams = (); + type ConnState = (); + type Action = action::Raw; + + async fn create_state(_ctx: &Ctx, (): Self::Input) -> Result { + Ok(()) + } + + async fn create(_ctx: &Ctx) -> Result { + Ok(Self) + } + + async fn run(self: Arc, _ctx: Ctx) -> Result<()> { + let _dropped = LostRunDropFlag; + LOST_RUN_STARTED.store(true, Ordering::SeqCst); + std::future::pending::<()>().await; + Ok(()) + } + } + + #[tokio::test] + async fn lost_generation_stops_the_run_task_after_its_event_loop_is_aborted() { + let (tx, rx) = unbounded_channel(); + let (start, ctx) = unit_start_with_ctx::(rx.into()); + ctx.inner().set_started(true); + let actor = tokio::spawn(run_actor::(start)); + for _ in 0..1000 { + if LOST_RUN_STARTED.load(Ordering::SeqCst) { + break; + } + tokio::task::yield_now().await; + } + assert!(LOST_RUN_STARTED.load(Ordering::SeqCst), "run should start"); + + ctx.inner().lost_signal().cancel(); + // Core aborts the adapter event loop after the lost deadline, detaching the run task. + actor.abort(); + for _ in 0..1000 { + if LOST_RUN_DROPPED.load(Ordering::SeqCst) { + break; + } + tokio::task::yield_now().await; + } + assert!( + LOST_RUN_DROPPED.load(Ordering::SeqCst), + "a lost generation's run task must stop even when its event loop is aborted" + ); + drop(tx); + } + + #[tokio::test] + async fn lost_generation_stops_pending_action_and_queue_replies() { + let (_tx, rx) = unbounded_channel(); + let (_start, ctx) = unit_start_with_ctx::(rx.into()); + let (action_tx, action_rx) = oneshot::channel(); + spawn_action_reply( + ctx.clone(), + action_tx.into(), + Box::pin(std::future::pending::>>()), + ); + let (queue_tx, queue_rx) = oneshot::channel(); + spawn_queue_reply( + ctx.clone(), + queue_tx.into(), + Box::pin(std::future::pending::>>>()), + ); + + ctx.inner().lost_signal().cancel(); + + let action = tokio::time::timeout(Duration::from_secs(1), action_rx) + .await + .expect("a lost generation's pending action must stop") + .expect("action reply should be sent"); + assert!(action.is_err()); + let queue = tokio::time::timeout(Duration::from_secs(1), queue_rx) + .await + .expect("a lost generation's pending queue send must stop") + .expect("queue reply should be sent"); + assert!(queue.is_err()); + } + struct PanickingRunActor; #[async_trait] diff --git a/rivetkit-typescript/packages/rivetkit-napi/index.d.ts b/rivetkit-typescript/packages/rivetkit-napi/index.d.ts index 1f83c8a6f7..3fa0e9abc7 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/index.d.ts +++ b/rivetkit-typescript/packages/rivetkit-napi/index.d.ts @@ -373,6 +373,7 @@ export declare class ActorContext { destroy(): void destroyRequested(): boolean generation(): number | null + isLost(): boolean waitForDestroyCompletion(): Promise setPreventSleep(preventSleep: boolean): void preventSleep(): boolean diff --git a/rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs b/rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs index e1cc6a99f7..2a98c8d944 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs +++ b/rivetkit-typescript/packages/rivetkit-napi/src/actor_context.rs @@ -612,6 +612,11 @@ impl ActorContext { self.inner.generation() } + #[napi] + pub fn is_lost(&self) -> bool { + self.inner.is_lost() + } + #[napi] pub async fn wait_for_destroy_completion(&self) { self.inner.wait_for_destroy_completion_public().await; diff --git a/rivetkit-typescript/packages/rivetkit-napi/src/napi_actor_events.rs b/rivetkit-typescript/packages/rivetkit-napi/src/napi_actor_events.rs index d841261f0d..c4e775c17f 100644 --- a/rivetkit-typescript/packages/rivetkit-napi/src/napi_actor_events.rs +++ b/rivetkit-typescript/packages/rivetkit-napi/src/napi_actor_events.rs @@ -244,6 +244,7 @@ async fn run_preamble( // Run database migrations before any user lifecycle hook so `c.db` is // usable from createState, onCreate, and createVars. if let Some(callback) = &bindings.on_migrate { + ensure_preamble_not_lost(ctx)?; with_timeout( "onMigrate", config.on_migrate_timeout, @@ -254,6 +255,7 @@ async fn run_preamble( if is_new { if let Some(callback) = &bindings.create_state { + ensure_preamble_not_lost(ctx)?; let bytes = with_timeout( "createState", config.create_state_timeout, @@ -263,6 +265,7 @@ async fn run_preamble( ctx.set_state_initial(bytes)?; } if let Some(callback) = &bindings.on_create { + ensure_preamble_not_lost(ctx)?; with_timeout( "onCreate", config.on_create_timeout, @@ -274,6 +277,7 @@ async fn run_preamble( } else if let Some(snapshot) = snapshot { ctx.set_state_initial(snapshot)?; } else if let Some(callback) = &bindings.create_state { + ensure_preamble_not_lost(ctx)?; let bytes = with_timeout( "createState", config.create_state_timeout, @@ -297,6 +301,7 @@ async fn run_preamble( } if let Some(callback) = &bindings.create_vars { + ensure_preamble_not_lost(ctx)?; with_timeout( "createVars", config.create_vars_timeout, @@ -306,6 +311,7 @@ async fn run_preamble( } if let Some(callback) = &bindings.on_wake { + ensure_preamble_not_lost(ctx)?; with_timeout( "onWake", config.on_wake_timeout, @@ -315,6 +321,7 @@ async fn run_preamble( } if let Some(callback) = &bindings.on_before_actor_start { + ensure_preamble_not_lost(ctx)?; with_timeout( "onBeforeActorStart", config.on_before_actor_start_timeout, @@ -323,6 +330,8 @@ async fn run_preamble( .await?; } + // The last startup hook may finish after the generation was lost. `run` must not start then. + ensure_preamble_not_lost(ctx)?; let run_handler = configure_run_handler(bindings, ctx); if bindings.run.is_some() { tokio::task::yield_now().await; @@ -585,6 +594,8 @@ pub(crate) async fn dispatch_event( } if let Some(callback) = create_conn_state { + // `onBeforeConnect` may finish after the generation was lost. + ensure_preamble_not_lost(&ctx)?; let state = with_timeout( "createConnState", create_conn_state_timeout, @@ -996,6 +1007,10 @@ fn spawn_run_handler( ) -> JoinHandle<()> { let run_handler_active = RunHandlerActiveGuard::new(ctx.inner().clone()); tokio::spawn(async move { + // A lost generation starts no new user code, including a restarted `run`. + if ctx.is_lost() { + return; + } let result = call_run(&callback, &ctx).await; match result { Ok(callback_accepted) => { @@ -1412,6 +1427,16 @@ fn actor_shutting_down() -> anyhow::Error { ActorLifecycle::Stopping.build() } +/// A generation lost during its startup preamble starts no further user hook. A hook that was +/// already running finishes, but nothing after it is dispatched. +fn ensure_preamble_not_lost(ctx: &ActorContext) -> Result<()> { + if ctx.is_lost() { + return Err(actor_shutting_down() + .context("actor generation was declared lost during its startup preamble")); + } + Ok(()) +} + fn missing_callback(name: &str) -> anyhow::Error { NapiInvalidState { state: format!("callback {name}"), diff --git a/rivetkit-typescript/packages/rivetkit-wasm/src/lib.rs b/rivetkit-typescript/packages/rivetkit-wasm/src/lib.rs index d6062e3974..66027675ed 100644 --- a/rivetkit-typescript/packages/rivetkit-wasm/src/lib.rs +++ b/rivetkit-typescript/packages/rivetkit-wasm/src/lib.rs @@ -13,7 +13,7 @@ use rivet_error::{ ActorSpecifier, MacroMarker, RivetError as RivetTransportError, RivetErrorKind, RivetErrorSchema, }; -use rivetkit_core::error::public_error_status_code; +use rivetkit_core::error::{ActorLifecycle, public_error_status_code}; use rivetkit_core::inspector::InspectorAuth; use rivetkit_core::{ ActorConfig, ActorConfigInput, ActorEvent, ActorFactory as CoreActorFactory, ActorHttpResponse, @@ -633,6 +633,10 @@ fn start_run_handler(callbacks: &WasmCallbacks, ctx: &WasmActorContext) -> Resul if ctx.inner.run_handler_active() { return Err(anyhow!("wasm run handler is already active")); } + // A lost generation starts no new user code, including a restarted `run`. + if ctx.inner.is_lost() { + return Ok(()); + } let ctx = ctx.clone(); ctx.inner.begin_run_handler(); spawn_local(async move { @@ -651,6 +655,17 @@ fn start_run_handler(callbacks: &WasmCallbacks, ctx: &WasmActorContext) -> Resul Ok(()) } +/// A generation lost during its startup preamble starts no further user hook. A hook that was +/// already running finishes, but nothing after it is dispatched. +fn ensure_preamble_not_lost(ctx: &WasmActorContext) -> Result<()> { + if ctx.inner.is_lost() { + return Err(ActorLifecycle::Stopping + .build() + .context("actor generation was declared lost during its startup preamble")); + } + Ok(()) +} + async fn run_preamble( callbacks: &WasmCallbacks, ctx: &WasmActorContext, @@ -659,6 +674,7 @@ async fn run_preamble( snapshot: Option>, ) -> Result<()> { if let Some(callback) = &callbacks.on_migrate { + ensure_preamble_not_lost(ctx)?; let payload = object(); set_anyhow(&payload, "ctx", JsValue::from(ctx.clone()))?; set_anyhow(&payload, "isNew", JsValue::from_bool(is_new))?; @@ -667,6 +683,7 @@ async fn run_preamble( if is_new { if let Some(callback) = &callbacks.create_state { + ensure_preamble_not_lost(ctx)?; let payload = object(); set_anyhow(&payload, "ctx", JsValue::from(ctx.clone()))?; if let Some(input) = input.as_ref() { @@ -676,6 +693,7 @@ async fn run_preamble( ctx.inner.set_state_initial(state); } if let Some(callback) = &callbacks.on_create { + ensure_preamble_not_lost(ctx)?; let payload = object(); set_anyhow(&payload, "ctx", JsValue::from(ctx.clone()))?; if let Some(input) = input.as_ref() { @@ -686,6 +704,7 @@ async fn run_preamble( } else if let Some(snapshot) = snapshot { ctx.inner.set_state_initial(snapshot); } else if let Some(callback) = &callbacks.create_state { + ensure_preamble_not_lost(ctx)?; let payload = object(); set_anyhow(&payload, "ctx", JsValue::from(ctx.clone()))?; if let Some(input) = input.as_ref() { @@ -699,23 +718,28 @@ async fn run_preamble( } if let Some(callback) = &callbacks.create_vars { + ensure_preamble_not_lost(ctx)?; let payload = object(); set_anyhow(&payload, "ctx", JsValue::from(ctx.clone()))?; call_callback(callback, &payload.into()).await?; } if let Some(callback) = &callbacks.on_wake { + ensure_preamble_not_lost(ctx)?; let payload = object(); set_anyhow(&payload, "ctx", JsValue::from(ctx.clone()))?; call_callback(callback, &payload.into()).await?; } if let Some(callback) = &callbacks.on_before_actor_start { + ensure_preamble_not_lost(ctx)?; let payload = object(); set_anyhow(&payload, "ctx", JsValue::from(ctx.clone()))?; call_callback(callback, &payload.into()).await?; } + // The last startup hook may finish after the generation was lost. `run` must not start then. + ensure_preamble_not_lost(ctx)?; Ok(()) } @@ -760,7 +784,12 @@ async fn dispatch_event(callbacks: &WasmCallbacks, ctx: &WasmActorContext, event } let mut output = call_callback_bytes(&callback, &payload.into()).await?; - if let Some(callback) = &on_before_action_response { + // The action may finish after its generation was lost, and a lost + // generation starts no new user callbacks. + if let Some(callback) = on_before_action_response + .as_ref() + .filter(|_| !ctx.is_lost()) + { let payload = object(); set_anyhow(&payload, "ctx", JsValue::from(ctx.clone()))?; set_anyhow(&payload, "name", JsValue::from_str(&name))?; @@ -1081,6 +1110,8 @@ async fn run_connection_preflight( } if let Some(callback) = &callbacks.create_conn_state { + // `onBeforeConnect` may finish after the generation was lost. + ensure_preamble_not_lost(ctx)?; let payload = object(); set_anyhow(&payload, "ctx", JsValue::from(ctx.clone()))?; set_anyhow( @@ -1286,6 +1317,11 @@ impl WasmActorContext { self.inner.generation() } + #[wasm_bindgen(js_name = isLost)] + pub fn is_lost(&self) -> bool { + self.inner.is_lost() + } + #[wasm_bindgen] pub fn name(&self) -> String { self.inner.name().to_owned() diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/lost-callbacks.test.ts b/rivetkit-typescript/packages/rivetkit/src/registry/lost-callbacks.test.ts new file mode 100644 index 0000000000..8681e5ea6c --- /dev/null +++ b/rivetkit-typescript/packages/rivetkit/src/registry/lost-callbacks.test.ts @@ -0,0 +1,64 @@ +import { EvictedError } from "@rivetkit/workflow-engine"; +import { describe, expect, test } from "vitest"; +import { ACTOR_CONTEXT_INTERNAL_SYMBOL } from "@/actor/config"; +import { throwIfGenerationLost } from "@/workflow/context"; +import { refuseCallbacksAfterLost } from "./native"; +import type { CoreRuntime } from "./runtime"; + +function fakeRuntime(isLost: () => boolean): CoreRuntime { + return { actorIsLost: () => isLost() } as unknown as CoreRuntime; +} + +describe("lost generation callback gate", () => { + test("refuses runtime callbacks once the generation is lost", async () => { + let lost = false; + const guarded = refuseCallbacksAfterLost( + fakeRuntime(() => lost), + { + createVars: async (_error: unknown, _payload: unknown) => + "createVars ran", + onSleep: async (_error: unknown, _payload: unknown) => + "cleanup ran", + onDestroy: async (_error: unknown, _payload: unknown) => + "destroy cleanup ran", + actions: { + increment: async (_error: unknown, _payload: unknown) => + "action ran", + }, + hasState: true, + onRequest: undefined, + }, + ); + const payload = { ctx: {} }; + + await expect(guarded.createVars(null, payload)).resolves.toBe( + "createVars ran", + ); + await expect(guarded.actions.increment(null, payload)).resolves.toBe( + "action ran", + ); + + lost = true; + await expect(guarded.createVars(null, payload)).rejects.toBeDefined(); + await expect( + guarded.actions.increment(null, payload), + ).rejects.toBeDefined(); + // Cleanup callbacks still run so the generation releases its runtime state. + await expect(guarded.onSleep(null, payload)).resolves.toBe( + "cleanup ran", + ); + await expect(guarded.onDestroy(null, payload)).resolves.toBe( + "destroy cleanup ran", + ); + expect(guarded.hasState).toBe(true); + expect(guarded.onRequest).toBeUndefined(); + }); + + test("workflow code stops before more user callbacks once lost", () => { + const actor = { isLost: false }; + const runCtx = { [ACTOR_CONTEXT_INTERNAL_SYMBOL]: actor }; + expect(() => throwIfGenerationLost(runCtx)).not.toThrow(); + actor.isLost = true; + expect(() => throwIfGenerationLost(runCtx)).toThrow(EvictedError); + }); +}); diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts b/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts index 05c834aeff..c4ebdb357d 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/napi-runtime.ts @@ -604,6 +604,10 @@ export class NapiCoreRuntime implements CoreRuntime { return asNativeActorContext(ctx).generation() ?? undefined; } + actorIsLost(ctx: ActorContextHandle): boolean { + return asNativeActorContext(ctx).isLost(); + } + runWithActorInvocationContext(ctx: ActorContextHandle, run: () => T): T { return this.#runAs(asNativeActorContext(ctx), run); } diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/native.ts b/rivetkit-typescript/packages/rivetkit/src/registry/native.ts index ce4a5af4bd..c773b82fb6 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/native.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/native.ts @@ -541,8 +541,11 @@ async function cleanupNativeSleepRuntimeState( void runtime .actorWaitForTrackedShutdownWorkUnbounded(ctx) .then(async () => { - await afterTrackedWorkDrained?.(); - clearNativeRuntimeState(runtime, ctx); + try { + await afterTrackedWorkDrained?.(); + } finally { + clearNativeRuntimeState(runtime, ctx); + } }) .catch((error) => { logger().warn({ @@ -553,10 +556,31 @@ async function cleanupNativeSleepRuntimeState( return; } - await afterTrackedWorkDrained?.(); + try { + await afterTrackedWorkDrained?.(); + } finally { + try { + await closeNativeDatabaseClient(runtime, ctx); + await closeNativeSqlDatabase(runtime, ctx); + } finally { + clearNativeRuntimeState(runtime, ctx); + } + } +} + +async function cleanupNativeLostRuntimeState( + runtime: CoreRuntime, + ctx: ActorContextHandle, +): Promise { + // Core has already abandoned a lost generation's user work, so this does not wait for it. + // A custom database client's `close` is user code, so a lost generation drops the client + // without calling it. The native database underneath is closed by the runtime. await closeNativeDatabaseClient(runtime, ctx); - await closeNativeSqlDatabase(runtime, ctx); - clearNativeRuntimeState(runtime, ctx); + try { + await closeNativeSqlDatabase(runtime, ctx); + } finally { + clearNativeRuntimeState(runtime, ctx); + } } function closeNativeSqlDatabase( @@ -585,6 +609,11 @@ async function closeNativeDatabaseClient( runtimeState.databaseClient = undefined; + // `close` is user code from a database provider. Cleanup paths await before reaching this, + // so check here: a lost generation drops the client without calling it. + if (callNativeSync(() => runtime.actorIsLost(ctx))) { + return; + } if ( entry.client && typeof entry.client === "object" && @@ -1301,6 +1330,76 @@ function wrapNativeCallback, Result>( }; } +/** + * Callbacks that still run for a lost generation. They release the generation's runtime + * resources and skip the user hook themselves. + */ +const LOST_GENERATION_CLEANUP_CALLBACKS = new Set(["onSleep", "onDestroy"]); + +/** + * Refuses to start a user callback for a generation that was declared lost. Call it right + * before the callback, after any other user code such as schema validation. + */ +function throwIfActorLost(runtime: CoreRuntime, ctx: ActorContextHandle): void { + if (callNativeSync(() => runtime.actorIsLost(ctx))) { + throw actorGenerationLostError(); + } +} + +function actorGenerationLostError(): RivetError { + return new RivetError( + "actor", + "stopping", + "Actor generation was declared lost; it starts no new user callbacks.", + ); +} + +/** + * Refuses every runtime callback for a generation that has been declared lost. Runtime + * callbacks can be queued or spawned before the loss and only start afterward, so the check + * runs when JS actually enters the callback. + */ +export function refuseCallbacksAfterLost>( + runtime: CoreRuntime, + callbacks: T, +): T { + const guard = (callback: (...args: unknown[]) => unknown) => { + return async (...args: unknown[]) => { + const ctx = (args[1] as { ctx?: ActorContextHandle } | undefined) + ?.ctx; + if ( + ctx !== undefined && + callNativeSync(() => runtime.actorIsLost(ctx)) + ) { + throw encodeNativeCallbackError(actorGenerationLostError()); + } + return await callback(...args); + }; + }; + const guarded: Record = {}; + for (const [name, value] of Object.entries(callbacks)) { + if (LOST_GENERATION_CLEANUP_CALLBACKS.has(name)) { + guarded[name] = value; + } else if (typeof value === "function") { + guarded[name] = guard(value as (...args: unknown[]) => unknown); + } else if ( + name === "actions" && + value !== null && + typeof value === "object" + ) { + guarded[name] = Object.fromEntries( + Object.entries(value).map(([actionName, handler]) => [ + actionName, + guard(handler as (...args: unknown[]) => unknown), + ]), + ); + } else { + guarded[name] = value; + } + } + return guarded as T; +} + function decodeArgs(value?: RuntimeBytes | null): unknown[] { const decoded = decodeValue(value); return normalizeArgs(decoded); @@ -1880,9 +1979,26 @@ class NativeQueueAdapter { this.#schemas = schemas; } + // Queue results arrive after awaits, and their schema validation is user code, so each + // validation checks that the generation was not lost first. + #validateBody(name: string, body: unknown): unknown { + throwIfActorLost(this.#runtime, this.#ctx); + return validateQueueBody(this.#schemas, name, body); + } + + #validateComplete(name: string, response: unknown): unknown { + throwIfActorLost(this.#runtime, this.#ctx); + return validateQueueComplete(this.#schemas, name, response); + } + + #wrap(message: Parameters[0]) { + throwIfActorLost(this.#runtime, this.#ctx); + return wrapQueueMessage(message, this.#schemas); + } + async send(name: string, body: unknown) { - const validatedBody = validateQueueBody(this.#schemas, name, body); - return wrapQueueMessage( + const validatedBody = this.#validateBody(name, body); + return this.#wrap( await callNative(() => this.#runtime.actorQueueSend( this.#ctx, @@ -1890,7 +2006,6 @@ class NativeQueueAdapter { encodeValue(validatedBody), ), ), - this.#schemas, ); } @@ -1948,9 +2063,7 @@ class NativeQueueAdapter { token, ), ); - const wrapped = messages.map((message) => - wrapQueueMessage(message, this.#schemas), - ); + const wrapped = messages.map((message) => this.#wrap(message)); return completable ? wrapped.map((message) => this.#makeCompletableMessage(message), @@ -1975,7 +2088,7 @@ class NativeQueueAdapter { ); try { - return wrapQueueMessage( + return this.#wrap( await callNative(() => this.#runtime.actorQueueWaitForNames( this.#ctx, @@ -1987,7 +2100,6 @@ class NativeQueueAdapter { token, ), ), - this.#schemas, ); } finally { cleanup?.(); @@ -2032,8 +2144,7 @@ class NativeQueueAdapter { message: { id: number | bigint; name: string }, response?: unknown, ): Promise { - const validatedResponse = validateQueueComplete( - this.#schemas, + const validatedResponse = this.#validateComplete( message.name, response, ); @@ -2056,7 +2167,7 @@ class NativeQueueAdapter { signal?: AbortSignal; }, ) { - const validatedBody = validateQueueBody(this.#schemas, name, body); + const validatedBody = this.#validateBody(name, body); const { token, cleanup } = await createCancellationTokenHandle( this.#runtime, options?.signal, @@ -2076,11 +2187,7 @@ class NativeQueueAdapter { ); return response === undefined || response === null ? undefined - : validateQueueComplete( - this.#schemas, - name, - decodeValue(response), - ); + : this.#validateComplete(name, decodeValue(response)); } finally { cleanup?.(); } @@ -2583,6 +2690,12 @@ class TrackedWebSocketHandleAdapter implements UniversalWebSocket { handler: TrackedWebSocketListener, event: any, ): void { + // Messages and opens start new user work, so a lost generation drops them. Each listener + // is checked because an earlier one can run while the generation is lost. Close and + // error still reach listeners, like the abort signal, so user code can wind down. + if ((type === "message" || type === "open") && this.#ctx.isLost) { + return; + } try { const result = handler(event); if (!this.#isPromiseLike(result)) { @@ -2967,6 +3080,11 @@ export class ActorContextHandleAdapter { return callNativeSync(() => this.#runtime.actorGeneration(this.#ctx)); } + /** Whether this generation was declared lost. A lost generation starts no new user code. */ + get isLost(): boolean { + return callNativeSync(() => this.#runtime.actorIsLost(this.#ctx)); + } + get name(): string { return callNativeSync(() => this.#runtime.actorName(this.#ctx)); } @@ -3100,6 +3218,10 @@ export class ActorContextHandleAdapter { } const actorId = this.actorId; + // `createClient` is user code from a database provider. + if (this.isLost) { + throw actorGenerationLostError(); + } const createdClient = await this.#databaseProvider.createClient({ actorId, kv: { @@ -3277,9 +3399,12 @@ export class ActorContextHandleAdapter { return; } - await this.#databaseProvider.onMigrate( - (await this.ensureDatabaseClient()) as never, - ); + const client = await this.ensureDatabaseClient(); + // Creating the client awaits, and the generation can be lost meanwhile. + if (this.isLost) { + throw actorGenerationLostError(); + } + await this.#databaseProvider.onMigrate(client as never); } async closeDatabase(): Promise { @@ -3614,7 +3739,8 @@ export class ActorContextHandleAdapter { this.#runtime.actorRequestSave(this.#ctx, { immediate: false }), ); - if (!this.#onStateChange) { + // State changes are flushed asynchronously and can land after the generation was lost. + if (!this.#onStateChange || this.isLost) { return; } @@ -3670,6 +3796,11 @@ class NativeWorkflowRuntimeAdapter { return this.#ctx.actorGeneration; } + /** Whether this generation was declared lost. Workflow code checks it before each user callback. */ + get isLost(): boolean { + return this.#ctx.isLost; + } + readonly id: string; readonly driver: { kvBatchGet: ( @@ -4113,19 +4244,29 @@ export function buildNativeFactory( ? new RunHandlerCoordinator(config.run) : undefined; const getNativeWorkflowInspector = (ctx: ActorContextHandle) => { + // Inspector workflow calls run user code (history, state, replay). A lost generation + // exposes no workflow inspector. + if (callNativeSync(() => runtime.actorIsLost(ctx))) { + return undefined; + } const actorId = callNativeSync(() => runtime.actorId(ctx)); const restart = () => callNativeSync(() => runtime.actorRestartRunHandler(ctx)); const actorGeneration = callNativeSync(() => runtime.actorGeneration(ctx), ); - return (runHandlerCoordinator?.getInspector( + const inspector = (runHandlerCoordinator?.getInspector( actorId, restart, actorGeneration, )?.workflow ?? getRunInspectorConfig(config.run, actorId)?.workflow) as | NativeWorkflowInspectorConfig | undefined; + // Building the inspector can run a user factory, and the generation can be lost meanwhile. + if (callNativeSync(() => runtime.actorIsLost(ctx))) { + return undefined; + } + return inspector; }; const onStateChange = typeof config.onStateChange === "function" @@ -4245,6 +4386,11 @@ export function buildNativeFactory( } } + // Reading connection params runs the user schema, and authentication awaited before this. + const inspectorConnParams = (conn: { params: unknown }) => { + throwIfActorLost(runtime, ctx); + return conn.params; + }; const workflowHistory = () => serializeWorkflowHistoryForJson( getNativeWorkflowInspector(ctx)?.getHistory() ?? null, @@ -4283,7 +4429,7 @@ export function buildNativeFactory( id: conn.id, details: { type: null, - params: conn.params, + params: inspectorConnParams(conn), stateEnabled: true, state: conn.state, subscriptions: 0, @@ -4543,7 +4689,7 @@ export function buildNativeFactory( id: conn.id, details: { type: null, - params: conn.params, + params: inspectorConnParams(conn), stateEnabled: true, state: conn.state, subscriptions: 0, @@ -4605,15 +4751,22 @@ export function buildNativeFactory( body.properties !== undefined ? [body.properties] : normalizeArgs(body.args); + // Authentication and body parsing may finish after the generation was lost. + if (callNativeSync(() => runtime.actorIsLost(ctx))) { + return jsonResponse( + { error: "actor generation was declared lost" }, + { status: 503 }, + ); + } try { - const output = await action( - actorCtx, - ...validateActionArgs( - schemaConfig.actionInputSchemas, - actionName, - args, - ), + const validatedArgs = validateActionArgs( + schemaConfig.actionInputSchemas, + actionName, + args, ); + // Schema validation is user code and can run while the generation is lost. + throwIfActorLost(runtime, ctx); + const output = await action(actorCtx, ...validatedArgs); return jsonResponse({ output }); } catch (error) { logger().error({ @@ -4755,7 +4908,14 @@ export function buildNativeFactory( await actorCtx.closeDatabase(); } await actorCtx.runDatabaseMigrations(); - if (typeof config.onMigrate === "function") { + // Migrations may finish after the generation was lost, and a lost + // generation starts no new user hooks. + if ( + typeof config.onMigrate === "function" && + !callNativeSync(() => + runtime.actorIsLost(ctx), + ) + ) { await config.onMigrate(actorCtx, isNew); } } catch (error) { @@ -4811,10 +4971,38 @@ export function buildNativeFactory( const { ctx } = unwrapTsfnPayload(error, payload); const actorCtx = makeActorCtx(ctx); const actorId = callNativeSync(() => runtime.actorId(ctx)); + const actorGeneration = callNativeSync(() => + runtime.actorGeneration(ctx), + ); + // The engine can declare this generation lost before or during `onSleep`. + // A lost generation releases its own runtime resources only: no user hook, + // no state save, and nothing keyed by actor id alone. + const isLost = () => + callNativeSync(() => runtime.actorIsLost(ctx)); + if (isLost()) { + try { + await cleanupNativeLostRuntimeState(runtime, ctx); + } finally { + runHandlerCoordinator?.destroy( + actorId, + actorGeneration, + { + disposeInspector: false, + }, + ); + await actorCtx.dispose(); + } + return; + } // TODO: Move this save hook into cleanupNativeSleepRuntimeState // so immediate and deferred sleep cleanup share one save-state // path instead of passing a callback through cleanup. const saveActorState = async () => { + // Cleanup waits on tracked work first, and the generation can be lost + // during that wait. + if (isLost()) { + return; + } if (runtime.kind === "wasm") { // Wasm cannot use the native context save helper here because // the runtime owns the serialized state handoff. @@ -4835,17 +5023,26 @@ export function buildNativeFactory( await saveActorState(); } finally { try { - await cleanupNativeSleepRuntimeState( - runtime, - ctx, - saveActorState, - ); + if (isLost()) { + await cleanupNativeLostRuntimeState(runtime, ctx); + } else { + await cleanupNativeSleepRuntimeState( + runtime, + ctx, + saveActorState, + ); + } } finally { runHandlerCoordinator?.destroy( actorId, - callNativeSync(() => runtime.actorGeneration(ctx)), + actorGeneration, + { + disposeInspector: !isLost(), + }, ); - disposeRunInspector(config.run, actorId); + if (!isLost()) { + disposeRunInspector(config.run, actorId); + } await actorCtx.dispose(); } } @@ -4856,15 +5053,22 @@ export function buildNativeFactory( const { ctx } = unwrapTsfnPayload(error, payload); const actorCtx = makeActorCtx(ctx); const actorId = callNativeSync(() => runtime.actorId(ctx)); + // A lost generation still releases its runtime state but starts no user hook. Loss + // is rechecked before each user callback because it can happen while one runs. + const isLost = () => + callNativeSync(() => runtime.actorIsLost(ctx)); // Close run control before user cleanup so replay cannot race actor // destruction. Recreating this actor id receives a fresh controller. runHandlerCoordinator?.destroy( actorId, callNativeSync(() => runtime.actorGeneration(ctx)), + { disposeInspector: !isLost() }, ); - disposeRunInspector(config.run, actorId); + if (!isLost()) { + disposeRunInspector(config.run, actorId); + } try { - if (typeof config.onDestroy === "function") { + if (typeof config.onDestroy === "function" && !isLost()) { await config.onDestroy(actorCtx); } } finally { @@ -4902,12 +5106,16 @@ export function buildNativeFactory( : undefined, ); try { + const validatedParams = validateConnParams( + schemaConfig.connParamsSchema, + decodeValue(params), + ); + // Schema validation is user code and can run while the + // generation is lost. + throwIfActorLost(runtime, ctx); await config.onBeforeConnect( actorCtx, - validateConnParams( - schemaConfig.connParamsSchema, - decodeValue(params), - ), + validatedParams, ); } finally { await actorCtx.dispose(); @@ -4955,15 +5163,25 @@ export function buildNativeFactory( () => actorCtx.assertCanMutateState(), ); try { - const nextConnState = hasStaticConnState - ? structuredClone(config.connState) - : await config.createConnState( + let nextConnState: unknown; + if (hasStaticConnState) { + nextConnState = structuredClone( + config.connState, + ); + } else { + const validatedParams = validateConnParams( + schemaConfig.connParamsSchema, + decodeValue(params), + ); + // Schema validation is user code and can run while the + // generation is lost. + throwIfActorLost(runtime, ctx); + nextConnState = + await config.createConnState( actorCtx, - validateConnParams( - schemaConfig.connParamsSchema, - decodeValue(params), - ), + validatedParams, ); + } connAdapter.initializeState(nextConnState); return encodeValue(nextConnState); } finally { @@ -5122,6 +5340,13 @@ export function buildNativeFactory( ) => { const { ctx, name, args, output } = unwrapTsfnPayload(error, payload); + // The action may finish after its generation was lost, and a lost + // generation starts no new user callbacks. + if ( + callNativeSync(() => runtime.actorIsLost(ctx)) + ) { + return output; + } const actorCtx = makeActorCtx(ctx); try { return encodeValue( @@ -5219,6 +5444,9 @@ export function buildNativeFactory( } }; try { + // Inspector routing awaits, and the generation can be lost + // meanwhile. Schema validation is user code. + throwIfActorLost(runtime, ctx); const connParams = validateConnParams( schemaConfig.connParamsSchema, rawConnParams @@ -5257,6 +5485,14 @@ export function buildNativeFactory( abortRequest, ); } + // Connection setup awaits, and the generation can be lost meanwhile. + if ( + callNativeSync(() => + runtime.actorIsLost(ctx), + ) + ) { + throw actorGenerationLostError(); + } const response = await config.onRequest( requestCtx, handlerRequest, @@ -5357,6 +5593,11 @@ export function buildNativeFactory( const { ctx } = unwrapTsfnPayload(error, payload); const actorId = callNativeSync(() => runtime.actorId(ctx)); const executeRun = async () => { + // The coordinator can hold `run` behind other run work until after the + // generation was lost. + if (callNativeSync(() => runtime.actorIsLost(ctx))) { + return; + } const actorCtx = makeActorCtx(ctx); try { await run(actorCtx); @@ -5449,14 +5690,18 @@ export function buildNativeFactory( : makeActorCtx(ctx, undefined, cancelToken); const runAction = async () => { try { + const validatedArgs = validateActionArgs( + schemaConfig.actionInputSchemas, + name, + decodeArgs(args), + ); + // Schema validation is user code and can run while the + // generation is lost. + throwIfActorLost(runtime, ctx); return encodeValue( await handler( actorCtx, - ...validateActionArgs( - schemaConfig.actionInputSchemas, - name, - decodeArgs(args), - ), + ...validatedArgs, ...(scheduledFire ? [scheduledFire] : []), @@ -5535,6 +5780,9 @@ export function buildNativeFactory( if (canPublish && !(await canPublish(actorCtx))) { throw forbiddenError(); } + // `canPublish` can finish after the generation was lost, and + // sending validates the body with user schema code. + throwIfActorLost(runtime, ctx); const decodedBody = decodeValue(body); if (wait) { @@ -5608,7 +5856,7 @@ export function buildNativeFactory( }; return runtime.createActorFactory( - callbacks, + refuseCallbacksAfterLost(runtime, callbacks), buildActorConfig(definition, registryConfig, runtime.kind), ); } diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/run-handler-coordinator.ts b/rivetkit-typescript/packages/rivetkit/src/registry/run-handler-coordinator.ts index 1cc8d391dd..623ac88409 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/run-handler-coordinator.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/run-handler-coordinator.ts @@ -61,9 +61,15 @@ function createRunState(actorGeneration: number | undefined): ActorRunState { }; } -function closeRunState(state: ActorRunState): void { +/** + * Closes run state. The inspector disposer is user code, so it is skipped for a generation + * that was declared lost. + */ +function closeRunState(state: ActorRunState, disposeInspector = true): void { state.closed = true; - state.inspector?.dispose?.(); + if (disposeInspector) { + state.inspector?.dispose?.(); + } state.inspector = undefined; notifyStateChanged(state); } @@ -153,7 +159,11 @@ export class RunHandlerCoordinator { * Closes the run state for `actorId`. With `actorGeneration`, only state owned by that * generation is closed, so an older generation's late cleanup cannot close a newer one. */ - destroy(actorId: string, actorGeneration?: number): void { + destroy( + actorId: string, + actorGeneration?: number, + options: { disposeInspector?: boolean } = {}, + ): void { const state = this.#states.get(actorId); if (!state) return; if ( @@ -164,7 +174,7 @@ export class RunHandlerCoordinator { return; } - closeRunState(state); + closeRunState(state, options.disposeInspector ?? true); this.#states.delete(actorId); } @@ -232,8 +242,10 @@ export class RunHandlerCoordinator { return detached; } // A newer generation replaces state left by an older one, which may still be - // marked active because its JS callback cannot be cancelled. - closeRunState(state); + // marked active because its JS callback cannot be cancelled. That generation did not + // clean up after itself, which only happens when it was lost, so its user inspector + // disposer is not called. + closeRunState(state, false); this.#states.delete(actorId); state = undefined; } diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts b/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts index dfa85e8ba4..5c8ab7ed8a 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/runtime.ts @@ -584,6 +584,7 @@ export interface CoreRuntime { ): Promise; actorId(ctx: ActorContextHandle): string; actorGeneration(ctx: ActorContextHandle): number | undefined; + actorIsLost(ctx: ActorContextHandle): boolean; /** * Runs one actor callback with `ctx` as the current invocation: operations * on retained handles for the same actor resolve to it, and its Core span diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts b/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts index 9384f17d17..463d26307c 100644 --- a/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts +++ b/rivetkit-typescript/packages/rivetkit/src/registry/wasm-runtime.ts @@ -542,6 +542,10 @@ export class WasmCoreRuntime implements CoreRuntime { return callHandle(asWasmActorContext(ctx), "generation") ?? undefined; } + actorIsLost(ctx: ActorContextHandle): boolean { + return callHandle(asWasmActorContext(ctx), "isLost"); + } + runWithActorInvocationContext( _ctx: ActorContextHandle, run: () => T, diff --git a/rivetkit-typescript/packages/rivetkit/src/workflow/context.ts b/rivetkit-typescript/packages/rivetkit/src/workflow/context.ts index a4323434af..d5b55f3a9b 100644 --- a/rivetkit-typescript/packages/rivetkit/src/workflow/context.ts +++ b/rivetkit-typescript/packages/rivetkit/src/workflow/context.ts @@ -12,6 +12,7 @@ import type { WorkflowContextInterface, WorkflowQueueMessage, } from "@rivetkit/workflow-engine"; +import { EvictedError } from "@rivetkit/workflow-engine"; import type { QueueFilterName, QueueNextBatchOptions, @@ -19,6 +20,7 @@ import type { QueueResultMessageForName, RunContext, } from "@/actor/config"; +import { ACTOR_CONTEXT_INTERNAL_SYMBOL } from "@/actor/config"; import type { AnyActorDefinition, BaseActorDefinition, @@ -282,6 +284,20 @@ export type WorkflowBranchConfig< // never appears on the public surface. const DEACTIVATE_STEP = Symbol("workflow.step.deactivate"); +/** + * Stops the workflow before it starts more user code once its actor generation was declared + * lost. The engine treats the eviction like a sleep and replays from history in the next + * generation. + */ +export function throwIfGenerationLost(runCtx: unknown): void { + const actor = ( + runCtx as { [ACTOR_CONTEXT_INTERNAL_SYMBOL]?: { isLost?: boolean } } + )[ACTOR_CONTEXT_INTERNAL_SYMBOL]; + if (actor?.isLost === true) { + throw new EvictedError(); + } +} + /** * The context handed to a workflow step (`step` / `tryStep` callbacks). This is * the only scope where actor data (state, vars, db, client) and side effects @@ -994,6 +1010,7 @@ export class WorkflowContext< TQueues >(this.#runCtx, () => this.#markGuardTriggered()); + throwIfGenerationLost(this.#runCtx); try { return await run(stepCtx); } finally { @@ -1044,6 +1061,7 @@ export class WorkflowContext< TEvents, TQueues >(this.#runCtx, () => this.#markGuardTriggered()); + throwIfGenerationLost(this.#runCtx); try { await rollback(stepCtx, output); } finally { @@ -1097,6 +1115,8 @@ export class WorkflowContext< TEvents, TQueues > { + // Every try, loop, join, and race body runs through a child context. + throwIfGenerationLost(this.#runCtx); return new WorkflowContext(ctx, this.#runCtx); } } diff --git a/rivetkit-typescript/packages/rivetkit/src/workflow/driver.ts b/rivetkit-typescript/packages/rivetkit/src/workflow/driver.ts index 4d7ebd4422..7cdc9022ff 100644 --- a/rivetkit-typescript/packages/rivetkit/src/workflow/driver.ts +++ b/rivetkit-typescript/packages/rivetkit/src/workflow/driver.ts @@ -10,6 +10,7 @@ import type { RunContext } from "@/actor/config"; import type { AnyStaticActorInstance } from "@/actor/definition"; import { makeWorkflowKey, workflowStoragePrefix } from "@/actor/keys"; import type { RawAccess } from "@/common/database/config"; +import { throwIfGenerationLost } from "./context"; const WORKFLOW_STORAGE_PREFIX = workflowStoragePrefix(); // Keep workflow flushes below depot's 320-dirty-page commit ceiling. The @@ -69,6 +70,26 @@ function normalizeSqlBlob(value: unknown): Uint8Array { throw new Error("workflow sqlite value was not a blob"); } +function guardRawAccessAgainstLost( + access: RawAccess, + runCtx: RunContext | undefined, +): RawAccess { + return { + ...access, + execute: ((...args: Parameters) => { + throwIfGenerationLost(runCtx); + return access.execute(...args); + }) as RawAccess["execute"], + transaction: (callback, options) => { + throwIfGenerationLost(runCtx); + return access.transaction( + (tx) => callback(guardRawAccessAgainstLost(tx, runCtx)), + options, + ); + }, + }; +} + function runtimeDbFromContext( runCtx?: RunContext, ): RawAccess | undefined { @@ -79,7 +100,10 @@ function runtimeDbFromContext( "execute" in db && "transaction" in db ) { - return db as RawAccess; + // The client can come from a user database provider. Workflow storage calls it after + // awaits, including inside its own transactions, so every call checks that the + // generation was not lost first. + return guardRawAccessAgainstLost(db as RawAccess, runCtx); } return undefined; } diff --git a/rivetkit-typescript/packages/rivetkit/src/workflow/mod.ts b/rivetkit-typescript/packages/rivetkit/src/workflow/mod.ts index 5e024b495c..0e8a52c96f 100644 --- a/rivetkit-typescript/packages/rivetkit/src/workflow/mod.ts +++ b/rivetkit-typescript/packages/rivetkit/src/workflow/mod.ts @@ -24,7 +24,7 @@ import { isActorAbortedError, RivetError } from "@/actor/errors"; import type { EventSchemaConfig, QueueSchemaConfig } from "@/actor/schema"; import type { AnyDatabaseProvider } from "@/common/database/config"; import { stringifyError } from "@/utils"; -import { WorkflowContext } from "./context"; +import { throwIfGenerationLost, WorkflowContext } from "./context"; import { ActorWorkflowControlDriver, ActorWorkflowDriver } from "./driver"; import { createWorkflowInspectorAdapter } from "./inspector"; @@ -265,7 +265,10 @@ export function workflow< const handle = runWorkflow( actor.id, - async (ctx) => await fn(new WorkflowContext(ctx, runCtx)), + async (ctx) => { + throwIfGenerationLost(runCtx); + return await fn(new WorkflowContext(ctx, runCtx)); + }, undefined, driver, { @@ -275,7 +278,15 @@ export function workflow< logger: runCtx.log as RunWorkflowOptions["logger"], onHistoryUpdated: workflowInspector.update, onError: onError - ? async (event) => await onError(runCtx, event) + ? async (event) => { + // A lost generation starts no new user hook. + if ( + (actor as { isLost?: boolean }).isLost === true + ) { + return; + } + await onError(runCtx, event); + } : undefined, }, ); diff --git a/rivetkit-typescript/packages/rivetkit/tests/database-state-transaction.test.ts b/rivetkit-typescript/packages/rivetkit/tests/database-state-transaction.test.ts index dbe79872f9..24578be914 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/database-state-transaction.test.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/database-state-transaction.test.ts @@ -32,6 +32,7 @@ describe("experimental database state transactions", () => { const context = new ActorContextHandleAdapter( { actorId: () => "actor-a", + actorIsLost: () => false, actorRuntimeState: () => runtimeState, } as never, {} as never, diff --git a/rivetkit-typescript/packages/rivetkit/tests/driver/serverless-lost-generation.test.ts b/rivetkit-typescript/packages/rivetkit/tests/driver/serverless-lost-generation.test.ts index 7f7864356e..75fec244a8 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/driver/serverless-lost-generation.test.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/driver/serverless-lost-generation.test.ts @@ -14,17 +14,13 @@ import { createClient } from "../../src/client/mod"; import { getOrStartSharedEngine, TOKEN } from "./shared-harness"; interface HookEvent { - hook: "createVars" | "onSleep" | "onSleepDone" | "clearRuntimeState"; + hook: "createVars" | "onSleep" | "clearRuntimeState"; boot: string | undefined; } // The registry runs in this process, so hooks and runtime cleanup report straight to the test. const hookEvents: HookEvent[] = []; -// Holds one generation's `onSleep` open so the next generation starts while it is still shutting -// down, which is the overlap seen in production. -let heldSleep: { boot: string; released: Promise } | undefined; - const lostVarsActor = actor({ state: {}, createVars: () => { @@ -34,10 +30,6 @@ const lostVarsActor = actor({ }, onSleep: async (c) => { hookEvents.push({ hook: "onSleep", boot: c.vars.boot }); - if (heldSleep?.boot === c.vars.boot) { - await heldSleep.released; - } - hookEvents.push({ hook: "onSleepDone", boot: c.vars.boot }); }, actions: { readBoot: (c) => c.vars.boot, @@ -247,7 +239,7 @@ async function withTimeout( } describe("Serverless Lost generation", () => { - test("a Lost generation's cleanup does not clear the next generation's state", async () => { + test("a Lost generation stops without user hooks and leaves the next generation's state intact", async () => { const engine = await getOrStartSharedEngine(); const namespace = `serverless-lost-${crypto.randomUUID()}`; const poolName = `serverless-lost-${crypto.randomUUID()}`; @@ -307,24 +299,25 @@ describe("Serverless Lost generation", () => { handle.readBoot(), ); - let releaseOldSleep = () => {}; - heldSleep = { - boot: firstBoot, - released: new Promise((resolve) => { - releaseOldSleep = resolve; - }), - }; - // Reset the engine's `/start` stream for the running generation. The engine treats - // that as the envoy being lost, sends `StopActor { reason: Lost }` to the still - // connected envoy, and the old generation starts its graceful sleep. + // that as the envoy being lost and sends `StopActor { reason: Lost }` to the still + // connected envoy, then starts the next generation on it when the actor is woken. expect(proxy.resetEngineConnections()).toBeGreaterThan(0); - await waitFor("old generation onSleep", 30_000, async () => - hasEvent("onSleep", firstBoot) ? true : undefined, + + // The lost generation releases its own runtime state without running user hooks. + const oldGenerationClear = await waitFor( + "lost generation runtime cleanup", + 30_000, + async () => + hookEvents.find( + (event) => + event.hook === "clearRuntimeState" && + event.boot === firstBoot, + ), ); + expect(oldGenerationClear.boot).toBe(firstBoot); + expect(hasEvent("onSleep", firstBoot)).toBe(false); - // Waking the actor starts the next generation on the same envoy while the old one is - // still inside `onSleep`. const secondBoot = await waitFor( "second generation", 30_000, @@ -333,22 +326,15 @@ describe("Serverless Lost generation", () => { return boot === firstBoot ? undefined : boot; }, ); - expect(hasEvent("onSleepDone", firstBoot)).toBe(false); - // Let the old generation finish and wait for the runtime-state clear its cleanup runs. - const eventsBeforeRelease = hookEvents.length; - releaseOldSleep(); - const oldGenerationClear = await waitFor( - "old generation runtime cleanup", - 30_000, - async () => - hookEvents - .slice(eventsBeforeRelease) - .find((event) => event.hook === "clearRuntimeState"), - ); - - // The old generation's cleanup must clear its own state, not the new generation's. - expect(oldGenerationClear.boot).toBe(firstBoot); + // The lost generation's cleanup must not have cleared the new generation's state. + expect( + hookEvents.some( + (event) => + event.hook === "clearRuntimeState" && + event.boot === secondBoot, + ), + ).toBe(false); await expect(withTimeout(handle.hasVars(), 10_000)).resolves.toBe( true, ); @@ -357,7 +343,6 @@ describe("Serverless Lost generation", () => { ); } finally { stopObservingClears(); - heldSleep = undefined; await proxy.close(); await new Promise((resolve) => { if (runner) { diff --git a/rivetkit-typescript/packages/rivetkit/tests/runtime-parity.test.ts b/rivetkit-typescript/packages/rivetkit/tests/runtime-parity.test.ts index afeef468f5..d905304d1d 100644 --- a/rivetkit-typescript/packages/rivetkit/tests/runtime-parity.test.ts +++ b/rivetkit-typescript/packages/rivetkit/tests/runtime-parity.test.ts @@ -200,6 +200,10 @@ class FakeActorContext { return null; } + isLost(): boolean { + return false; + } + restartRunHandler(): void { this.scenario.runRestarts += 1; }