rt: share vtable between waker and waker ref (#2045)

The `Waker::will_wake` compares both a data pointer and a vtable to
decide if wakers are equivalent. To avoid false negatives during
comparison, use the same vtable for a waker stored in `WakerRef`.

tokio/src/task/harness.rs

@@ -300,10 +300,6 @@ where
         self.drop_waker();
     }
 
-    pub(super) fn wake_by_local_ref(&self) {
-        self.wake_by_ref();
-    }
-
     pub(super) fn wake_by_ref(&self) {
         if self.header().state.transition_to_notified() {
             unsafe {

tokio/src/task/tests/task.rs

@@ -641,3 +641,21 @@ fn shutdown_from_task_after_notified() {
 
     assert_ready_err!(handle.poll());
 }
+
+#[test]
+fn waker_ref_will_wake_clone() {
+    use std::task::Poll::Ready;
+
+    let (task, handle) = task::joinable(poll_fn(|cx| {
+        let waker = cx.waker().clone();
+        assert!(cx.waker().will_wake(&waker));
+        Ready(())
+    }));
+    let mut handle = spawn(handle);
+
+    let mock = mock().bind(&task).release_local();
+    let mock = &mut || Some(From::from(&mock));
+
+    assert_none!(task.run(mock));
+    assert_ready_ok!(handle.poll());
+}

tokio/src/task/waker.rs

@@ -3,11 +3,12 @@ use crate::task::{Header, Schedule};
 
 use std::future::Future;
 use std::marker::PhantomData;
+use std::mem::ManuallyDrop;
 use std::ops;
 use std::task::{RawWaker, RawWakerVTable, Waker};
 
 pub(super) struct WakerRef<'a, S: 'static> {
-    waker: Waker,
+    waker: ManuallyDrop<Waker>,
     _p: PhantomData<(&'a Header, S)>,
 }
 
@@ -18,16 +19,15 @@ where
     T: Future,
     S: Schedule,
 {
-    let ptr = meta as *const _ as *const ();
+    // `Waker::will_wake` uses the VTABLE pointer as part of the check. This
-
+    // means that `will_wake` will always return false when using the current
-    let vtable = &RawWakerVTable::new(
+    // task's waker. (discussion at rust-lang/rust#66281).
-        clone_waker::<T, S>,
+    //
-        wake_unreachable,
+    // To fix this, we use a single vtable. Since we pass in a reference at this
-        wake_by_local_ref::<T, S>,
+    // point and not an *owned* waker, we must ensure that `drop` is never
-        noop,
+    // called on this waker instance. This is done by wrapping it with
-    );
+    // `ManuallyDrop` and then never calling drop.
-
+    let waker = unsafe { ManuallyDrop::new(Waker::from_raw(raw_waker::<T, S>(meta))) };
-    let waker = unsafe { Waker::from_raw(RawWaker::new(ptr, vtable)) };
 
     WakerRef {
         waker,
@@ -50,15 +50,7 @@ where
 {
     let meta = ptr as *const Header;
     (*meta).state.ref_inc();
-
+    raw_waker::<T, S>(meta)
-    let vtable = &RawWakerVTable::new(
-        clone_waker::<T, S>,
-        wake_by_val::<T, S>,
-        wake_by_ref::<T, S>,
-        drop_waker::<T, S>,
-    );
-
-    RawWaker::new(ptr, vtable)
 }
 
 unsafe fn drop_waker<T, S>(ptr: *const ())
@@ -70,11 +62,6 @@ where
     harness.drop_waker();
 }
 
-// `wake()` cannot be called on the ref variaant.
-unsafe fn wake_unreachable(_data: *const ()) {
-    unreachable!();
-}
-
 unsafe fn wake_by_val<T, S>(ptr: *const ())
 where
     T: Future,
@@ -84,16 +71,6 @@ where
     harness.wake_by_val();
 }
 
-// This function can only be called when on the runtime.
-unsafe fn wake_by_local_ref<T, S>(ptr: *const ())
-where
-    T: Future,
-    S: Schedule,
-{
-    let harness = Harness::<T, S>::from_raw(ptr as *mut _);
-    harness.wake_by_local_ref();
-}
-
 // Wake without consuming the waker
 unsafe fn wake_by_ref<T, S>(ptr: *const ())
 where
@@ -104,4 +81,17 @@ where
     harness.wake_by_ref();
 }
 
-unsafe fn noop(_ptr: *const ()) {}
+fn raw_waker<T, S>(meta: *const Header) -> RawWaker
+where
+    T: Future,
+    S: Schedule,
+{
+    let ptr = meta as *const ();
+    let vtable = &RawWakerVTable::new(
+        clone_waker::<T, S>,
+        wake_by_val::<T, S>,
+        wake_by_ref::<T, S>,
+        drop_waker::<T, S>,
+    );
+    RawWaker::new(ptr, vtable)
+}