tokio-rs/tokio
1
0
Fork 0
mirror of https://github.com/tokio-rs/tokio.git synced 2026-08-21 00:00:10 +02:00
Code Issues Packages Projects Releases Wiki Activity

coop: add cooperative and poll_proceed (#7405)

Browse Source
This commit is contained in:
Pepijn Van Eeckhoudt
2025-07-11 08:07:44 +00:00
committed by GitHub
parent 888ee60e41
commit 611b793356
2 changed files with 137 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tokio/src/task/coop/mod.rs
+136 -26
View File
@@ -250,14 +250,27 @@ cfg_coop! {
use pin_project_lite::pin_project;
use std::cell::Cell;
use std::future::Future;
use std::marker::PhantomData;
use std::pin::Pin;
use std::task::{ready, Context, Poll};
/// Value returned by the [`poll_proceed`] method.
#[derive(Debug)]
#[must_use]
pub(crate) struct RestoreOnPending(Cell<Budget>);
pub struct RestoreOnPending(Cell<Budget>, PhantomData<*mut ()>);
impl RestoreOnPending {
pub(crate) fn made_progress(&self) {
fn new(budget: Budget) -> Self {
RestoreOnPending(
Cell::new(budget),
PhantomData,
)
}
/// Signals that the task that obtained this `RestoreOnPending` was able to make
/// progress. This prevents the task budget from being restored to the value
/// it had prior to obtaining this instance when it is dropped.
pub fn made_progress(&self) {
self.0.set(Budget::unconstrained());
}
}
@@ -275,27 +288,102 @@ cfg_coop! {
}
}
/// Returns `Poll::Pending` if the current task has exceeded its budget and should yield.
/// Decrements the task budget and returns [`Poll::Pending`] if the budget is depleted.
/// This indicates that the task should yield to the scheduler. Otherwise, returns
/// [`RestoreOnPending`] which can be used to commit the budget consumption.
///
/// When you call this method, the current budget is decremented. However, to ensure that
/// progress is made every time a task is polled, the budget is automatically restored to its
/// former value if the returned `RestoreOnPending` is dropped. It is the caller's
/// responsibility to call `RestoreOnPending::made_progress` if it made progress, to ensure
/// that the budget empties appropriately.
/// The returned [`RestoreOnPending`] will revert the budget to its former
/// value when dropped unless [`RestoreOnPending::made_progress`]
/// is called. It is the caller's responsibility to do so when it _was_ able to
/// make progress after the call to [`poll_proceed`].
/// Restoring the budget automatically ensures the task can try to make progress in some other
/// way.
///
/// Note that `RestoreOnPending` restores the budget **as it was before `poll_proceed`**.
/// Therefore, if the budget is _further_ adjusted between when `poll_proceed` returns and
/// `RestRestoreOnPending` is dropped, those adjustments are erased unless the caller indicates
/// Note that [`RestoreOnPending`] restores the budget **as it was before [`poll_proceed`]**.
/// Therefore, if the budget is _further_ adjusted between when [`poll_proceed`] returns and
/// [`RestoreOnPending`] is dropped, those adjustments are erased unless the caller indicates
/// that progress was made.
///
/// # Examples
///
/// This example shows a simple countdown latch that uses [`poll_proceed`] to participate in
/// cooperative scheduling.
///
/// ```
/// use std::future::{Future};
/// use std::pin::Pin;
/// use std::task::{ready, Context, Poll, Waker};
/// use tokio::task::coop;
///
/// struct CountdownLatch<T> {
/// counter: usize,
/// value: Option<T>,
/// waker: Option<Waker>
/// }
///
/// impl<T> CountdownLatch<T> {
/// fn new(value: T, count: usize) -> Self {
/// CountdownLatch {
/// counter: count,
/// value: Some(value),
/// waker: None
/// }
/// }
/// fn count_down(&mut self) {
/// if self.counter <= 0 {
/// return;
/// }
///
/// self.counter -= 1;
/// if self.counter == 0 {
/// if let Some(w) = self.waker.take() {
/// w.wake();
/// }
/// }
/// }
/// }
///
/// impl<T> Future for CountdownLatch<T> {
/// type Output = T;
///
/// fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
/// // `poll_proceed` checks with the runtime if this task is still allowed to proceed
/// // with performing work.
/// // If not, `Pending` is returned and `ready!` ensures this function returns.
/// // If we are allowed to proceed, coop now represents the budget consumption
/// let coop = ready!(coop::poll_proceed(cx));
///
/// // Get a mutable reference to the CountdownLatch
/// let this = Pin::get_mut(self);
///
/// // Next we check if the latch is ready to release its value
/// if this.counter == 0 {
/// let t = this.value.take();
/// // The latch made progress so call `made_progress` to ensure the budget
/// // is not reverted.
/// coop.made_progress();
/// Poll::Ready(t.unwrap())
/// } else {
/// // If the latch is not ready so return pending and simply drop `coop`.
/// // This will restore the budget making it available again to perform any
/// // other work.
/// this.waker = Some(cx.waker().clone());
/// Poll::Pending
/// }
/// }
/// }
///
/// impl<T> Unpin for CountdownLatch<T> {}
/// ```
#[inline]
pub(crate) fn poll_proceed(cx: &mut Context<'_>) -> Poll<RestoreOnPending> {
pub fn poll_proceed(cx: &mut Context<'_>) -> Poll<RestoreOnPending> {
context::budget(|cell| {
let mut budget = cell.get();
let decrement = budget.decrement();
if decrement.success {
let restore = RestoreOnPending(Cell::new(cell.get()));
let restore = RestoreOnPending::new(cell.get());
cell.set(budget);
// avoid double counting
@@ -308,7 +396,7 @@ cfg_coop! {
register_waker(cx);
Poll::Pending
}
}).unwrap_or(Poll::Ready(RestoreOnPending(Cell::new(Budget::unconstrained()))))
}).unwrap_or(Poll::Ready(RestoreOnPending::new(Budget::unconstrained())))
}
/// Returns `Poll::Ready` if the current task has budget to consume, and `Poll::Pending` otherwise.
@@ -380,15 +468,9 @@ cfg_coop! {
}
pin_project! {
/// Future wrapper to ensure cooperative scheduling.
///
/// When being polled `poll_proceed` is called before the inner future is polled to check
/// if the inner future has exceeded its budget. If the inner future resolves, this will
/// automatically call `RestoreOnPending::made_progress` before resolving this future with
/// the result of the inner one. If polling the inner future is pending, polling this future
/// type will also return a `Poll::Pending`.
/// Future wrapper to ensure cooperative scheduling created by [`cooperative`].
#[must_use = "futures do nothing unless polled"]
pub(crate) struct Coop<F: Future> {
pub struct Coop<F: Future> {
#[pin]
pub(crate) fut: F,
}
@@ -409,11 +491,39 @@ cfg_coop! {
}
}
/// Run a future with a budget constraint for cooperative scheduling.
/// If the future exceeds its budget while being polled, control is yielded back to the
/// runtime.
/// Creates a wrapper future that makes the inner future cooperate with the Tokio scheduler.
///
/// When polled, the wrapper will first call [`poll_proceed`] to consume task budget, and
/// immediately yield if the budget has been depleted. If budget was available, the inner future
/// is polled. The budget consumption will be made final using [`RestoreOnPending::made_progress`]
/// if the inner future resolves to its final value.
///
/// # Examples
///
/// When you call `recv` on the `Receiver` of a [`tokio::sync::mpsc`](crate::sync::mpsc)
/// channel, task budget will automatically be consumed when the next value is returned.
/// This makes tasks that use Tokio mpsc channels automatically cooperative.
///
/// If you're using [`futures::channel::mpsc`](https://docs.rs/futures/latest/futures/channel/mpsc/index.html)
/// instead, automatic task budget consumption will not happen. This example shows how can use
/// `cooperative` to make `futures::channel::mpsc` channels cooperate with the scheduler in the
/// same way Tokio channels do.
///
/// ```
/// use tokio::task::coop::cooperative;
/// use futures::channel::mpsc::Receiver;
/// use futures::stream::StreamExt;
///
/// async fn receive_next<T>(receiver: &mut Receiver<T>) -> Option<T> {
/// // Use `StreamExt::next` to obtain a `Future` that resolves to the next value
/// let recv_future = receiver.next();
/// // Wrap it a cooperative wrapper
/// let coop_future = cooperative(recv_future);
/// // And await
/// coop_future.await
/// }
#[inline]
pub(crate) fn cooperative<F: Future>(fut: F) -> Coop<F> {
pub fn cooperative<F: Future>(fut: F) -> Coop<F> {
Coop { fut }
}
}
tokio/tests/async_send_sync.rs
+1
View File
@@ -454,6 +454,7 @@ assert_value!(tokio::task::JoinSet<NN>: !Send & !Sync & Unpin);
assert_value!(tokio::task::JoinSet<YN>: Send & Sync & Unpin);
assert_value!(tokio::task::JoinSet<YY>: Send & Sync & Unpin);
assert_value!(tokio::task::LocalSet: !Send & !Sync & Unpin);
assert_value!(tokio::task::coop::RestoreOnPending: !Send & !Sync & Unpin);
async_assert_fn!(tokio::sync::Barrier::wait(_): Send & Sync & !Unpin);
async_assert_fn!(tokio::sync::Mutex<NN>::lock(_): !Send & !Sync & !Unpin);
async_assert_fn!(tokio::sync::Mutex<NN>::lock_owned(_): !Send & !Sync & !Unpin);