Add LocalWaker support

[tvallotton]

Proposal

Problem statement

The current implementation of Waker is Send + Sync, which means that it can be freely moved between threads. While this is necessary for work stealing, it also imposes a performance penalty on single threaded and thread per core (TPC) runtimes. This runtime cost is contrary to Rust's zero cost abstraction philosophy. Furthermore, this missing API has led some async runtimes to implement unsound wakers, both by neglect and intentionally.

Motivation, use-cases

Local wakers leverage Rust's compile-time thread safety to avoid unnecessary thread synchronization, offering a compile-time guarantee that a waker has not been moved to another thread. This allows async runtimes to specialize their waker implementations and skip unnecessary atomic operations, thereby improving performance.

It is noteworthy to mention that while this is especially useful for TPC runtimes, work stealing runtimes may also benefit from performing this specialization. So this should not be considered a niche use case.

Solution sketches

Construction

Constructing a local waker is done in the same way as a shared waker. You just use the from_raw function, which takes a RawWaker.

use std::task::LocalWaker; 
let waker = unsafe { LocalWaker::from_raw(raw_waker) }

Alternatively, the LocalWake trait may be used analogously to Wake.

use std::task::LocalWake;

thread_local! {
    /// A queue containing all woken tasks
    static WOKEN_TASKS: RefCell<VecDeque<Task>>; 
}

struct Task(Box<dyn Future<Output = ()>>);

impl LocalWake for Task {
    fn wake(self: Rc<Self>) {
        WOKEN_TASKS.with(|woken_tasks| {
            woken_tasks.borrow_mut().push_back(self)
        })
    }
}

The safety requirements for constructing a LocalWaker from a RawWaker would be the same as a Waker, except for thread safety.

Usage

A local waker can be accessed with the local_waker method on Context and woken just like a regular waker. All contexts will return a valid local_waker, regardless of whether they are specialized for this case or not.

let local_waker: &LocalWaker = cx.local_waker(); 
local_waker.wake(); 

ContextBuilder

In order to construct a specialized Context, the ConstextBuilder type must be used. This type may be extended in the future to allow for more functionality for Context.

use std::task::{Context, LocalWaker, Waker, ContextBuilder}; 

let waker: Waker = /* ... */; 
let local_waker: LocalWaker = /* ... */;

let context = ContextBuilder::new()
    .waker(&waker)
    .local_waker(&local_waker)
    .build()
    .expect("at least one waker must be set"); 

Then it can be accessed with the local_waker method on Context.

let local_waker: &LocalWaker = cx.local_waker(); 

If a LocalWaker is not set on a context, this one would still return a valid LocalWaker, because a local waker can be trivially constructed from the context's waker.

If a runtime does not intend to support thread safe wakers, they should not provide a Waker to ContextBuilder, and it will construct a Context that panics in the call to waker().

Links and related work

• IRLO thread about LocalWaker
• Context made !Send + !Sync with this use case in mind
• Monoio's Waker being unsound
• Glommio's Waker being unsound
• Withoutboat's comments on the waker API v1
• Withoutboat's comments on the waker API v2

What happens now?

This issue is part of the libs-api team API change proposal process. Once this issue is filed the libs-api team will review open proposals in its weekly meeting. You should receive feedback within a week or two.

[withoutboats]

(NOT A CONTRIBUTION)

A local waker can be accessed with the local_waker method on Context and woken just like a regular waker. All contexts will return a valid local_waker, regardless of weather they are specialized for this case or not.

Should clarify what you mean by this: a context without a local waker set will just return its Waker as a LocalWaker, because every Waker is a valid LocalWaker (they are basically just local wakers that implement Send).

It'd be helpful probably to sketch out what the definition of Context would look like in this proposal to support all of these operations:

1. Returning the Waker for .waker() if the waker is set, and panicking if not.
2. Returning the LocalWaker for .local_waker() if the local waker is set, and returning the Waker is not.

Should be possible but it may be a bit tricky.

---

let context = ContextBuilder::new()
    .waker(&waker)
    .local_waker(&local_waker)
    .build()
    .expect("at least one waker must be set"); 

