neon/libs/utils/src/yielding_loop.rs

use tokio_util::sync::CancellationToken;

#[derive(thiserror::Error, Debug)]
pub enum YieldingLoopError {
    #[error("Cancelled")]
    Cancelled,
}

/// Helper for long synchronous loops, e.g. over all tenants in the system.  Periodically
/// yields to avoid blocking the executor, and after resuming checks the provided
/// cancellation token to drop out promptly on shutdown.
#[inline(always)]
pub async fn yielding_loop<I, T, F>(
    interval: usize,
    cancel: &CancellationToken,
    iter: I,
    mut visitor: F,
) -> Result<(), YieldingLoopError>
where
    I: Iterator<Item = T>,
    F: FnMut(T),
{
    for (i, item) in iter.enumerate() {
        visitor(item);

        if i + 1 % interval == 0 {
            tokio::task::yield_now().await;
            if cancel.is_cancelled() {
                return Err(YieldingLoopError::Cancelled);
            }
        }
    }

    Ok(())
}