Preserve same vtable pointer when cloning raw waker, to fix Waker::wi…

…ll_wake
rust-lang · Feb 26, 2024 · c4c5b3d · c4c5b3d
1 parent 793b45f
commit c4c5b3d
Showing 1 changed file with 16 additions and 16 deletions.
diff --git a/library/alloc/src/task.rs b/library/alloc/src/task.rs
@@ -138,10 +138,7 @@ fn raw_waker<W: Wake + Send + Sync + 'static>(waker: Arc<W>) -> RawWaker {
     // Increment the reference count of the arc to clone it.
     unsafe fn clone_waker<W: Wake + Send + Sync + 'static>(waker: *const ()) -> RawWaker {
         unsafe { Arc::increment_strong_count(waker as *const W) };
-        RawWaker::new(
-            waker as *const (),
-            &RawWakerVTable::new(clone_waker::<W>, wake::<W>, wake_by_ref::<W>, drop_waker::<W>),
-        )
+        RawWaker::new(waker, vtable::<W>())
     }
 
     // Wake by value, moving the Arc into the Wake::wake function
@@ -161,10 +158,15 @@ fn raw_waker<W: Wake + Send + Sync + 'static>(waker: Arc<W>) -> RawWaker {
         unsafe { Arc::decrement_strong_count(waker as *const W) };
     }
 
-    RawWaker::new(
-        Arc::into_raw(waker) as *const (),
-        &RawWakerVTable::new(clone_waker::<W>, wake::<W>, wake_by_ref::<W>, drop_waker::<W>),
-    )
+    fn vtable<W: Wake + Send + Sync + 'static>() -> &'static RawWakerVTable {
+        // Perform this const promotion in a single place used by both initial
+        // construction of the RawWaker below and by clone_waker above. This way
+        // the vtable pointer comparison performed by Waker::will_wake sees that
+        // a cloned waker will wake the original one and vice versa.
+        &RawWakerVTable::new(clone_waker::<W>, wake::<W>, wake_by_ref::<W>, drop_waker::<W>)
+    }
+
+    RawWaker::new(Arc::into_raw(waker) as *const (), vtable::<W>())
 }
 
 /// An analogous trait to `Wake` but used to construct a `LocalWaker`. This API
@@ -306,10 +308,7 @@ fn local_raw_waker<W: LocalWake + 'static>(waker: Rc<W>) -> RawWaker {
     // Increment the reference count of the Rc to clone it.
     unsafe fn clone_waker<W: LocalWake + 'static>(waker: *const ()) -> RawWaker {
         unsafe { Rc::increment_strong_count(waker as *const W) };
-        RawWaker::new(
-            waker as *const (),
-            &RawWakerVTable::new(clone_waker::<W>, wake::<W>, wake_by_ref::<W>, drop_waker::<W>),
-        )
+        RawWaker::new(waker, vtable::<W>())
     }
 
     // Wake by value, moving the Rc into the LocalWake::wake function
@@ -329,8 +328,9 @@ fn local_raw_waker<W: LocalWake + 'static>(waker: Rc<W>) -> RawWaker {
         unsafe { Rc::decrement_strong_count(waker as *const W) };
     }
 
-    RawWaker::new(
-        Rc::into_raw(waker) as *const (),
-        &RawWakerVTable::new(clone_waker::<W>, wake::<W>, wake_by_ref::<W>, drop_waker::<W>),
-    )
+    fn vtable<W: LocalWake + 'static>() -> &'static RawWakerVTable {
+        &RawWakerVTable::new(clone_waker::<W>, wake::<W>, wake_by_ref::<W>, drop_waker::<W>)
+    }
+
+    RawWaker::new(Rc::into_raw(waker) as *const (), vtable::<W>())
 }