diff --git a/lib/wasix/src/runtime/task_manager/tokio.rs b/lib/wasix/src/runtime/task_manager/tokio.rs index b9f2a4c098c..6118717b7b8 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(()) }) } @@ -210,3 +206,35 @@ impl VirtualTaskManager for TokioTaskManager { .unwrap_or(8)) } } + +// Used by [`VirtualTaskManager::sleep_now`] to abort a sleep task when drop. +#[derive(Default)] +struct SleepNow { + abort_handle: Option, +} + +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 Drop for SleepNow { + fn drop(&mut self) { + if let Some(h) = self.abort_handle.as_ref() { + h.abort() + } + } +}