Zero-allocation futures

Get rid of `Arc` in the `poll` methods, replacing with a
`UnparkHandle`. Implement a custom trait object that can clone itself
by limiting the maximum size of the object. Upon `park`, instead of
giving `Arc` references we clone the contents of the `UnparkHandle`,
this generalizes the behaviour of cloning the `Arc`.

Author: leoyvens

src/executor.rs

@@ -7,4 +7,4 @@
 //!
 //! [online]: https://tokio.rs/docs/going-deeper/tasks/
 
-pub use task_impl::{Spawn, spawn, Unpark, Executor, Run};
+pub use task_impl::{Spawn, spawn, Unpark, Executor, Run, UnparkHandle};

src/task_impl/mod.rs

@@ -12,17 +12,19 @@ use future::BoxFuture;
 
 mod unpark_mutex;
 use self::unpark_mutex::UnparkMutex;
-
+mod unpark_handle;
+use self::unpark_handle::UnparkObj;
 mod task_rc;
 mod data;
 #[allow(deprecated)]
 #[cfg(feature = "with-deprecated")]
 pub use self::task_rc::TaskRc;
 pub use self::data::LocalKey;
+pub use self::unpark_handle::UnparkHandle;
 
 struct BorrowedTask<'a> {
     id: usize,
-    unpark: &'a Arc<Unpark>,
+    unpark: UnparkHandle<'a>,
     map: &'a data::LocalMap,
     events: Events,
 }
@@ -81,7 +83,7 @@ fn with<F: FnOnce(&BorrowedTask) -> R, R>(f: F) -> R {
 #[derive(Clone)]
 pub struct Task {
     id: usize,
-    unpark: Arc<Unpark>,
+    unpark: UnparkObj,
     events: Events,
 }
 
@@ -121,7 +123,7 @@ pub fn park() -> Task {
         Task {
             id: task.id,
             events: task.events.clone(),
-            unpark: task.unpark.clone(),
+            unpark: UnparkObj::from(task.unpark),
         }
     })
 }
