diff --git a/crates/shim/examples/empty_shim.rs b/crates/shim/examples/empty_shim.rs
index 0e85bdb..1dd2411 100644
--- a/crates/shim/examples/empty_shim.rs
+++ b/crates/shim/examples/empty_shim.rs
@@ -17,11 +17,11 @@
 use containerd_shim as shim;
 
 use log::info;
-use shim::{api, TtrpcContext, TtrpcResult};
+use shim::{api, ExitSignal, TtrpcContext, TtrpcResult};
 
 #[derive(Clone)]
 struct Service {
-    exit: shim::ExitSignal,
+    exit: ExitSignal,
 }
 
 impl shim::Shim for Service {
@@ -33,9 +33,10 @@ impl shim::Shim for Service {
         _namespace: &str,
         _publisher: shim::RemotePublisher,
         _config: &mut shim::Config,
-        exit: shim::ExitSignal,
     ) -> Self {
-        Service { exit }
+        Service {
+            exit: ExitSignal::default(),
+        }
     }
 
     fn start_shim(&mut self, opts: shim::StartOpts) -> Result<String, shim::Error> {
@@ -43,6 +44,10 @@ impl shim::Shim for Service {
         Ok(address)
     }
 
+    fn wait(&mut self) {
+        self.exit.wait();
+    }
+
     fn get_task_service(&self) -> Self::T {
         self.clone()
     }
diff --git a/crates/shim/src/lib.rs b/crates/shim/src/lib.rs
index 20e8a97..6a68abc 100644
--- a/crates/shim/src/lib.rs
+++ b/crates/shim/src/lib.rs
@@ -109,7 +109,7 @@ impl ExitSignal {
     }
 
     /// Wait for the exit signal to be set.
-    fn wait(&self) {
+    pub fn wait(&self) {
         let (lock, cvar) = &*self.0;
         let mut started = lock.lock().unwrap();
         while !*started {
@@ -129,13 +129,7 @@ pub trait Shim {
     type T: Task + Send + Sync;
 
     /// Create a new instance of Shim.
-    fn new(
-        id: &str,
-        namespace: &str,
-        publisher: RemotePublisher,
-        config: &mut Config,
-        exit: ExitSignal,
-    ) -> Self;
+    fn new(id: &str, namespace: &str, publisher: RemotePublisher, config: &mut Config) -> Self;
 
     /// Start shim will be called by containerd when launching new shim instance.
     ///
@@ -149,6 +143,9 @@ pub trait Shim {
         Ok(DeleteResponse::default())
     }
 
+    /// Wait for the shim to exit.
+    fn wait(&mut self);
+
     /// Get the task service object.
     fn get_task_service(&self) -> Self::T;
 }
@@ -176,15 +173,8 @@ where
     let publisher = publisher::RemotePublisher::new(&ttrpc_address)?;
 
     // Create shim instance
-    let exit_signal = ExitSignal::default();
     let mut config = Config::default();
-    let mut shim = T::new(
-        id,
-        &flags.namespace,
-        publisher,
-        &mut config,
-        exit_signal.clone(),
-    );
+    let mut shim = T::new(id, &flags.namespace, publisher, &mut config);
 
     if !config.no_sub_reaper {
         reap::set_subreaper()?;
@@ -237,7 +227,7 @@ where
             server.start()?;
 
             info!("Shim successfully started, waiting for exit signal...");
-            exit_signal.wait();
+            shim.wait();
 
             info!("Shutting down shim instance");
             server.shutdown();