tokio-interval-set 0.1.0

//! # tokio-interval-set
//!
//! A lightweight utility for running a closure repeatedly on a fixed time
//! interval, backed by [`tokio`]'s time and task system.
//!
//! ## Example
//!
//! ```rust
//! use std::sync::Arc;
//! use std::time::Duration;
//! use tokio::sync::Mutex;
//! use tokio::time::interval;
//! use tokio_interval_set::set_interval;
//!
//! let rt = tokio::runtime::Builder::new_current_thread()
//!     .enable_time()
//!     .build()
//!     .unwrap();
//!
//! rt.block_on(async {
//!     let count = Arc::new(Mutex::new(0u64));
//!     let count_cloned = Arc::clone(&count);
//!
//!     let handle = set_interval(interval(Duration::from_millis(100)), move || {
//!         let count = Arc::clone(&count_cloned);
//!         async move {
//!             let mut c = count.lock().await;
//!             *c += 1;
//!         }
//!     });
//!
//!     tokio::time::sleep(Duration::from_millis(350)).await;
//!     handle.clear(); // stop the interval
//!
//!     let final_count = *count.lock().await;
//!     assert!(final_count >= 3, "expected at least 3 ticks, got {}", final_count);
//! });
//! ```

use tokio::{select, sync::oneshot, time::Interval};

/// A handle that controls a running interval task.
///
/// When the handle is **dropped** (without calling [`clear()`](IntervalHandle::clear)),
/// the associated interval task will be **cancelled** on the next tick boundary.
///
/// Call [`clear()`](IntervalHandle::clear) to synchronously wait for the
/// interval task to shut down (by sending a shutdown signal through the channel).
///
/// # Panics
///
/// [`clear()`](IntervalHandle::clear) will **panic** if the spawned interval
/// task has already panicked (detected by a closed channel).
#[derive(Debug)]
pub struct IntervalHandle {
    shutdown: oneshot::Sender<()>,
}

impl IntervalHandle {
    /// Stops the interval loop.
    ///
    /// Sends a shutdown signal to the spawned task. The loop will break on the
    /// next iteration, after any in-flight call to `func` completes.
    ///
    /// # Panics
    ///
    /// Panics with the message `"Interval was panic"` if the spawned task has
    /// already panicked (i.e., the oneshot receiver was dropped without receiving
    /// the shutdown signal).
    ///
    /// # Example
    ///
    /// ```rust
    /// use std::sync::Arc;
    /// use std::time::Duration;
    /// use tokio::sync::Mutex;
    /// use tokio::time::interval;
    /// use tokio_interval_set::set_interval;
    ///
    /// let rt = tokio::runtime::Builder::new_current_thread()
    ///     .enable_time()
    ///     .build()
    ///     .unwrap();
    ///
    /// rt.block_on(async {
    ///     let count = Arc::new(Mutex::new(0u64));
    ///     let count_clone = Arc::clone(&count);
    ///     let handle = set_interval(
    ///         interval(Duration::from_secs(1)),
    ///         move || {
    ///             let count = Arc::clone(&count_clone);
    ///             async move {
    ///                 let mut c = count.lock().await;
    ///                 *c += 1;
    ///             }
    ///         },
    ///     );
    ///
    ///     tokio::time::sleep(Duration::from_secs(3)).await;
    ///     handle.clear(); // ✅ clean shutdown
    ///     assert_eq!(*count.lock().await, 3);
    /// });
    /// ```
    pub fn clear(self) {
        if self.shutdown.is_closed() {
            panic!("Interval was panic");
        }
        self.shutdown.send(()).unwrap();
    }
}

