Add Waker::wake_by_ref and make Waker::wake consume the Waker

Author: cramertj

src/libcore/task/wake.rs

@@ -72,10 +72,18 @@ pub struct RawWakerVTable {
     /// This function will be called when `wake` is called on the [`Waker`].
     /// It must wake up the task associated with this [`RawWaker`].
     ///
-    /// The implemention of this function must not consume the provided data
-    /// pointer.
+    /// The implementation of this function must make sure to release any
+    /// resources that are associated with this instance of a [`RawWaker`] and
+    /// associated task.
     wake: unsafe fn(*const ()),
 
+    /// This function will be called when `wake_by_ref` is called on the [`Waker`].
+    /// It must wake up the task associated with this [`RawWaker`].
+    ///
+    /// This function is similar to `wake`, but must not consume the provided data
+    /// pointer.
+    wake_by_ref: unsafe fn(*const ()),
+
     /// This function gets called when a [`RawWaker`] gets dropped.
     ///
     /// The implementation of this function must make sure to release any
@@ -85,8 +93,8 @@ pub struct RawWakerVTable {
 }
 
 impl RawWakerVTable {
-    /// Creates a new `RawWakerVTable` from the provided `clone`, `wake`, and
-    /// `drop` functions.
+    /// Creates a new `RawWakerVTable` from the provided `clone`, `wake`,
+    /// `wake_by_ref`, and `drop` functions.
     ///
     /// # `clone`
     ///
@@ -103,7 +111,16 @@ impl RawWakerVTable {
     /// This function will be called when `wake` is called on the [`Waker`].
     /// It must wake up the task associated with this [`RawWaker`].
     ///
-    /// The implemention of this function must not consume the provided data
+    /// The implementation of this function must make sure to release any
+    /// resources that are associated with this instance of a [`RawWaker`] and
+    /// associated task.
+    ///
+    /// # `wake_by_ref`
+    ///
+    /// This function will be called when `wake_by_ref` is called on the [`Waker`].
+    /// It must wake up the task associated with this [`RawWaker`].
+    ///
+    /// This function is similar to `wake`, but must not consume the provided data
     /// pointer.
     ///
     /// # `drop`
@@ -120,11 +137,13 @@ impl RawWakerVTable {
     pub const fn new(
         clone: unsafe fn(*const ()) -> RawWaker,
         wake: unsafe fn(*const ()),
+        wake_by_ref: unsafe fn(*const ()),
         drop: unsafe fn(*const ()),
     ) -> Self {
         Self {
             clone,
             wake,
+            wake_by_ref,
             drop,
         }
     }
@@ -187,14 +206,33 @@ unsafe impl Sync for Waker {}
 impl Waker {
     /// Wake up the task associated with this `Waker`.
     #[inline]
-    pub fn wake(&self) {
+    pub fn wake(self) {
         // The actual wakeup call is delegated through a virtual function call
         // to the implementation which is defined by the executor.
+        let wake = self.waker.vtable.wake;
+        let data = self.waker.data;
+
+        // Don't call `drop` -- the waker will be consumed by `wake`.
+        crate::mem::forget(self);
 
         // SAFETY: This is safe because `Waker::new_unchecked` is the only way
         // to initialize `wake` and `data` requiring the user to acknowledge
         // that the contract of `RawWaker` is upheld.
-        unsafe { (self.waker.vtable.wake)(self.waker.data) }
+        unsafe { (wake)(data) };
+    }
+
+    /// Wake up the task associated with this `Waker` without consuming the `Waker`.
+    ///
+    /// This is similar to `wake`, but may be slightly less efficient in the case
+    /// where an owned `Waker` is available. This method should be preferred to
+    /// calling `waker.clone().wake()`.
+    #[inline]
+    pub fn wake_by_ref(&self) {
+        // The actual wakeup call is delegated through a virtual function call
+        // to the implementation which is defined by the executor.
+
+        // SAFETY: see `wake`
+        unsafe { (self.waker.vtable.wake_by_ref)(self.waker.data) }
     }
 
     /// Returns `true` if this `Waker` and another `Waker` have awoken the same task.

src/test/run-pass/async-await.rs

@@ -19,7 +19,10 @@ struct Counter {
 }
 
 impl ArcWake for Counter {
-    fn wake(arc_self: &Arc<Self>) {
+    fn wake(self: Arc<Self>) {
+        Self::wake_by_ref(&self)
+    }
+    fn wake_by_ref(arc_self: &Arc<Self>) {
         arc_self.wakes.fetch_add(1, atomic::Ordering::SeqCst);
     }
 }
@@ -34,7 +37,7 @@ impl Future for WakeOnceThenComplete {
         if self.0 {
             Poll::Ready(())
         } else {
-            cx.waker().wake();
+            cx.waker().wake_by_ref();
             self.0 = true;
             Poll::Pending
         }

src/test/run-pass/auxiliary/arc_wake.rs

@@ -12,25 +12,30 @@ macro_rules! waker_vtable {
         &RawWakerVTable::new(
             clone_arc_raw::<$ty>,
             wake_arc_raw::<$ty>,
+            wake_by_ref_arc_raw::<$ty>,
             drop_arc_raw::<$ty>,
         )
     };
 }
 
 pub trait ArcWake {
-    fn wake(arc_self: &Arc<Self>);
+    fn wake(self: Arc<Self>);
+
+    fn wake_by_ref(arc_self: &Arc<Self>) {
+        arc_self.clone().wake()
+    }
 
     fn into_waker(wake: Arc<Self>) -> Waker where Self: Sized
     {
-        let ptr = Arc::into_raw(wake) as *const();
+        let ptr = Arc::into_raw(wake) as *const ();
 
         unsafe {
             Waker::new_unchecked(RawWaker::new(ptr, waker_vtable!(Self)))
         }
     }
 }
 
-unsafe fn increase_refcount<T: ArcWake>(data: *const()) {
+unsafe fn increase_refcount<T: ArcWake>(data: *const ()) {
     // Retain Arc by creating a copy
     let arc: Arc<T> = Arc::from_raw(data as *const T);
     let arc_clone = arc.clone();
@@ -39,18 +44,23 @@ unsafe fn increase_refcount<T: ArcWake>(data: *const()) {
     let _ = Arc::into_raw(arc_clone);
 }
 
-unsafe fn clone_arc_raw<T: ArcWake>(data: *const()) -> RawWaker {
+unsafe fn clone_arc_raw<T: ArcWake>(data: *const ()) -> RawWaker {
     increase_refcount::<T>(data);
     RawWaker::new(data, waker_vtable!(T))
 }
 
-unsafe fn drop_arc_raw<T: ArcWake>(data: *const()) {
+unsafe fn drop_arc_raw<T: ArcWake>(data: *const ()) {
     // Drop Arc
     let _: Arc<T> = Arc::from_raw(data as *const T);
 }
 
-unsafe fn wake_arc_raw<T: ArcWake>(data: *const()) {
+unsafe fn wake_arc_raw<T: ArcWake>(data: *const ()) {
+    let arc: Arc<T> = Arc::from_raw(data as *const T);
+    ArcWake::wake(arc);
+}
+
+unsafe fn wake_by_ref_arc_raw<T: ArcWake>(data: *const ()) {
     let arc: Arc<T> = Arc::from_raw(data as *const T);
-    ArcWake::wake(&arc);
+    ArcWake::wake_by_ref(&arc);
     let _ = Arc::into_raw(arc);
 }

src/test/run-pass/futures-api.rs

@@ -20,7 +20,10 @@ struct Counter {
 }
 
 impl ArcWake for Counter {
-    fn wake(arc_self: &Arc<Self>) {
+    fn wake(self: Arc<Self>) {
+        Self::wake_by_ref(&self)
+    }
+    fn wake_by_ref(arc_self: &Arc<Self>) {
         arc_self.wakes.fetch_add(1, atomic::Ordering::SeqCst);
     }
 }
@@ -32,8 +35,8 @@ impl Future for MyFuture {
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         // Wake twice
         let waker = cx.waker();
-        waker.wake();
-        waker.wake();
+        waker.wake_by_ref();
+        waker.wake_by_ref();
         Poll::Ready(())
     }
 }