feat(container-runner): drain child before SIGTERM on engine pause

container-runner/src/actor.rs

@@ -17,7 +17,8 @@
 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 @@
 		// 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 @@
 			.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");
 		}
@@ -262,10 +288,10 @@
 
 	/// Engine-initiated sleep. `no_sleep` suppresses idle sleep, but the
 	/// engine can still sleep an actor (dashboard, crash policy, eviction
 	/// ahead of instance retirement); leaving the child running would orphan
 	/// it on an instance the engine considers vacated.
 	async fn on_sleep(self: Arc<Self>, ctx: Ctx<Self>) -> Result<()> {
-		self.stop_child(ctx.actor_id(), "actor sleeping").await;
+		self.drain_then_stop_child(ctx.actor_id(), "actor sleeping").await;
 		Ok(())
 	}
 

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<String>,
 	/// 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<Duration> = LazyLock::new(|| {
 	let secs = std::env::var("RIVET_SIGTERM_BUDGET_SECS")
 		.ok()
 		.and_then(|value| value.parse::<u64>().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<Duration> = LazyLock::new(|| {
+	let secs = std::env::var("RIVET_DRAIN_GRACE_SECS")
+		.ok()
+		.and_then(|value| value.parse::<u64>().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<RunnerConfig> {
 	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<Arc<ChildProcess>> = 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<u16> {

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 @@
 		);
 	}
 
+	// Collect actors with a stop in the raw batch before dedup, so a replayed
+	// (skipped) stop is still re-acked instead of being replayed forever.
+	let stopped_actors: Vec<(String, u32)> = 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 @@
 			}
 		}
 	}
+
+	// Ack stops immediately since their actors are removed before the periodic
+	// tick. Scope to just the stopped actors instead of a full-state ack, and do
+	// not clear dedup; the tick handles full re-acks, recovery, and clearing.
+	if !stopped_actors.is_empty() {
+		send_stop_command_acks(ctx, &stopped_actors).await;
+	}
 }
 
-pub async fn send_command_ack(ctx: &mut EnvoyContext) {
-	let mut last_command_checkpoints: Vec<protocol::ActorCheckpoint> = 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 @@
 	// 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<protocol::ActorCheckpoint> {
+	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<protocol::ActorCheckpoint>,
+) -> bool {
+	ws_send(
+		&ctx.shared,
+		protocol::ToRivet::ToRivetAckCommands(protocol::ToRivetAckCommands {
+			last_command_checkpoints,
+		}),
+	)
+	.await
+}