feat(container-runner): drain engine sleeps unless the idle timer fired

container-runner/src/actor.rs

@@ -12,7 +12,7 @@
 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, ActorState};
 use crate::{
 	children, drain_grace, effective_stop_grace, exit_token, idle_timeout, idle_timeout_with_jitter,
@@ -26,14 +26,16 @@
 	LazyLock::new(scc::HashMap::new);
 
 /// One-shot idle lifecycle for a generation, in a single atomic. The startup timer
-/// sleeps the actor only while it is `ARMED`; a request moves it to `REQUESTED` so the
-/// timer will not sleep.
+/// sleeps the actor only while it is `ARMED`; a request moves it to `REQUESTED` (the
+/// timer will not sleep), and the timer firing with no request moves it to
+/// `IDLE_SLEEPING`, which `on_sleep` reads to skip the drain.
 const IDLE_ARMED: u8 = 0;
 const IDLE_REQUESTED: u8 = 1;
+const IDLE_SLEEPING: u8 = 2;
 
 pub struct GameServer {
 	child: TokioMutex<Option<Arc<ChildProcess>>>,
-	/// One-shot idle state: `IDLE_ARMED` / `IDLE_REQUESTED`.
+	/// One-shot idle state: `IDLE_ARMED` / `IDLE_REQUESTED` / `IDLE_SLEEPING`.
 	idle_state: AtomicU8,
 	/// 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`].
@@ -108,7 +110,14 @@
 				_ = tokio::time::sleep(delay) => {}
 				_ = abort.cancelled() => return,
 			}
-			if this.idle_state.load(Ordering::Relaxed) == IDLE_REQUESTED {
+			// Sleep only if still armed. If a request raced in, the CAS fails and we do
+			// nothing; on success the state records this as an idle-timer sleep so
+			// `on_sleep` skips the drain.
+			if this
+				.idle_state
+				.compare_exchange(IDLE_ARMED, IDLE_SLEEPING, Ordering::SeqCst, Ordering::SeqCst)
+				.is_err()
+			{
 				return;
 			}
 			tracing::info!(actor_id = %actor_id, ?delay, "no request within idle timeout, sleeping");
@@ -362,12 +371,12 @@
 		crate::proxy::ws_proxy(child_port, path, ws).await
 	}
 
-	/// Engine-initiated sleep. `no_sleep` blocks only idle sleep; the engine can
-	/// still sleep an actor (dashboard, crash policy, eviction), so we stop the child.
-	/// In idle-timeout mode the sleep is (treated as) an idle sleep, so it skips the
-	/// drain and stops promptly; otherwise it drains for in-flight work first.
+	/// Engine-initiated sleep. `no_sleep` blocks only idle sleep; the engine can still
+	/// sleep an actor (dashboard, crash policy, eviction). A no-request idle-timer sleep
+	/// has nothing to drain, so it stops promptly; every other sleep (an active actor,
+	/// or idle timeout disabled) drains in-flight work first.
 	async fn on_sleep(self: Arc<Self>, ctx: Ctx<Self>) -> Result<()> {
-		if idle_timeout().is_some() {
+		if self.idle_state.load(Ordering::SeqCst) == IDLE_SLEEPING {
 			self.stop_child(ctx.actor_id(), "actor sleeping (idle)").await;
 		} else {
 			self.drain_then_stop_child(ctx.actor_id(), "actor sleeping").await;