From e2b44df850260fefc3541a9be465f44a266626de Mon Sep 17 00:00:00 2001
From: kuangmingfu <kuang.mingfu@zte.com.cn>
Date: Tue, 30 Jul 2024 19:02:31 +0800
Subject: [PATCH] add exec delete timeout

Signed-off-by: kuangmingfu <kuang.mingfu@zte.com.cn>
---
 crates/runc-shim/src/main.rs        |  3 ++-
 crates/runc-shim/src/runc.rs        |  3 +++
 crates/shim/src/asynchronous/mod.rs | 18 +++++++++++++++++-
 3 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/crates/runc-shim/src/main.rs b/crates/runc-shim/src/main.rs
index 2998f2f8..492d3422 100644
--- a/crates/runc-shim/src/main.rs
+++ b/crates/runc-shim/src/main.rs
@@ -15,8 +15,9 @@
 */
 
 use std::env;
-use tokio::runtime::Builder;
+
 use containerd_shim::{asynchronous::run, parse};
+use tokio::runtime::Builder;
 
 mod cgroup_memory;
 mod common;
diff --git a/crates/runc-shim/src/runc.rs b/crates/runc-shim/src/runc.rs
index 92bb0f23..fe2a9bed 100644
--- a/crates/runc-shim/src/runc.rs
+++ b/crates/runc-shim/src/runc.rs
@@ -457,6 +457,9 @@ impl ProcessLifecycle<ExecProcess> for RuncExecLifecycle {
 
     async fn delete(&self, p: &mut ExecProcess) -> Result<()> {
         self.exit_signal.signal();
+        self.exit_signal
+            .wait_for_exit(tokio::time::Duration::from_secs(2))
+            .await;
         let exec_pid_path = Path::new(self.bundle.as_str()).join(format!("{}.pid", p.id));
         remove_file(exec_pid_path).await.unwrap_or_default();
         Ok(())
diff --git a/crates/shim/src/asynchronous/mod.rs b/crates/shim/src/asynchronous/mod.rs
index 4a145310..27d720ef 100644
--- a/crates/shim/src/asynchronous/mod.rs
+++ b/crates/shim/src/asynchronous/mod.rs
@@ -47,7 +47,7 @@ use nix::{
     unistd::Pid,
 };
 use signal_hook_tokio::Signals;
-use tokio::{io::AsyncWriteExt, sync::Notify};
+use tokio::{self, io::AsyncWriteExt, sync::Notify, time::Duration};
 
 use crate::{
     args,
@@ -237,6 +237,22 @@ impl ExitSignal {
             notified.await;
         }
     }
+
+    /// Wait for the exit signal to be set or return after a timeout.
+    pub async fn wait_for_exit(&self, timeout_duration: Duration) {
+        let timeout_task = async {
+            tokio::time::sleep(timeout_duration).await;
+        };
+
+        let exit_task = async {
+            self.wait().await;
+        };
+
+        tokio::select! {
+            _ = timeout_task => {},
+            _ = exit_task => {},
+        }
+    }
 }
 
 /// Spawn is a helper func to launch shim process asynchronously.