diff --git a/src/libcore/task/wake.rs b/src/libcore/task/wake.rs index 979a0792608c1..006cbbb6ce6bd 100644 --- a/src/libcore/task/wake.rs +++ b/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; - // SAFETY: This is safe because `Waker::new_unchecked` is the only way + // Don't call `drop` -- the waker will be consumed by `wake`. + crate::mem::forget(self); + + // SAFETY: This is safe because `Waker::from_raw` 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. @@ -215,7 +253,7 @@ impl Waker { /// in [`RawWaker`]'s and [`RawWakerVTable`]'s documentation is not upheld. /// Therefore this method is unsafe. #[inline] - pub unsafe fn new_unchecked(waker: RawWaker) -> Waker { + pub unsafe fn from_raw(waker: RawWaker) -> Waker { Waker { waker, } @@ -226,7 +264,7 @@ impl Clone for Waker { #[inline] fn clone(&self) -> Self { Waker { - // SAFETY: This is safe because `Waker::new_unchecked` is the only way + // SAFETY: This is safe because `Waker::from_raw` is the only way // to initialize `clone` and `data` requiring the user to acknowledge // that the contract of [`RawWaker`] is upheld. waker: unsafe { (self.waker.vtable.clone)(self.waker.data) }, @@ -237,7 +275,7 @@ impl Clone for Waker { impl Drop for Waker { #[inline] fn drop(&mut self) { - // SAFETY: This is safe because `Waker::new_unchecked` is the only way + // SAFETY: This is safe because `Waker::from_raw` is the only way // to initialize `drop` and `data` requiring the user to acknowledge // that the contract of `RawWaker` is upheld. unsafe { (self.waker.vtable.drop)(self.waker.data) } diff --git a/src/test/run-pass/async-await.rs b/src/test/run-pass/async-await.rs index 4f5f7724ad076..518452aefc152 100644 --- a/src/test/run-pass/async-await.rs +++ b/src/test/run-pass/async-await.rs @@ -19,7 +19,10 @@ struct Counter { } impl ArcWake for Counter { - fn wake(arc_self: &Arc) { + fn wake(self: Arc) { + Self::wake_by_ref(&self) + } + fn wake_by_ref(arc_self: &Arc) { 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 } diff --git a/src/test/run-pass/auxiliary/arc_wake.rs b/src/test/run-pass/auxiliary/arc_wake.rs index 74ec56f55171a..93e074e7ee55c 100644 --- a/src/test/run-pass/auxiliary/arc_wake.rs +++ b/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); + fn wake(self: Arc); + + fn wake_by_ref(arc_self: &Arc) { + arc_self.clone().wake() + } fn into_waker(wake: Arc) -> 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))) + Waker::from_raw(RawWaker::new(ptr, waker_vtable!(Self))) } } } -unsafe fn increase_refcount(data: *const()) { +unsafe fn increase_refcount(data: *const ()) { // Retain Arc by creating a copy let arc: Arc = Arc::from_raw(data as *const T); let arc_clone = arc.clone(); @@ -39,18 +44,23 @@ unsafe fn increase_refcount(data: *const()) { let _ = Arc::into_raw(arc_clone); } -unsafe fn clone_arc_raw(data: *const()) -> RawWaker { +unsafe fn clone_arc_raw(data: *const ()) -> RawWaker { increase_refcount::(data); RawWaker::new(data, waker_vtable!(T)) } -unsafe fn drop_arc_raw(data: *const()) { +unsafe fn drop_arc_raw(data: *const ()) { // Drop Arc let _: Arc = Arc::from_raw(data as *const T); } -unsafe fn wake_arc_raw(data: *const()) { +unsafe fn wake_arc_raw(data: *const ()) { + let arc: Arc = Arc::from_raw(data as *const T); + ArcWake::wake(arc); +} + +unsafe fn wake_by_ref_arc_raw(data: *const ()) { let arc: Arc = Arc::from_raw(data as *const T); - ArcWake::wake(&arc); + ArcWake::wake_by_ref(&arc); let _ = Arc::into_raw(arc); } diff --git a/src/test/run-pass/futures-api.rs b/src/test/run-pass/futures-api.rs index 5d0b0db510f41..6094f15569bb6 100644 --- a/src/test/run-pass/futures-api.rs +++ b/src/test/run-pass/futures-api.rs @@ -20,7 +20,10 @@ struct Counter { } impl ArcWake for Counter { - fn wake(arc_self: &Arc) { + fn wake(self: Arc) { + Self::wake_by_ref(&self) + } + fn wake_by_ref(arc_self: &Arc) { 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 { // Wake twice let waker = cx.waker(); - waker.wake(); - waker.wake(); + waker.wake_by_ref(); + waker.wake_by_ref(); Poll::Ready(()) } }