twizzler_async/
throttle.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
use std::{
    cell::Cell,
    task::{Context, Poll},
};

use scoped_tls_hkt::scoped_thread_local;

scoped_thread_local! {
    static BUDGET: Cell<u32>
}

pub(crate) fn setup<T>(poll: impl FnOnce() -> T) -> T {
    BUDGET.set(&Cell::new(200), poll)
}

#[allow(dead_code)]
pub(crate) fn poll(cx: &mut Context<'_>) -> Poll<()> {
    if BUDGET.is_set() && BUDGET.with(|b| b.replace(b.get().saturating_sub(1))) == 0 {
        cx.waker().wake_by_ref();
        return Poll::Pending;
    }
    Poll::Ready(())
}