diff --git a/src/runtime/mod.rs b/src/runtime/mod.rs index f6550640..f535e69e 100644 --- a/src/runtime/mod.rs +++ b/src/runtime/mod.rs @@ -2,6 +2,7 @@ use crate::driver::Driver; use std::future::Future; use std::io; +use std::mem::ManuallyDrop; use std::os::unix::io::{AsRawFd, RawFd}; use tokio::io::unix::AsyncFd; use tokio::task::LocalSet; @@ -20,19 +21,10 @@ pub struct Runtime { uring_fd: RawFd, /// LocalSet for !Send tasks - local: LocalSet, + local: ManuallyDrop, /// Tokio runtime, always current-thread - rt: tokio::runtime::Runtime, - - /// This is here for drop order reasons. - /// - /// We can't unset the driver in the runtime drop method, because the inner runtime needs to - /// be dropped first so that there are no tasks running. - /// - /// The rust drop order rules guarantee that the inner runtime will be dropped first, and - /// this last. - _driver_guard: RuntimeContextGuard, + rt: ManuallyDrop, } /// Spawns a new asynchronous task, returning a [`JoinHandle`] for it. @@ -77,7 +69,9 @@ impl Runtime { .enable_all() .build()?; - let local = LocalSet::new(); + let rt = ManuallyDrop::new(rt); + + let local = ManuallyDrop::new(LocalSet::new()); let driver = Driver::new(b)?; @@ -89,7 +83,6 @@ impl Runtime { uring_fd: driver_fd, local, rt, - _driver_guard: RuntimeContextGuard, }) } @@ -124,10 +117,16 @@ impl Runtime { } } -struct RuntimeContextGuard; - -impl Drop for RuntimeContextGuard { +impl Drop for Runtime { fn drop(&mut self) { + // drop tasks + unsafe { + ManuallyDrop::drop(&mut self.local); + ManuallyDrop::drop(&mut self.rt); + } + + // once tasks are dropped, we can unset the driver + // this will block until all completions are received CONTEXT.with(|rc| rc.unset_driver()) } }