tokio_interval_set/lib.rs
1//! # tokio-interval-set
2//!
3//! A lightweight utility for running a closure repeatedly on a fixed time
4//! interval, backed by [`tokio`]'s time and task system.
5//!
6//! ## Example
7//!
8//! ```rust
9//! use std::sync::Arc;
10//! use std::time::Duration;
11//! use tokio::sync::Mutex;
12//! use tokio::time::interval;
13//! use tokio_interval_set::set_interval;
14//!
15//! let rt = tokio::runtime::Builder::new_current_thread()
16//! .enable_time()
17//! .build()
18//! .unwrap();
19//!
20//! rt.block_on(async {
21//! let count = Arc::new(Mutex::new(0u64));
22//! let count_cloned = Arc::clone(&count);
23//!
24//! let handle = set_interval(interval(Duration::from_millis(100)), move || {
25//! let count = Arc::clone(&count_cloned);
26//! async move {
27//! let mut c = count.lock().await;
28//! *c += 1;
29//! }
30//! });
31//!
32//! tokio::time::sleep(Duration::from_millis(350)).await;
33//! handle.clear(); // stop the interval
34//!
35//! let final_count = *count.lock().await;
36//! assert!(final_count >= 3, "expected at least 3 ticks, got {}", final_count);
37//! });
38//! ```
39
40use tokio::{select, sync::oneshot, time::Interval};
41
42/// A handle that controls a running interval task.
43///
44/// When the handle is **dropped** (without calling [`clear()`](IntervalHandle::clear)),
45/// the associated interval task will be **cancelled** on the next tick boundary.
46///
47/// Call [`clear()`](IntervalHandle::clear) to synchronously wait for the
48/// interval task to shut down (by sending a shutdown signal through the channel).
49///
50/// # Panics
51///
52/// [`clear()`](IntervalHandle::clear) will **panic** if the spawned interval
53/// task has already panicked (detected by a closed channel).
54#[derive(Debug)]
55pub struct IntervalHandle {
56 shutdown: oneshot::Sender<()>,
57}
58
59impl IntervalHandle {
60 /// Stops the interval loop.
61 ///
62 /// Sends a shutdown signal to the spawned task. The loop will break on the
63 /// next iteration, after any in-flight call to `func` completes.
64 ///
65 /// # Panics
66 ///
67 /// Panics with the message `"Interval was panic"` if the spawned task has
68 /// already panicked (i.e., the oneshot receiver was dropped without receiving
69 /// the shutdown signal).
70 ///
71 /// # Example
72 ///
73 /// ```rust
74 /// use std::sync::Arc;
75 /// use std::time::Duration;
76 /// use tokio::sync::Mutex;
77 /// use tokio::time::interval;
78 /// use tokio_interval_set::set_interval;
79 ///
80 /// let rt = tokio::runtime::Builder::new_current_thread()
81 /// .enable_time()
82 /// .build()
83 /// .unwrap();
84 ///
85 /// rt.block_on(async {
86 /// let count = Arc::new(Mutex::new(0u64));
87 /// let count_clone = Arc::clone(&count);
88 /// let handle = set_interval(
89 /// interval(Duration::from_secs(1)),
90 /// move || {
91 /// let count = Arc::clone(&count_clone);
92 /// async move {
93 /// let mut c = count.lock().await;
94 /// *c += 1;
95 /// }
96 /// },
97 /// );
98 ///
99 /// tokio::time::sleep(Duration::from_secs(3)).await;
100 /// handle.clear(); // ✅ clean shutdown
101 /// assert_eq!(*count.lock().await, 3);
102 /// });
103 /// ```
104 pub fn clear(self) {
105 if self.shutdown.is_closed() {
106 panic!("Interval was panic");
107 }
108 self.shutdown.send(()).unwrap();
109 }
110}
111
112/// Runs `func` repeatedly on the given `interval`.
113///
114/// The function returns an [`IntervalHandle`] that can be used to stop the loop
115/// via [`clear()`](IntervalHandle::clear). If the handle is **dropped** without
116/// calling `clear()`, the underlying task will be cancelled when the interval
117/// next ticks (the oneshot sender is dropped, causing the receiver to wake up
118/// with a closed error).
119///
120/// # Type Parameters
121///
122/// * `F` — The closure type. Must be [`Fn`] (callable multiple times),
123/// [`Send`] + [`Sync`] + `'static`.
124/// * `Fut` — The future returned by the closure. Must be [`Send`] + `'static`.
125///
126/// # Arguments
127///
128/// * `interval` — A [`tokio::time::Interval`] that controls the tick timing.
129/// * `func` — A closure that returns a future. Called on every tick.
130///
131/// # Panics
132///
133/// If the spawned tokio task panics (e.g., because `func` panics), the
134/// [`IntervalHandle::clear()`] method will panic when called.
135///
136/// # Example
137///
138/// ```rust
139/// use std::sync::Arc;
140/// use std::time::Duration;
141/// use tokio::sync::Mutex;
142/// use tokio::time::interval;
143/// use tokio_interval_set::set_interval;
144///
145/// let rt = tokio::runtime::Builder::new_current_thread()
146/// .enable_time()
147/// .build()
148/// .unwrap();
149///
150/// rt.block_on(async {
151/// let values = Arc::new(Mutex::new(Vec::new()));
152/// let values_clone = Arc::clone(&values);
153///
154/// let handle = set_interval(
155/// interval(Duration::from_millis(100)),
156/// move || {
157/// let values = Arc::clone(&values_clone);
158/// async move {
159/// let mut v = values.lock().await;
160/// let idx = v.len();
161/// v.push(idx);
162/// }
163/// },
164/// );
165///
166/// tokio::time::sleep(Duration::from_millis(250)).await;
167/// handle.clear();
168///
169/// let v = values.lock().await;
170/// assert!(v.len() >= 2, "expected at least 2 items, got {}", v.len());
171/// });
172/// ```
173pub fn set_interval<F, Fut>(mut interval: Interval, func: F) -> IntervalHandle
174where
175 F: Fn() -> Fut + Sync + Send + 'static,
176 Fut: Future<Output = ()> + Send,
177{
178 let (tx, mut rx) = oneshot::channel();
179
180 tokio::spawn(async move {
181 loop {
182 select! {
183 _ = interval.tick() => {
184 func().await;
185 },
186 _ = &mut rx => break,
187 }
188 }
189 });
190
191 IntervalHandle { shutdown: tx }
192}
193
194// ---------------------------------------------------------------------------
195// Tests
196// ---------------------------------------------------------------------------
197#[cfg(test)]
198mod tests {
199 use std::{sync::Arc, time::Duration};
200
201 use tokio::{sync::Mutex, time::interval};
202
203 use super::*;
204
205 /// Verify that the callback fires the expected number of times.
206 ///
207 /// We check for a range (`EXPECTED` or `EXPECTED + 1`) to tolerate small
208 /// timing jitter inherent in async test execution.
209 #[test]
210 fn tick_count() {
211 let rt = tokio::runtime::Builder::new_current_thread()
212 .enable_time()
213 .build()
214 .unwrap();
215
216 const INTERVAL_MS: u64 = 50;
217 const EXPECTED_TICKS: u64 = 5;
218 // Sleep for slightly more than `INTERVAL_MS * EXPECTED_TICKS` so we
219 // definitely observe the Nth tick, but accept one extra if timing is
220 // slightly off.
221 const SLEEP_MS: u64 = INTERVAL_MS * EXPECTED_TICKS + INTERVAL_MS / 2;
222
223 rt.block_on(async {
224 let count = Arc::new(Mutex::new(0u64));
225 let count_cloned = Arc::clone(&count);
226
227 let _handle = set_interval(
228 interval(Duration::from_millis(INTERVAL_MS)),
229 move || {
230 let count = Arc::clone(&count_cloned);
231 async move {
232 let mut c = count.lock().await;
233 *c += 1;
234 }
235 },
236 );
237
238 tokio::time::sleep(Duration::from_millis(SLEEP_MS)).await;
239 let final_count = *count.lock().await;
240 assert!(
241 final_count == EXPECTED_TICKS || final_count == EXPECTED_TICKS + 1,
242 "expected {EXPECTED_TICKS} or {} ticks, got {final_count}",
243 EXPECTED_TICKS + 1
244 );
245 });
246 }
247
248 /// Verify that calling `clear()` stops the interval loop immediately
249 /// (no more ticks after the signal is received).
250 #[test]
251 fn clear_stops_interval() {
252 let rt = tokio::runtime::Builder::new_current_thread()
253 .enable_time()
254 .build()
255 .unwrap();
256
257 const INTERVAL_MS: u64 = 50;
258
259 rt.block_on(async {
260 let count = Arc::new(Mutex::new(0u64));
261 let count_cloned = Arc::clone(&count);
262
263 let handle = set_interval(
264 interval(Duration::from_millis(INTERVAL_MS)),
265 move || {
266 let count = Arc::clone(&count_cloned);
267 async move {
268 let mut c = count.lock().await;
269 *c += 1;
270 }
271 },
272 );
273
274 // Let it tick a couple of times
275 tokio::time::sleep(Duration::from_millis(120)).await;
276 let after_two = *count.lock().await;
277 assert!(
278 after_two >= 2,
279 "expected at least 2 ticks before clear, got {after_two}"
280 );
281
282 // Stop the interval
283 handle.clear();
284
285 // Capture the count *right after* clear
286 let count_after_clear = *count.lock().await;
287
288 // Wait significantly longer — if the loop were still running,
289 // the count would increase noticeably.
290 tokio::time::sleep(Duration::from_millis(300)).await;
291 let final_count = *count.lock().await;
292
293 assert_eq!(
294 final_count, count_after_clear,
295 "count increased after clear(): was {count_after_clear}, now {final_count}"
296 );
297 });
298 }
299
300 /// Verify that `clear()` panics when the spawned task has already panicked.
301 #[test]
302 #[should_panic(expected = "Interval was panic")]
303 fn clear_panics_on_panicked_interval() {
304 let rt = tokio::runtime::Builder::new_current_thread()
305 .enable_time()
306 .build()
307 .unwrap();
308
309 rt.block_on(async {
310 let handle = set_interval(interval(Duration::from_millis(10)), move || {
311 async move {
312 panic!("intentional panic in callback");
313 }
314 });
315
316 // Wait long enough for the spawned task to panic
317 tokio::time::sleep(Duration::from_millis(100)).await;
318
319 // The receiver was dropped when the task panicked, so
320 // is_closed() returns true → clear() should panic.
321 handle.clear();
322 });
323 }
324
325 /// Verify that dropping the handle (without calling clear) also stops
326 /// the interval (the oneshot sender is dropped → receiver gets closed).
327 #[test]
328 fn drop_handle_stops_interval() {
329 let rt = tokio::runtime::Builder::new_current_thread()
330 .enable_time()
331 .build()
332 .unwrap();
333
334 const INTERVAL_MS: u64 = 50;
335
336 rt.block_on(async {
337 let count = Arc::new(Mutex::new(0u64));
338 let count_cloned = Arc::clone(&count);
339
340 let handle = set_interval(
341 interval(Duration::from_millis(INTERVAL_MS)),
342 move || {
343 let count = Arc::clone(&count_cloned);
344 async move {
345 let mut c = count.lock().await;
346 *c += 1;
347 }
348 },
349 );
350
351 // Let it tick a couple of times
352 tokio::time::sleep(Duration::from_millis(120)).await;
353 assert!(*count.lock().await >= 2);
354
355 // Drop the handle (equivalent to letting it go out of scope)
356 drop(handle);
357
358 let count_after_drop = *count.lock().await;
359
360 // Wait — if the loop were alive the count would increase
361 tokio::time::sleep(Duration::from_millis(300)).await;
362 let final_count = *count.lock().await;
363
364 assert_eq!(
365 final_count, count_after_drop,
366 "count increased after handle dropped: was {count_after_drop}, now {final_count}"
367 );
368 });
369 }
370}