Maybe build should just panic if no waker is set, instead of returning a Result? Not clear what anyone would do except call expect on this result.

[tvallotton]

Yeah, currenty Waker is transparent to RawWaker, so a transmute from a &Waker to a &LocalWaker should probably be fine, and it would require no clones (assuming LocalWaker would too be transparent to RawWaker).

[withoutboats]

(NOT A CONTRIBUTION)

possibly LocalWaker could also just be defined as struct LocalWaker(Waker) and you could store the waker as a LocalWaker in Context. Would be a bit counterintuitive but would avoid transmutes.

[pitaj]

Agreed, please outline the exact function signatures for WakerBuilder and any added or changed functions and fields for Context. I am wondering if a builder is really warranted here. Do we plan on adding even more options to Context?

[pitaj]

Do we want to allow constructing a Context with only a LocalWaker (no Waker)? In that case, calling .waker() would panic? Can we somehow make that a monomorphization error instead?

[tvallotton]

Note I proposed a ContextBuilder type only, that for the moment will have:

• local_waker: to set the local waker on context
• waker: to set the thread safe waker on context.
• build: to build the context.

The whole purpose of Context was to leave room for future extensibility, so I believe a builder is warranted.

[tvallotton]

Do we want to allow constructing a Context with only a LocalWaker (no Waker)? In that case, calling .waker() would panic? Can we somehow make that a monomorphization error instead?

I doubt it, contexts would be entirely dynamic with respect to their Waker support. Now, it would be possible to add this optional support later without it being a breaking change.

[pitaj]

I'm a little confused.

I doubt it, contexts would be entirely dynamic with respect to their Waker support.

You're saying you doubt it would be possible to make this a monomorphization error?

Just to be clear, you're not proposing to make Waker optional here, just leaving it open for the future.

[pitaj]

I think something like try_local_waker() -> Option<&LocalWaker> could be useful as well. I'm wondering if maybe local_waker() should just return an option. Then we don't have to worry about transmutes or anything. Let the user fall back to waker() themselves or unwrap to panic if they always expect a local waker to exist.

[withoutboats]

(NOT A CONTRIBUTION)

You're saying you doubt it would be possible to make this a monomorphization error?

Just to be clear, you're not proposing to make Waker optional here, just leaving it open for the future.

Via context builder, @tvallotton's proposal does make it possible to construct a Context without a waker. Calling waker on that Context would panic.

The alternative is to not provide a way to construct a Context without a waker. If Rust does that, executors that don't support being awoken from other threads will just have to define a Waker that panics when you wakes it. This seems strictly worse to me.

I think something like try_local_waker() -> Option<&LocalWaker> could be useful as well. I'm wondering if maybe local_waker() should just return an option. Then we don't have to worry about transmutes or anything. Let the user fall back to waker() themselves or unwrap to panic if they always expect a local waker to exist.

There's no reason for this. It's always safe to wake the task from the same thread via the waker mechanism.

The point of this API is that reactors that wake from the same thread that they were polled on can use the local_waker mechanism instead of waker, and they will always just work. Since these reactors often store the waker somewhere, it's important that they be able to get the same type (LocalWaker) regardless of if they are actually waking via a thread safe mechanism or not. Reactors written that way can be run on a single threaded executor or a multithreaded executor and will be totally agnostic to it.

[pitaj]

I was just thinking that if someone is trying to enhance performance by using local wakers, knowing if you actually got a local waker or if it just fell back to the non-local waker would be useful.

[tvallotton]

I think it is safe to say that a 100% of the people that ask for a LocalWaker will conform with a Waker given that all wakers are local wakers. Its preferable to have std perform the transmute than the ecosystem.

[withoutboats]

(NOT A CONTRIBUTION)

I was just thinking that if someone is trying to enhance performance by using local wakers, knowing if you actually got a local waker or if it just fell back to the non-local waker would be useful.

They're not the same parties. The person who wants the enhanced performance is the one who selects the executor, which isn't the same one who writes the reactor and deals with wakers, it's the one who calls spawn.