Skip to content

Commit

Permalink
Add 'static bound to waker_ref
Browse files Browse the repository at this point in the history
  • Loading branch information
taiki-e committed Feb 25, 2024
1 parent 0fa6a1a commit 6e6d35a
Show file tree
Hide file tree
Showing 4 changed files with 103 additions and 11 deletions.
12 changes: 6 additions & 6 deletions futures-task/src/waker.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ use alloc::sync::Arc;
use core::mem;
use core::task::{RawWaker, RawWakerVTable, Waker};

pub(super) fn waker_vtable<W: ArcWake>() -> &'static RawWakerVTable {
pub(super) fn waker_vtable<W: ArcWake + 'static>() -> &'static RawWakerVTable {
&RawWakerVTable::new(
clone_arc_raw::<W>,
wake_arc_raw::<W>,
Expand All @@ -28,31 +28,31 @@ where
// FIXME: panics on Arc::clone / refcount changes could wreak havoc on the
// code here. We should guard against this by aborting.

unsafe fn increase_refcount<T: ArcWake>(data: *const ()) {
unsafe fn increase_refcount<T: ArcWake + 'static>(data: *const ()) {
// Retain Arc, but don't touch refcount by wrapping in ManuallyDrop
let arc = mem::ManuallyDrop::new(unsafe { Arc::<T>::from_raw(data.cast::<T>()) });
// Now increase refcount, but don't drop new refcount either
let _arc_clone: mem::ManuallyDrop<_> = arc.clone();
}

// used by `waker_ref`
unsafe fn clone_arc_raw<T: ArcWake>(data: *const ()) -> RawWaker {
unsafe fn clone_arc_raw<T: ArcWake + 'static>(data: *const ()) -> RawWaker {
unsafe { increase_refcount::<T>(data) }
RawWaker::new(data, waker_vtable::<T>())
}

unsafe fn wake_arc_raw<T: ArcWake>(data: *const ()) {
unsafe fn wake_arc_raw<T: ArcWake + 'static>(data: *const ()) {
let arc: Arc<T> = unsafe { Arc::from_raw(data.cast::<T>()) };
ArcWake::wake(arc);
}

// used by `waker_ref`
unsafe fn wake_by_ref_arc_raw<T: ArcWake>(data: *const ()) {
unsafe fn wake_by_ref_arc_raw<T: ArcWake + 'static>(data: *const ()) {
// Retain Arc, but don't touch refcount by wrapping in ManuallyDrop
let arc = mem::ManuallyDrop::new(unsafe { Arc::<T>::from_raw(data.cast::<T>()) });
ArcWake::wake_by_ref(&arc);
}

unsafe fn drop_arc_raw<T: ArcWake>(data: *const ()) {
unsafe fn drop_arc_raw<T: ArcWake + 'static>(data: *const ()) {
drop(unsafe { Arc::<T>::from_raw(data.cast::<T>()) })
}
2 changes: 1 addition & 1 deletion futures-task/src/waker_ref.rs
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,7 @@ impl Deref for WakerRef<'_> {
#[inline]
pub fn waker_ref<W>(wake: &Arc<W>) -> WakerRef<'_>
where
W: ArcWake,
W: ArcWake + 'static,
{
// simply copy the pointer instead of using Arc::into_raw,
// as we don't actually keep a refcount by using ManuallyDrop.<
Expand Down
3 changes: 2 additions & 1 deletion futures-util/src/stream/futures_unordered/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -510,7 +510,8 @@ impl<Fut: Future> Stream for FuturesUnordered<Fut> {
// We are only interested in whether the future is awoken before it
// finishes polling, so reset the flag here.
task.woken.store(false, Relaxed);
let waker = Task::waker_ref(task);
// SAFETY: see the comments of Bomb and this block.
let waker = unsafe { Task::waker_ref(task) };
let mut cx = Context::from_waker(&waker);

// Safety: We won't move the future ever again
Expand Down
97 changes: 94 additions & 3 deletions futures-util/src/stream/futures_unordered/task.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ use core::sync::atomic::{AtomicBool, AtomicPtr};

use super::abort::abort;
use super::ReadyToRunQueue;
use crate::task::{waker_ref, ArcWake, WakerRef};
use crate::task::ArcWake;

pub(super) struct Task<Fut> {
// The future
Expand Down Expand Up @@ -77,8 +77,8 @@ impl<Fut> ArcWake for Task<Fut> {

impl<Fut> Task<Fut> {
/// Returns a waker reference for this task without cloning the Arc.
pub(super) fn waker_ref(this: &Arc<Self>) -> WakerRef<'_> {
waker_ref(this)
pub(super) unsafe fn waker_ref(this: &Arc<Self>) -> waker_ref::WakerRef<'_> {
unsafe { waker_ref::waker_ref(this) }
}

/// Spins until `next_all` is no longer set to `pending_next_all`.
Expand Down Expand Up @@ -123,3 +123,94 @@ impl<Fut> Drop for Task<Fut> {
}
}
}

mod waker_ref {
use alloc::sync::Arc;
use core::marker::PhantomData;
use core::mem;
use core::mem::ManuallyDrop;
use core::ops::Deref;
use core::task::{RawWaker, RawWakerVTable, Waker};
use futures_task::ArcWake;

pub(crate) struct WakerRef<'a> {
waker: ManuallyDrop<Waker>,
_marker: PhantomData<&'a ()>,
}

impl WakerRef<'_> {
#[inline]
fn new_unowned(waker: ManuallyDrop<Waker>) -> Self {
Self { waker, _marker: PhantomData }
}
}

impl Deref for WakerRef<'_> {
type Target = Waker;

#[inline]
fn deref(&self) -> &Waker {
&self.waker
}
}

/// Copy of `future_task::waker_ref` without `W: 'static` bound.
///
/// # Safety
///
/// The caller must guarantee that use-after-free will not occur.
#[inline]
pub(crate) unsafe fn waker_ref<W>(wake: &Arc<W>) -> WakerRef<'_>
where
W: ArcWake,
{
// simply copy the pointer instead of using Arc::into_raw,
// as we don't actually keep a refcount by using ManuallyDrop.<
let ptr = Arc::as_ptr(wake).cast::<()>();

let waker =
ManuallyDrop::new(unsafe { Waker::from_raw(RawWaker::new(ptr, waker_vtable::<W>())) });
WakerRef::new_unowned(waker)
}

fn waker_vtable<W: ArcWake>() -> &'static RawWakerVTable {
&RawWakerVTable::new(
clone_arc_raw::<W>,
wake_arc_raw::<W>,
wake_by_ref_arc_raw::<W>,
drop_arc_raw::<W>,
)
}

// FIXME: panics on Arc::clone / refcount changes could wreak havoc on the
// code here. We should guard against this by aborting.

unsafe fn increase_refcount<T: ArcWake>(data: *const ()) {
// Retain Arc, but don't touch refcount by wrapping in ManuallyDrop
let arc = mem::ManuallyDrop::new(unsafe { Arc::<T>::from_raw(data.cast::<T>()) });
// Now increase refcount, but don't drop new refcount either
let _arc_clone: mem::ManuallyDrop<_> = arc.clone();
}

// used by `waker_ref`
unsafe fn clone_arc_raw<T: ArcWake>(data: *const ()) -> RawWaker {
unsafe { increase_refcount::<T>(data) }
RawWaker::new(data, waker_vtable::<T>())
}

unsafe fn wake_arc_raw<T: ArcWake>(data: *const ()) {
let arc: Arc<T> = unsafe { Arc::from_raw(data.cast::<T>()) };
ArcWake::wake(arc);
}

// used by `waker_ref`
unsafe fn wake_by_ref_arc_raw<T: ArcWake>(data: *const ()) {
// Retain Arc, but don't touch refcount by wrapping in ManuallyDrop
let arc = mem::ManuallyDrop::new(unsafe { Arc::<T>::from_raw(data.cast::<T>()) });
ArcWake::wake_by_ref(&arc);
}

unsafe fn drop_arc_raw<T: ArcWake>(data: *const ()) {
drop(unsafe { Arc::<T>::from_raw(data.cast::<T>()) })
}
}

0 comments on commit 6e6d35a

Please sign in to comment.