/// Runs `func` repeatedly on the given `interval`.
///
/// The function returns an [`IntervalHandle`] that can be used to stop the loop
/// via [`clear()`](IntervalHandle::clear). If the handle is **dropped** without
/// calling `clear()`, the underlying task will be cancelled when the interval
/// next ticks (the oneshot sender is dropped, causing the receiver to wake up
/// with a closed error).
///
/// # Type Parameters
///
/// * `F` — The closure type. Must be [`Fn`] (callable multiple times),
///   [`Send`] + [`Sync`] + `'static`.
/// * `Fut` — The future returned by the closure. Must be [`Send`] + `'static`.
///
/// # Arguments
///
/// * `interval` — A [`tokio::time::Interval`] that controls the tick timing.
/// * `func` — A closure that returns a future. Called on every tick.
///
/// # Panics
///
/// If the spawned tokio task panics (e.g., because `func` panics), the
/// [`IntervalHandle::clear()`] method will panic when called.
///
/// # Example
///
/// ```rust
/// use std::sync::Arc;
/// use std::time::Duration;
/// use tokio::sync::Mutex;
/// use tokio::time::interval;
/// use tokio_interval_set::set_interval;
///
/// let rt = tokio::runtime::Builder::new_current_thread()
///     .enable_time()
///     .build()
///     .unwrap();
///
/// rt.block_on(async {
///     let values = Arc::new(Mutex::new(Vec::new()));
///     let values_clone = Arc::clone(&values);
///
///     let handle = set_interval(
///         interval(Duration::from_millis(100)),
///         move || {
///             let values = Arc::clone(&values_clone);
///             async move {
///                 let mut v = values.lock().await;
///                 let idx = v.len();
///                 v.push(idx);
///             }
///         },
///     );
///
///     tokio::time::sleep(Duration::from_millis(250)).await;
///     handle.clear();
///
///     let v = values.lock().await;
///     assert!(v.len() >= 2, "expected at least 2 items, got {}", v.len());
/// });
/// ```
pub fn set_interval<F, Fut>(mut interval: Interval, func: F) -> IntervalHandle
where
    F: Fn() -> Fut + Sync + Send + 'static,
    Fut: Future<Output = ()> + Send,
{
    let (tx, mut rx) = oneshot::channel();

    tokio::spawn(async move {
        loop {
            select! {
                _ = interval.tick() => {
                    func().await;
                },
                _ = &mut rx => break,
            }
        }
    });

    IntervalHandle { shutdown: tx }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------
#[cfg(test)]
mod tests {
    use std::{sync::Arc, time::Duration};

    use tokio::{sync::Mutex, time::interval};

    use super::*;

    /// Verify that the callback fires the expected number of times.
    ///
    /// We check for a range (`EXPECTED` or `EXPECTED + 1`) to tolerate small
    /// timing jitter inherent in async test execution.
    #[test]
    fn tick_count() {
        let rt = tokio::runtime::Builder::new_current_thread()
            .enable_time()
            .build()
            .unwrap();

        const INTERVAL_MS: u64 = 50;
        const EXPECTED_TICKS: u64 = 5;
        // Sleep for slightly more than `INTERVAL_MS * EXPECTED_TICKS` so we
        // definitely observe the Nth tick, but accept one extra if timing is
        // slightly off.
        const SLEEP_MS: u64 = INTERVAL_MS * EXPECTED_TICKS + INTERVAL_MS / 2;

        rt.block_on(async {
            let count = Arc::new(Mutex::new(0u64));
            let count_cloned = Arc::clone(&count);

            let _handle = set_interval(
                interval(Duration::from_millis(INTERVAL_MS)),
                move || {
                    let count = Arc::clone(&count_cloned);
                    async move {
                        let mut c = count.lock().await;
                        *c += 1;
                    }
                },
            );

            tokio::time::sleep(Duration::from_millis(SLEEP_MS)).await;
            let final_count = *count.lock().await;
            assert!(
                final_count == EXPECTED_TICKS || final_count == EXPECTED_TICKS + 1,
                "expected {EXPECTED_TICKS} or {} ticks, got {final_count}",
                EXPECTED_TICKS + 1
            );
        });
    }

    /// Verify that calling `clear()` stops the interval loop immediately
    /// (no more ticks after the signal is received).
    #[test]
    fn clear_stops_interval() {
        let rt = tokio::runtime::Builder::new_current_thread()
            .enable_time()
            .build()
            .unwrap();

        const INTERVAL_MS: u64 = 50;

        rt.block_on(async {
            let count = Arc::new(Mutex::new(0u64));
            let count_cloned = Arc::clone(&count);

            let handle = set_interval(
                interval(Duration::from_millis(INTERVAL_MS)),
                move || {
                    let count = Arc::clone(&count_cloned);
                    async move {
                        let mut c = count.lock().await;
                        *c += 1;
                    }
                },
            );

            // Let it tick a couple of times
            tokio::time::sleep(Duration::from_millis(120)).await;
            let after_two = *count.lock().await;
            assert!(
                after_two >= 2,
                "expected at least 2 ticks before clear, got {after_two}"
            );

            // Stop the interval
            handle.clear();

            // Capture the count *right after* clear
            let count_after_clear = *count.lock().await;

            // Wait significantly longer — if the loop were still running,
            // the count would increase noticeably.
            tokio::time::sleep(Duration::from_millis(300)).await;
            let final_count = *count.lock().await;

            assert_eq!(
                final_count, count_after_clear,
                "count increased after clear(): was {count_after_clear}, now {final_count}"
            );
        });
    }

    /// Verify that `clear()` panics when the spawned task has already panicked.
    #[test]
    #[should_panic(expected = "Interval was panic")]
    fn clear_panics_on_panicked_interval() {
        let rt = tokio::runtime::Builder::new_current_thread()
            .enable_time()
            .build()
            .unwrap();

        rt.block_on(async {
            let handle = set_interval(interval(Duration::from_millis(10)), move || {
                async move {
                    panic!("intentional panic in callback");
                }
            });

            // Wait long enough for the spawned task to panic
            tokio::time::sleep(Duration::from_millis(100)).await;

            // The receiver was dropped when the task panicked, so
            // is_closed() returns true → clear() should panic.
            handle.clear();
        });
    }

    /// Verify that dropping the handle (without calling clear) also stops
    /// the interval (the oneshot sender is dropped → receiver gets closed).
    #[test]
    fn drop_handle_stops_interval() {
        let rt = tokio::runtime::Builder::new_current_thread()
            .enable_time()
            .build()
            .unwrap();

        const INTERVAL_MS: u64 = 50;

        rt.block_on(async {
            let count = Arc::new(Mutex::new(0u64));
            let count_cloned = Arc::clone(&count);

            let handle = set_interval(
                interval(Duration::from_millis(INTERVAL_MS)),
                move || {
                    let count = Arc::clone(&count_cloned);
                    async move {
                        let mut c = count.lock().await;
                        *c += 1;
                    }
                },
            );

            // Let it tick a couple of times
            tokio::time::sleep(Duration::from_millis(120)).await;
            assert!(*count.lock().await >= 2);

            // Drop the handle (equivalent to letting it go out of scope)
            drop(handle);

            let count_after_drop = *count.lock().await;

            // Wait — if the loop were alive the count would increase
            tokio::time::sleep(Duration::from_millis(300)).await;
            let final_count = *count.lock().await;

            assert_eq!(
                final_count, count_after_drop,
                "count increased after handle dropped: was {count_after_drop}, now {final_count}"
            );
        });
    }
}