diff --git a/container-runner/src/actor.rs b/container-runner/src/actor.rs index 79dc1dd7b7..800618ca49 100644 --- a/container-runner/src/actor.rs +++ b/container-runner/src/actor.rs @@ -17,7 +17,8 @@ use tokio::sync::Mutex as TokioMutex; use crate::child::{ChildProcess, SpawnSpec, log_prefix}; use crate::input::ActorInput; use crate::{ - children, effective_stop_grace, release_child_port, reserve_child_port, runner_config, + children, drain_grace, effective_stop_grace, exit_token, release_child_port, + reserve_child_port, runner_config, }; /// Live actor contexts on this instance, keyed by actor id. Lets the process @@ -54,6 +55,31 @@ impl GameServer { // self-exit could otherwise lose. tracing::info!(actor_id = %actor_id, reason, "actor stopped, keeping instance warm"); } + + /// Engine pause path (sleep, lost, going-away). Give the child up to + /// `DRAIN_GRACE` to finish its in-flight work and exit on its own before we + /// force a stop; the child is not signalled during the window. A natural + /// child exit ends the wait immediately, and a platform SIGTERM (which + /// cancels the exit token) cuts it short so the reclaim's SIGTERM→SIGKILL + /// budget is honored. + async fn drain_then_stop_child(&self, actor_id: &str, reason: &str) { + let child = self.child.lock().await.clone(); + if let Some(child) = child { + if !child.has_exited() { + let prefix = log_prefix(actor_id, child.key.as_deref()); + println!( + "{prefix} runner: draining child for up to {:?} before stopping", + drain_grace() + ); + tokio::select! { + _ = child.wait_exit() => {} + _ = tokio::time::sleep(drain_grace()) => {} + _ = exit_token().cancelled() => {} + } + } + } + self.stop_child(actor_id, reason).await; + } } #[async_trait] @@ -183,7 +209,7 @@ impl Actor for GameServer { .is_err() { // Unreachable given the duplicate-start check above; defensive. - child.stop(cfg.stop_grace).await; + child.stop(effective_stop_grace()).await; release_child_port(child_port).await; anyhow::bail!("a child for actor {actor_id} is already registered"); } @@ -265,7 +291,7 @@ impl Actor for GameServer { /// ahead of instance retirement); leaving the child running would orphan /// it on an instance the engine considers vacated. async fn on_sleep(self: Arc, ctx: Ctx) -> Result<()> { - self.stop_child(ctx.actor_id(), "actor sleeping").await; + self.drain_then_stop_child(ctx.actor_id(), "actor sleeping").await; Ok(()) } diff --git a/container-runner/src/main.rs b/container-runner/src/main.rs index 4594598431..04f5f492ba 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; @@ -64,8 +65,6 @@ pub struct RunnerConfig { pub command_template: Vec, /// Default local port for the child when `input.port` is absent. pub default_child_port: u16, - /// SIGTERM→SIGKILL grace period on stop. - pub stop_grace: Duration, /// How long to wait for the child's port to open before failing the start. pub readiness_timeout: Duration, } @@ -103,29 +102,42 @@ 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) -} +/// How long an engine-initiated pause (sleep, lost, going-away) lets the child +/// keep serving and finish its own work before we force a stop. The child is not +/// signalled during this window; it either exits on its own or is SIGTERM'd at +/// the end. Defaults to 15 minutes. It must fit inside the engine's per-runner +/// `drain_grace_period` or a platform reclaim cuts it short. +static DRAIN_GRACE: LazyLock = LazyLock::new(|| { + let secs = std::env::var("RIVET_DRAIN_GRACE_SECS") + .ok() + .and_then(|value| value.parse::().ok()) + .unwrap_or(900); + Duration::from_secs(secs) +}); -fn signal_child_stop_grace() -> Duration { - SIGTERM_BUDGET.mul_f64(0.4) +/// The drain window for an engine pause. See [`DRAIN_GRACE`]. +pub fn drain_grace() -> Duration { + *DRAIN_GRACE } -const SIGNAL_SWEEP_GRACE: Duration = Duration::from_secs(1); +/// The cancellation token that fires on a platform shutdown signal. Used to cut +/// a drain wait short so the platform's SIGTERM→SIGKILL budget is honored. +pub fn exit_token() -> &'static CancellationToken { + &EXIT +} pub fn runner_config() -> Arc { RUNNER_CONFIG @@ -197,16 +209,12 @@ pub async fn release_child_port(port: u16) { RESERVED_PORTS.remove_async(&port).await; } -/// The SIGTERM→SIGKILL grace to give a child right now: the configured -/// `--stop-grace-secs` normally, capped to the platform budget while the -/// process is shutting down due to an OS signal. +/// The SIGTERM→SIGKILL window for the child: always `SIGTERM_BUDGET`, whatever +/// triggered the stop, so the child never gets a different kill deadline on one +/// path than another. The wait *before* SIGTERM (letting the child exit on its +/// own during an engine pause) is separate; see [`drain_grace`]. pub fn effective_stop_grace() -> Duration { - let grace = runner_config().stop_grace; - if SIGNAL_SHUTDOWN.load(Ordering::Acquire) { - grace.min(signal_child_stop_grace()) - } else { - grace - } + *SIGTERM_BUDGET } /// End the process. Only the platform shutdown signal drives this now: actors @@ -286,10 +294,6 @@ struct Args { #[arg(long, env = "RIVET_SERVERLESS_BASE_PATH", default_value = "/api/rivet")] base_path: String, - /// SIGTERM→SIGKILL grace period (seconds) when stopping the child. - #[arg(long, env = "RIVET_STOP_GRACE_SECS", default_value_t = 10)] - stop_grace_secs: u64, - /// How long (seconds) to wait for the child's port to open before failing start. #[arg(long, env = "RIVET_READINESS_TIMEOUT_SECS", default_value_t = 30)] readiness_timeout_secs: u64, @@ -335,12 +339,10 @@ async fn async_main() -> Result<()> { .or_else(|| env_u16("PORT")) .unwrap_or(8080); - let stop_grace = Duration::from_secs(args.stop_grace_secs); RUNNER_CONFIG .set(Arc::new(RunnerConfig { command_template: args.command.clone(), default_child_port: args.child_port, - stop_grace, readiness_timeout: Duration::from_secs(args.readiness_timeout_secs), })) .map_err(|_| anyhow::anyhow!("runner config already set"))?; @@ -351,9 +353,12 @@ async fn async_main() -> Result<()> { ActorConfig { // Game servers hold live in-memory state; never idle-sleep the actor. no_sleep: true, - // The destroy grace deadline must outlast the child's SIGTERM→SIGKILL - // window or core aborts `on_destroy` mid-stop and leaks the child. - sleep_grace_period: stop_grace + Duration::from_secs(10), + // Core force-aborts the stop hook after this deadline. `on_sleep` + // legitimately runs for the whole drain window plus the SIGTERM + // budget, so the deadline must outlast both or core would abort the + // drain mid-flight and leak the child. The extra 5s keeps the + // deadline from racing a hook that finishes right on time. + sleep_grace_period: *DRAIN_GRACE + *SIGTERM_BUDGET + Duration::from_secs(5), sleep_grace_period_overridden: true, ..Default::default() }, @@ -395,11 +400,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,15 +424,17 @@ 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; + stop_all_children(*SIGTERM_BUDGET).await; runtime.shutdown().await; } serve_shutdown.cancel(); @@ -442,7 +451,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 +466,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); +}