TLA  Line data    Source code

       1                 : use tokio_util::sync::CancellationToken;
       2                 : 
       3 UBC           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 CBC        3327 : pub async fn yielding_loop<I, T, F>(
      14            3327 :     interval: usize,
      15            3327 :     cancel: &CancellationToken,
      16            3327 :     iter: I,
      17            3327 :     mut visitor: F,
      18            3327 : ) -> Result<(), YieldingLoopError>
      19            3327 : where
      20            3327 :     I: Iterator<Item = T>,
      21            3327 :     F: FnMut(T),
      22            3327 : {
      23            3327 :     for (i, item) in iter.enumerate() {
      24             841 :         visitor(item);
      25             841 : 
      26             841 :         if i + 1 % interval == 0 {
      27 UBC           0 :             tokio::task::yield_now().await;
      28               0 :             if cancel.is_cancelled() {
      29               0 :                 return Err(YieldingLoopError::Cancelled);
      30               0 :             }
      31 CBC         841 :         }
      32                 :     }
      33                 : 
      34            3327 :     Ok(())
      35            3327 : }