ACP: Add Waker::with_arc_wake

[orlp]

Proposal

Problem statement

The Rust standard library provides the Wake trait to allow you to create Wakers without unsafe code. This is nice, but there is no way to create a Waker from an &Arc<W> without cloning the Arc first. However, this is unnecessary, this clone could happen on the first time the Waker is cloned instead (and thus may be avoided if the Waker isn't cloned at all).

Motivating examples or use cases

I have some waker: Arc<W> inside a struct I re-use every time I wish to block_on using this struct. The code might look something like this:

pub fn block_on<F: Future>(&self, fut: F) -> <F as Future>::Output {
    let mut fut = pin!(fut);
    let mut cx = Context::from_waker(self.waker.clone());
    loop {
        match fut.as_mut().poll(&mut cx) {
            Poll::Ready(output) => return output,
            Poll::Pending => self.waker.block_until_notified(),
        }
    }
}

Note how we have to call self.waker.clone() every time, regardless of whether we end up blocking or not.

Solution sketch

We could add the function Waker::with_arc_wake which uses a closure to guarantee the &Waker it provides
does not outlive the arc_wake, without having to call clone:

/// Creates a temporary [`Waker`] using the given [`Wake`] implementation.
fn with_arc_wake<R, W: Wake + Send + Sync + 'static, F: FnOnce(&Waker) -> R>(
    arc_wake: &Arc<W>,
    f: F,
) -> R {
    // SAFETY: we create a clone of the Arc without incrementing the refcount.
    // This is safe because we never drop this clone and the lifetime of
    // arc_wake guarantees the Arc will not die prematurely either.
    let tmp = ManuallyDrop::new(Waker::from(unsafe { core::ptr::from_ref(arc_wake).read() }));
    f(&tmp)
}

With this function the above block_on could be written as such:

pub fn block_on<F: Future>(&self, fut: F) -> <F as Future>::Output {
    let mut fut = pin!(fut);
    Waker::with_arc_wake(&self.waker, |waker| {
        let mut cx = Context::from_waker(waker);
        loop {
            match fut.as_mut().poll(&mut cx) {
                Poll::Ready(output) => return output,
                Poll::Pending => self.waker.block_until_notified(),
            }
        }
    })
}

No more premature cloning.

Alternatives

The futures crate uses the function waker_ref returning a WakerRef<'a> to achieve the same effect. While more flexible than with_arc_wake it does add an extra type to the standard library which is quite a bit of extra surface to document, open questions, etc. In all use-cases I could think of with_arc_wake achieves the same result with a smaller API.

One could of course also implement this yourself using unsafe code, but as the point of the Wake trait is to provide a safe alternative to RawWaker, it would be nice if using it optimally also doesn't involve unsafe code.

Links and related work

https://docs.rs/futures/latest/futures/task/fn.waker_ref.html
https://docs.rs/futures/latest/futures/task/struct.WakerRef.html

[traviscross]

cc @rust-lang/wg-async

[yoshuawuyts]

@orlp thanks for filing this issue! Can you say a little more about your use case here? Why are you storing an Arc<W> inside of your struct, rather than a concrete Waker?

I'm asking this because Context::from_waker takes an &Waker. If you can have your type store an owned Waker, you should be able to reuse that between multiple block_on calls without needing to clone the underlying Arc on each invocation.

Example

If you have a type that is defined as follows:

use std::task::Waker;

struct MyType {
    waker: Waker,
}

I believe you should be able to define the block_on method like this (playground):

  pub fn block_on<F: Future>(&self, fut: F) -> <F as Future>::Output {
      let mut fut = pin!(fut);
-    let mut cx = Context::from_waker(self.waker.clone());
+    let mut cx = Context::from_waker(&self.waker);
      loop {
          match fut.as_mut().poll(&mut cx) {
              Poll::Ready(output) => return output,
              Poll::Pending => self.waker.block_until_notified(),
          }
      }
  }

[Darksonn]

The futures crate uses the function waker_ref returning a WakerRef<'a> to achieve the same effect. While more flexible than with_arc_wake it does add an extra type to the standard library which is quite a bit of extra surface to document, open questions, etc. In all use-cases I could think of with_arc_wake achieves the same result with a smaller API.

Personally I prefer the API provided by the futures crate. One way we could avoid defining an extra type is this:

fn waker_ref<W: Wake + Send + Sync + 'static>(arc_wake: &Arc<W>) -> impl Deref<Target = Waker>

[orlp]

@yoshuawuyts You dodged the difficult part in your playground example in this line:

                Poll::Pending => todo!(),

that line used to say

              Poll::Pending => self.waker.block_until_notified(),

However, Waker has no block_until_notified method, that was part of the concrete W inside the Arc<W>. So storing a Waker would not suffice.

You are right in that I could store both a Waker and an Arc<W>, one being the type-erased version of the other. But considering they are essentially exactly the same object, this feels rather wasteful.

[orlp]

@Darksonn Do we have impl Deref in return types anywhere else in the standard library? Or impl Trait in general? I honestly kind of assumed it was forbidden in the standard library, and that returning concrete types should always be preferred. I at least can't think of any public stdlib API right now which has impl Trait in the return type.

[Darksonn]

I mean, I would also be okay with returning an explicit WakerRef.

[yoshuawuyts]

However, Waker has no block_until_notified method, that was part of the concrete W inside the Arc<W>. So storing a Waker would not suffice.

@orlp Ah, I see - thank you for clarifying. Follow up question: why does your type W: Wake implement the method block_until_notified? Why isn't this logic for example contained in the block_on method?

I reckon that the reason for that is probably obvious to you, but it isn't yet to me. I'm asking additional questions so I can fully understand your use case here, and be able to explain why existing alternatives aren't sufficient and we need an API like this to solve it.

[orlp]

Why isn't this logic for example contained in the block_on method?

@yoshuawuyts Because block_until_notified needs to access the same shared inner state inside the Waker to realize when it has been woken. Here's a complete self-contained example you can play with on the playground.