From aba7eba189e02f1d8f00f772ca08bd1ed84ad862 Mon Sep 17 00:00:00 2001
From: Kevin Wang <wy721@qq.com>
Date: Tue, 29 Aug 2023 07:37:48 +0000
Subject: [PATCH] pRuntime: Fix sidevm stucks at Restart

---
 crates/phala-scheduler/src/task_scheduler.rs | 4 ++++
 crates/sidevm/host-runtime/src/run.rs        | 4 +++-
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/crates/phala-scheduler/src/task_scheduler.rs b/crates/phala-scheduler/src/task_scheduler.rs
index 22eac10366..234542e6ca 100644
--- a/crates/phala-scheduler/src/task_scheduler.rs
+++ b/crates/phala-scheduler/src/task_scheduler.rs
@@ -37,6 +37,10 @@ impl<TaskId: TaskIdType> TaskScheduler<TaskId> {
         self.inner.lock().unwrap().poll_resume(cx, task_id, weight)
     }
 
+    pub fn reset(&self, task_id: &TaskId) {
+        self.exit(task_id)
+    }
+
     pub fn exit(&self, task_id: &TaskId) {
         self.inner.lock().unwrap().exit(task_id)
     }
diff --git a/crates/sidevm/host-runtime/src/run.rs b/crates/sidevm/host-runtime/src/run.rs
index 3125a4464b..f6d74aa07e 100644
--- a/crates/sidevm/host-runtime/src/run.rs
+++ b/crates/sidevm/host-runtime/src/run.rs
@@ -24,7 +24,8 @@ pub struct WasmRun {
 
 impl Drop for WasmRun {
     fn drop(&mut self) {
-        self.env.cleanup()
+        self.env.cleanup();
+        self.scheduler.exit(&self.id);
     }
 }
 
@@ -73,6 +74,7 @@ impl WasmRun {
         env.set_instance(instance);
         env.set_gas_per_breath(gas_per_breath);
         env.set_weight(weight);
+        scheduler.reset(&id);
         Ok((
             WasmRun {
                 env: env.clone(),