Add spawn_local (#176)

Author: sarsko

Cargo.toml

@@ -36,6 +36,7 @@ regex = "1.5.5"
 tempfile = "3.2.0"
 test-log = { version = "0.2.8", default-features = false, features = ["trace"] }
 tracing-subscriber = { version = "0.3.9", features = ["env-filter"] }
+trybuild = "1.0"
 pin-project = "1.1.3"
 
 [lib]

src/future/mod.rs

@@ -18,11 +18,10 @@ use std::task::{Context, Poll, Waker};
 
 pub mod batch_semaphore;
 
-/// Spawn a new async task that the executor will run to completion.
-pub fn spawn<T, F>(fut: F) -> JoinHandle<T>
+fn spawn_inner<F>(fut: F) -> JoinHandle<F::Output>
 where
-    F: Future<Output = T> + Send + 'static,
-    T: Send + 'static,
+    F: Future + 'static,
+    F::Output: 'static,
 {
     let stack_size = ExecutionState::with(|s| s.config.stack_size);
     let inner = Arc::new(std::sync::Mutex::new(JoinHandleInner::default()));
@@ -33,6 +32,25 @@ where
     JoinHandle { task_id, inner }
 }
 
+/// Spawn a new async task that the executor will run to completion.
+pub fn spawn<F>(fut: F) -> JoinHandle<F::Output>
+where
+    F: Future + Send + 'static,
+    F::Output: Send + 'static,
+{
+    spawn_inner(fut)
+}
+
+/// Spawn a new async task that the executor will run to completion.
+/// This is just `spawn` without the `Send` bound, and it mirrors `spawn_local` from Tokio.
+pub fn spawn_local<F>(fut: F) -> JoinHandle<F::Output>
+where
+    F: Future + 'static,
+    F::Output: 'static,
+{
+    spawn_inner(fut)
+}
+
 /// An owned permission to join on an async task (await its termination).
 #[derive(Debug)]
 pub struct JoinHandle<T> {
@@ -120,27 +138,28 @@ impl<T> Future for JoinHandle<T> {
 // contains a mutex-wrapped field that stores the value and the waker for the task
 // waiting on the join handle. When `poll` returns `Poll::Ready`, the `Wrapper` stores
 // the result in the `result` field and wakes the `waker`.
-struct Wrapper<T, F> {
+struct Wrapper<F: Future> {
     future: Pin<Box<F>>,
-    inner: std::sync::Arc<std::sync::Mutex<JoinHandleInner<T>>>,
+    inner: std::sync::Arc<std::sync::Mutex<JoinHandleInner<F::Output>>>,
 }
 
-impl<T, F> Wrapper<T, F>
+impl<F> Wrapper<F>
 where
-    F: Future<Output = T> + Send + 'static,
+    F: Future + 'static,
+    F::Output: 'static,
 {
-    fn new(future: F, inner: std::sync::Arc<std::sync::Mutex<JoinHandleInner<T>>>) -> Self {
+    fn new(future: F, inner: std::sync::Arc<std::sync::Mutex<JoinHandleInner<F::Output>>>) -> Self {
         Self {
             future: Box::pin(future),
             inner,
         }
     }
 }
 
-impl<T, F> Future for Wrapper<T, F>
+impl<F> Future for Wrapper<F>
 where
-    F: Future<Output = T> + Send + 'static,
-    T: Send + 'static,
+    F: Future + 'static,
+    F::Output: 'static,
 {
     type Output = ();
 

src/runtime/execution.rs

@@ -372,7 +372,7 @@ impl ExecutionState {
     /// if it wants to give the new task a chance to run immediately.
     pub(crate) fn spawn_future<F>(future: F, stack_size: usize, name: Option<String>) -> TaskId
     where
-        F: Future<Output = ()> + Send + 'static,
+        F: Future<Output = ()> + 'static,
     {
         let task_id = Self::with(|state| {
             let schedule_len = state.current_schedule.len();

src/runtime/task/mod.rs

@@ -206,7 +206,7 @@ impl Task {
         parent_task_id: Option<TaskId>,
     ) -> Self
     where
-        F: FnOnce() + Send + 'static,
+        F: FnOnce() + 'static,
     {
         assert!(id.0 < clock.time.len());
         let mut continuation = ContinuationPool::acquire(stack_size);
@@ -288,7 +288,7 @@ impl Task {
         parent_task_id: Option<TaskId>,
     ) -> Self
     where
-        F: Future<Output = ()> + Send + 'static,
+        F: Future<Output = ()> + 'static,
     {
         let mut future = Box::pin(future);
 

src/runtime/thread/continuation.rs

@@ -32,11 +32,11 @@ pub(crate) struct Continuation {
 /// A cell to pass functions into continuations
 #[allow(clippy::type_complexity)]
 #[derive(Clone)]
-struct ContinuationFunction(Rc<Cell<Option<Box<dyn FnOnce() + Send>>>>);
+struct ContinuationFunction(Rc<Cell<Option<Box<dyn FnOnce()>>>>);
 
 // Safety: we arrange for the `function` field of `Continuation` to only be accessed by one thread
 // at a time: Shuttle tests are single threaded, and continuations are never shared across threads
-// by the ContinuationPool, which is thread-local. The function itself already implements `Send`.
+// by the ContinuationPool, which is thread-local.
 unsafe impl Send for ContinuationFunction {}
 
 /// Inputs that we can pass to a continuation.
@@ -104,7 +104,7 @@ impl Continuation {
 
     /// Provide a new function for the continuation to execute. The continuation must
     /// be in reusable state.
-    pub fn initialize(&mut self, fun: Box<dyn FnOnce() + Send>) {
+    pub fn initialize(&mut self, fun: Box<dyn FnOnce()>) {
         debug_assert_eq!(
             self.state,
             ContinuationState::NotReady,

tests/basic/thread.rs

@@ -1,7 +1,8 @@
 use shuttle::current::{get_name_for_task, me};
 use shuttle::sync::{Barrier, Condvar, Mutex};
-use shuttle::{check_dfs, check_random, thread};
+use shuttle::{check_dfs, check_random, future, thread};
 use std::collections::HashSet;
+use std::rc::Rc;
 use std::sync::atomic::{AtomicBool, AtomicU8, Ordering};
 use std::sync::Arc;
 use test_log::test;
@@ -642,3 +643,14 @@ fn thread_unpark_after_spurious_wakeup() {
         None,
     )
 }
+
+#[test]
+fn spawn_local_sanity() {
+    check_dfs(
+        || {
+            let rc = Rc::new(0);
+            shuttle::future::block_on(future::spawn_local(async { drop(rc) })).unwrap()
+        },
+        None,
+    );
+}

tests/mod.rs

@@ -5,6 +5,12 @@ mod data;
 mod demo;
 mod future;
 
+#[test]
+fn ui() {
+    let t = trybuild::TestCases::new();
+    t.compile_fail("tests/ui/*.rs");
+}
+
 use shuttle::scheduler::{ReplayScheduler, Scheduler};
 use shuttle::{check_random_with_seed, replay_from_file, Config, FailurePersistence, Runner};
 use std::panic::{self, RefUnwindSafe, UnwindSafe};

tests/ui/spawn_not_send.rs

@@ -0,0 +1,13 @@
+use shuttle::future;
+use shuttle::check_dfs;
+use std::rc::Rc;
+
+fn main() {
+    check_dfs(
+        || {
+            let rc = Rc::new(0);
+            shuttle::future::block_on(future::spawn(async { drop(rc) })).unwrap()
+        },
+        None,
+    );
+}

tests/ui/spawn_not_send.stderr

@@ -0,0 +1,20 @@
+error: future cannot be sent between threads safely
+ --> tests/ui/spawn_not_send.rs:9:53
+  |
+9 |             shuttle::future::block_on(future::spawn(async { drop(rc) })).unwrap()
+  |                                                     ^^^^^^^^^^^^^^^^^^ future created by async block is not `Send`
+  |
+  = help: within `{async block@$DIR/tests/ui/spawn_not_send.rs:9:53: 9:58}`, the trait `Send` is not implemented for `Rc<i32>`
+note: captured value is not `Send`
+ --> tests/ui/spawn_not_send.rs:9:66
+  |
+9 |             shuttle::future::block_on(future::spawn(async { drop(rc) })).unwrap()
+  |                                                                  ^^ has type `Rc<i32>` which is not `Send`
+note: required by a bound in `shuttle::future::spawn`
+ --> src/future/mod.rs
+  |
+  | pub fn spawn<F>(fut: F) -> JoinHandle<F::Output>
+  |        ----- required by a bound in this function
+  | where
+  |     F: Future + Send + 'static,
+  |                 ^^^^ required by this bound in `spawn`