diff --git a/container-runner/src/main.rs b/container-runner/src/main.rs index 4594598431..9059e10bec 100644 --- a/container-runner/src/main.rs +++ b/container-runner/src/main.rs @@ -29,6 +29,7 @@ use std::time::Duration; use anyhow::Result; use clap::Parser; +use futures_util::future::join_all; use rivetkit::serverless_http::{self, ListenerConfig}; use rivetkit::{ActorConfig, EngineSpawnMode, Registry, ServeConfig}; use tokio_util::sync::CancellationToken; @@ -103,30 +104,19 @@ static SIGNAL_SHUTDOWN: AtomicBool = AtomicBool::new(false); static PLATFORM_RECLAIM: AtomicBool = AtomicBool::new(false); /// How long the platform gives this container between SIGTERM and SIGKILL. -/// Defaults to 10 seconds; keep this in sync with the platform's actual budget -/// via RIVET_SIGTERM_BUDGET_SECS. -/// The signal-path teardown splits the budget: ~60% for the engine drain -/// (whose per-actor stops SIGTERM children with a grace capped at ~40%), 1s -/// for the straggler sweep, and the rest as margin. +/// Defaults to 9s, one second under the common ~10s platform budget so the whole +/// teardown lands before SIGKILL; keep it in sync with the platform's actual +/// budget via RIVET_SIGTERM_BUDGET_SECS. On the signal path the engine drain and +/// the child kills run concurrently, each bounded by this full budget. static SIGTERM_BUDGET: LazyLock = LazyLock::new(|| { let secs = std::env::var("RIVET_SIGTERM_BUDGET_SECS") .ok() .and_then(|value| value.parse::().ok()) - .unwrap_or(10) + .unwrap_or(9) .max(3); Duration::from_secs(secs) }); -fn signal_drain_timeout() -> Duration { - SIGTERM_BUDGET.mul_f64(0.6) -} - -fn signal_child_stop_grace() -> Duration { - SIGTERM_BUDGET.mul_f64(0.4) -} - -const SIGNAL_SWEEP_GRACE: Duration = Duration::from_secs(1); - pub fn runner_config() -> Arc { RUNNER_CONFIG .get() @@ -203,7 +193,7 @@ pub async fn release_child_port(port: u16) { pub fn effective_stop_grace() -> Duration { let grace = runner_config().stop_grace; if SIGNAL_SHUTDOWN.load(Ordering::Acquire) { - grace.min(signal_child_stop_grace()) + grace.min(*SIGTERM_BUDGET) } else { grace } @@ -395,11 +385,13 @@ async fn async_main() -> Result<()> { // currently unreachable and kept only as a fallback for a future // actor-driven exit. // - // Signal (platform is reclaiming the instance): tell the engine FIRST so - // it can start re-placing actors immediately. Its per-actor stops run our - // on_destroy hooks, which SIGTERM children with the capped signal grace. - // The drain is bounded so an unreachable engine cannot eat the whole - // platform budget; the sweep then catches any child whose hooks never ran. + // Signal (platform is reclaiming the instance): kill our children AND + // notify the engine at the SAME time, each bounded by the full SIGTERM + // budget. The direct sweep is what guarantees children die within budget + // rather than waiting on an engine round-trip to run our on_destroy hooks; + // notifying the engine in parallel just lets it start re-placing actors + // immediately. Bounding the drain means an unreachable engine cannot eat + // the budget the children need. // // Fallback actor-driven exit (unreachable today): no platform deadline. // Children are already reaped by the hooks (the sweep is a no-op backstop), @@ -417,13 +409,15 @@ async fn async_main() -> Result<()> { ) .await; } - if tokio::time::timeout(signal_drain_timeout(), runtime.shutdown()) - .await - .is_err() - { - tracing::warn!("engine drain exceeded the signal budget, sweeping children directly"); - } - stop_all_children(SIGNAL_SWEEP_GRACE).await; + let drain = async { + if tokio::time::timeout(*SIGTERM_BUDGET, runtime.shutdown()) + .await + .is_err() + { + tracing::warn!("engine drain exceeded the signal budget"); + } + }; + tokio::join!(drain, stop_all_children(*SIGTERM_BUDGET)); } else { stop_all_children(stop_grace).await; runtime.shutdown().await; @@ -442,7 +436,9 @@ async fn async_main() -> Result<()> { /// Stop every child still in the registry. Actor `on_destroy` normally reaps /// its own child first; this is the belt-and-suspenders sweep for the signal -/// path so children are never orphaned. +/// path so children are never orphaned. Children are stopped concurrently so +/// each gets the full `grace` within the SIGTERM budget instead of queueing +/// behind the others. async fn stop_all_children(grace: Duration) { let mut children: Vec> = Vec::new(); CHILDREN @@ -455,10 +451,11 @@ async fn stop_all_children(grace: Duration) { return; } println!("runner: shutdown, stopping {} child(ren)", children.len()); - for child in children { + join_all(children.into_iter().map(|child| async move { child.stop(grace).await; release_child_port(child.child_port).await; - } + })) + .await; } fn env_u16(key: &str) -> Option { diff --git a/engine/sdks/rust/envoy-client/src/commands.rs b/engine/sdks/rust/envoy-client/src/commands.rs index 342fa4d022..ffe86f993b 100644 --- a/engine/sdks/rust/envoy-client/src/commands.rs +++ b/engine/sdks/rust/envoy-client/src/commands.rs @@ -1,3 +1,5 @@ +use std::collections::HashMap; + use rivet_envoy_protocol as protocol; use crate::actor::create_actor; @@ -16,6 +18,14 @@ pub async fn handle_commands(ctx: &mut EnvoyContext, commands: Vec = commands + .iter() + .filter(|c| matches!(c.inner, protocol::Command::CommandStopActor(_))) + .map(|c| (c.checkpoint.actor_id.clone(), c.checkpoint.generation)) + .collect(); + for command_wrapper in commands { let checkpoint = command_wrapper.checkpoint; let dedup_key = (checkpoint.actor_id.clone(), checkpoint.generation); @@ -80,35 +90,59 @@ pub async fn handle_commands(ctx: &mut EnvoyContext, commands: Vec = Vec::new(); +/// Ack only the given actors' latest processed command index. Used for the +/// immediate stop ack. Does not clear dedup (see the race note in +/// `send_command_ack`); a failed send is retried by the replayed batch or tick. +async fn send_stop_command_acks(ctx: &EnvoyContext, actors: &[(String, u32)]) { + let mut highest: HashMap<(String, u32), i64> = HashMap::new(); + for key in actors { + if let Some(&index) = ctx.processed_command_idx.get(key) { + highest.insert(key.clone(), index); + } + } + if highest.is_empty() { + return; + } + + send_ack_checkpoints(ctx, checkpoints_from(highest)).await; +} + +pub async fn send_command_ack(ctx: &mut EnvoyContext) { + // Merge live actors and the dedup map, highest index per actor-generation. + // Live actors are re-acked every tick (recovers an ack accepted locally but + // never committed by the server); the dedup map covers stops whose actor was + // already removed and is cleared once acked. + let mut highest: HashMap<(String, u32), i64> = HashMap::new(); for (actor_id, generations) in &ctx.actors { for (generation, entry) in generations { - if entry.last_command_idx < 0 { - continue; + if entry.last_command_idx >= 0 { + highest.insert((actor_id.clone(), *generation), entry.last_command_idx); } - last_command_checkpoints.push(protocol::ActorCheckpoint { - actor_id: actor_id.clone(), - generation: *generation, - index: entry.last_command_idx, - }); } } + for ((actor_id, generation), &index) in &ctx.processed_command_idx { + highest + .entry((actor_id.clone(), *generation)) + .and_modify(|existing| *existing = (*existing).max(index)) + .or_insert(index); + } - if last_command_checkpoints.is_empty() { + if highest.is_empty() { return; } - let send_failed = ws_send( - &ctx.shared, - protocol::ToRivet::ToRivetAckCommands(protocol::ToRivetAckCommands { - last_command_checkpoints: last_command_checkpoints.clone(), - }), - ) - .await; + let last_command_checkpoints = checkpoints_from(highest); + let send_failed = send_ack_checkpoints(ctx, last_command_checkpoints.clone()).await; // Skip the dedup clear if the ack never left this process. Otherwise // `pegboard-envoy` would replay the commands on reconnect with no dedup @@ -127,8 +161,35 @@ pub async fn send_command_ack(ctx: &mut EnvoyContext) { // window is narrow (the gap between OS-accepted bytes and the FDB // commit), but a strictly correct fix needs an ack-of-ack from // `pegboard-envoy` so we only clear after positive confirmation. + // This now also applies to removed actors whose stops are acked here: a + // short-lived actor can be resurrected in the same window. Same fix. for cp in &last_command_checkpoints { ctx.processed_command_idx .remove(&(cp.actor_id.clone(), cp.generation)); } } + +fn checkpoints_from(highest: HashMap<(String, u32), i64>) -> Vec { + highest + .into_iter() + .map(|((actor_id, generation), index)| protocol::ActorCheckpoint { + actor_id, + generation, + index, + }) + .collect() +} + +/// Send an ack for the given checkpoints. Returns whether the send failed. +async fn send_ack_checkpoints( + ctx: &EnvoyContext, + last_command_checkpoints: Vec, +) -> bool { + ws_send( + &ctx.shared, + protocol::ToRivet::ToRivetAckCommands(protocol::ToRivetAckCommands { + last_command_checkpoints, + }), + ) + .await +} diff --git a/engine/sdks/rust/envoy-client/tests/command_dedup.rs b/engine/sdks/rust/envoy-client/tests/command_dedup.rs index 0f48ef5dfe..79b445cba8 100644 --- a/engine/sdks/rust/envoy-client/tests/command_dedup.rs +++ b/engine/sdks/rust/envoy-client/tests/command_dedup.rs @@ -3,7 +3,7 @@ use std::sync::Arc; use rivet_envoy_client::actor::ToActor; use rivet_envoy_client::async_counter::AsyncCounter; -use rivet_envoy_client::commands::handle_commands; +use rivet_envoy_client::commands::{handle_commands, send_command_ack}; use rivet_envoy_client::config::{ BoxFuture, EnvoyCallbacks, EnvoyConfig, HttpRequest, HttpResponse, WebSocketHandler, WebSocketSender, @@ -18,6 +18,7 @@ use rivet_envoy_client::sqlite::{ use rivet_envoy_client::utils::{BufferMap, RemoteSqliteIndeterminateResultError}; use rivet_envoy_protocol as protocol; use tokio::sync::mpsc; +use vbare::OwnedVersionedData; struct IdleCallbacks; @@ -271,3 +272,153 @@ async fn replayed_command_is_dropped_after_remote_sql_lost_response() { handle_commands(&mut ctx, vec![stop_command("actor-replay", 1, 5)]).await; assert!(actor_rx.try_recv().is_err()); } + +fn decode_ack_checkpoints(msg: WsTxMessage) -> Vec { + let WsTxMessage::Send(bytes) = msg else { + panic!("expected a websocket send, got a close"); + }; + let message = protocol::versioned::ToRivet::deserialize(&bytes, protocol::PROTOCOL_VERSION) + .expect("failed to decode ToRivet message"); + match message { + protocol::ToRivet::ToRivetAckCommands(val) => val.last_command_checkpoints, + _ => panic!("expected ToRivetAckCommands"), + } +} + +#[tokio::test] +async fn stop_command_is_acked_immediately() { + let mut ctx = new_envoy_context(); + let (actor_tx, mut actor_rx) = mpsc::unbounded_channel::(); + ctx.insert_actor( + "actor-a".to_string(), + 1, + actor_tx, + Arc::new(AsyncCounter::new()), + "actor-a".to_string(), + -1, + ); + + let (ws_tx, mut ws_rx) = mpsc::unbounded_channel(); + *ctx.shared.ws_tx.lock().await = Some(ws_tx); + + handle_commands(&mut ctx, vec![stop_command("actor-a", 1, 5)]).await; + assert!(matches!( + actor_rx.try_recv(), + Ok(ToActor::Stop { command_idx: 5, .. }) + )); + + // The stop must be acked right away rather than waiting for the periodic + // tick, otherwise the actor entry is gone before the next ack. + let checkpoints = + decode_ack_checkpoints(ws_rx.try_recv().expect("stop should trigger an immediate ack")); + assert_eq!(checkpoints.len(), 1); + assert_eq!(checkpoints[0].actor_id, "actor-a"); + assert_eq!(checkpoints[0].generation, 1); + assert_eq!(checkpoints[0].index, 5); + + // A successful immediate ack must still retain the dedup entry. Only the + // periodic tick clears it, so a replay can re-ack if the server never + // committed this ack. + assert_eq!( + ctx.processed_command_idx.get(&("actor-a".to_string(), 1)), + Some(&5) + ); +} + +#[tokio::test] +async fn stop_ack_retried_via_replay_after_failed_send() { + let mut ctx = new_envoy_context(); + let (actor_tx, mut actor_rx) = mpsc::unbounded_channel::(); + ctx.insert_actor( + "actor-a".to_string(), + 1, + actor_tx, + Arc::new(AsyncCounter::new()), + "actor-a".to_string(), + -1, + ); + + // No websocket is connected, so the immediate ack send fails. The processed + // index must be retained so a later replay can re-ack it. + handle_commands(&mut ctx, vec![stop_command("actor-a", 1, 5)]).await; + assert!(matches!( + actor_rx.try_recv(), + Ok(ToActor::Stop { command_idx: 5, .. }) + )); + assert_eq!( + ctx.processed_command_idx.get(&("actor-a".to_string(), 1)), + Some(&5) + ); + + // Remove the actor, as happens once it emits its Stopped event. The re-ack + // must still work with no live actor, sourced from the dedup map. + ctx.remove_actor("actor-a", 1); + + // Reconnect and replay the same stop. Dedup skips reprocessing, but the batch + // still carries a stop, so the retained checkpoint is re-acked. + let (ws_tx, mut ws_rx) = mpsc::unbounded_channel(); + *ctx.shared.ws_tx.lock().await = Some(ws_tx); + handle_commands(&mut ctx, vec![stop_command("actor-a", 1, 5)]).await; + assert!( + actor_rx.try_recv().is_err(), + "replayed stop must not be reprocessed" + ); + + let checkpoints = + decode_ack_checkpoints(ws_rx.try_recv().expect("replayed stop should re-ack")); + assert_eq!(checkpoints.len(), 1); + assert_eq!(checkpoints[0].index, 5); +} + +#[tokio::test] +async fn unknown_actor_stop_is_acked() { + let mut ctx = new_envoy_context(); + let (ws_tx, mut ws_rx) = mpsc::unbounded_channel(); + *ctx.shared.ws_tx.lock().await = Some(ws_tx); + + // No actor inserted: models a stop replayed to a process that never started + // it (e.g. after restart). It must still be acked to stop the replay. + handle_commands(&mut ctx, vec![stop_command("actor-gone", 3, 9)]).await; + + let checkpoints = decode_ack_checkpoints( + ws_rx + .try_recv() + .expect("unknown-actor stop should still ack"), + ); + assert_eq!(checkpoints.len(), 1); + assert_eq!(checkpoints[0].actor_id, "actor-gone"); + assert_eq!(checkpoints[0].generation, 3); + assert_eq!(checkpoints[0].index, 9); +} + +#[tokio::test] +async fn live_actor_is_reacked_on_each_tick() { + let mut ctx = new_envoy_context(); + let (actor_tx, _actor_rx) = mpsc::unbounded_channel::(); + // A live actor whose latest command index is 3. + ctx.insert_actor( + "actor-a".to_string(), + 1, + actor_tx, + Arc::new(AsyncCounter::new()), + "actor-a".to_string(), + 3, + ); + + let (ws_tx, mut ws_rx) = mpsc::unbounded_channel(); + *ctx.shared.ws_tx.lock().await = Some(ws_tx); + + // The first tick acks index 3 and clears the dedup map. The second tick must + // still re-ack from the live-actor scan, recovering an ack the server may + // never have committed. + send_command_ack(&mut ctx).await; + let first = decode_ack_checkpoints(ws_rx.try_recv().expect("first tick should ack")); + assert_eq!(first.len(), 1); + assert_eq!(first[0].index, 3); + + send_command_ack(&mut ctx).await; + let second = decode_ack_checkpoints(ws_rx.try_recv().expect("second tick should re-ack")); + assert_eq!(second.len(), 1); + assert_eq!(second[0].actor_id, "actor-a"); + assert_eq!(second[0].index, 3); +}