1 files changed, 24 insertions, 34 deletions

diff --git a/tokio/src/task/waker.rs b/tokio/src/task/waker.rs
index e0e1f36c..9892f1be 100644
--- a/tokio/src/task/waker.rs
+++ b/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 ();
-
-    let vtable = &RawWakerVTable::new(
-        clone_waker::<T, S>,
-        wake_unreachable,
-        wake_by_local_ref::<T, S>,
-        noop,
-    );
-
-    let waker = unsafe { Waker::from_raw(RawWaker::new(ptr, vtable)) };
+    // `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
+    // task's waker. (discussion at rust-lang/rust#66281).
+    //
+    // To fix this, we use a single vtable. Since we pass in a reference at this
+    // point and not an *owned* waker, we must ensure that `drop` is never
+    // 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))) };
 
     WakerRef {
         waker,
@@ -50,15 +50,7 @@ where
 {
     let meta = ptr as *const Header;
     (*meta).state.ref_inc();
-
-    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)
+    raw_waker::<T, S>(meta)
 }
 
 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,24 +71,27 @@ 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 ())
+// Wake without consuming the waker
+unsafe fn wake_by_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();
+    harness.wake_by_ref();
 }
 
-// Wake without consuming the waker
-unsafe fn wake_by_ref<T, S>(ptr: *const ())
+fn raw_waker<T, S>(meta: *const Header) -> RawWaker
 where
     T: Future,
     S: Schedule,
 {
-    let harness = Harness::<T, S>::from_raw(ptr as *mut _);
-    harness.wake_by_ref();
+    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)
 }
-
-unsafe fn noop(_ptr: *const ()) {}