A question of ordering: is it better for cooperation to take a "yield before" shape or a "yield after" shape?
// yield before
async fn cooperative<F: IntoFuture>(inner: F) -> F::Output {
let inner = pin!(inner.into_future());
poll_fn(move |cx| {
// fn Context::poll_cooperate(cx) -> Poll<Cooperation>;
let mut coop = ready!(cx.poll_cooperate(cx));
let output = ready!(inner.poll(cx));
// fn Cooperation::record_progress(self);
coop.record_progress();
Ready(output)
}).await
}
// yield after
async fn cooperative<F: Future>(mut inner: F) -> F::Output {
let output = inner.await;
// fn Context::yield_cooperatively() -> bool;
poll_fn(|cx| {
if cx.yield_cooperatively() { Ready(()) } else { Pending }
).await;
output
}
Having written that down, I think the former is probably preferable; yield after imposes an extra state (made progress but yielding cooperatively), and you can shim from yield before to yield after but not the other way around.
async fn yield_cooperatively() {
poll_fn(|cx| {
let coop = ready!(cx.poll_cooperation());
coop.record_progress();
Ready(())
}).await
}
// useful (abusable?) for back-edge yielding, e.g.
while let Some(task) = tasks.next() {
make_progress(task); // sync
yield_cooperarively().await;
}
// kinda like yield_now(), but yielding less frequently
I'd still suggest some minor tweaks, though:
- I have a slight preference for
record_progress over made_progress, because the former is an active verb, whereas the latter "feels" more like a fn(&self) -> bool query due to being past tense.
-
Cooperation::record_progress should probably take self by value, because it's only intended to attempt and then record a single instance of progress after acquiring a cooperation handle. If you want to make more progress, you should poll for another.
-
Cooperation should capture the lifetime of the context, since it's not intended to hold onto an instance across multiple polls.
- Doing so is almost certainly an error, so we should make it an error to write
poll_fn(|cx| cx.poll_cooperation()). I originally saw this was possible and thought that it was a correct way to implement cooperation, and thus that record_progress needed to return Poll as well for "yield after" behavior for repeated progress on the same Cooperation.
- In the extreme,
Cooperation could even just be &dyn Cooperate rather than going through a manual vtable like Waker. Though it's probably realistic to consider that it might want to hold onto some state which the Context doesn't have.
Before stabilization, it'll be strongly desirable to provide a reasonably usable definition of when you should guard a poll with cooperation. The starting position is probably roughly "any nonguaranteed progress by means other than polling another future" and let people use their best judgement from there.
The simplest possible implementation would I think be
use core::task::{Poll, Waker};
use core::marker::PhantomData;
use core::mem::align_of;
pub struct Context<'a> {
waker: &'a Waker,
cooperation: &'a mut (dyn coop::Cooperate + 'a),
// plus the PhantomData markers
}
#[derive(Default)]
#[non_exhaustive]
pub struct ExtraContext<'a> {
pub cooperation: Option<&'a mut (dyn coop::Cooperate + 'a)>,
_marker: PhantomData<Context<'a>>,
}
impl<'a> Context<'a> {
pub fn from_waker_and(waker: &'a Waker, extra: ExtraContext<'a>) -> Self {
let cooperation = extra.cooperation.unwrap_or_else(|| unsafe {
// SAFETY: coop::Never is a unit struct ZST; it's
// `Box::leak(Box::new(coop::Never))` without alloc
&mut *(align_of::<coop::Never>() as *mut coop::Never)
});
Self { waker, cooperation, .. }
}
pub fn with_waker<'new>(&'new mut self, waker: &'new Waker) -> Context<'new> {
Context { waker, cooperation: self.cooperation, .. }
}
pub fn poll_cooperate(&mut self) -> Poll<coop::Cooperation<'_>> {
if self.cooperation.should_yield() {
Poll::Pending
} else {
Poll::Ready(coop::Cooperation(self.cooperation))
}
}
}
pub mod coop {
pub trait Cooperate {
fn should_yield(&self) -> bool;
fn record_progress(&mut self);
}
pub struct Never;
impl Cooperate for Never {
fn should_yield(&self) -> bool { false }
fn record_progress(&mut self) {}
}
pub struct Cooperation<'a>(pub(super) &'a mut (dyn Cooperate + 'a));
impl Cooperation<'_> {
pub fn record_progress(self) {
self.0.record_progress()
}
}
}
If we do much more than this, I'd want to see some small example of where more flexibility might be desirable. I've failed to construct an artificial scenario where this isn't sufficient.
The reason Waker uses a manual vtable is that it needs to be cloneable and owned. Cooperation I think doesn't, and in fact wants to not last beyond a single poll's context.
I do personally think that ultimately task::Context should be a Provider such that runtimes can dynamically thread extra context through which futures can then dynamically use if available. Cooperative yielding is actually a good application of this IMHO; if a task doesn't cooperate, it's an unfortunate pessimization but it's not going to result in incorrect results.
Rough "could look like" sketch
pub struct Context<'a> {
provider: &'a dyn Provider + 'a,
waker: &'a Waker,
}
impl<'_> Provider for Context<'_> {
fn provide<'a>(&'a self, demand: &mut Demand<'a>) {
demand.provide_ref(waker);
demand.provide_value_with(|| waker.clone());
self.provider.provide(demand);
}
}
impl<'a> Context<'a> {
pub fn from_waker(waker: &'a Waker) -> Self {
Self::from_waker_and_provider(waker, &NullProvider)
}
#[track_caller]
pub fn from_provider(provider: &'a dyn Provider + 'a) -> Self {
let waker = request_ref(provider)
.expect("context provider should provide a waker");
Self::from_waker_and_provider(waker, provider)
}
pub fn from_waker_and_provider(waker: &'a Waker, provider: &'a dyn Provider + 'a) -> Self {
Context { provider, waker }
}
pub fn with_waker<'b>(&mut self, waker: &'b Waker) -> Context<'b>
where
'a: 'b,
{
Context::from_waker_and_provider(waker, self.provider)
}
}
impl<'a> Context<'a> {
pub fn waker(&self) -> &'a Waker { self.waker }
pub fn poll_cooperate(&self) -> Poll<impl FnOnce() + '_> {
let coop = request_ref::<dyn Cooperate>();
if let Some(coop) = &coop && coop.should_yield() {
Pending
} else {
Ready(move || if let Some(coop) = coop {
coop.record_progress();
})
}
}
}
pub trait Cooperate {
fn should_yield(&self) -> bool;
fn record_progress(&self);
}
// NB: there's no mut-capable version of Provider currently.
// Provider could support demanding &mut 'static, it just doesn't currently.
// If there were, I'd use it, so recording progress could be mut.
// And that way, the budget can be owned
The waker is kept outside the provider for two main reasons:
- Roughly every future will access the waker, so keeping
cx.waker() nondynamic is probably beneficial.
- Creating a new context with a replaced waker is common (e.g. for
join/select combinators tracking which future got awoken). The waker reference being separated allows easy replacement, whereas it being behind Provider would require more work to chain Providers up the context stack.
- If futures want to add on things accessible via the provider API, they can still wrap and delegate up the context stack themselves, e.g. with
Context::with_waker_and_provider(&new_waker, cx).
The !Send + !Sync now comes from containing a dyn trait object. I relaxed the invariance to covariance because cx.with_waker(cx.waker()) basically does the same as covariance. (It could just drop any invariant context, but doing so is at best a footgun, since replacing wakers is a reasonably common thing to do. The current use of Context::from_waker to do so already drops any extra context we might attach in the future, which is unfortunate since that means any currently existing code which does that won't play well with that extra context.)
But even if arbitrary context can be provided, it still makes sense for std to provide dedicated types/accessors for widely applicable context, such as for cooperative yielding.
Where it becomes more interesting is when non-std context is required (e.g. Tokio IO futures which require being in a scoped Tokio context with an active reactor/runtime). Then you do want some marker with which to record your needed additional context. With the Provider API, that maps somewhat nicely as a const generic set of TypeId (for the lifetime-erased type tags, not the demanded type, which could be a lifetime-covariant type), but still carries a lot of questions on allowing providing (and propagating) "too much" guaranteed context.