use std::sync::atomic::AtomicUsize;
pub struct EveryNQueryYielder {
counter: AtomicUsize,
n: usize,
}
impl EveryNQueryYielder {
pub const fn new(n: usize) -> Self {
Self { counter: AtomicUsize::new(0), n: n + 1 }
}
pub async fn tick(&self) {
let count =
self.counter.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
if count.is_multiple_of(self.n) {
tokio::task::yield_now().await;
}
}
}
pub enum Yielder {
Never,
EveryNQuery(EveryNQueryYielder),
}
impl Yielder {
pub const fn every_n_query(n: usize) -> Self {
Self::EveryNQuery(EveryNQueryYielder::new(n))
}
pub const fn never() -> Self { Self::Never }
pub async fn tick(&self) {
match self {
Self::Never => {}
Self::EveryNQuery(yielder) => yielder.tick().await,
}
}
}