From 57602c595bd262ea1a512ecfd01a476869311a84 Mon Sep 17 00:00:00 2001 From: ABCxFF <79597906+abcxff@users.noreply.github.com> Date: Mon, 5 Oct 2026 20:29:07 -0400 Subject: [PATCH] fix(envoy-client): keep the ping silence deadline bounded across reconnects --- docs-internal/engine/sleep-sequence.md | 2 +- engine/sdks/rust/envoy-client/src/actor.rs | 1 + .../rust/envoy-client/src/connection/mod.rs | 37 +++-- .../envoy-client/src/connection/native.rs | 11 +- .../rust/envoy-client/src/connection/wasm.rs | 12 +- engine/sdks/rust/envoy-client/src/context.rs | 72 ++++++++ engine/sdks/rust/envoy-client/src/envoy.rs | 111 +++++++++---- engine/sdks/rust/envoy-client/src/events.rs | 1 + engine/sdks/rust/envoy-client/src/sqlite.rs | 1 + .../rust/envoy-client/tests/command_dedup.rs | 92 ++++++++-- .../envoy-client/tests/engine_liveness.rs | 77 +++++++++ .../tests/ping_silence_reconnect.rs | 157 +++++++++++++++++- .../tests/sqlite_vfs_disconnect_stall.rs | 1 + .../envoy-client/tests/wait_actors_drained.rs | 1 + .../tests/ws_tx_lock_contention.rs | 1 + .../packages/rivetkit-core/src/testing.rs | 1 + .../packages/rivetkit-core/tests/context.rs | 3 + .../tests/migrate_kv_to_sqlite.rs | 1 + .../rivetkit-core/tests/serverless.rs | 1 + .../packages/rivetkit-core/tests/sleep.rs | 1 + .../packages/rivetkit-core/tests/sqlite.rs | 1 + .../packages/rivetkit-core/tests/task.rs | 1 + .../packages/rivetkit/tests/client.rs | 1 + 23 files changed, 524 insertions(+), 63 deletions(-) create mode 100644 engine/sdks/rust/envoy-client/tests/engine_liveness.rs diff --git a/docs-internal/engine/sleep-sequence.md b/docs-internal/engine/sleep-sequence.md index 75d557c5d8..59522b45b9 100644 --- a/docs-internal/engine/sleep-sequence.md +++ b/docs-internal/engine/sleep-sequence.md @@ -57,7 +57,7 @@ When the grace deadline elapses before `can_finalize_sleep()` returns true: - 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. -- envoy-client declares all its actors lost once no engine ping has arrived for `envoy_lost_threshold` minus a 3s margin. The engine refreshes its liveness timestamp just before each ping and reallocates actors once that timestamp is older than the threshold, so the envoy normally stops first, including on half-open connections that never report a close. This is best effort: a ping delayed in transit by more than the margin can still let the engine give up first. +- envoy-client declares all its actors lost once no engine ping has arrived for `envoy_lost_threshold` minus a 3s margin. The engine refreshes its liveness timestamp just before each ping and reallocates actors once that timestamp is older than the threshold, so the envoy normally stops first, including on half-open connections that never report a close. This is best effort: a ping delayed in transit by more than the margin can still let the engine give up first. The deadline counts from `EngineLiveness::baseline` on a monotonic clock. After a reconnect the engine refreshes its liveness timestamp only when it claims the connection, after sending `ToEnvoyInit`, so the baseline keeps the last proven refresh (a ping, or the install time of a connection claimed by commands) until a ping or WebSocket commands arrive on the new connection, then counts from when the client started opening that connection. Serverless start payloads injected over HTTP do not count as a claim. A generation created after the baseline counts from its creation instead, since the engine only assigns it over a claimed connection, and only generations past their own deadline are stopped. ## Guarding lifecycle requests diff --git a/engine/sdks/rust/envoy-client/src/actor.rs b/engine/sdks/rust/envoy-client/src/actor.rs index 097d1ae0ca..cbdc1c36cf 100644 --- a/engine/sdks/rust/envoy-client/src/actor.rs +++ b/engine/sdks/rust/envoy-client/src/actor.rs @@ -1695,6 +1695,7 @@ mod tests { protocol_metadata: Arc::new(tokio::sync::Mutex::new(None)), shutting_down: std::sync::atomic::AtomicBool::new(false), last_ping_ts: std::sync::atomic::AtomicI64::new(0), + engine_liveness: Default::default(), stopped_tx: tokio::sync::watch::channel(true).0, }); (shared, envoy_rx) diff --git a/engine/sdks/rust/envoy-client/src/connection/mod.rs b/engine/sdks/rust/envoy-client/src/connection/mod.rs index f5579882ce..37498e34ba 100644 --- a/engine/sdks/rust/envoy-client/src/connection/mod.rs +++ b/engine/sdks/rust/envoy-client/src/connection/mod.rs @@ -49,7 +49,14 @@ pub(crate) async fn install_connection( let _ = message.written.send(result); } }); - install_connection_with_http(shared, tx, http_tx, http_byte_budget).await + install_connection_with_http( + shared, + tx, + http_tx, + http_byte_budget, + crate::time::Instant::now(), + ) + .await } pub(crate) async fn install_connection_with_http( @@ -57,6 +64,9 @@ pub(crate) async fn install_connection_with_http( tx: tokio::sync::mpsc::UnboundedSender, http_tx: tokio::sync::mpsc::Sender, http_byte_budget: std::sync::Arc, + // When the client started opening this connection. The engine can only claim the connection + // after that, so it is a safe lower bound for the claim even when the install itself is late. + connect_started_at: crate::time::Instant, ) -> u64 { let mut guard = shared.ws_tx.lock().await; let mut http_guard = shared.http_ws_tx.lock().await; @@ -64,9 +74,9 @@ pub(crate) async fn install_connection_with_http( .next_connection_session .fetch_add(1, Ordering::AcqRel) .saturating_add(1); - // A ping from an earlier connection says nothing about this one. Clearing it keeps the engine - // ping silence check disabled until this connection receives its first ping. - shared.last_ping_ts.store(0, Ordering::Release); + shared + .engine_liveness + .connection_installed(connect_started_at); shared.connection_session.store(session, Ordering::Release); *guard = Some(tx); *http_guard = Some(HttpConnectionTx { @@ -240,18 +250,20 @@ async fn forward_to_envoy(shared: &SharedContext, session: u64, message: protoco match message { protocol::ToEnvoy::ToEnvoyPing(ping) => { - let previous_ping_ts = shared + shared .last_ping_ts - .swap(crate::time::now_millis(), Ordering::AcqRel); + .store(crate::time::now_millis(), Ordering::Release); + let newly_claimed = shared + .engine_liveness + .ping_received(crate::time::Instant::now()); ws_send( shared, protocol::ToRivet::ToRivetPong(protocol::ToRivetPong { ts: ping.ts }), ) .await; - // The envoy loop does not arm the engine ping silence check while no ping is - // recorded. Wake it on the first ping so the check arms without waiting for an - // unrelated event. - if previous_ping_ts == 0 { + // The first ping on a connection moves the engine ping silence baseline. The envoy + // loop computes the silence wait only when a turn starts, so wake it. + if newly_claimed { let _ = crate::envoy::send_to_envoy_tx( shared, ToEnvoyMessage::ConnMessage { @@ -262,6 +274,11 @@ async fn forward_to_envoy(shared: &SharedContext, session: u64, message: protoco } } other => { + if matches!(other, protocol::ToEnvoy::ToEnvoyCommands(_)) { + // Recorded before forwarding so the envoy loop sees the claim when it handles the + // commands. + shared.engine_liveness.commands_received(); + } let _ = crate::envoy::send_to_envoy_tx( shared, ToEnvoyMessage::ConnMessage { diff --git a/engine/sdks/rust/envoy-client/src/connection/native.rs b/engine/sdks/rust/envoy-client/src/connection/native.rs index 111dd52fde..719a1665b4 100644 --- a/engine/sdks/rust/envoy-client/src/connection/native.rs +++ b/engine/sdks/rust/envoy-client/src/connection/native.rs @@ -125,6 +125,7 @@ async fn single_connection( .body(()) .map_err(|e| anyhow::anyhow!("failed to build ws request: {e}"))?; + let connect_started_at = crate::time::Instant::now(); let (ws_stream, _) = tokio_tungstenite::connect_async_with_config(request, Some(websocket_config()), false) .await?; @@ -136,8 +137,14 @@ async fn single_connection( let (http_ws_tx, mut http_ws_rx) = mpsc::channel::(super::HTTP_WS_MESSAGE_CAPACITY); let http_byte_budget = Arc::new(Semaphore::new(super::HTTP_WS_BYTE_CAPACITY)); - let session = - super::install_connection_with_http(shared, ws_tx, http_ws_tx, http_byte_budget).await; + let session = super::install_connection_with_http( + shared, + ws_tx, + http_ws_tx, + http_byte_budget, + connect_started_at, + ) + .await; tracing::info!( endpoint = %shared.config.endpoint, diff --git a/engine/sdks/rust/envoy-client/src/connection/wasm.rs b/engine/sdks/rust/envoy-client/src/connection/wasm.rs index 24d4917a71..a5aef7a81e 100644 --- a/engine/sdks/rust/envoy-client/src/connection/wasm.rs +++ b/engine/sdks/rust/envoy-client/src/connection/wasm.rs @@ -104,6 +104,7 @@ mod imp { ) -> anyhow::Result<(u64, Option)> { let url = super::super::ws_url(shared); let protocols = protocols(&shared.config.token); + let connect_started_at = crate::time::Instant::now(); let ws = WebSocket::new_with_str_sequence(&url, protocols.as_ref()) .map_err(|error| anyhow::anyhow!("failed to create websocket: {}", js_error(error)))?; ws.set_binary_type(BinaryType::Arraybuffer); @@ -177,9 +178,14 @@ mod imp { let (http_ws_tx, mut http_ws_rx) = mpsc::channel::(super::super::HTTP_WS_MESSAGE_CAPACITY); let http_byte_budget = Arc::new(Semaphore::new(super::super::HTTP_WS_BYTE_CAPACITY)); - let session = - super::super::install_connection_with_http(shared, ws_tx, http_ws_tx, http_byte_budget) - .await; + let session = super::super::install_connection_with_http( + shared, + ws_tx, + http_ws_tx, + http_byte_budget, + connect_started_at, + ) + .await; tracing::info!( endpoint = %shared.config.endpoint, diff --git a/engine/sdks/rust/envoy-client/src/context.rs b/engine/sdks/rust/envoy-client/src/context.rs index f37c9d8afb..d4e0cd9a3c 100644 --- a/engine/sdks/rust/envoy-client/src/context.rs +++ b/engine/sdks/rust/envoy-client/src/context.rs @@ -4,6 +4,7 @@ use std::sync::Mutex as StdMutex; use std::sync::atomic::{AtomicBool, AtomicI64, AtomicU64}; use crate::async_counter::AsyncCounter; +use crate::time::Instant; use rivet_envoy_protocol as protocol; use tokio::sync::{Mutex, Notify, Semaphore, mpsc, oneshot, watch}; @@ -43,6 +44,9 @@ pub struct SharedContext { /// `EnvoyHandle::is_ping_healthy` to surface a dead engine link to upstream health checks. /// Zero means no ping has been received yet. pub last_ping_ts: AtomicI64, + /// What this envoy can rely on about the engine's liveness clock. Drives the engine ping + /// silence check. Only the connection task updates it. + pub engine_liveness: EngineLiveness, // Latched signal fired by `envoy_loop` after its cleanup block completes. // Waiters observing `true` are guaranteed that the loop has exited and // every pending KV/SQLite request has been resolved (with `EnvoyShutdownError` @@ -68,3 +72,71 @@ pub struct HttpWsTxMessage { pub _byte_permit: tokio::sync::OwnedSemaphorePermit, pub written: oneshot::Sender>, } + +/// Tracks the earliest time the engine could have last refreshed this envoy's liveness timestamp. +/// +/// The engine declares the envoy's actors lost once that timestamp is older than +/// `envoy_lost_threshold`. It refreshes the timestamp when it claims a connection, which happens +/// after it sends `ToEnvoyInit`, and before each ping. The engine sends commands and pings on a +/// connection only after claiming it, so either one proves the claim. Until then the engine may +/// still be counting from a ping on an earlier connection, for example when the claim fails. +/// Times are monotonic so wall clock jumps cannot trigger or delay the check. +#[derive(Default)] +pub struct EngineLiveness(parking_lot::Mutex); + +#[derive(Default, Clone, Copy)] +struct EngineLivenessState { + /// Latest time the engine is known to have refreshed its liveness timestamp, on this or an + /// earlier connection: a ping, or the install time of a connection the engine claimed. + last_refresh: Option, + /// When the client started opening the current connection. + connection_installed_at: Option, + /// Whether the engine has claimed the current connection. + connection_claimed: bool, +} + +impl EngineLiveness { + /// A new connection was installed. `connect_started_at` is when the client started opening it, + /// which precedes any claim even when the engine claims it before the client finishes the + /// handshake. The engine has not proven a claim yet, so the last refresh is kept until it does. + pub fn connection_installed(&self, connect_started_at: Instant) { + let mut state = self.0.lock(); + state.connection_installed_at = Some(connect_started_at); + state.connection_claimed = false; + } + + /// A ping arrived on the current connection. Returns true when this is the first evidence that + /// the engine claimed the connection. + pub fn ping_received(&self, now: Instant) -> bool { + let mut state = self.0.lock(); + state.last_refresh = Some(now); + let newly_claimed = !state.connection_claimed; + state.connection_claimed = true; + newly_claimed + } + + /// Commands arrived on the current connection's WebSocket. Returns true when this is the first + /// evidence that the engine claimed the connection. Serverless start payloads injected over + /// HTTP do not count. + pub fn commands_received(&self) -> bool { + let mut state = self.0.lock(); + let newly_claimed = !state.connection_claimed; + state.connection_claimed = true; + // The claim refreshed the engine's timestamp after the install, so the install time is a + // safe lower bound for it. + state.last_refresh = match (state.last_refresh, state.connection_installed_at) { + (Some(refresh), Some(installed)) => Some(refresh.max(installed)), + (refresh, installed) => refresh.or(installed), + }; + newly_claimed + } + + /// The time the engine ping silence deadline counts from. Until the engine claims a + /// connection it may still be counting from the last refresh, so a reconnect never moves the + /// baseline forward on its own. An envoy the engine never refreshed counts from the install, + /// which is earlier than any timestamp the engine could hold for it. + pub fn baseline(&self) -> Option { + let state = *self.0.lock(); + state.last_refresh.or(state.connection_installed_at) + } +} diff --git a/engine/sdks/rust/envoy-client/src/envoy.rs b/engine/sdks/rust/envoy-client/src/envoy.rs index a18dded331..e0521323d5 100644 --- a/engine/sdks/rust/envoy-client/src/envoy.rs +++ b/engine/sdks/rust/envoy-client/src/envoy.rs @@ -104,6 +104,8 @@ pub struct ActorEntry { pub name: String, pub last_command_idx: i64, pub received_stop: bool, + /// When this generation was added. The engine ping silence deadline never counts from earlier. + pub created_at: crate::time::Instant, } #[derive(Default)] @@ -280,6 +282,7 @@ impl EnvoyContext { name, last_command_idx, received_stop: false, + created_at: crate::time::Instant::now(), }, ); self.shared @@ -440,6 +443,7 @@ fn start_envoy_sync_inner(config: EnvoyConfig) -> EnvoyHandle { protocol_metadata: Arc::new(tokio::sync::Mutex::new(None)), shutting_down: std::sync::atomic::AtomicBool::new(false), last_ping_ts: std::sync::atomic::AtomicI64::new(0), + engine_liveness: Default::default(), stopped_tx, }); @@ -619,7 +623,7 @@ async fn envoy_loop( } } => { branch = "lost_timeout"; - declare_actors_lost(&mut ctx, "stopping all actors due to envoy lost threshold"); + declare_actors_lost(&mut ctx, "envoy_lost_threshold"); lost_timeout = None; } _ = async { @@ -629,12 +633,9 @@ async fn envoy_loop( } }, if !ctx.actors.is_empty() => { branch = "engine_ping_silence"; - // A ping may have arrived while this branch slept. - if engine_ping_silence_expired(&ctx, crate::time::now_millis()) { - declare_actors_lost( - &mut ctx, - "stopping all actors because the engine stopped pinging and is about to declare them lost", - ); + // A ping may have arrived while this branch slept, so this rechecks each deadline. + declare_silent_actors_lost(&mut ctx, crate::time::Instant::now()); + if ctx.actors.is_empty() { lost_timeout = None; } } @@ -772,7 +773,7 @@ async fn handle_conn_message( } protocol::ToEnvoy::ToEnvoyPing(_) => { // The connection task records pings. It forwards only the first ping of a connection, - // so this loop takes a turn and arms the engine ping silence check. + // so this loop takes a turn and recomputes the engine ping silence wait. } } @@ -801,36 +802,90 @@ fn envoy_lost_threshold_ms(ctx: &EnvoyContext) -> i64 { .unwrap_or(10_000) } -/// When the envoy stops trusting its actors after the last engine ping. The engine refreshes its -/// envoy liveness timestamp right before sending each ping and declares the envoy's actors lost -/// once that timestamp is older than `envoy_lost_threshold`, then may start them elsewhere. The -/// envoy receives each ping after that refresh, so giving up a margin before the threshold normally -/// stops the actors before the engine can reallocate them. A ping delayed in transit by more than -/// the margin can still let the engine give up first. This also covers half-open connections that -/// never report a close. -pub fn engine_ping_silence_deadline_ms(ctx: &EnvoyContext) -> Option { - let last_ping_ts = ctx.shared.last_ping_ts.load(Ordering::Acquire); - if last_ping_ts == 0 { - return None; - } +/// When the envoy stops trusting a generation created at `created_at`. The engine declares the +/// envoy's actors lost once its liveness timestamp is older than `envoy_lost_threshold`, then may +/// start them elsewhere. The deadline counts from `EngineLiveness::baseline`, the latest time the +/// envoy knows the engine refreshed that timestamp, and gives up a margin before the threshold, so +/// the envoy normally stops its actors first. This also covers half-open connections that never +/// report a close. A ping delayed in transit by more than the margin can still let the engine give +/// up first. +/// +/// A generation created after the baseline counts from its creation instead. The engine only +/// assigns it over a connection it has claimed, which refreshed its timestamp no earlier than the +/// creation, so judging it against a stale baseline would stop a freshly started actor at once. +pub fn engine_ping_silence_deadline( + ctx: &EnvoyContext, + created_at: crate::time::Instant, +) -> Option { + let baseline = ctx.shared.engine_liveness.baseline()?; let threshold = envoy_lost_threshold_ms(ctx); let margin = LOST_THRESHOLD_SAFETY_MARGIN_MS.min(threshold / 2); - Some(last_ping_ts + threshold - margin) + Some( + baseline.max(created_at) + + std::time::Duration::from_millis((threshold - margin).max(0) as u64), + ) } -pub fn engine_ping_silence_expired(ctx: &EnvoyContext, now_ms: i64) -> bool { - engine_ping_silence_deadline_ms(ctx).is_some_and(|deadline| now_ms >= deadline) +/// The earliest engine ping silence deadline among this envoy's actors. +fn next_engine_ping_silence_deadline(ctx: &EnvoyContext) -> Option { + ctx.actors + .values() + .flat_map(|generations| generations.values()) + .filter_map(|entry| engine_ping_silence_deadline(ctx, entry.created_at)) + .min() } fn engine_ping_silence_wait(ctx: &EnvoyContext) -> Option { - let deadline = engine_ping_silence_deadline_ms(ctx)?; - let remaining = (deadline - crate::time::now_millis()).max(0); - Some(std::time::Duration::from_millis(remaining as u64)) + let deadline = next_engine_ping_silence_deadline(ctx)?; + Some(deadline.saturating_duration_since(crate::time::Instant::now())) +} + +/// Stops the generations whose engine ping silence deadline has passed. When every generation +/// expired, this is the same as `declare_actors_lost`. +pub fn declare_silent_actors_lost(ctx: &mut EnvoyContext, now: crate::time::Instant) { + let mut total = 0; + let mut expired = Vec::new(); + for (actor_id, generations) in &ctx.actors { + for (generation, entry) in generations { + total += 1; + if engine_ping_silence_deadline(ctx, entry.created_at) + .is_some_and(|deadline| now >= deadline) + { + expired.push((actor_id.clone(), *generation)); + } + } + } + if expired.is_empty() { + return; + } + if expired.len() == total { + declare_actors_lost(ctx, "engine_ping_silence"); + return; + } + + tracing::warn!( + actor_count = expired.len(), + reason = "engine_ping_silence", + "declaring envoy actors lost" + ); + for (actor_id, generation) in expired { + if let Some(entry) = ctx + .actors + .get(&actor_id) + .and_then(|generations| generations.get(&generation)) + { + entry.lost.cancel(); + if !entry.handle.is_closed() { + let _ = entry.handle.send(ToActor::Lost); + } + } + ctx.remove_actor(&actor_id, generation); + } } /// Stops every actor on this envoy because the engine has given up on them. The lost token is /// cancelled first so each generation stops writing storage before its task sees the message. -pub fn declare_actors_lost(ctx: &mut EnvoyContext, message: &'static str) { +pub fn declare_actors_lost(ctx: &mut EnvoyContext, reason: &'static str) { for (_id, request) in ctx.kv_requests.drain() { METRICS.kv_requests_inflight.dec(); let _ = request @@ -845,7 +900,7 @@ pub fn declare_actors_lost(ctx: &mut EnvoyContext, message: &'static str) { } tracing::warn!( actor_count = ctx.actors.len(), - reason = message, + reason, "declaring envoy actors lost" ); for (_actor_id, gens) in &ctx.actors { diff --git a/engine/sdks/rust/envoy-client/src/events.rs b/engine/sdks/rust/envoy-client/src/events.rs index b2dfa3c39f..f48c2db351 100644 --- a/engine/sdks/rust/envoy-client/src/events.rs +++ b/engine/sdks/rust/envoy-client/src/events.rs @@ -228,6 +228,7 @@ mod tests { protocol_metadata: Arc::new(tokio::sync::Mutex::new(None)), shutting_down: std::sync::atomic::AtomicBool::new(false), last_ping_ts: std::sync::atomic::AtomicI64::new(0), + engine_liveness: Default::default(), stopped_tx: tokio::sync::watch::channel(true).0, }); let handle = EnvoyHandle { diff --git a/engine/sdks/rust/envoy-client/src/sqlite.rs b/engine/sdks/rust/envoy-client/src/sqlite.rs index 4afc1203f4..f3692c213b 100644 --- a/engine/sdks/rust/envoy-client/src/sqlite.rs +++ b/engine/sdks/rust/envoy-client/src/sqlite.rs @@ -683,6 +683,7 @@ mod tests { protocol_metadata: Arc::new(tokio::sync::Mutex::new(None)), shutting_down: std::sync::atomic::AtomicBool::new(false), last_ping_ts: std::sync::atomic::AtomicI64::new(0), + engine_liveness: Default::default(), stopped_tx: tokio::sync::watch::channel(true).0, }); diff --git a/engine/sdks/rust/envoy-client/tests/command_dedup.rs b/engine/sdks/rust/envoy-client/tests/command_dedup.rs index 073ff983e2..5d1ea062dd 100644 --- a/engine/sdks/rust/envoy-client/tests/command_dedup.rs +++ b/engine/sdks/rust/envoy-client/tests/command_dedup.rs @@ -10,7 +10,7 @@ use rivet_envoy_client::config::{ }; use rivet_envoy_client::context::{SharedContext, WsTxMessage}; use rivet_envoy_client::envoy::{ - EnvoyContext, declare_actors_lost, engine_ping_silence_deadline_ms, engine_ping_silence_expired, + EnvoyContext, declare_actors_lost, declare_silent_actors_lost, engine_ping_silence_deadline, }; use rivet_envoy_client::handle::EnvoyHandle; use rivet_envoy_client::sqlite::{ @@ -107,6 +107,7 @@ fn new_envoy_context() -> EnvoyContext { protocol_metadata: Arc::new(tokio::sync::Mutex::new(None)), shutting_down: std::sync::atomic::AtomicBool::new(false), last_ping_ts: std::sync::atomic::AtomicI64::new(0), + engine_liveness: Default::default(), stopped_tx: tokio::sync::watch::channel(true).0, }); EnvoyContext { @@ -570,16 +571,16 @@ async fn queued_kv_write_is_dropped_once_its_generation_is_lost() { assert!(format!("{error:#}").contains("declared lost")); } -/// The engine declares actors lost `envoy_lost_threshold` after its last ping refresh. The envoy -/// gives up a safety margin earlier, measured from the last ping it received, so its actors stop -/// before the engine can start them elsewhere. +/// The engine declares actors lost `envoy_lost_threshold` after its last liveness refresh. The +/// envoy gives up a safety margin earlier, counting from the latest time it knows the engine +/// refreshed, so its actors stop before the engine can start them elsewhere. #[tokio::test] async fn engine_ping_silence_expires_before_the_engine_lost_threshold() { let ctx = new_envoy_context(); - assert_eq!( - engine_ping_silence_deadline_ms(&ctx), - None, - "no deadline before the first engine ping" + let created_at = std::time::Instant::now(); + assert!( + engine_ping_silence_deadline(&ctx, created_at).is_none(), + "no deadline before the first connection" ); *ctx.shared.protocol_metadata.lock().await = Some(protocol::ProtocolMetadata { @@ -587,22 +588,79 @@ async fn engine_ping_silence_expires_before_the_engine_lost_threshold() { actor_stop_threshold: 1_800_000, max_response_payload_size: 20_971_520, }); - let last_ping_ts = 1_700_000_000_000; - ctx.shared - .last_ping_ts - .store(last_ping_ts, std::sync::atomic::Ordering::Release); + ctx.shared.engine_liveness.connection_installed(created_at); + let pinged_at = created_at + std::time::Duration::from_secs(1); + ctx.shared.engine_liveness.ping_received(pinged_at); - let deadline = engine_ping_silence_deadline_ms(&ctx).expect("deadline after a ping"); + let deadline = engine_ping_silence_deadline(&ctx, created_at).expect("deadline after a ping"); assert!( - deadline < last_ping_ts + 15_000, + deadline < pinged_at + std::time::Duration::from_secs(15), "the envoy must give up before the engine's threshold" ); assert!( - deadline > last_ping_ts + 3_000, + deadline > pinged_at + std::time::Duration::from_secs(3), "the envoy must tolerate several missed ping intervals" ); - assert!(!engine_ping_silence_expired(&ctx, deadline - 1)); - assert!(engine_ping_silence_expired(&ctx, deadline)); + + // A generation created after the baseline counts from its creation instead. + let started_late = pinged_at + std::time::Duration::from_secs(20); + assert_eq!( + engine_ping_silence_deadline(&ctx, started_late), + Some(started_late + (deadline - pinged_at)), + ); +} + +/// Only generations whose own deadline passed are stopped. An actor started after a stale +/// baseline, for example by a serverless start while the WebSocket is down, keeps running. +#[tokio::test] +async fn ping_silence_stops_only_generations_past_their_own_deadline() { + let mut ctx = new_envoy_context(); + let (actor_tx, mut actor_rx) = mpsc::unbounded_channel::(); + let old_lost = tokio_util::sync::CancellationToken::new(); + let new_lost = tokio_util::sync::CancellationToken::new(); + let installed_at = std::time::Instant::now(); + ctx.shared + .engine_liveness + .connection_installed(installed_at); + ctx.shared.engine_liveness.ping_received(installed_at); + for (actor_id, lost) in [ + ("actor-old", old_lost.clone()), + ("actor-new", new_lost.clone()), + ] { + ctx.insert_actor_with_lost_signal( + actor_id.to_string(), + 1, + actor_tx.clone(), + lost, + Arc::new(AsyncCounter::new()), + actor_id.to_string(), + -1, + ); + } + let old_deadline = engine_ping_silence_deadline(&ctx, installed_at).expect("deadline"); + // Treat the old actor as created at the baseline and the new one as created after it. + ctx.actors + .get_mut("actor-old") + .and_then(|generations| generations.get_mut(&1)) + .expect("old actor") + .created_at = installed_at; + ctx.actors + .get_mut("actor-new") + .and_then(|generations| generations.get_mut(&1)) + .expect("new actor") + .created_at = old_deadline; + + declare_silent_actors_lost(&mut ctx, old_deadline); + + assert!(old_lost.is_cancelled(), "the expired generation is stopped"); + assert!( + !new_lost.is_cancelled(), + "the newer generation keeps running" + ); + assert!(ctx.get_actor("actor-old", Some(1)).is_none()); + assert!(ctx.get_actor("actor-new", Some(1)).is_some()); + assert!(matches!(actor_rx.try_recv(), Ok(ToActor::Lost))); + assert!(actor_rx.try_recv().is_err()); } #[tokio::test] diff --git a/engine/sdks/rust/envoy-client/tests/engine_liveness.rs b/engine/sdks/rust/envoy-client/tests/engine_liveness.rs new file mode 100644 index 0000000000..745be55c3a --- /dev/null +++ b/engine/sdks/rust/envoy-client/tests/engine_liveness.rs @@ -0,0 +1,77 @@ +//! The engine ping silence baseline must never be later than the engine's liveness timestamp, +//! must not count a ping from an earlier connection against actors started on a claimed +//! connection, must stay bounded before the first ping of a connection, and must not move forward +//! on a reconnect the engine never claims. + +use std::time::{Duration, Instant}; + +use rivet_envoy_client::context::EngineLiveness; + +fn at(base: Instant, ms: u64) -> Instant { + base + Duration::from_millis(ms) +} + +#[test] +fn no_connection_has_no_baseline() { + assert_eq!(EngineLiveness::default().baseline(), None); +} + +#[test] +fn fresh_connection_counts_from_install_before_the_first_ping() { + let t = Instant::now(); + let liveness = EngineLiveness::default(); + liveness.connection_installed(at(t, 100)); + assert_eq!(liveness.baseline(), Some(at(t, 100))); +} + +#[test] +fn reconnect_claimed_by_commands_ignores_the_previous_connection_ping() { + // The reported sequence: ping, disconnect, reconnect, missed start commands before the + // first ping on the new connection. + let t = Instant::now(); + let liveness = EngineLiveness::default(); + liveness.connection_installed(at(t, 0)); + liveness.ping_received(at(t, 1_000)); + liveness.connection_installed(at(t, 8_000)); + assert!(liveness.commands_received(), "commands prove the claim"); + assert_eq!(liveness.baseline(), Some(at(t, 8_000))); +} + +#[test] +fn unclaimed_reconnect_keeps_the_previous_ping() { + // The engine refreshes its timestamp only when it claims the connection. Until then, and + // for good if the claim fails, it is still counting from the previous ping. + let t = Instant::now(); + let liveness = EngineLiveness::default(); + liveness.connection_installed(at(t, 0)); + liveness.ping_received(at(t, 1_000)); + liveness.connection_installed(at(t, 8_000)); + assert_eq!(liveness.baseline(), Some(at(t, 1_000))); +} + +#[test] +fn first_ping_on_a_connection_claims_it() { + let t = Instant::now(); + let liveness = EngineLiveness::default(); + liveness.connection_installed(at(t, 0)); + assert!(liveness.ping_received(at(t, 500))); + assert!( + !liveness.ping_received(at(t, 3_500)), + "only the first ping claims" + ); + assert_eq!(liveness.baseline(), Some(at(t, 3_500))); +} + +#[test] +fn unclaimed_reconnect_keeps_a_claim_proven_only_by_commands() { + // The engine refreshed its timestamp when it claimed the first connection, which only + // commands proved. A reconnect it never claims must keep counting from that claim. + let t = Instant::now(); + let liveness = EngineLiveness::default(); + liveness.connection_installed(at(t, 0)); + assert!(liveness.commands_received()); + liveness.connection_installed(at(t, 8_000)); + assert_eq!(liveness.baseline(), Some(at(t, 0))); + assert!(liveness.commands_received(), "commands prove the new claim"); + assert_eq!(liveness.baseline(), Some(at(t, 8_000))); +} diff --git a/engine/sdks/rust/envoy-client/tests/ping_silence_reconnect.rs b/engine/sdks/rust/envoy-client/tests/ping_silence_reconnect.rs index 1648e688e0..25b2566fad 100644 --- a/engine/sdks/rust/envoy-client/tests/ping_silence_reconnect.rs +++ b/engine/sdks/rust/envoy-client/tests/ping_silence_reconnect.rs @@ -116,6 +116,17 @@ async fn recv(ws: &mut WebSocketStream) -> protocol::ToRivet { } } +/// Waits for the actor to report stopped. Unacknowledged events, such as an earlier running +/// state, are replayed on a new connection and skipped here. +async fn recv_actor_stopped(ws: &mut WebSocketStream) -> protocol::ActorState { + loop { + let state = recv_actor_state(ws).await; + if matches!(state, protocol::ActorState::ActorStateStopped(_)) { + return state; + } + } +} + async fn recv_actor_state(ws: &mut WebSocketStream) -> protocol::ActorState { loop { if let protocol::ToRivet::ToRivetEvents(events) = recv(ws).await { @@ -158,8 +169,7 @@ fn start_actor() -> protocol::ToEnvoy { }]) } -#[tokio::test] -async fn ping_from_previous_connection_does_not_stop_actor_started_after_reconnect() { +async fn start_envoy() -> (TcpListener, EnvoyHandle) { let listener = TcpListener::bind("127.0.0.1:0") .await .expect("bind fake engine"); @@ -175,6 +185,28 @@ async fn ping_from_previous_connection_does_not_stop_actor_started_after_reconne debug_latency_ms: None, callbacks: Arc::new(IdleCallbacks), }); + (listener, handle) +} + +async fn ping(ws: &mut WebSocketStream) { + send( + ws, + protocol::ToEnvoy::ToEnvoyPing(protocol::ToEnvoyPing { ts: 0 }), + ) + .await; + while !matches!(recv(ws).await, protocol::ToRivet::ToRivetPong(_)) {} +} + +fn assert_stopped_with_error(state: protocol::ActorState) { + let protocol::ActorState::ActorStateStopped(stopped) = state else { + panic!("expected a stopped actor, got {state:?}"); + }; + assert_eq!(stopped.code, protocol::StopCode::Error); +} + +#[tokio::test] +async fn ping_from_previous_connection_does_not_stop_actor_started_after_reconnect() { + let (listener, handle) = start_envoy().await; // The first connection is pinged once, then the engine drops it. let mut first = accept(&listener).await; @@ -229,3 +261,124 @@ async fn ping_from_previous_connection_does_not_stop_actor_started_after_reconne handle.shutdown(true); } + +/// The engine claims a connection and sends commands, then the link goes half-open before its +/// first ping. The silence check must still stop the actor, counting from the new connection. +#[tokio::test] +async fn half_open_connection_before_first_ping_still_stops_actors() { + let (listener, handle) = start_envoy().await; + + let mut ws = accept(&listener).await; + let installed_at = Instant::now(); + send(&mut ws, init()).await; + send(&mut ws, start_actor()).await; + assert!(matches!( + recv_actor_state(&mut ws).await, + protocol::ActorState::ActorStateRunning + )); + + // No ping ever arrives and the connection never reports a close. + let state = tokio::time::timeout(Duration::from_secs(5), recv_actor_state(&mut ws)) + .await + .expect("actor should stop even though the connection was never pinged"); + assert!( + installed_at.elapsed() >= SILENCE_DEADLINE - Duration::from_millis(100), + "actor stopped {:?} after the connection opened, before the silence deadline", + installed_at.elapsed() + ); + assert_stopped_with_error(state); + + handle.shutdown(true); +} + +/// After a reconnect the engine sends Init before it claims the connection, and only the claim +/// refreshes its liveness timestamp. If the claim never happens, the engine is still counting from +/// the last ping on the previous connection, so the envoy must stop its actors on that deadline. +#[tokio::test] +async fn unclaimed_reconnect_stops_actors_on_the_previous_ping_deadline() { + let (listener, handle) = start_envoy().await; + + let mut first = accept(&listener).await; + send(&mut first, init()).await; + send(&mut first, start_actor()).await; + assert!(matches!( + recv_actor_state(&mut first).await, + protocol::ActorState::ActorStateRunning + )); + ping(&mut first).await; + let old_ping_at = Instant::now(); + drop(first); + + // The envoy reconnects and receives Init, but the engine never claims the connection: no + // commands and no ping follow. + let mut second = accept(&listener).await; + send(&mut second, init()).await; + + let state = tokio::time::timeout(Duration::from_secs(5), recv_actor_stopped(&mut second)) + .await + .expect("actor should stop on the previous ping's deadline"); + assert!( + old_ping_at.elapsed() < SILENCE_DEADLINE + Duration::from_millis(500), + "actor stopped {:?} after the last ping, later than the engine's view allows", + old_ping_at.elapsed() + ); + assert_stopped_with_error(state); + + handle.shutdown(true); +} + +/// A serverless start injected over HTTP while the WebSocket is not yet claimed must not be judged +/// against the stale baseline. The actor runs and stops only at its own silence deadline. +#[tokio::test] +async fn serverless_start_on_unclaimed_connection_counts_from_its_creation() { + let (listener, handle) = start_envoy().await; + + let mut first = accept(&listener).await; + send(&mut first, init()).await; + ping(&mut first).await; + let old_ping_at = Instant::now(); + drop(first); + + // The new connection receives Init but is never claimed. + let mut second = accept(&listener).await; + send(&mut second, init()).await; + tokio::time::sleep_until(old_ping_at + SILENCE_DEADLINE + Duration::from_millis(200)).await; + + let mut payload = protocol::PROTOCOL_VERSION.to_le_bytes().to_vec(); + payload.extend( + protocol::versioned::ToEnvoy::wrap_latest(start_actor()) + .serialize(protocol::PROTOCOL_VERSION) + .expect("encode serverless start"), + ); + handle + .start_serverless_actor(&payload) + .await + .expect("inject serverless start"); + let started_at = Instant::now(); + assert!(matches!( + recv_actor_state(&mut second).await, + protocol::ActorState::ActorStateRunning + )); + + let early = tokio::time::timeout( + SILENCE_DEADLINE - Duration::from_millis(300), + recv_actor_stopped(&mut second), + ) + .await; + assert!( + early.is_err(), + "actor stopped against the stale baseline: {early:?}" + ); + + let state = tokio::time::timeout(Duration::from_secs(5), recv_actor_stopped(&mut second)) + .await + .expect("actor should stop at its own silence deadline"); + assert!( + started_at.elapsed() >= SILENCE_DEADLINE - Duration::from_millis(100), + "actor stopped {:?} after it started, before its silence deadline", + started_at.elapsed() + ); + assert_stopped_with_error(state); + + handle.shutdown(true); +} diff --git a/engine/sdks/rust/envoy-client/tests/sqlite_vfs_disconnect_stall.rs b/engine/sdks/rust/envoy-client/tests/sqlite_vfs_disconnect_stall.rs index 12b7c5be46..e0d28f4a67 100644 --- a/engine/sdks/rust/envoy-client/tests/sqlite_vfs_disconnect_stall.rs +++ b/engine/sdks/rust/envoy-client/tests/sqlite_vfs_disconnect_stall.rs @@ -133,6 +133,7 @@ fn new_envoy_context() -> EnvoyContext { protocol_metadata: Arc::new(tokio::sync::Mutex::new(None)), shutting_down: std::sync::atomic::AtomicBool::new(false), last_ping_ts: std::sync::atomic::AtomicI64::new(0), + engine_liveness: Default::default(), stopped_tx: tokio::sync::watch::channel(true).0, }); EnvoyContext { diff --git a/engine/sdks/rust/envoy-client/tests/wait_actors_drained.rs b/engine/sdks/rust/envoy-client/tests/wait_actors_drained.rs index 265548bca7..8ef03efc66 100644 --- a/engine/sdks/rust/envoy-client/tests/wait_actors_drained.rs +++ b/engine/sdks/rust/envoy-client/tests/wait_actors_drained.rs @@ -102,6 +102,7 @@ fn build_shared(stopped: bool) -> Arc { protocol_metadata: Arc::new(tokio::sync::Mutex::new(None)), shutting_down: std::sync::atomic::AtomicBool::new(false), last_ping_ts: std::sync::atomic::AtomicI64::new(0), + engine_liveness: Default::default(), stopped_tx: tokio::sync::watch::channel(stopped).0, }) } diff --git a/engine/sdks/rust/envoy-client/tests/ws_tx_lock_contention.rs b/engine/sdks/rust/envoy-client/tests/ws_tx_lock_contention.rs index 4e377f52d2..a7dffaf7d8 100644 --- a/engine/sdks/rust/envoy-client/tests/ws_tx_lock_contention.rs +++ b/engine/sdks/rust/envoy-client/tests/ws_tx_lock_contention.rs @@ -140,6 +140,7 @@ fn new_shared_context() -> Arc { protocol_metadata: Arc::new(tokio::sync::Mutex::new(None)), shutting_down: std::sync::atomic::AtomicBool::new(false), last_ping_ts: std::sync::atomic::AtomicI64::new(0), + engine_liveness: Default::default(), stopped_tx: tokio::sync::watch::channel(true).0, }) } diff --git a/rivetkit-rust/packages/rivetkit-core/src/testing.rs b/rivetkit-rust/packages/rivetkit-core/src/testing.rs index b9d0600cb5..3662f10987 100644 --- a/rivetkit-rust/packages/rivetkit-core/src/testing.rs +++ b/rivetkit-rust/packages/rivetkit-core/src/testing.rs @@ -208,6 +208,7 @@ fn test_envoy_handle(endpoint: String) -> (EnvoyHandle, mpsc::UnboundedReceiver< protocol_metadata: Arc::new(tokio::sync::Mutex::new(None)), shutting_down: std::sync::atomic::AtomicBool::new(false), last_ping_ts: std::sync::atomic::AtomicI64::new(i64::MAX), + engine_liveness: Default::default(), stopped_tx: tokio::sync::watch::channel(true).0, }); (EnvoyHandle::from_shared(shared), envoy_rx) diff --git a/rivetkit-rust/packages/rivetkit-core/tests/context.rs b/rivetkit-rust/packages/rivetkit-core/tests/context.rs index 6e8e070d12..3c175f5ae7 100644 --- a/rivetkit-rust/packages/rivetkit-core/tests/context.rs +++ b/rivetkit-rust/packages/rivetkit-core/tests/context.rs @@ -182,6 +182,7 @@ fn test_envoy_handle() -> (TestEnvoyHandle, mpsc::UnboundedReceiver (EnvoyHandle, mpsc::UnboundedReceiver) protocol_metadata: Arc::new(AsyncMutex::new(None)), shutting_down: AtomicBool::new(false), last_ping_ts: std::sync::atomic::AtomicI64::new(i64::MAX), + engine_liveness: Default::default(), stopped_tx: tokio::sync::watch::channel(true).0, }); diff --git a/rivetkit-rust/packages/rivetkit-core/tests/serverless.rs b/rivetkit-rust/packages/rivetkit-core/tests/serverless.rs index 1a25f4c92d..d3ec0d803b 100644 --- a/rivetkit-rust/packages/rivetkit-core/tests/serverless.rs +++ b/rivetkit-rust/packages/rivetkit-core/tests/serverless.rs @@ -290,6 +290,7 @@ mod moved_tests { // Zero means no engine ping has been received yet, which // `is_ping_healthy` treats as unhealthy. last_ping_ts: AtomicI64::new(0), + engine_liveness: Default::default(), stopped_tx: tokio::sync::watch::channel(true).0, }); diff --git a/rivetkit-rust/packages/rivetkit-core/tests/sleep.rs b/rivetkit-rust/packages/rivetkit-core/tests/sleep.rs index 821bd2df8a..28dac2b596 100644 --- a/rivetkit-rust/packages/rivetkit-core/tests/sleep.rs +++ b/rivetkit-rust/packages/rivetkit-core/tests/sleep.rs @@ -660,6 +660,7 @@ mod moved_tests { protocol_metadata: Arc::new(tokio::sync::Mutex::new(None)), shutting_down: AtomicBool::new(false), last_ping_ts: std::sync::atomic::AtomicI64::new(i64::MAX), + engine_liveness: Default::default(), stopped_tx: tokio::sync::watch::channel(true).0, }); diff --git a/rivetkit-rust/packages/rivetkit-core/tests/sqlite.rs b/rivetkit-rust/packages/rivetkit-core/tests/sqlite.rs index eb7c99a9b8..04a720c450 100644 --- a/rivetkit-rust/packages/rivetkit-core/tests/sqlite.rs +++ b/rivetkit-rust/packages/rivetkit-core/tests/sqlite.rs @@ -185,6 +185,7 @@ fn test_envoy_handle_with_shared() -> ( protocol_metadata: Arc::new(AsyncMutex::new(None)), shutting_down: AtomicBool::new(false), last_ping_ts: std::sync::atomic::AtomicI64::new(i64::MAX), + engine_liveness: Default::default(), stopped_tx: tokio::sync::watch::channel(true).0, }); diff --git a/rivetkit-rust/packages/rivetkit-core/tests/task.rs b/rivetkit-rust/packages/rivetkit-core/tests/task.rs index a7b9c8c199..4b4d75f80c 100644 --- a/rivetkit-rust/packages/rivetkit-core/tests/task.rs +++ b/rivetkit-rust/packages/rivetkit-core/tests/task.rs @@ -443,6 +443,7 @@ pub(crate) mod moved_tests { protocol_metadata: Arc::new(tokio::sync::Mutex::new(None)), shutting_down: AtomicBool::new(false), last_ping_ts: std::sync::atomic::AtomicI64::new(i64::MAX), + engine_liveness: Default::default(), stopped_tx: tokio::sync::watch::channel(true).0, }); diff --git a/rivetkit-rust/packages/rivetkit/tests/client.rs b/rivetkit-rust/packages/rivetkit/tests/client.rs index abe046cd30..30a2401a91 100644 --- a/rivetkit-rust/packages/rivetkit/tests/client.rs +++ b/rivetkit-rust/packages/rivetkit/tests/client.rs @@ -338,6 +338,7 @@ fn test_envoy_handle(endpoint: String) -> EnvoyHandle { protocol_metadata: Arc::new(tokio::sync::Mutex::new(None)), shutting_down: AtomicBool::new(false), last_ping_ts: std::sync::atomic::AtomicI64::new(0), + engine_liveness: Default::default(), stopped_tx: tokio::sync::watch::channel(true).0, });