From 85385811db1dba91ab6a0360ddb1feba44a057e0 Mon Sep 17 00:00:00 2001 From: xuanyuan300 Date: Sat, 9 Dec 2023 22:14:27 +0800 Subject: [PATCH 1/3] Optimize sleep_now for TokioTaskManager Signed-off-by: xuanyuan300 --- lib/wasix/src/runtime/task_manager/tokio.rs | 51 +++++++++++++++++---- 1 file changed, 41 insertions(+), 10 deletions(-) diff --git a/lib/wasix/src/runtime/task_manager/tokio.rs b/lib/wasix/src/runtime/task_manager/tokio.rs index b9f2a4c098c..67e7ce916f0 100644 --- a/lib/wasix/src/runtime/task_manager/tokio.rs +++ b/lib/wasix/src/runtime/task_manager/tokio.rs @@ -117,17 +117,13 @@ impl<'g> Drop for TokioRuntimeGuard<'g> { impl VirtualTaskManager for TokioTaskManager { /// See [`VirtualTaskManager::sleep_now`]. fn sleep_now(&self, time: Duration) -> Pin + Send + Sync>> { - let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel(); - self.rt.handle().spawn(async move { - if time == Duration::ZERO { - tokio::task::yield_now().await; - } else { - tokio::time::sleep(time).await; - } - tx.send(()).ok(); - }); + let handle = self.runtime_handle(); Box::pin(async move { - rx.recv().await; + SleepNow::default() + .enter(handle, time) + .await + .ok() + .unwrap_or_else(|| ()) }) } @@ -210,3 +206,38 @@ impl VirtualTaskManager for TokioTaskManager { .unwrap_or(8)) } } + +// Used by [`VirtualTaskManager::sleep_now`] to abort a sleep task when drop. +struct SleepNow { + abort_handle: Option<::tokio::task::AbortHandle>, +} + +impl SleepNow { + async fn enter( + &mut self, + handle: ::tokio::runtime::Handle, + time: Duration, + ) -> Result<(), ::tokio::task::JoinError> { + let handle = handle.spawn(async move { + if time == Duration::ZERO { + ::tokio::task::yield_now().await; + } else { + ::tokio::time::sleep(time).await; + } + }); + self.abort_handle = Some(handle.abort_handle()); + handle.await + } +} + +impl Default for SleepNow { + fn default() -> Self { + Self { abort_handle: None } + } +} + +impl Drop for SleepNow { + fn drop(&mut self) { + self.abort_handle.as_ref().map(|h| h.abort()); + } +} From 2bac4bc84ff2a764d0697cc404f32f965a748d98 Mon Sep 17 00:00:00 2001 From: xuanyuan300 Date: Sat, 9 Dec 2023 22:22:23 +0800 Subject: [PATCH 2/3] style fmt Signed-off-by: xuanyuan300 --- lib/wasix/src/runtime/task_manager/tokio.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/lib/wasix/src/runtime/task_manager/tokio.rs b/lib/wasix/src/runtime/task_manager/tokio.rs index 67e7ce916f0..ceee20435d9 100644 --- a/lib/wasix/src/runtime/task_manager/tokio.rs +++ b/lib/wasix/src/runtime/task_manager/tokio.rs @@ -209,20 +209,20 @@ impl VirtualTaskManager for TokioTaskManager { // Used by [`VirtualTaskManager::sleep_now`] to abort a sleep task when drop. struct SleepNow { - abort_handle: Option<::tokio::task::AbortHandle>, + abort_handle: Option, } impl SleepNow { async fn enter( &mut self, - handle: ::tokio::runtime::Handle, + handle: tokio::runtime::Handle, time: Duration, - ) -> Result<(), ::tokio::task::JoinError> { + ) -> Result<(), tokio::task::JoinError> { let handle = handle.spawn(async move { if time == Duration::ZERO { - ::tokio::task::yield_now().await; + tokio::task::yield_now().await; } else { - ::tokio::time::sleep(time).await; + tokio::time::sleep(time).await; } }); self.abort_handle = Some(handle.abort_handle()); From 7b3f09f17c427ceac738006042bcdbdbfe909970 Mon Sep 17 00:00:00 2001 From: xuanyuan300 Date: Sat, 9 Dec 2023 23:11:51 +0800 Subject: [PATCH 3/3] fix cargo clippy Signed-off-by: xuanyuan300 --- lib/wasix/src/runtime/task_manager/tokio.rs | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/lib/wasix/src/runtime/task_manager/tokio.rs b/lib/wasix/src/runtime/task_manager/tokio.rs index ceee20435d9..6118717b7b8 100644 --- a/lib/wasix/src/runtime/task_manager/tokio.rs +++ b/lib/wasix/src/runtime/task_manager/tokio.rs @@ -123,7 +123,7 @@ impl VirtualTaskManager for TokioTaskManager { .enter(handle, time) .await .ok() - .unwrap_or_else(|| ()) + .unwrap_or(()) }) } @@ -208,6 +208,7 @@ impl VirtualTaskManager for TokioTaskManager { } // Used by [`VirtualTaskManager::sleep_now`] to abort a sleep task when drop. +#[derive(Default)] struct SleepNow { abort_handle: Option, } @@ -230,14 +231,10 @@ impl SleepNow { } } -impl Default for SleepNow { - fn default() -> Self { - Self { abort_handle: None } - } -} - impl Drop for SleepNow { fn drop(&mut self) { - self.abort_handle.as_ref().map(|h| h.abort()); + if let Some(h) = self.abort_handle.as_ref() { + h.abort() + } } }