diff --git a/Cargo.lock b/Cargo.lock index a2ade3bdd1..6be5252f96 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5317,6 +5317,7 @@ version = "2.3.7" dependencies = [ "anyhow", "async-trait", + "ciborium", "clap", "futures-util", "nix 0.30.1", diff --git a/container-runner/Cargo.toml b/container-runner/Cargo.toml index e01f51bc66..5ecd8d93d1 100644 --- a/container-runner/Cargo.toml +++ b/container-runner/Cargo.toml @@ -28,3 +28,6 @@ tokio-tungstenite.workspace = true tokio-util = { workspace = true, features = ["rt"] } tracing.workspace = true tracing-subscriber.workspace = true + +[dev-dependencies] +ciborium.workspace = true diff --git a/container-runner/src/actor.rs b/container-runner/src/actor.rs index 47594f52fb..859fa30683 100644 --- a/container-runner/src/actor.rs +++ b/container-runner/src/actor.rs @@ -13,10 +13,10 @@ 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::input::{ActorInput, ActorState}; use crate::{ - children, drain_grace, effective_stop_grace, exit_token, idle_timeout, release_child_port, - request_exit, reserve_child_port, runner_config, + children, drain_grace, effective_stop_grace, exit_token, idle_timeout, reject_second_start, + release_child_port, request_exit, reserve_child_port, runner_config, }; /// Live actor contexts keyed by actor id, so the shutdown path can report actors @@ -29,6 +29,9 @@ pub struct GameServer { /// Set on the first request, disarming the one-shot startup idle timer. Only /// meaningful when [`idle_timeout`] is enabled. got_request: AtomicBool, + /// Set when `on_start` detected a repeat start and skipped spawning a child, so + /// `run` sleeps the actor instead of running. See [`reject_second_start`]. + reject_start: AtomicBool, } impl GameServer { @@ -106,13 +109,23 @@ impl GameServer { } }); } + + /// Record that the actor received a request: disarms the one-shot idle timer, + /// and in idle mode marks the real start on the first request so an idle-slept + /// actor that never served one can wake without tripping the second-start guard. + fn note_request(&self, ctx: &Ctx) { + let first = !self.got_request.swap(true, Ordering::Relaxed); + if first && reject_second_start() && idle_timeout().is_some() { + mark_started_once(ctx); + } + } } #[async_trait] impl Actor for GameServer { - // The launch spec is the persisted state: a woken actor restores the same - // spec without the engine re-sending input. - type State = ActorInput; + // The persisted state (launch spec plus `started_once`) is restored on wake, so a + // woken actor keeps its spec without the engine re-sending input. + type State = ActorState; type Input = ActorInput; type Actions = (); type Events = (); @@ -122,13 +135,17 @@ impl Actor for GameServer { type Action = action::Raw; async fn create_state(_ctx: &Ctx, input: Self::Input) -> Result { - Ok(input) + Ok(ActorState { + input, + started_once: false, + }) } async fn create(_ctx: &Ctx) -> Result { Ok(Self { child: TokioMutex::new(None), got_request: AtomicBool::new(false), + reject_start: AtomicBool::new(false), }) } @@ -160,9 +177,19 @@ impl Actor for GameServer { } } + // Second-start guard: if this actor already did its real start (a persisted + // flag that survives sleep), do not run again. Skip spawning a child; `run` + // sleeps the actor back down. + if reject_second_start() && ctx.state().started_once { + tracing::warn!(actor_id = %actor_id, "actor tried a second-start"); + self.reject_start.store(true, Ordering::Relaxed); + return Ok(()); + } + // Copy the launch spec out of the state guard before any await. let (input_port, mut parts, env) = { - let input = ctx.state(); + let state = ctx.state(); + let input = &state.input; // input.command overrides the CLI template; input.args are appended. let mut parts = input .command @@ -236,6 +263,11 @@ impl Actor for GameServer { // hook to remove the entry, so registering earlier would leak it. register_ctx(&actor_id, &ctx).await; *self.child.lock().await = Some(child); + // Non-idle: the real start is complete, so record it. Idle mode defers this to + // the first request (see `note_request`) so an idle-slept actor can wake. + if reject_second_start() && idle_timeout().is_none() { + mark_started_once(&ctx); + } self.arm_idle_timeout(&ctx, actor_id); Ok(()) } @@ -244,6 +276,15 @@ impl Actor for GameServer { /// first, so winning the `remove` race means the exit was unexpected: a clean /// exit destroys the actor, any other reports an errored stop (a crash). async fn run(self: Arc, ctx: Ctx) -> Result<()> { + // A rejected repeat start spawned no child; sleep the actor so it does not run + // again. `ctx.sleep()` is valid here because startup has completed. + if self.reject_start.load(Ordering::Relaxed) { + if let Err(err) = ctx.sleep() { + tracing::debug!(error = ?err, actor_id = %ctx.actor_id(), "reject-start sleep failed"); + } + return Ok(()); + } + let Some(child) = self.child.lock().await.clone() else { anyhow::bail!("run: child process was never spawned"); }; @@ -272,7 +313,7 @@ impl Actor for GameServer { } async fn on_fetch(self: Arc, ctx: Ctx, req: Request) -> Result { - self.got_request.store(true, Ordering::Relaxed); + self.note_request(&ctx); let child_port = self .child .lock() @@ -289,7 +330,7 @@ impl Actor for GameServer { ws: WebSocket, req: Request, ) -> Result<()> { - self.got_request.store(true, Ordering::Relaxed); + self.note_request(&ctx); let child_port = self .child .lock() @@ -324,6 +365,16 @@ impl Actor for GameServer { } } +/// Persist `started_once` so a later start is treated as a repeat. Idempotent: a +/// no-op when already set. The read guard is released before the write. +fn mark_started_once(ctx: &Ctx) { + if ctx.state().started_once { + return; + } + ctx.state_mut().started_once = true; + ctx.request_save(); +} + /// Register an actor context for crash-on-shutdown reporting. Overwrites any /// stale entry left by a prior generation with the same id. async fn register_ctx(actor_id: &str, ctx: &Ctx) { diff --git a/container-runner/src/input.rs b/container-runner/src/input.rs index c39fb6d0bd..ca70f1c62d 100644 --- a/container-runner/src/input.rs +++ b/container-runner/src/input.rs @@ -1,16 +1,29 @@ -//! The actor input payload describing how to launch the child game server. +//! The actor input payload and persisted state for the child game server. //! -//! The command, args, env, and port are carried in the actor's create-time `input` -//! (CBOR per RivetKit); anything omitted falls back to the CLI template -//! (`rivet-container-runner -- `). This is also the actor's persisted -//! state, so a woken actor restores the same launch spec. +//! [`ActorInput`] carries the command, args, env, and port from the engine's +//! create-time `input` (CBOR per RivetKit); anything omitted falls back to the CLI +//! template (`rivet-container-runner -- `). [`ActorState`] is what +//! persists across sleep: the launch spec plus lifecycle bookkeeping. use serde::{Deserialize, Serialize}; use std::collections::HashMap; +/// Persisted actor state, restored on wake. The launch spec is flattened in so an +/// actor persisted before `started_once` existed (state was a bare [`ActorInput`]) +/// still decodes. +#[derive(Debug, Default, Serialize, Deserialize)] +pub struct ActorState { + #[serde(flatten)] + pub input: ActorInput, + /// Set once the actor has performed its real start. The reject-second-start guard + /// self-sleeps a repeat start when this is already set. See `RIVET_REJECT_SECOND_START`. + #[serde(default)] + pub started_once: bool, +} + /// Shape of the actor `input` payload. Unknown fields are ignored, not rejected: -/// this is also the persisted state, and a strict decode would break waking actors -/// after a rollback to a binary predating a new field. +/// it nests in the persisted [`ActorState`], and a strict decode would break waking +/// actors after a rollback to a binary predating a new field. #[derive(Debug, Default, Serialize, Deserialize)] pub struct ActorInput { /// Overrides the CLI command template entirely (program + fixed args). diff --git a/container-runner/src/main.rs b/container-runner/src/main.rs index 2e854eae5a..9d7d7b6fbd 100644 --- a/container-runner/src/main.rs +++ b/container-runner/src/main.rs @@ -133,6 +133,20 @@ pub fn idle_timeout() -> Option { *IDLE_TIMEOUT } +/// When set, an actor that starts a second time self-sleeps instead of running +/// again; its persisted `started_once` records the first real start. Configured via +/// RIVET_REJECT_SECOND_START (truthy `1`/`true`/`yes`/`on`). Off by default. +static REJECT_SECOND_START: LazyLock = LazyLock::new(|| { + std::env::var("RIVET_REJECT_SECOND_START") + .map(|value| matches!(value.trim().to_ascii_lowercase().as_str(), "1" | "true" | "yes" | "on")) + .unwrap_or(false) +}); + +/// Whether the reject-second-start guard is enabled. See [`REJECT_SECOND_START`]. +pub fn reject_second_start() -> bool { + *REJECT_SECOND_START +} + /// Token that fires on a platform shutdown signal. Cuts a drain wait short so the /// platform's SIGTERM→SIGKILL budget is honored. pub fn exit_token() -> &'static CancellationToken { diff --git a/container-runner/tests/inline/input.rs b/container-runner/tests/inline/input.rs index 2dc78209fd..30e3dc3038 100644 --- a/container-runner/tests/inline/input.rs +++ b/container-runner/tests/inline/input.rs @@ -46,3 +46,44 @@ fn default_matches_empty() { assert!(default.env.is_empty()); assert!(default.port.is_none()); } + +fn cbor_round_trip(value: &T) -> anyhow::Result +where + T: serde::Serialize, + U: serde::de::DeserializeOwned, +{ + let mut buf = Vec::new(); + ciborium::into_writer(value, &mut buf)?; + Ok(ciborium::from_reader(&buf[..])?) +} + +#[test] +fn actor_state_cbor_round_trips() { + // State persists as CBOR (ciborium), so the flattened input must survive a CBOR + // round trip, not just the JSON the other tests use. + let state = ActorState { + input: ActorInput { + port: Some(7777), + args: vec!["-x".to_string()], + ..Default::default() + }, + started_once: true, + }; + let decoded: ActorState = cbor_round_trip(&state).unwrap(); + assert_eq!(decoded.input.port, Some(7777)); + assert_eq!(decoded.input.args, vec!["-x".to_string()]); + assert!(decoded.started_once); +} + +#[test] +fn legacy_bare_input_state_decodes_into_actor_state() { + // State written before `started_once` existed was a bare ActorInput. The + // flattened input must still decode, defaulting `started_once` to false. + let legacy = ActorInput { + port: Some(7777), + ..Default::default() + }; + let decoded: ActorState = cbor_round_trip(&legacy).unwrap(); + assert_eq!(decoded.input.port, Some(7777)); + assert!(!decoded.started_once); +}