steamkit/util/
scheduled_function.rs

1//! A function that runs on a fixed schedule (immediately, then every `delay`),
2//! on a tokio task. Stopping it cancels the task.
3
4use std::sync::atomic::{AtomicBool, Ordering};
5use std::sync::{Arc, Mutex, RwLock};
6use std::time::Duration;
7
8use tokio_util::sync::CancellationToken;
9
10use crate::util::debug_log::DebugLog;
11
12/// The mutable state of a scheduled function.
13#[derive(Debug)]
14struct ScheduledFunctionState {
15    /// The period; `None` means "run once".
16    delay: RwLock<Option<Duration>>,
17    /// Whether it's running.
18    running: AtomicBool,
19    /// Cancels the running task.
20    cancel: Mutex<Option<CancellationToken>>,
21}
22
23/// A function that runs on a fixed schedule.
24pub(crate) struct ScheduledFunction {
25    /// The function to run.
26    func: Arc<dyn Fn() + Send + Sync>,
27    /// The mutable state.
28    state: Arc<ScheduledFunctionState>,
29}
30
31impl std::fmt::Debug for ScheduledFunction {
32    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
33        formatter
34            .debug_struct("ScheduledFunction")
35            .field("delay", &self.delay())
36            .field("running", &self.state.running.load(Ordering::SeqCst))
37            .finish()
38    }
39}
40
41impl ScheduledFunction {
42    /// Creates a scheduled function that runs once when started.
43    pub(crate) fn new(func: impl Fn() + Send + Sync + 'static) -> ScheduledFunction {
44        ScheduledFunction::with_delay(func, None)
45    }
46
47    /// Creates a scheduled function that runs immediately and then every
48    /// `delay`. `None` means it runs once.
49    pub(crate) fn with_delay(
50        func: impl Fn() + Send + Sync + 'static,
51        delay: Option<Duration>,
52    ) -> ScheduledFunction {
53        ScheduledFunction {
54            func: Arc::new(func),
55            state: Arc::new(ScheduledFunctionState {
56                delay: RwLock::new(delay),
57                running: AtomicBool::new(false),
58                cancel: Mutex::new(None),
59            }),
60        }
61    }
62
63    /// The period; `None` means "run once".
64    pub(crate) fn delay(&self) -> Option<Duration> {
65        *self.state.delay.read().unwrap_or_else(|error| error.into_inner())
66    }
67
68    /// Sets the period. Takes effect the next time it's started.
69    pub(crate) fn set_delay(&self, delay: Option<Duration>) {
70        *self.state.delay.write().unwrap_or_else(|error| error.into_inner()) = delay;
71    }
72
73    /// Starts it: runs the function now and then every [`Self::delay`]. Does
74    /// nothing if it's already running.
75    pub(crate) fn start(&self) {
76        if self.state.running.swap(true, Ordering::SeqCst) {
77            return;
78        }
79
80        let Ok(handle) = tokio::runtime::Handle::try_current() else {
81            self.state.running.store(false, Ordering::SeqCst);
82            DebugLog::write_line(
83                "ScheduledFunction",
84                "Unable to start a scheduled function outside of a tokio runtime.",
85            );
86            return;
87        };
88
89        let token = CancellationToken::new();
90
91        {
92            let mut cancel = self.state.cancel.lock().unwrap_or_else(|error| error.into_inner());
93            *cancel = Some(token.clone());
94        }
95
96        let delay = self.delay();
97        let func = self.func.clone();
98
99        handle.spawn(async move {
100            match delay {
101                Some(period) if !period.is_zero() => {
102                    let mut interval = tokio::time::interval(period);
103                    interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
104
105                    loop {
106                        tokio::select! {
107                            _ = token.cancelled() => break,
108                            _ = interval.tick() => func(),
109                        }
110                    }
111                }
112                _ => {
113                    tokio::select! {
114                        biased;
115                        _ = token.cancelled() => {}
116                        _ = std::future::ready(()) => func(),
117                    }
118                }
119            }
120        });
121    }
122
123    /// Stops it. Does nothing if it isn't running.
124    pub(crate) fn stop(&self) {
125        if !self.state.running.swap(false, Ordering::SeqCst) {
126            return;
127        }
128
129        let token = {
130            let mut cancel = self.state.cancel.lock().unwrap_or_else(|error| error.into_inner());
131            cancel.take()
132        };
133
134        if let Some(token) = token {
135            token.cancel();
136        }
137    }
138}
139
140impl Drop for ScheduledFunction {
141    fn drop(&mut self) {
142        self.stop();
143    }
144}
steamkit/util/scheduled_function.rs

steamkit/util/
scheduled_function.rs
