Add task::consume_budget util for cooperative scheduling

[psarna]

Motivation

For cpu-only computations that do not use any Tokio resources, budgeting does not really kick in in order to yield and prevent other tasks from starvation. One way to take part in budgeting is to use a Tokio resource, but some workloads simply don't need any, as they just need to compute a lot and occasionally yield in order to keep latencies low (e.g. a Wasm engine).

Solution

The new mechanism - task::consume_budget, performs a budget check and consumes a unit of it as if a Tokio resource was used (except it wasn't, so it's cheaper) and yields only if the task exceeded the budget. That allows cpu-intenstive computations to define yield points that yield only conditionally, since unconditional yielding to runtime is a potentially heavy operation.

[carllerche]

I think this is probably a good feature to add. I would probably not name it "maybe_yield_now()". Instead, name it something to reflect that it represents conceptually tracking work.

task::did_work().await
task::consume_work_budget().await

^^ Don't love either of those names, but yeah.

[psarna]

Right, the name should reflect that there's a side effect - burning budget. I'll figure out a new name and send it in v2, along with the tests

[psarna]

v2:

• renamed the utility to task::consume_budget
• added a test case
• simplified the code by using Poll:map instead of an explicit match; possible, because coop::poll_proceed already wakes up the task, so it doesn't need to be done twice

[carllerche]

@jonhoo I would love to get your thoughts on naming and the feature in general.

[hawkw]

Overall, I think something like this is probably a good idea. I had some minor nitpicks and suggestions, but they're not hard blockers.

If we want to avoid a very long bikeshed over this API, it occurs to me that we could just go ahead and ship it as an unstable feature, so people can start experimenting with it, and then we can come back and stabilize the API separately?

[hawkw]

one note about the documentation for this is that this is the first actual mention of the cooperative scheduling budget in the API documentation. I think it is probably worth adding a high-level summary of the task budget somewhere in the tokio::task documentation and having this API documentation reference that? Currently, the task budget is basically only documented in this blog post which seems maybe not ideal...

[hawkw]

i wonder if we should pre-emptively mark this as an unstable feature, so that we can ship it sooner in the next release while reserving the ability to bikeshed the API for a bit longer?

[psarna]

Could you point me to a code example showing how to mark a feature as unstable?

[psarna]

nvm, just found #4499, I'll follow it as a guideline

[Darksonn]

If I view the generated docs, then it is not listed as unstable. You should add that to the doc(cfg(..)).

#[cfg_attr(docsrs, doc(cfg(all(tokio_unstable, feature = "rt"))))]

[psarna]

done, thanks!

[hawkw]

nit, take it or leave it: i think we could probably implement this a bit more simply using poll_fn:

Suggested change

	struct ConsumeBudget {
	status: Poll<()>,
	}
	impl Future for ConsumeBudget {
	type Output = ();
	fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
	if self.status.is_ready() {
	return self.status;
	}
	self.status = crate::coop::poll_proceed(cx).map(|restore| {
	restore.made_progress();
	});
	self.status
	}
	}
	ConsumeBudget { status: Poll::Pending }.await
	let mut status = Poll::Pending;
	crate::future::poll_fn(move |cx| {
	if status.is_ready() {
	return status;
	}
	status = crate::coop::poll_proceed(cx).map(|restore| {
	restore.made_progress();
	});
	status
	}).await

i believe this should work identically, and avoids a little bit of boilerplate by not having to impl Future for a type.

[psarna]

Neat! Looks much better, thanks

[hawkw]

nit: how about

Suggested change

	/// runtime *if* the task went out ouf budget.
	/// runtime *if* the task's coop budget was exhausted.

[hawkw]

nit: maybe

Suggested change

	/// The task will only yield if it ran out of its coop budget.
	/// It can be used in order to insert optional yield points into long
	/// The task will only yield if its entire coop budget has been exhausted.
	/// This function can be used in order to insert optional yield points into long

[hawkw]

nit:

Suggested change

	/// without redundantly yielding to runtime each time.
	/// without redundantly yielding to the runtime each time.

[hawkw]

also, it could be nice to have an example in the documentation, maybe showing this in a long-running loop or something?

[jonhoo]

We can always make this a no-op down the line should we decide that coop isn't a thing we want, and make it forward to an external lib if that's where coop ends up living, so I'm 👍 on this in general. In fact, a method like this existed in the original implementation in #2160, initially called cooperate, and later renamed to coop::proceed. It was never made publicly available though, and was removed as part of the coop simplification in #2498. I'm happy to see it coming back. I do think we should likely also expose poll_proceed, since coop is low enough level that it's fairly likely a pollable API will be needed.

[psarna]

v3:

• marked the feature as unstable
• applied the wording suggestions in comments (thanks!)
• added a simple example to the doc
• refactored the implementation to use poll_fn instead, as recommended (thanks again!)

I also retested everything with the additional RUSTFLAGS='--cfg tokio_unstable' flag, so that my new test would also get executed. Is it the correct way to go?

[psarna]

This failure in CI/test tokio full run looks unrelated to my patch - it errors out on some macro tests that I haven't touched and I'm able to reproduce the same failures on current master branch.

[Darksonn]

Those errors are due to the previous release of rustc. Please rebase or merge master into your branch to fix them.

[psarna]

got it, rebased, thanks

[psarna]

(ping)

[psarna]

(rebased on top of newest master)

[psarna]

(monthly rebase)

[Darksonn]

There are some other coop APIs that I wonder if they are more important to add. (Mainly, a public version of our poll_proceed.) However, this function does seem useful too, and it is marked as unstable, so it seems fine to just go ahead and merge this.

[Darksonn]

Sorry for taking so long.

[psarna]

Thanks! No worries, I often take active part in review processes that take years to finish, this one was actually quite painless and swift (: