rt: add support for non-send closures for thread (un)parking

Add support for non `Send`+`Sync` closures for thread parking and unparking
callbacks when using a `LocalRuntime`. Since a `LocalRuntime` will always
run its tasks on the same thread, its safe to accept a non `Send`+`Sync`
closure.

Signed-off-by: Sanskar Jaiswal <[email protected]>

Author: aryan9600

tokio/src/runtime/builder.rs

@@ -13,6 +13,8 @@ use std::io;
 use std::thread::ThreadId;
 use std::time::Duration;
 
+// use super::LocalOptions;
+
 /// Builds Tokio Runtime with custom configuration values.
 ///
 /// Methods can be chained in order to set the configuration values. The
@@ -923,7 +925,7 @@ impl Builder {
     #[cfg_attr(docsrs, doc(cfg(tokio_unstable)))]
     pub fn build_local(&mut self, options: LocalOptions) -> io::Result<LocalRuntime> {
         match &self.kind {
-            Kind::CurrentThread => self.build_current_thread_local_runtime(),
+            Kind::CurrentThread => self.build_current_thread_local_runtime(options),
             #[cfg(feature = "rt-multi-thread")]
             Kind::MultiThread => panic!("multi_thread is not supported for LocalRuntime"),
         }
@@ -1435,11 +1437,16 @@ impl Builder {
     }
 
     #[cfg(tokio_unstable)]
-    fn build_current_thread_local_runtime(&mut self) -> io::Result<LocalRuntime> {
+    fn build_current_thread_local_runtime(
+        &mut self,
+        opts: LocalOptions,
+    ) -> io::Result<LocalRuntime> {
         use crate::runtime::local_runtime::LocalRuntimeScheduler;
 
         let tid = std::thread::current().id();
 
+        self.before_park = opts.before_park;
+        self.after_unpark = opts.after_unpark;
         let (scheduler, handle, blocking_pool) =
             self.build_current_thread_runtime_components(Some(tid))?;
 

tokio/src/runtime/local_runtime/options.rs

@@ -1,18 +1,144 @@
 use std::marker::PhantomData;
 
+use crate::runtime::Callback;
+
 /// [`LocalRuntime`]-only config options
 ///
-/// Currently, there are no such options, but in the future, things like `!Send + !Sync` hooks may
-/// be added.
-///
 /// Use `LocalOptions::default()` to create the default set of options. This type is used with
 /// [`Builder::build_local`].
 ///
+/// When using [`Builder::build_local`], this overrides any pre-configured options set on the
+/// [`Builder`].
+///
 /// [`Builder::build_local`]: crate::runtime::Builder::build_local
 /// [`LocalRuntime`]: crate::runtime::LocalRuntime
-#[derive(Default, Debug)]
+/// [`Builder`]: crate::runtime::Builder
+#[derive(Default)]
 #[non_exhaustive]
+#[allow(missing_debug_implementations)]
 pub struct LocalOptions {
     /// Marker used to make this !Send and !Sync.
     _phantom: PhantomData<*mut u8>,
+
+    /// Callback for a worker parking itself
+    pub(crate) before_park: Option<Callback>,
+
+    /// Callback for a worker unparking itself
+    pub(crate) after_unpark: Option<Callback>,
+}
+
+impl LocalOptions {
+    /// Executes function `f` just before a thread is parked (goes idle).
+    /// `f` is called within the Tokio context, so functions like [`tokio::spawn`](crate::spawn)
+    /// can be called, and may result in this thread being unparked immediately.
+    ///
+    /// This can be used to start work only when the executor is idle, or for bookkeeping
+    /// and monitoring purposes.
+    ///
+    /// This differs from the `Builder::on_thread_park` method in that it accepts a non Send + Sync
+    /// closure.
+    ///
+    /// Note: There can only be one park callback for a runtime; calling this function
+    /// more than once replaces the last callback defined, rather than adding to it.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use tokio::runtime::{Builder, LocalOptions};
+    /// # pub fn main() {
+    /// let (tx, rx) = std::sync::mpsc::channel();
+    /// let mut opts = LocalOptions::default();
+    /// opts.on_thread_park(move || match rx.recv() {
+    ///    Ok(x) => println!("Received from channel: {}", x),
+    ///    Err(e) => println!("Error receiving from channel: {}", e),
+    /// });
+    ///
+    /// let runtime = Builder::new_current_thread()
+    ///     .enable_time()
+    ///     .build_local(opts)
+    ///     .unwrap();
+    ///
+    /// runtime.block_on(async {
+    ///     tokio::task::spawn_local(async move {
+    ///         tx.send(42).unwrap();
+    ///     });
+    ///     tokio::time::sleep(std::time::Duration::from_millis(1)).await;
+    /// })
+    /// }
+    /// ```
+    pub fn on_thread_park<F>(&mut self, f: F) -> &mut Self
+    where
+        F: Fn() + 'static,
+    {
+        self.before_park = Some(std::sync::Arc::new(to_send_sync(f)));
+        self
+    }
+
+    /// Executes function `f` just after a thread unparks (starts executing tasks).
+    ///
+    /// This is intended for bookkeeping and monitoring use cases; note that work
+    /// in this callback will increase latencies when the application has allowed one or
+    /// more runtime threads to go idle.
+    ///
+    /// This differs from the `Builder::on_thread_unpark` method in that it accepts a non Send + Sync
+    /// closure.
+    ///
+    /// Note: There can only be one unpark callback for a runtime; calling this function
+    /// more than once replaces the last callback defined, rather than adding to it.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # use tokio::runtime::{Builder, LocalOptions};
+    /// # pub fn main() {
+    /// let (tx, rx) = std::sync::mpsc::channel();
+    /// let mut opts = LocalOptions::default();
+    /// opts.on_thread_unpark(move || match rx.recv() {
+    ///    Ok(x) => println!("Received from channel: {}", x),
+    ///    Err(e) => println!("Error receiving from channel: {}", e),
+    /// });
+    ///
+    /// let runtime = Builder::new_current_thread()
+    ///     .enable_time()
+    ///     .build_local(opts)
+    ///     .unwrap();
+    ///
+    /// runtime.block_on(async {
+    ///     tokio::task::spawn_local(async move {
+    ///         tx.send(42).unwrap();
+    ///     });
+    ///     tokio::time::sleep(std::time::Duration::from_millis(1)).await;
+    /// })
+    /// }
+    /// ```
+    pub fn on_thread_unpark<F>(&mut self, f: F) -> &mut Self
+    where
+        F: Fn() + 'static,
+    {
+        self.after_unpark = Some(std::sync::Arc::new(to_send_sync(f)));
+        self
+    }
+}
+
+// A wrapper type to allow non-Send + Sync closures to be used in a Send + Sync context.
+// This is specifically used for executing callbacks when using a `LocalRuntime`.
+struct UnsafeSendSync<T>(T);
+
+// SAFETY: This type is only used in a context where it is guaranteed that the closure will not be
+// sent across threads.
+unsafe impl<T> Send for UnsafeSendSync<T> {}
+unsafe impl<T> Sync for UnsafeSendSync<T> {}
+
+impl<T: Fn()> UnsafeSendSync<T> {
+    fn call(&self) {
+        (self.0)()
+    }
+}
+
+fn to_send_sync<F>(f: F) -> impl Fn() + Send + Sync
+where
+    F: Fn(),
+{
+    let f = UnsafeSendSync(f);
+    move || f.call()
 }

tokio/tests/rt_local.rs

@@ -6,7 +6,7 @@ use tokio::task::spawn_local;
 
 #[test]
 fn test_spawn_local_in_runtime() {
-    let rt = rt();
+    let rt = rt(LocalOptions::default());
 
     let res = rt.block_on(async move {
         let (tx, rx) = tokio::sync::oneshot::channel();
@@ -22,9 +22,34 @@ fn test_spawn_local_in_runtime() {
     assert_eq!(res, 5);
 }
 
+#[test]
+fn test_on_thread_park_in_runtime() {
+    let mut opts = LocalOptions::default();
+    let called = std::rc::Rc::new(std::cell::RefCell::new(false));
+    let cc = called.clone();
+    opts.on_thread_park(move || {
+        *cc.borrow_mut() = true;
+    });
+    let rt = rt(opts);
+
+    rt.block_on(async move {
+        let (tx, rx) = tokio::sync::oneshot::channel();
+
+        spawn_local(async {
+            tokio::task::yield_now().await;
+            tx.send(5).unwrap();
+        });
+
+        // this is not really required execpt to ensure on_thread_park is called
+        rx.await.unwrap()
+    });
+
+    assert!(*called.borrow());
+}
+
 #[test]
 fn test_spawn_from_handle() {
-    let rt = rt();
+    let rt = rt(LocalOptions::default());
 
     let (tx, rx) = tokio::sync::oneshot::channel();
 
@@ -40,7 +65,7 @@ fn test_spawn_from_handle() {
 
 #[test]
 fn test_spawn_local_on_runtime_object() {
-    let rt = rt();
+    let rt = rt(LocalOptions::default());
 
     let (tx, rx) = tokio::sync::oneshot::channel();
 
@@ -56,7 +81,7 @@ fn test_spawn_local_on_runtime_object() {
 
 #[test]
 fn test_spawn_local_from_guard() {
-    let rt = rt();
+    let rt = rt(LocalOptions::default());
 
     let (tx, rx) = tokio::sync::oneshot::channel();
 
@@ -78,7 +103,7 @@ fn test_spawn_from_guard_other_thread() {
     let (tx, rx) = std::sync::mpsc::channel();
 
     std::thread::spawn(move || {
-        let rt = rt();
+        let rt = rt(LocalOptions::default());
         let handle = rt.handle().clone();
 
         tx.send(handle).unwrap();
@@ -98,7 +123,7 @@ fn test_spawn_local_from_guard_other_thread() {
     let (tx, rx) = std::sync::mpsc::channel();
 
     std::thread::spawn(move || {
-        let rt = rt();
+        let rt = rt(LocalOptions::default());
         let handle = rt.handle().clone();
 
         tx.send(handle).unwrap();
@@ -111,9 +136,9 @@ fn test_spawn_local_from_guard_other_thread() {
     spawn_local(async {});
 }
 
-fn rt() -> tokio::runtime::LocalRuntime {
+fn rt(opts: LocalOptions) -> tokio::runtime::LocalRuntime {
     tokio::runtime::Builder::new_current_thread()
         .enable_all()
-        .build_local(LocalOptions::default())
+        .build_local(opts)
         .unwrap()
 }