runtime-rs: Block WaitSandbox until sandbox exits

Rework sandbox waiting so the WaitSandbox path blocks on sandbox
lifetime rather than directly borrowing the hypervisor wait call.

Once stop has been observed, the cached exit result is returned to
later waiters. While the sandbox is still alive, waiters subscribe to
the internal stop notifier and sleep until shutdown or VM exit records
the final result.

Together with the preceding support commits, this keeps the overall
behaviour identical to the original WaitSandbox fix while making the
dependency chain explicit.

Signed-off-by: Alex Lyn <alex.lyn@antgroup.com>

src/runtime-rs/crates/runtimes/virt_container/src/sandbox.rs

@@ -181,6 +181,20 @@ impl VirtSandbox {
         self.hypervisor.clone()
     }
 
+    async fn record_stop(&self, exit_status: u32, exited_at: SystemTime) {
+        let mut inner = self.inner.write().await;
+        if inner.state == SandboxState::Stopped {
+            return;
+        }
+
+        inner.state = SandboxState::Stopped;
+        inner.exit_info = Some(SandboxExitInfo {
+            exit_status,
+            exited_at: Some(exited_at),
+        });
+        let _ = self.exit_notify_tx.send(true);
+    }
+
     #[instrument]
     async fn prepare_for_start_sandbox(
         &self,
@@ -758,6 +772,22 @@ impl Sandbox for VirtSandbox {
         self.hypervisor.start_vm(10_000).await.context("start vm")?;
         info!(sl!(), "start vm");
 
+        let sandbox = self.clone();
+        // wait for vm exit in background, and record the exit status and time when vm exited.
+        tokio::spawn(async move {
+            match sandbox.hypervisor.wait_vm().await {
+                Ok(exit_code) => {
+                    sandbox
+                        .record_stop(exit_code as u32, SystemTime::now())
+                        .await;
+                }
+                Err(err) => {
+                    warn!(sl!(), "failed waiting for sandbox VM exit: {:?}", err);
+                    sandbox.record_stop(255, SystemTime::now()).await;
+                }
+            }
+        });
+
         // execute pre-start hook functions, including Prestart Hooks and CreateRuntime Hooks
         let (prestart_hooks, create_runtime_hooks) =
             if let Some(hooks) = sandbox_config.hooks.as_ref() {
@@ -944,6 +974,22 @@ impl Sandbox for VirtSandbox {
             .await
             .context("start template vm")?;
         info!(sl!(), "vm started from template");
+
+        let sandbox = self.clone();
+        tokio::spawn(async move {
+            match sandbox.hypervisor.wait_vm().await {
+                Ok(exit_code) => {
+                    sandbox
+                        .record_stop(exit_code as u32, SystemTime::now())
+                        .await;
+                }
+                Err(err) => {
+                    warn!(sl!(), "failed waiting for sandbox VM exit: {:?}", err);
+                    sandbox.record_stop(255, SystemTime::now()).await;
+                }
+            }
+        });
+
         Ok(())
     }
 
@@ -962,23 +1008,47 @@ impl Sandbox for VirtSandbox {
 
     async fn wait(&self) -> Result<SandboxExitInfo> {
         info!(sl!(), "wait sandbox");
-        let exit_code = self.hypervisor.wait_vm().await.context("wait vm")?;
-        Ok(SandboxExitInfo {
-            exit_status: exit_code as u32,
-            exited_at: Some(std::time::SystemTime::now()),
-        })
+        {
+            let inner = self.inner.read().await;
+            if inner.state == SandboxState::Stopped {
+                return Ok(inner.exit_info.clone().unwrap_or_default());
+            }
+        }
+
+        let mut exit_notify_rx = self.exit_notify_tx.subscribe();
+        while !*exit_notify_rx.borrow() {
+            exit_notify_rx
+                .changed()
+                .await
+                .context("wait for sandbox stop notification")?;
+        }
+
+        let inner = self.inner.read().await;
+        Ok(inner.exit_info.clone().unwrap_or_default())
     }
 
     async fn stop(&self) -> Result<()> {
-        let mut sandbox_inner = self.inner.write().await;
+        let state = {
+            let sandbox_inner = self.inner.read().await;
+            sandbox_inner.state
+        };
 
-        if sandbox_inner.state != SandboxState::Stopped {
-            info!(sl!(), "begin stop sandbox");
-            self.hypervisor.stop_vm().await.context("stop vm")?;
-            sandbox_inner.state = SandboxState::Stopped;
-            info!(sl!(), "sandbox stopped");
+        if state == SandboxState::Stopped {
+            return Ok(());
         }
 
+        info!(sl!(), "begin stop sandbox");
+        if state == SandboxState::Init {
+            let _ = self.hypervisor.stop_vm().await;
+            self.record_stop(0, SystemTime::now()).await;
+            info!(sl!(), "sandbox stopped during Init");
+            return Ok(());
+        }
+
+        self.hypervisor.stop_vm().await.context("stop vm")?;
+        self.wait().await.context("wait for vm exit after stop")?;
+        info!(sl!(), "sandbox stopped");
+
         Ok(())
     }