From 6b8f04bfa2e0fca4312093a013f7e61b5f2c8ef5 Mon Sep 17 00:00:00 2001
From: Roy Lin <roylin@a3s.box>
Date: Wed, 17 Jun 2026 09:46:47 +0800
Subject: [PATCH] fix(pool): close warm-pool shutdown-vs-replenish push race
 (orphaned-VM leak)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

From the concurrency audit. A freshly-booted warm VM could be pushed into the
idle pool AFTER drain_idle already cleared it, leaking the VM (mount/IP/volume)
since there is no Drop reaper.

The replenish task checked `*shutdown_rx.borrow()` and only THEN acquired the
idle lock to push — with an await in between. drain_all does signal_shutdown()
then drain_idle() (which locks idle, drains, and runs exactly once). If the
unlocked check observed shutdown=false and drain_idle completed before the task
acquired the lock, the VM was pushed into the already-drained pool and never
reclaimed.

Fix: acquire the idle lock FIRST, then re-check shutdown UNDER it before pushing
— making the check-and-push atomic against drain_idle (which holds the same lock
while draining, always after signal_shutdown). On shutdown, destroy the VM
instead of pushing. The same latent leak existed in `release()` (returns a VM to
the pool with no shutdown check); guard it the same way.

Validated: full a3s-box-runtime pool/lib suite (55) green, fmt + clippy clean.
A minimal check-under-lock change; the interleaving itself needs a real-VM
harness to exercise (like the kill-identity / durable-write fixes).
---
 src/runtime/src/pool/warm_pool.rs | 21 +++++++++++++++++++--
 1 file changed, 19 insertions(+), 2 deletions(-)

diff --git a/src/runtime/src/pool/warm_pool.rs b/src/runtime/src/pool/warm_pool.rs
index 36634fae..665ec963 100644
--- a/src/runtime/src/pool/warm_pool.rs
+++ b/src/runtime/src/pool/warm_pool.rs
@@ -259,6 +259,17 @@ impl WarmPool {
     pub async fn release(&self, vm: VmManager) -> Result<()> {
         let mut idle = self.idle.lock().await;
 
+        // Don't return a VM to a pool that is shutting down: drain_idle has (or
+        // soon will have) cleared `idle` and won't run again, so a push here leaks
+        // the VM (no Drop reaper). Checked under the idle lock so it is atomic with
+        // a concurrent drain_idle. Destroy the VM instead.
+        if *self.shutdown_rx.borrow() {
+            drop(idle);
+            let mut vm = vm;
+            vm.destroy().await?;
+            return Ok(());
+        }
+
         if idle.len() >= self.config.max_size {
             // Pool is full — destroy the VM
             drop(idle); // Release lock before async destroy
@@ -809,8 +820,15 @@ impl WarmPool {
                                         // booting, drain_idle has already cleared
                                         // `idle` and will not run again, so a VM
                                         // pushed now leaks (no Drop reaper). Destroy
-                                        // it instead.
+                                        // it instead. Acquire the idle lock FIRST and
+                                        // re-check shutdown UNDER it: drain_idle drains
+                                        // while holding this same lock (always after
+                                        // signal_shutdown), so the check-and-push is
+                                        // atomic against it — closing the TOCTOU window
+                                        // that an unlocked `borrow()` check left open.
+                                        let mut pool = idle.lock().await;
                                         if *shutdown_rx.borrow() {
+                                            drop(pool);
                                             tracing::debug!(
                                                 box_id = %box_id,
                                                 "Pool shutting down mid-replenish; destroying freshly-booted VM"
@@ -818,7 +836,6 @@ impl WarmPool {
                                             let _ = vm.destroy_with_timeout(2000).await;
                                             continue;
                                         }
-                                        let mut pool = idle.lock().await;
                                         pool.push(WarmVm {
                                             vm,
                                             created_at: Instant::now(),