@@ -333,8 +335,8 @@ impl<F: Future> Spawn<F> {
     /// scheduled to receive a notification when poll can be called again.
     /// Otherwise if `Ready` or `Err` is returned, the `Spawn` task can be
     /// safely destroyed.
-    pub fn poll_future(&mut self, unpark: Arc<Unpark>) -> Poll<F::Item, F::Error> {
-        self.enter(&unpark, |f| f.poll())
+    pub fn poll_future(&mut self, unpark: UnparkHandle) -> Poll<F::Item, F::Error> {
+        self.enter(unpark, |f| f.poll())
     }
 
     /// Waits for the internal future to complete, blocking this thread's
@@ -344,9 +346,11 @@ impl<F: Future> Spawn<F> {
     /// to complete. When a future cannot make progress it will use
     /// `thread::park` to block the current thread.
     pub fn wait_future(&mut self) -> Result<F::Item, F::Error> {
-        let unpark = Arc::new(ThreadUnpark::new(thread::current()));
+        let unpark = Arc::new(ThreadUnpark::new());
+        let mut unpark2 = unpark.clone();
+        let handle = UnparkHandle::new(&mut unpark2);
         loop {
-            match try!(self.poll_future(unpark.clone())) {
+            match try!(self.poll_future(handle)) {
                 Async::NotReady => unpark.park(),
                 Async::Ready(e) => return Ok(e),
             }
@@ -391,17 +395,19 @@ impl<F: Future> Spawn<F> {
 
 impl<S: Stream> Spawn<S> {
     /// Like `poll_future`, except polls the underlying stream.
-    pub fn poll_stream(&mut self, unpark: Arc<Unpark>)
+    pub fn poll_stream(&mut self, unpark: UnparkHandle)
                        -> Poll<Option<S::Item>, S::Error> {
-        self.enter(&unpark, |stream| stream.poll())
+        self.enter(unpark, |stream| stream.poll())
     }
 
     /// Like `wait_future`, except only waits for the next element to arrive on
     /// the underlying stream.
     pub fn wait_stream(&mut self) -> Option<Result<S::Item, S::Error>> {
-        let unpark = Arc::new(ThreadUnpark::new(thread::current()));
+        let unpark = Arc::new(ThreadUnpark::new());
+        let mut unpark2 = unpark.clone();
+        let handle = UnparkHandle::new(&mut unpark2);
         loop {
-            match self.poll_stream(unpark.clone()) {
+            match self.poll_stream(handle) {
                 Ok(Async::NotReady) => unpark.park(),
                 Ok(Async::Ready(Some(e))) => return Some(Ok(e)),
                 Ok(Async::Ready(None)) => return None,
@@ -417,7 +423,7 @@ impl<S: Sink> Spawn<S> {
     /// If the underlying operation returns `NotReady` then the `unpark` value
     /// passed in will receive a notification when the operation is ready to be
     /// attempted again.
-    pub fn start_send(&mut self, value: S::SinkItem, unpark: &Arc<Unpark>)
+    pub fn start_send(&mut self, value: S::SinkItem, unpark: UnparkHandle)
                        -> StartSend<S::SinkItem, S::SinkError> {
         self.enter(unpark, |sink| sink.start_send(value))
     }
@@ -427,7 +433,7 @@ impl<S: Sink> Spawn<S> {
     /// If the underlying operation returns `NotReady` then the `unpark` value
     /// passed in will receive a notification when the operation is ready to be
     /// attempted again.
-    pub fn poll_flush(&mut self, unpark: &Arc<Unpark>)
+    pub fn poll_flush(&mut self, unpark: UnparkHandle)
                        -> Poll<(), S::SinkError> {
         self.enter(unpark, |sink| sink.poll_complete())
     }
@@ -439,10 +445,11 @@ impl<S: Sink> Spawn<S> {
     /// be blocked until it's able to send the value.
     pub fn wait_send(&mut self, mut value: S::SinkItem)
                      -> Result<(), S::SinkError> {
-        let unpark = Arc::new(ThreadUnpark::new(thread::current()));
-        let unpark2 = unpark.clone() as Arc<Unpark>;
+        let unpark = Arc::new(ThreadUnpark::new());
+        let mut unpark2 = unpark.clone();
+        let handle = UnparkHandle::new(&mut unpark2);
         loop {
-            value = match try!(self.start_send(value, &unpark2)) {
+            value = match try!(self.start_send(value, handle)) {
                 AsyncSink::NotReady(v) => v,
                 AsyncSink::Ready => return Ok(()),
             };
@@ -459,10 +466,11 @@ impl<S: Sink> Spawn<S> {
     /// The thread will be blocked until `poll_complete` returns that it's
     /// ready.
     pub fn wait_flush(&mut self) -> Result<(), S::SinkError> {
-        let unpark = Arc::new(ThreadUnpark::new(thread::current()));
-        let unpark2 = unpark.clone() as Arc<Unpark>;
+        let unpark = Arc::new(ThreadUnpark::new());
+        let mut unpark2 = unpark.clone();
+        let handle = UnparkHandle::new(&mut unpark2);
         loop {
-            if try!(self.poll_flush(&unpark2)).is_ready() {
+            if try!(self.poll_flush(handle)).is_ready() {
                 return Ok(())
             }
             unpark.park();
@@ -471,7 +479,7 @@ impl<S: Sink> Spawn<S> {
 }
 
 impl<T> Spawn<T> {
-    fn enter<F, R>(&mut self, unpark: &Arc<Unpark>, f: F) -> R
+    fn enter<F, R>(&mut self, unpark: UnparkHandle, f: F) -> R
         where F: FnOnce(&mut T) -> R
     {
         let task = BorrowedTask {
@@ -501,7 +509,7 @@ impl<T: fmt::Debug> fmt::Debug for Spawn<T> {
 /// `Spawn::poll_stream` functions. It's transitively used as part of the
 /// `Task::unpark` method to internally deliver notifications of readiness of a
 /// future to move forward.
-pub trait Unpark: Send + Sync {
+pub trait Unpark: Send {
     /// Indicates that an associated future and/or task are ready to make
     /// progress.
     ///
@@ -510,6 +518,12 @@ pub trait Unpark: Send + Sync {
     fn unpark(&self);
 }
 
+impl<T: Unpark + Sync> Unpark for Arc<T> {
+    fn unpark(&self) {
+        (**self).unpark();
+    }
+}
+
 /// A trait representing requests to poll futures.
 ///
 /// This trait is an argument to the `Spawn::execute` which is used to run a
@@ -526,9 +540,9 @@ struct ThreadUnpark {
 }
 
 impl ThreadUnpark {
-    fn new(thread: thread::Thread) -> ThreadUnpark {
+    fn new() -> ThreadUnpark {
         ThreadUnpark {
-            thread: thread,
+            thread: thread::current(),
             ready: AtomicBool::new(false),
         }
     }
@@ -563,15 +577,14 @@ impl Run {
     /// Actually run the task (invoking `poll` on its future) on the current
     /// thread.
     pub fn run(self) {
-        let Run { mut spawn, inner } = self;
-
+        let Run { mut spawn, mut inner } = self;
         // SAFETY: the ownership of this `Run` object is evidence that
         // we are in the `POLLING`/`REPOLL` state for the mutex.
         unsafe {
             inner.mutex.start_poll();
 
             loop {
-                match spawn.poll_future(inner.clone()) {
+                match spawn.poll_future(UnparkHandle::new(&mut inner)) {
                     Ok(Async::NotReady) => {}
                     Ok(Async::Ready(())) |
                     Err(()) => return inner.mutex.complete(),

src/task_impl/unpark_handle.rs

@@ -0,0 +1,155 @@
+use core::ptr;
+use core::mem::{forget, size_of};
+use core::marker::PhantomData;
+use super::Unpark;
+
+/// Maximum size in bytes that will fit in a `UnparkObject`.
+/// TODO: What should this value be?
+/// Should we expose this?
+/// We probably want to say that this value may increase but never decrease in a 1.x release.
+const MAX_OBJ_BYTES: usize = 64;
+
+/// A VTable that knows how to clone because the data has a maximum size.
+#[derive(Copy)]
+struct UnparkVtable {
+    unpark: fn(*const ()),
+    clone_to_byte_buffer: fn(*const ()) -> [u8; MAX_OBJ_BYTES],
+    drop_in_place: unsafe fn(*mut ()),
+}
+
+impl Clone for UnparkVtable {
+    fn clone(&self) -> Self {
+        UnparkVtable { ..*self }
+    }
+}
+
+impl UnparkVtable {
+    fn new<T: Unpark + Clone>() -> UnparkVtable {
+        assert!(size_of::<T>() <= MAX_OBJ_BYTES);
+        UnparkVtable {
+            unpark: Self::call_unpark::<T>,
+            clone_to_byte_buffer: Self::clone_to_byte_buffer::<T>,
+            drop_in_place: Self::drop_in_place::<T>,
+        }
+    }
+
+    fn call_unpark<T: Unpark>(data: *const ()) {
+        let downcasted = unsafe { &*(data as *const _ as *const T) };
+        downcasted.unpark()
+    }
+
+    /// Returns array with bytes of the cloned data. Make sure data is shorter than MAX_OBJ_BYTES.
+    /// The caller owns the new data and is responsible for dropping it with `drop_in_place<T>`.
+    fn clone_to_byte_buffer<T: Clone>(data: *const ()) -> [u8; MAX_OBJ_BYTES] {
+        let downcasted = unsafe { &*(data as *const _ as *const T) };
+        let cloned = downcasted.clone();
+        let mut buffer = [0; MAX_OBJ_BYTES];
+        // View cloned and buffer as raw bytes.
+        let cloned_ptr = &cloned as *const _ as *const u8;
+        let buffer_ptr = &mut buffer as *mut _ as *mut u8;
+        // Copy from cloned to the buffer and forget cloned.
+        // Semantically, the buffer now owns cloned.
+        unsafe {
+            ptr::copy_nonoverlapping(cloned_ptr, buffer_ptr, size_of::<T>());
+        }
+        forget(cloned);
+        buffer
+    }
+
+    /// Make sure the original value is forgotten to avoid double free.
+    unsafe fn drop_in_place<T>(data: *mut ()) {
+        ptr::drop_in_place(&mut *(data as *mut _ as *mut T));
+    }
+}
+
+/// `UnparkHandle` is used as an argument to methods like `poll_future`.
+///  Each `Task` handle generated by `task::park` will contain a fresh clone
+///  of the `unpark` argument provided to `UnparkHandle::new`.
+///  Depending on your use case, you may want to wrap your `unpark` value in
+///  an `Arc` before passing it to the `UnparkHandle`.
+///
+/// # Deciding whether to use an `Arc`
+/// If your `unpark` is not `Clone` or has lifetime parameters then you must use an `Arc`.
+/// If you use inner mutability in your `unpark`, then you should carefully
+/// consider what happens when it is cloned and what is the behaviour you want.
+/// Inner mutability aside, the only difference between using or not an `Arc` should be performance.
+/// An `Arc` will cost an allocation upfront and updates an atomic ref count on park,
+/// while no `Arc` has no upfront cost but will cost a clone on park.
+/// The best strategy depends on how often your futures park and how costly it is to clone `unpark`.
+#[derive(Copy, Clone)]
+// Safe to copy because `UnparkHandle` never mutates itself.
+#[allow(missing_debug_implementations)]
+pub struct UnparkHandle<'a> {
+    // Trait object that can clone `data` into an 'UnparkObj` to be put in a `Task`.
+    data: *const (),
+    data_lifetime: PhantomData<&'a ()>,
+    vtable: UnparkVtable,
+}
+
+impl<'a> UnparkHandle<'a> {
+    /// Constructs a `UnparkHandle`.
+    /// A `Task` handle returned by `park` will contain a clone of the `unpark`
+    /// argument provided here. You may want to wrap `unpark` in an `Arc`.
+    ///
+    /// # Panic
+    /// Panics if the size of 'T' is too large. If you get a panic try wrapping the argument
+    /// in an `Arc`.
+
+    // Take unpark as &mut data to make sure it dosen't change under our feet due to inner mutability.
+    // T must be static so we don't clone around references to a stack or something.
+    pub fn new<T: Unpark + Clone + 'static>(unpark: &mut T) -> UnparkHandle<'a> {
+        if size_of::<T>() > MAX_OBJ_BYTES {
+            // TODO: Panicking here seems reasonable and could be a compile time error when we
+            // get a const sytem in Rust. But what about libraries that pass a user supplied type as T?
+            // Should we expose MAX_OBJ_BYTES? Offer a version that return an error?
+            panic!("The size of T is {} bytes which is too large. Try wrapping the unpark argument in an Arc.");
+        }
+
+        UnparkHandle {
+            data: unpark as *const _ as *const (),
+            data_lifetime: PhantomData,
+            vtable: UnparkVtable::new::<T>(),
+        }
+    }
+}
+
+/// A custom trait object that takes ownership of the data as a slice of bytes.
+pub struct UnparkObj {
+    data: [u8; MAX_OBJ_BYTES],
+    vtable: UnparkVtable,
+}
+
+impl Drop for UnparkObj {
+    fn drop(&mut self) {
+        unsafe {
+            (self.vtable.drop_in_place)(&mut self.data as *mut _ as *mut ());
+        }
+    }
+}
+
+impl UnparkObj {
+    fn new(data: *const (), vtable: UnparkVtable) -> Self {
+        UnparkObj {
+            data: (vtable.clone_to_byte_buffer)(data),
+            vtable: vtable,
+        }
+    }
+}
+
+impl Clone for UnparkObj {
+    fn clone(&self) -> Self {
+        UnparkObj::new(&self.data as *const _ as *const (), self.vtable)
+    }
+}
+
+impl<'a> From<UnparkHandle<'a>> for UnparkObj {
+    fn from(handle: UnparkHandle) -> UnparkObj {
+        UnparkObj::new(handle.data, handle.vtable)
+    }
+}
+
+impl Unpark for UnparkObj {
+    fn unpark(&self) {
+        (self.vtable.unpark)(&self.data as *const _ as *const ())
+    }
+}

tests/all.rs

@@ -323,7 +323,7 @@ fn select2() {
         let b = b.map(move |v| { btx.send(v).unwrap(); v });
         let d = d.map(move |v| { dtx.send(v).unwrap(); v });
         let f = b.select(d);
-        drop(executor::spawn(f).poll_future(support::unpark_noop()));
+        drop(executor::spawn(f).poll_future(unpark_noop()));
         assert!(drx.recv().is_err());
         assert!(brx.recv().is_err());
     }