Skip to content

Commit

Permalink
Preserve same vtable pointer when cloning raw waker, to fix Waker::wi…
Browse files Browse the repository at this point in the history
…ll_wake
  • Loading branch information
dtolnay committed Feb 26, 2024
1 parent 793b45f commit c4c5b3d
Showing 1 changed file with 16 additions and 16 deletions.
32 changes: 16 additions & 16 deletions library/alloc/src/task.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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>())
}

0 comments on commit c4c5b3d

Please sign in to comment.