Line data    Source code

       1              : use tokio_util::sync::CancellationToken;
       2              : 
       3            0 : #[derive(thiserror::Error, Debug)]
       4              : pub enum YieldingLoopError {
       5              :     #[error("Cancelled")]
       6              :     Cancelled,
       7              : }
       8              : 
       9              : /// Helper for long synchronous loops, e.g. over all tenants in the system.  Periodically
      10              : /// yields to avoid blocking the executor, and after resuming checks the provided
      11              : /// cancellation token to drop out promptly on shutdown.
      12              : #[inline(always)]
      13            0 : pub async fn yielding_loop<I, T, F>(
      14            0 :     interval: usize,
      15            0 :     cancel: &CancellationToken,
      16            0 :     iter: I,
      17            0 :     mut visitor: F,
      18            0 : ) -> Result<(), YieldingLoopError>
      19            0 : where
      20            0 :     I: Iterator<Item = T>,
      21            0 :     F: FnMut(T),
      22            0 : {
      23            0 :     for (i, item) in iter.enumerate() {
      24            0 :         visitor(item);
      25            0 : 
      26            0 :         if i + 1 % interval == 0 {
      27            0 :             tokio::task::yield_now().await;
      28            0 :             if cancel.is_cancelled() {
      29            0 :                 return Err(YieldingLoopError::Cancelled);
      30            0 :             }
      31            0 :         }
      32              :     }
      33              : 
      34            0 :     Ok(())
      35            0 : }