driver: refactor the driver to be stored in a shared set of handles (tokio-rs#196)

This opens up a lot of further cleanup opportunities, especially around
shutdown logic. It also opens up the opportunity for us to eventually
provide a public Handle API.

I am trying to make it so that we don't need to orchestrate our entire
shutdown logic around the thread locals rather than purely the CONTEXT
thread local, which is slightly broken to begin with (I believe that
there are currently issues around creating two runtimes at the same
time).

Now the CONTEXT behaves more like in tokio. When you enter a runtime
context via block_on, we set the context. When we leave that call, we
unset the context.

This also allows us to have a bit more flexibility with things like
sending ops between multiple runtimes as well, so you can now in theory
await ops created from another runtime.

Part of the goal of this is to add new options around things like fixed
or provided buffers. This potentially allows us to work towards things
like transferring buffers between uring runtimes.

More importantly however, it is a crucial first step towards having a
handle to a uring runtime.

More refactoring to come.

Author: Noah-Kennedy

src/buf/fixed/pool.rs

@@ -1,6 +1,8 @@
 use super::handle::CheckedOutBuf;
 use super::{FixedBuf, FixedBuffers};
 
+use crate::runtime::driver::WeakHandle;
+use crate::runtime::CONTEXT;
 use libc::{iovec, UIO_MAXIOV};
 use std::cell::RefCell;
 use std::cmp;
@@ -53,12 +55,12 @@ use std::slice;
 /// # let (memlock_limit, _) = getrlimit(Resource::RLIMIT_MEMLOCK)?;
 /// # let BUF_SIZE_LARGE = memlock_limit as usize / 8;
 /// # let BUF_SIZE_SMALL = memlock_limit as usize / 16;
-/// let pool = FixedBufPool::new(
-///     iter::once(Vec::with_capacity(BUF_SIZE_LARGE))
-///         .chain(iter::repeat_with(|| Vec::with_capacity(BUF_SIZE_SMALL)).take(2))
-/// );
-///
 /// tokio_uring::start(async {
+///     let pool = FixedBufPool::new(
+///          iter::once(Vec::with_capacity(BUF_SIZE_LARGE))
+///              .chain(iter::repeat_with(|| Vec::with_capacity(BUF_SIZE_SMALL)).take(2))
+///      );
+///
 ///     pool.register()?;
 ///
 ///     let buf = pool.try_next(BUF_SIZE_LARGE).unwrap();
@@ -82,6 +84,7 @@ use std::slice;
 #[derive(Clone)]
 pub struct FixedBufPool {
     inner: Rc<RefCell<Inner>>,
+    driver: WeakHandle,
 }
 
 impl FixedBufPool {
@@ -149,6 +152,12 @@ impl FixedBufPool {
     pub fn new(bufs: impl IntoIterator<Item = Vec<u8>>) -> Self {
         FixedBufPool {
             inner: Rc::new(RefCell::new(Inner::new(bufs.into_iter()))),
+            driver: CONTEXT.with(|x| {
+                x.handle()
+                    .as_ref()
+                    .expect("Not in a runtime context")
+                    .into()
+            }),
         }
     }
 
@@ -172,7 +181,10 @@ impl FixedBufPool {
     /// of the `tokio-uring` runtime this call is made in, the function returns
     /// an error.
     pub fn register(&self) -> io::Result<()> {
-        crate::io::register_buffers(Rc::clone(&self.inner) as _)
+        self.driver
+            .upgrade()
+            .expect("Runtime context is no longer present")
+            .register_buffers(Rc::clone(&self.inner) as _)
     }
 
     /// Unregisters this collection of buffers.
@@ -191,7 +203,10 @@ impl FixedBufPool {
     /// an error. Calling `unregister` when no `FixedBufPool` is currently
     /// registered on this runtime also returns an error.
     pub fn unregister(&self) -> io::Result<()> {
-        crate::io::unregister_buffers(Rc::clone(&self.inner) as _)
+        self.driver
+            .upgrade()
+            .expect("Runtime context is no longer present")
+            .unregister_buffers(Rc::clone(&self.inner) as _)
     }
 
     /// Returns a buffer of requested capacity from this pool

src/buf/fixed/registry.rs

@@ -1,6 +1,8 @@
 use super::handle::CheckedOutBuf;
 use super::{FixedBuf, FixedBuffers};
 
+use crate::runtime::driver::WeakHandle;
+use crate::runtime::CONTEXT;
 use libc::{iovec, UIO_MAXIOV};
 use std::cell::RefCell;
 use std::cmp;
@@ -35,6 +37,7 @@ use std::slice;
 #[derive(Clone)]
 pub struct FixedBufRegistry {
     inner: Rc<RefCell<Inner>>,
+    driver: WeakHandle,
 }
 
 impl FixedBufRegistry {
@@ -90,11 +93,10 @@ impl FixedBufRegistry {
     /// # let (memlock_limit, _) = getrlimit(Resource::RLIMIT_MEMLOCK)?;
     /// # let NUM_BUFFERS = std::cmp::max(memlock_limit as usize / 4096 / 8, 1);
     /// # let BUF_SIZE = 4096;
-    /// let registry = FixedBufRegistry::new(
-    ///     iter::repeat_with(|| Vec::with_capacity(BUF_SIZE)).take(NUM_BUFFERS)
-    /// );
-    ///
     /// tokio_uring::start(async {
+    ///     let registry = FixedBufRegistry::new(
+    ///         iter::repeat_with(|| Vec::with_capacity(BUF_SIZE)).take(NUM_BUFFERS)
+    ///     );
     ///     registry.register()?;
     ///     // ...
     ///     Ok(())
@@ -104,6 +106,12 @@ impl FixedBufRegistry {
     pub fn new(bufs: impl IntoIterator<Item = Vec<u8>>) -> Self {
         FixedBufRegistry {
             inner: Rc::new(RefCell::new(Inner::new(bufs.into_iter()))),
+            driver: CONTEXT.with(|x| {
+                x.handle()
+                    .as_ref()
+                    .expect("Not in a runtime context")
+                    .into()
+            }),
         }
     }
 
@@ -127,7 +135,10 @@ impl FixedBufRegistry {
     /// of the `tokio-uring` runtime this call is made in, the function returns
     /// an error.
     pub fn register(&self) -> io::Result<()> {
-        crate::io::register_buffers(Rc::clone(&self.inner) as _)
+        self.driver
+            .upgrade()
+            .expect("Runtime context is no longer present")
+            .register_buffers(Rc::clone(&self.inner) as _)
     }
 
     /// Unregisters this collection of buffers.
@@ -146,7 +157,10 @@ impl FixedBufRegistry {
     /// an error. Calling `unregister` when no `FixedBufRegistry` is currently
     /// registered on this runtime also returns an error.
     pub fn unregister(&self) -> io::Result<()> {
-        crate::io::unregister_buffers(Rc::clone(&self.inner) as _)
+        self.driver
+            .upgrade()
+            .expect("Runtime context is no longer present")
+            .unregister_buffers(Rc::clone(&self.inner) as _)
     }
 
     /// Returns a buffer identified by the specified index for use by the

src/io/accept.rs

@@ -1,6 +1,7 @@
 use crate::io::{SharedFd, Socket};
 use crate::runtime::driver::op;
 use crate::runtime::driver::op::{Completable, Op};
+use crate::runtime::CONTEXT;
 use std::net::SocketAddr;
 use std::{boxed::Box, io};
 
@@ -17,21 +18,23 @@ impl Op<Accept> {
             unsafe { std::mem::zeroed() },
             std::mem::size_of::<libc::sockaddr_storage>() as libc::socklen_t,
         ));
-        Op::submit_with(
-            Accept {
-                fd: fd.clone(),
-                socketaddr,
-            },
-            |accept| {
-                opcode::Accept::new(
-                    types::Fd(accept.fd.raw_fd()),
-                    &mut accept.socketaddr.0 as *mut _ as *mut _,
-                    &mut accept.socketaddr.1,
-                )
-                .flags(libc::O_CLOEXEC)
-                .build()
-            },
-        )
+        CONTEXT.with(|x| {
+            x.handle().expect("Not in a runtime context").submit_op(
+                Accept {
+                    fd: fd.clone(),
+                    socketaddr,
+                },
+                |accept| {
+                    opcode::Accept::new(
+                        types::Fd(accept.fd.raw_fd()),
+                        &mut accept.socketaddr.0 as *mut _ as *mut _,
+                        &mut accept.socketaddr.1,
+                    )
+                    .flags(libc::O_CLOEXEC)
+                    .build()
+                },
+            )
+        })
     }
 }
 

src/io/close.rs

@@ -1,5 +1,6 @@
 use crate::runtime::driver::op;
 use crate::runtime::driver::op::{Completable, Op};
+use crate::runtime::CONTEXT;
 use std::io;
 use std::os::unix::io::RawFd;
 
@@ -11,8 +12,12 @@ impl Op<Close> {
     pub(crate) fn close(fd: RawFd) -> io::Result<Op<Close>> {
         use io_uring::{opcode, types};
 
-        Op::submit_with(Close { fd }, |close| {
-            opcode::Close::new(types::Fd(close.fd)).build()
+        CONTEXT.with(|x| {
+            x.handle()
+                .expect("Not in a runtime context")
+                .submit_op(Close { fd }, |close| {
+                    opcode::Close::new(types::Fd(close.fd)).build()
+                })
         })
     }
 }

src/io/connect.rs

@@ -1,5 +1,6 @@
 use crate::io::SharedFd;
 use crate::runtime::driver::op::{Completable, CqeResult, Op};
+use crate::runtime::CONTEXT;
 use socket2::SockAddr;
 use std::io;
 
@@ -16,20 +17,22 @@ impl Op<Connect> {
     pub(crate) fn connect(fd: &SharedFd, socket_addr: SockAddr) -> io::Result<Op<Connect>> {
         use io_uring::{opcode, types};
 
-        Op::submit_with(
-            Connect {
-                fd: fd.clone(),
-                socket_addr: Box::new(socket_addr),
-            },
-            |connect| {
-                opcode::Connect::new(
-                    types::Fd(connect.fd.raw_fd()),
-                    connect.socket_addr.as_ptr(),
-                    connect.socket_addr.len(),
-                )
-                .build()
-            },
-        )
+        CONTEXT.with(|x| {
+            x.handle().expect("Not in a runtime context").submit_op(
+                Connect {
+                    fd: fd.clone(),
+                    socket_addr: Box::new(socket_addr),
+                },
+                |connect| {
+                    opcode::Connect::new(
+                        types::Fd(connect.fd.raw_fd()),
+                        connect.socket_addr.as_ptr(),
+                        connect.socket_addr.len(),
+                    )
+                    .build()
+                },
+            )
+        })
     }
 }
 

src/io/fsync.rs

@@ -2,6 +2,7 @@ use std::io;
 
 use crate::io::SharedFd;
 use crate::runtime::driver::op::{Completable, CqeResult, Op};
+use crate::runtime::CONTEXT;
 use io_uring::{opcode, types};
 
 pub(crate) struct Fsync {
@@ -10,16 +11,25 @@ pub(crate) struct Fsync {
 
 impl Op<Fsync> {
     pub(crate) fn fsync(fd: &SharedFd) -> io::Result<Op<Fsync>> {
-        Op::submit_with(Fsync { fd: fd.clone() }, |fsync| {
-            opcode::Fsync::new(types::Fd(fsync.fd.raw_fd())).build()
+        CONTEXT.with(|x| {
+            x.handle()
+                .expect("Not in a runtime context")
+                .submit_op(Fsync { fd: fd.clone() }, |fsync| {
+                    opcode::Fsync::new(types::Fd(fsync.fd.raw_fd())).build()
+                })
         })
     }
 
     pub(crate) fn datasync(fd: &SharedFd) -> io::Result<Op<Fsync>> {
-        Op::submit_with(Fsync { fd: fd.clone() }, |fsync| {
-            opcode::Fsync::new(types::Fd(fsync.fd.raw_fd()))
-                .flags(types::FsyncFlags::DATASYNC)
-                .build()
+        CONTEXT.with(|x| {
+            x.handle().expect("Not in a runtime context").submit_op(
+                Fsync { fd: fd.clone() },
+                |fsync| {
+                    opcode::Fsync::new(types::Fd(fsync.fd.raw_fd()))
+                        .flags(types::FsyncFlags::DATASYNC)
+                        .build()
+                },
+            )
         })
     }
 }

src/io/mod.rs

@@ -18,9 +18,6 @@ mod read_fixed;
 
 mod readv;
 
-mod register_buffers;
-pub(crate) use register_buffers::{register_buffers, unregister_buffers};
-
 mod recv_from;
 
 mod rename_at;

src/io/noop.rs

@@ -1,4 +1,5 @@
 use crate::runtime::driver::op::{Completable, CqeResult, Op};
+use crate::runtime::CONTEXT;
 use std::io;
 
 /// No operation. Just posts a completion event, nothing else.
@@ -10,7 +11,11 @@ impl Op<NoOp> {
     pub fn no_op() -> io::Result<Op<NoOp>> {
         use io_uring::opcode;
 
-        Op::submit_with(NoOp {}, |_| opcode::Nop::new().build())
+        CONTEXT.with(|x| {
+            x.handle()
+                .expect("Not in a runtime context")
+                .submit_op(NoOp {}, |_| opcode::Nop::new().build())
+        })
     }
 }
 

src/io/open.rs

@@ -2,6 +2,7 @@ use crate::fs::{File, OpenOptions};
 use crate::io::SharedFd;
 
 use crate::runtime::driver::op::{Completable, CqeResult, Op};
+use crate::runtime::CONTEXT;
 use std::ffi::CString;
 use std::io;
 use std::path::Path;
@@ -23,16 +24,20 @@ impl Op<Open> {
             | options.creation_mode()?
             | (options.custom_flags & !libc::O_ACCMODE);
 
-        Op::submit_with(Open { path, flags }, |open| {
-            // Get a reference to the memory. The string will be held by the
-            // operation state and will not be accessed again until the operation
-            // completes.
-            let p_ref = open.path.as_c_str().as_ptr();
+        CONTEXT.with(|x| {
+            x.handle()
+                .expect("Not in a runtime context")
+                .submit_op(Open { path, flags }, |open| {
+                    // Get a reference to the memory. The string will be held by the
+                    // operation state and will not be accessed again until the operation
+                    // completes.
+                    let p_ref = open.path.as_c_str().as_ptr();
 
-            opcode::OpenAt::new(types::Fd(libc::AT_FDCWD), p_ref)
-                .flags(flags)
-                .mode(options.mode)
-                .build()
+                    opcode::OpenAt::new(types::Fd(libc::AT_FDCWD), p_ref)
+                        .flags(flags)
+                        .mode(options.mode)
+                        .build()
+                })
         })
     }
 }