MetaCrate/crates/libremetaverse/src/repeat.rs

use std::time::Duration;

use libremetaverse_types::compat::{CancellationToken, Task};

use crate::Error;

pub struct Repeat;

impl Repeat {
    pub async fn interval_with_time_span_action_cancellation_token_boolean(
        poll_interval: Duration,
        action: Box<dyn Fn() + Send + Sync>,
        token: CancellationToken,
        immediately: Option<bool>,
    ) -> Result<(), Error> {
        if poll_interval.is_zero() {
            return Err(Error::Argument);
        }
        if immediately.unwrap_or(false) {
            action();
        }
        loop {
            tokio::select! {
                () = tokio::time::sleep(poll_interval) => action(),
                () = token.cancelled() => return Ok(()),
            }
        }
    }

    pub async fn interval_with_time_span_func_cancellation_token_boolean(
        poll_interval: Duration,
        async_action: Box<dyn Fn() -> Task<()> + Send + Sync>,
        token: CancellationToken,
        immediately: Option<bool>,
    ) -> Result<(), Error> {
        if poll_interval.is_zero() {
            return Err(Error::Argument);
        }
        if immediately.unwrap_or(false) {
            async_action().await?;
        }
        loop {
            tokio::select! {
                () = tokio::time::sleep(poll_interval) => async_action().await?,
                () = token.cancelled() => return Ok(()),
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;
    use std::sync::atomic::{AtomicUsize, Ordering};

    use libremetaverse_types::compat::CancellationTokenSource;

    use super::*;

    #[tokio::test]
    async fn repeats_immediately_and_stops_on_cancellation() {
        let source = CancellationTokenSource::new();
        let token = source.token();
        let count = Arc::new(AtomicUsize::new(0));
        let action_count = Arc::clone(&count);
        let cancel = source.clone();
        Repeat::interval_with_time_span_action_cancellation_token_boolean(
            Duration::from_millis(1),
            Box::new(move || {
                if action_count.fetch_add(1, Ordering::Relaxed) >= 2 {
                    cancel.cancel();
                }
            }),
            token,
            Some(true),
        )
        .await
        .unwrap();
        assert_eq!(count.load(Ordering::Relaxed), 3);
    }
}