From 63aea5e48c15907063cc8328ac27323f5bd2a2aa Mon Sep 17 00:00:00 2001
From: Jorge Prendes <jorge.prendes@gmail.com>
Date: Thu, 27 Feb 2025 16:17:23 +0000
Subject: [PATCH] use async tasks to wait for container exit

Signed-off-by: Jorge Prendes <jorge.prendes@gmail.com>
---
 .../src/sandbox/async_utils.rs                |  9 ++++++
 .../src/sandbox/shim/local.rs                 | 28 ++++++++-----------
 .../src/sys/unix/container/instance.rs        |  9 +++---
 crates/containerd-shim-wasm/src/testing.rs    |  2 +-
 4 files changed, 26 insertions(+), 22 deletions(-)

diff --git a/crates/containerd-shim-wasm/src/sandbox/async_utils.rs b/crates/containerd-shim-wasm/src/sandbox/async_utils.rs
index 644682973..154b86d95 100644
--- a/crates/containerd-shim-wasm/src/sandbox/async_utils.rs
+++ b/crates/containerd-shim-wasm/src/sandbox/async_utils.rs
@@ -4,6 +4,7 @@ use std::future::Future;
 use std::sync::LazyLock;
 use std::time::Duration;
 
+use tokio::task::JoinHandle;
 use tokio::time::timeout;
 
 // A thread local runtime that can be used to run futures to completion.
@@ -31,6 +32,14 @@ pub trait AmbientRuntime: Future {
     {
         timeout(t, self).await.ok()
     }
+
+    fn spawn(self) -> JoinHandle<Self::Output>
+    where
+        Self: Sized + Send + 'static,
+        Self::Output: Send + 'static,
+    {
+        RUNTIME.spawn(self)
+    }
 }
 
 impl<F: Future> AmbientRuntime for F {}
diff --git a/crates/containerd-shim-wasm/src/sandbox/shim/local.rs b/crates/containerd-shim-wasm/src/sandbox/shim/local.rs
index 15f1861bf..92670aca5 100644
--- a/crates/containerd-shim-wasm/src/sandbox/shim/local.rs
+++ b/crates/containerd-shim-wasm/src/sandbox/shim/local.rs
@@ -7,7 +7,7 @@ use std::thread;
 #[cfg(feature = "opentelemetry")]
 use std::time::Duration;
 
-use anyhow::{Context as AnyhowContext, ensure};
+use anyhow::ensure;
 use containerd_shim::api::{
     ConnectRequest, ConnectResponse, CreateTaskRequest, CreateTaskResponse, DeleteRequest, Empty,
     KillRequest, ShutdownRequest, StartRequest, StartResponse, StateRequest, StateResponse,
@@ -252,21 +252,17 @@ impl<T: Instance + Send + Sync, E: EventSender> Local<T, E> {
 
         let id = req.id().to_string();
 
-        thread::Builder::new()
-            .name(format!("{id}-wait"))
-            .spawn(move || {
-                let (exit_code, timestamp) = i.wait().block_on();
-                events.send(TaskExit {
-                    container_id: id.clone(),
-                    exit_status: exit_code,
-                    exited_at: Some(timestamp.to_timestamp()).into(),
-                    pid,
-                    id,
-                    ..Default::default()
-                });
-            })
-            .context("could not spawn thread to wait exit")
-            .map_err(Error::from)?;
+        async move {
+            let (exit_code, timestamp) = i.wait().await;
+            events.send(TaskExit {
+                container_id: id.clone(),
+                exit_status: exit_code,
+                exited_at: Some(timestamp.to_timestamp()).into(),
+                pid,
+                id,
+                ..Default::default()
+            });
+        }.spawn();
 
         debug!("started: {:?}", req);
 
diff --git a/crates/containerd-shim-wasm/src/sys/unix/container/instance.rs b/crates/containerd-shim-wasm/src/sys/unix/container/instance.rs
index 5af5d9612..ab3701e32 100644
--- a/crates/containerd-shim-wasm/src/sys/unix/container/instance.rs
+++ b/crates/containerd-shim-wasm/src/sys/unix/container/instance.rs
@@ -1,6 +1,5 @@
 use std::marker::PhantomData;
 use std::path::Path;
-use std::thread;
 
 use chrono::{DateTime, Utc};
 use libcontainer::container::builder::ContainerBuilder;
@@ -101,11 +100,11 @@ impl<E: Engine + Default> SandboxInstance for Instance<E> {
         self.container.start()?;
 
         let exit_code = self.exit_code.clone();
-        thread::spawn(move || {
-            // move the exit code guard into this thread
+        async move {
+            // move the exit code guard into this task
             let _guard = guard;
 
-            let status = match pidfd.wait().block_on() {
+            let status = match pidfd.wait().await {
                 Ok(WaitStatus::Exited(_, status)) => status,
                 Ok(WaitStatus::Signaled(_, sig, _)) => 128 + sig as i32,
                 Ok(res) => {
@@ -118,7 +117,7 @@ impl<E: Engine + Default> SandboxInstance for Instance<E> {
                 }
             } as u32;
             let _ = exit_code.set((status, Utc::now()));
-        });
+        }.spawn();
 
         Ok(pid as _)
     }
diff --git a/crates/containerd-shim-wasm/src/testing.rs b/crates/containerd-shim-wasm/src/testing.rs
index 625ec59b8..00fa93b77 100644
--- a/crates/containerd-shim-wasm/src/testing.rs
+++ b/crates/containerd-shim-wasm/src/testing.rs
@@ -19,7 +19,7 @@ use oci_spec::runtime::{
     get_default_namespaces,
 };
 
-use crate::sandbox::async_utils::AmbientRuntime;
+use crate::sandbox::async_utils::AmbientRuntime as _;
 use crate::sandbox::{Instance, InstanceConfig};
 
 pub const TEST_NAMESPACE: &str = "runwasi-test";