[runtime] Add spawn_blocking_ref to Spawner trait (#957)

Author: patrick-ogrady

runtime/src/deterministic.rs

@@ -814,6 +814,45 @@ impl crate::Spawner for Context {
         handle
     }
 
+    fn spawn_blocking_ref<F, T>(&mut self, dedicated: bool) -> impl FnOnce(F) -> Handle<T> + 'static
+    where
+        F: FnOnce() -> T + Send + 'static,
+        T: Send + 'static,
+    {
+        // Ensure a context only spawns one task
+        assert!(!self.spawned, "already spawned");
+        self.spawned = true;
+
+        // Get metrics
+        let label = if dedicated {
+            Label::blocking_dedicated(self.name.clone())
+        } else {
+            Label::blocking_shared(self.name.clone())
+        };
+        self.executor
+            .metrics
+            .tasks_spawned
+            .get_or_create(&label)
+            .inc();
+        let gauge = self
+            .executor
+            .metrics
+            .tasks_running
+            .get_or_create(&label)
+            .clone();
+
+        // Set up the task
+        let executor = self.executor.clone();
+        move |f: F| {
+            let (f, handle) = Handle::init_blocking(f, gauge, false);
+
+            // Spawn the task
+            let f = async move { f() };
+            Tasks::register_work(&executor.tasks, label, Box::pin(f));
+            handle
+        }
+    }
+
     fn stop(&self, value: i32) {
         self.executor.auditor.event(b"stop", |hasher| {
             hasher.update(value.to_be_bytes());

runtime/src/lib.rs

@@ -123,11 +123,7 @@ pub trait Spawner: Clone + Send + Sync + 'static {
 
     /// Enqueue a task to be executed (without consuming the context).
     ///
-    /// Unlike a future, a spawned task will start executing immediately (even if the caller
-    /// does not await the handle).
-    ///
-    /// In some cases, it may be useful to spawn a task without consuming the context (e.g. starting
-    /// an actor that already has a reference to context).
+    /// The semantics are the same as [Spawner::spawn].
     ///
     /// # Warning
     ///
@@ -161,6 +157,22 @@ pub trait Spawner: Clone + Send + Sync + 'static {
         F: FnOnce(Self) -> T + Send + 'static,
         T: Send + 'static;
 
+    /// Enqueue a blocking task to be executed (without consuming the context).
+    ///
+    /// The semantics are the same as [Spawner::spawn_blocking].
+    ///
+    /// # Warning
+    ///
+    /// If this function is used to spawn multiple tasks from the same context,
+    /// the runtime will panic to prevent accidental misuse.
+    fn spawn_blocking_ref<F, T>(
+        &mut self,
+        dedicated: bool,
+    ) -> impl FnOnce(F) -> Handle<T> + 'static
+    where
+        F: FnOnce() -> T + Send + 'static,
+        T: Send + 'static;
+
     /// Signals the runtime to stop execution and that all outstanding tasks
     /// should perform any required cleanup and exit. This method is idempotent and
     /// can be called multiple times.
@@ -933,6 +945,32 @@ mod tests {
         });
     }
 
+    fn test_spawn_blocking_ref<R: Runner>(runner: R, dedicated: bool)
+    where
+        R::Context: Spawner,
+    {
+        runner.start(|mut context| async move {
+            let spawn = context.spawn_blocking_ref(dedicated);
+            let handle = spawn(|| 42);
+            let result = handle.await;
+            assert!(matches!(result, Ok(42)));
+        });
+    }
+
+    fn test_spawn_blocking_ref_duplicate<R: Runner>(runner: R, dedicated: bool)
+    where
+        R::Context: Spawner,
+    {
+        runner.start(|mut context| async move {
+            let spawn = context.spawn_blocking_ref(dedicated);
+            let result = spawn(|| 42).await;
+            assert!(matches!(result, Ok(42)));
+
+            // Ensure context is consumed
+            context.spawn_blocking(dedicated, |_| 42);
+        });
+    }
+
     fn test_spawn_blocking_abort<R: Runner>(runner: R, dedicated: bool)
     where
         R::Context: Spawner,
@@ -1157,6 +1195,23 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_deterministic_spawn_blocking_ref() {
+        for dedicated in [false, true] {
+            let executor = deterministic::Runner::default();
+            test_spawn_blocking_ref(executor, dedicated);
+        }
+    }
+
+    #[test]
+    #[should_panic]
+    fn test_deterministic_spawn_blocking_ref_duplicate() {
+        for dedicated in [false, true] {
+            let executor = deterministic::Runner::default();
+            test_spawn_blocking_ref_duplicate(executor, dedicated);
+        }
+    }
+
     #[test]
     fn test_deterministic_metrics() {
         let executor = deterministic::Runner::default();
@@ -1311,6 +1366,23 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_tokio_spawn_blocking_ref() {
+        for dedicated in [false, true] {
+            let executor = tokio::Runner::default();
+            test_spawn_blocking_ref(executor, dedicated);
+        }
+    }
+
+    #[test]
+    #[should_panic]
+    fn test_tokio_spawn_blocking_ref_duplicate() {
+        for dedicated in [false, true] {
+            let executor = tokio::Runner::default();
+            test_spawn_blocking_ref_duplicate(executor, dedicated);
+        }
+    }
+
     #[test]
     fn test_tokio_metrics() {
         let executor = tokio::Runner::default();

runtime/src/tokio/runtime.rs

@@ -429,7 +429,7 @@ impl crate::Spawner for Context {
             .get_or_create(&label)
             .clone();
 
-        // Initialize the blocking task using the new function
+        // Set up the task
         let executor = self.executor.clone();
         let (f, handle) = Handle::init_blocking(|| f(self), gauge, executor.cfg.catch_panics);
 
@@ -442,6 +442,48 @@ impl crate::Spawner for Context {
         handle
     }
 
+    fn spawn_blocking_ref<F, T>(&mut self, dedicated: bool) -> impl FnOnce(F) -> Handle<T> + 'static
+    where
+        F: FnOnce() -> T + Send + 'static,
+        T: Send + 'static,
+    {
+        // Ensure a context only spawns one task
+        assert!(!self.spawned, "already spawned");
+        self.spawned = true;
+
+        // Get metrics
+        let label = if dedicated {
+            Label::blocking_dedicated(self.name.clone())
+        } else {
+            Label::blocking_shared(self.name.clone())
+        };
+        self.executor
+            .metrics
+            .tasks_spawned
+            .get_or_create(&label)
+            .inc();
+        let gauge = self
+            .executor
+            .metrics
+            .tasks_running
+            .get_or_create(&label)
+            .clone();
+
+        // Set up the task
+        let executor = self.executor.clone();
+        move |f: F| {
+            let (f, handle) = Handle::init_blocking(f, gauge, executor.cfg.catch_panics);
+
+            // Spawn the blocking task
+            if dedicated {
+                std::thread::spawn(f);
+            } else {
+                executor.runtime.spawn_blocking(f);
+            }
+            handle
+        }
+    }
+
     fn stop(&self, value: i32) {
         self.executor.signaler.lock().unwrap().signal(value);
     }