feat(container-runner): exit process when the last child stops

container-runner/src/actor.rs

@@ -14,10 +14,11 @@
 use rivetkit::{Actor, ActorKeySegment, Ctx, Request, Response, WebSocket, action};
 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,
+	request_exit, reserve_child_port, runner_config,
 };
 
 /// Live actor contexts on this instance, keyed by actor id. Lets the process
@@ -47,12 +48,45 @@
 			release_child_port(child.child_port).await;
 		}
 
-		// The instance stays alive and warm after its last actor stops, ready to
-		// host the next placement. It is reaped by the platform's own shutdown
-		// signal, not by self-exit. This keeps the serverless container long
-		// lived enough for the log agent to drain its stderr, which a fast
-		// self-exit could otherwise lose.
-		tracing::info!(actor_id = %actor_id, reason, "actor stopped, keeping instance warm");
+		// Exit the whole process once the last child on this instance stops. The
+		// runner is PID 1, so `request_exit` cancels `EXIT`, which wakes `main`
+		// to run the graceful envoy close and then return, stopping the container
+		// so the platform reaps it. Guarded on an empty registry so a multi-actor
+		// instance does not tear down siblings still hosting a child.
+		if children().is_empty() {
+			request_exit(actor_id, reason);
+		} else {
+			tracing::info!(
+				actor_id = %actor_id,
+				reason,
+				"actor stopped, other actors still running on this instance"
+			);
+		}
+	}
+
+	/// 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;
 	}
 }
 
@@ -183,7 +217,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 +296,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,22 +209,18 @@ 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
-/// stopping or failing to start no longer exit the instance, so it stays warm
-/// and reusable and its logs have time to drain. The runner is PID 1 in the
-/// image, so exiting stops the container and the platform reaps the instance.
+/// End the process. Driven by a platform shutdown signal or by the last child on
+/// this instance stopping (see `stop_child`). The runner is PID 1 in the image,
+/// so cancelling `EXIT` wakes `main` to run the graceful envoy close and return,
+/// which stops the container and lets the platform reap the instance.
 pub fn request_exit(actor_id: &str, reason: &str) {
 	tracing::info!(actor_id = %actor_id, reason, "shutting down container");
 	EXIT.cancel();
@@ -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()
 		},
@@ -389,21 +394,20 @@ async fn async_main() -> Result<()> {
 	));
 	tracing::info!(port, "container-runner serverless front door listening");
 
-	// Wait for an exit request, then tear down. Only the signal path is live
-	// today: nothing calls `request_exit` except `spawn_signal_handler`, which
-	// sets `SIGNAL_SHUTDOWN` before cancelling `EXIT`, so the `else` branch is
-	// currently unreachable and kept only as a fallback for a future
-	// actor-driven exit.
+	// Wait for an exit request, then tear down. Two shapes depending on why:
 	//
-	// 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),
-	// and the runtime drains unbounded so the /start SSE flushes cleanly.
+	// Actor-driven exit (the last child stopped, so `stop_child` cancelled
+	// `EXIT`): no platform deadline. The child is already reaped by the hook
+	// that triggered the exit (the sweep is a no-op backstop), and the runtime
+	// drains unbounded so the /start SSE flushes cleanly.
 	EXIT.cancelled().await;
 	if SIGNAL_SHUTDOWN.load(Ordering::Acquire) {
 		// A platform SIGTERM reclaims this instance. Report every actor as crashed
@@ -417,15 +421,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 +448,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 +463,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> {