Add 'static bound to waker_ref

Author: taiki-e

futures-task/src/waker.rs

@@ -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>,
@@ -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>()) })
 }

futures-task/src/waker_ref.rs

@@ -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.<

futures-util/src/stream/futures_unordered/mod.rs

@@ -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

futures-util/src/stream/futures_unordered/task.rs

@@ -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
@@ -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`.
@@ -123,3 +123,92 @@ 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();
+    }
+
+    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);
+    }
+
+    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>()) })
+    }
+}