协作式调度预算

async fn drop_all<I: Stream + Unpin>(mut input: I) {
    while let Some(_) = input.next().await {}
}

pub(crate) struct Budget(Option<u8>);

impl Budget {
    // ...
    const fn initial() -> Budget {
        Budget(Some(128))
    }

    /// Returns an unconstrained budget. Operations will not be limited.
    pub(crate) const fn unconstrained() -> Budget {
        Budget(None)
    }

    fn has_remaining(self) -> bool {
        self.0.is_none_or(|budget| budget > 0)
    }
}

pub(crate) fn budget<R>(f: impl FnOnce() -> R) -> R {
    with_budget(Budget::initial(), f)
}

fn with_budget<R>(budget: Budget, f: impl FnOnce() -> R) -> R {
    struct ResetGuard {
        prev: Budget,
    }

    impl Drop for ResetGuard {
        fn drop(&mut self) {
            let _ = context::budget(|cell| {
                cell.set(self.prev);
            });
        }
    }

    #[allow(unused_variables)]
    let maybe_guard = context::budget(|cell| {
        let prev = cell.get();
        cell.set(budget);

        ResetGuard { prev }
    });

    // The function is called regardless even if the budget is not successfully
    // set due to the thread-local being destroyed.
    f()
}

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::new(cell.get());
            cell.set(budget);

            // avoid double counting
            if decrement.hit_zero {
                inc_budget_forced_yield_count();
            }

            Poll::Ready(restore)
        } else {
            register_waker(cx);
            Poll::Pending
        }
    }).unwrap_or(Poll::Ready(RestoreOnPending::new(Budget::unconstrained())))
}

// Keep track of task budget
let coop = ready!(crate::task::coop::poll_proceed(cx));

// ...
match rx_fields.list.pop(&self.inner.tx) {
    Some(Read::Value(value)) => {
        self.inner.semaphore.add_permit();
        coop.made_progress();
        return Ready(Some(value));
    }
    // ...
}

// Don't wake the task immediately, as that would push it right back
// onto the run queue and it could be polled again before other tasks
// or the IO/timer driver get a chance to run. Instead, hand the waker
// to the scheduler, which wakes deferred tasks only after it has run
// out of ready tasks and polled the driver. When polled from outside
// a Tokio runtime, the waker is woken immediately.
context::defer(cx.waker());

// Records the message count after every poll of the inner future.
struct Probe<F> {
  inner: Pin<Box<F>>,
  count: Arc<AtomicUsize>,
  log: Vec<usize>,
}

impl<F: Future<Output = ()>> Future for Probe<F> {
  type Output = Vec<usize>;
  fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Vec<usize>> {
    let r = self.inner.as_mut().poll(cx);
    let n = self.count.load(Relaxed);
    self.log.push(n);
    match r {
      Poll::Ready(()) => Poll::Ready(std::mem::take(&mut self.log)),
      Poll::Pending => Poll::Pending,
    }
  }
}

async fn run(unconstrained: bool) {
  let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
  for i in 0..1000 {
    tx.send(i).unwrap();
  }
  drop(tx);
  let count = Arc::new(AtomicUsize::new(0));
  let c = count.clone();
  let consume = async move {
    while rx.recv().await.is_some() {
      c.fetch_add(1, Relaxed);
    }
  };
  let log = if unconstrained {
    tokio::spawn(tokio::task::unconstrained(Probe { inner: Box::pin(consume), count, log: vec![] })).await
  } else {
    tokio::spawn(Probe { inner: Box::pin(consume), count, log: vec![] }).await
  };
  println!("unconstrained = {unconstrained}: count after each poll {:?}", log.unwrap());
}

#[tokio::main(flavor = "current_thread")]
async fn main() {
  run(false).await;
  run(true).await;
}

unconstrained = false: count after each poll [128, 256, 384, 512, 640, 768, 896, 1000]
unconstrained = true: count after each poll [1000]

pub async fn consume_budget() {
    let mut status = std::task::Poll::Pending;

    std::future::poll_fn(move |cx| {
        std::task::ready!(crate::trace::trace_leaf(cx));
        if status.is_ready() {
            return status;
        }
        status = crate::task::coop::poll_proceed(cx).map(|restore| {
            restore.made_progress();
        });
        status
    })
    .await
}

#[tokio::main(flavor = "current_thread")]
async fn main() {
  let (_keep, mut empty) = tokio::sync::mpsc::unbounded_channel::<u32>();
  let (tx, mut full) = tokio::sync::mpsc::unbounded_channel();
  for i in 0..500u32 {
    tx.send(i).unwrap();
  }
  let count = Arc::new(AtomicUsize::new(0));
  let c = count.clone();
  let a = async move {
    while c.load(Relaxed) < 500 {
      tokio::select! {
        biased;
        _ = empty.recv() => {}
        Some(_) = full.recv() => { c.fetch_add(1, Relaxed); }
      }
    }
  };
  let log = tokio::spawn(Probe { inner: Box::pin(a), count, log: vec![] }).await.unwrap();
  println!("count after each poll: {log:?}");
}

count after each poll: [128, 256, 384, 500]