Skip to main content

self_update_extras/
throttle.rs

1//! Throttle wrapper: limits how often update checks run.
2
3use self_update::Status;
4use self_update::errors::{Error, Result};
5use self_update::update::{Release, ReleaseUpdate, UpdateStatus};
6use std::time::Duration;
7
8/// Default minimum time between update checks.
9const DEFAULT_THROTTLE_WINDOW: Duration = Duration::from_secs(15 * 60);
10
11/// Builder for a throttled [`Update`].
12///
13/// Configure the inner [`ReleaseUpdate`] backend and the throttle window,
14/// then call [`build`](Self::build) to produce a `Box<dyn ReleaseUpdate>`
15/// that skips update checks within the configured window.
16#[derive(Default)]
17pub struct UpdateBuilder {
18    release_update: Option<Box<dyn ReleaseUpdate>>,
19    throttle_window: Option<Duration>,
20}
21
22impl UpdateBuilder {
23    /// Initialize a new builder.
24    pub fn new() -> Self {
25        Default::default()
26    }
27
28    /// Set the release update implementation to wrap. Required.
29    pub fn release_update(&mut self, release_update: Box<dyn ReleaseUpdate>) -> &mut Self {
30        self.release_update = Some(release_update);
31        self
32    }
33
34    /// Set the minimum time between update checks. Defaults to 15 minutes.
35    pub fn throttle_window(&mut self, window: Duration) -> &mut Self {
36        self.throttle_window = Some(window);
37        self
38    }
39
40    /// Confirm config and create a ready-to-use `Update`.
41    ///
42    /// * Errors:
43    ///     * Config - `release_update` was not provided
44    pub fn build(&mut self) -> Result<Box<dyn ReleaseUpdate>> {
45        let inner = self
46            .release_update
47            .take()
48            .ok_or_else(|| Error::Config("`release_update` required".to_owned()))?;
49
50        Ok(Box::new(Update {
51            inner,
52            throttle_window: self.throttle_window.unwrap_or(DEFAULT_THROTTLE_WINDOW),
53        }))
54    }
55}
56
57/// Wraps a [`ReleaseUpdate`] and throttles how often update checks run.
58///
59/// The underlying update check is skipped when one was already performed
60/// within `throttle_window`. The time of the last check is tracked via a
61/// throttle file in the system temp directory.
62///
63/// # Throttle behavior
64///
65/// The throttle file is touched only after a *successful* update check.
66/// Failed checks do not reset the throttle window, allowing retries without
67/// waiting for the full window to elapse.
68pub struct Update {
69    inner: Box<dyn ReleaseUpdate>,
70    throttle_window: Duration,
71}
72
73impl Update {
74    /// Initialize a new `Update` builder.
75    pub fn configure() -> UpdateBuilder {
76        UpdateBuilder::new()
77    }
78
79    /// The minimum time that must elapse between update checks.
80    pub fn throttle_window(&self) -> Duration {
81        self.throttle_window
82    }
83
84    /// Returns true if enough time has elapsed to perform an update check.
85    fn should_check(&self) -> bool {
86        let path = self.throttle_path();
87        match std::fs::metadata(&path) {
88            Ok(meta) => {
89                if let Ok(mtime) = meta.modified()
90                    && let Ok(elapsed) = std::time::SystemTime::now().duration_since(mtime)
91                {
92                    return elapsed > self.throttle_window;
93                }
94                true
95            }
96            Err(_) => true,
97        }
98    }
99
100    /// Path to the throttle file in the system temp directory.
101    fn throttle_path(&self) -> std::path::PathBuf {
102        std::env::temp_dir().join(format!("{}.throttle", &self.inner.bin_name()))
103    }
104
105    /// Update the throttle file's mtime, creating it if necessary.
106    fn touch_throttle(&self) {
107        let _ = std::fs::OpenOptions::new()
108            .create(true)
109            .truncate(true)
110            .write(true)
111            .open(self.throttle_path());
112    }
113}
114
115impl ReleaseUpdate for Update {
116    fn get_latest_release(&self) -> Result<Release> {
117        self.inner.get_latest_release()
118    }
119
120    fn get_latest_releases(&self, current_version: &str) -> Result<Vec<Release>> {
121        self.inner.get_latest_releases(current_version)
122    }
123
124    fn get_release_version(&self, ver: &str) -> Result<Release> {
125        self.inner.get_release_version(ver)
126    }
127
128    fn current_version(&self) -> String {
129        self.inner.current_version()
130    }
131
132    fn target(&self) -> String {
133        self.inner.target()
134    }
135
136    fn target_version(&self) -> Option<String> {
137        self.inner.target_version()
138    }
139
140    fn bin_name(&self) -> String {
141        self.inner.bin_name()
142    }
143
144    fn bin_install_path(&self) -> std::path::PathBuf {
145        self.inner.bin_install_path()
146    }
147
148    fn bin_path_in_archive(&self) -> String {
149        self.inner.bin_path_in_archive()
150    }
151
152    fn show_download_progress(&self) -> bool {
153        self.inner.show_download_progress()
154    }
155
156    fn show_output(&self) -> bool {
157        self.inner.show_output()
158    }
159
160    fn no_confirm(&self) -> bool {
161        self.inner.no_confirm()
162    }
163
164    fn progress_template(&self) -> String {
165        self.inner.progress_template()
166    }
167
168    fn progress_chars(&self) -> String {
169        self.inner.progress_chars()
170    }
171
172    fn auth_token(&self) -> Option<String> {
173        self.inner.auth_token()
174    }
175
176    fn update(&self) -> Result<Status> {
177        let current_version = self.current_version();
178        self.update_extended()
179            .map(|s| s.into_status(current_version))
180    }
181
182    fn update_extended(&self) -> Result<UpdateStatus> {
183        if !self.should_check() {
184            return Ok(UpdateStatus::UpToDate);
185        }
186
187        let result = self.inner.update_extended();
188        self.touch_throttle();
189        result
190    }
191}
192
193#[cfg(test)]
194mod tests {
195    use super::*;
196    use crate::test_support::MockRelease;
197
198    fn remove(path: &std::path::Path) {
199        let _ = std::fs::remove_file(path);
200    }
201
202    /// Build a concrete `Update` for white-box tests of the private throttle
203    /// mechanics. `build()` returns a `Box<dyn ReleaseUpdate>`, which hides
204    /// those methods, so tests construct the type directly.
205    fn concrete(mock: MockRelease, window: Duration) -> Update {
206        Update {
207            inner: Box::new(mock),
208            throttle_window: window,
209        }
210    }
211
212    // Each test uses a distinct bin name, so their throttle files never
213    // collide and the tests can run in parallel.
214
215    #[test]
216    fn should_check_is_true_without_a_throttle_file() {
217        let updater = concrete(
218            MockRelease::new("mock-throttle-missing"),
219            DEFAULT_THROTTLE_WINDOW,
220        );
221        remove(&updater.throttle_path());
222
223        assert!(updater.should_check());
224    }
225
226    #[test]
227    fn should_check_is_false_with_a_recent_throttle_file() {
228        let updater = concrete(
229            MockRelease::new("mock-throttle-recent"),
230            DEFAULT_THROTTLE_WINDOW,
231        );
232        let path = updater.throttle_path();
233        remove(&path);
234
235        updater.touch_throttle();
236
237        assert!(!updater.should_check());
238        remove(&path);
239    }
240
241    #[test]
242    fn should_check_is_true_with_an_expired_throttle_file() {
243        let updater = concrete(
244            MockRelease::new("mock-throttle-expired"),
245            Duration::from_secs(60),
246        );
247        let path = updater.throttle_path();
248        std::fs::write(&path, "").unwrap();
249
250        let expired = std::time::SystemTime::now()
251            .duration_since(std::time::UNIX_EPOCH)
252            .unwrap()
253            - Duration::from_secs(120);
254        let expired = filetime::FileTime::from_unix_time(expired.as_secs() as i64, 0);
255        filetime::set_file_mtime(&path, expired).unwrap();
256
257        assert!(updater.should_check());
258        remove(&path);
259    }
260
261    #[test]
262    fn touch_throttle_creates_the_throttle_file() {
263        let updater = concrete(
264            MockRelease::new("mock-throttle-touch"),
265            DEFAULT_THROTTLE_WINDOW,
266        );
267        let path = updater.throttle_path();
268        remove(&path);
269
270        updater.touch_throttle();
271
272        assert!(path.exists());
273        remove(&path);
274    }
275
276    #[test]
277    fn update_runs_inner_and_touches_throttle_when_allowed() {
278        let mock = MockRelease::new("mock-throttle-run");
279        let calls = mock.call_counter();
280        let updater = concrete(mock, DEFAULT_THROTTLE_WINDOW);
281        let path = updater.throttle_path();
282        remove(&path);
283
284        let status = updater.update().unwrap();
285
286        assert!(matches!(status, Status::UpToDate(_)));
287        assert_eq!(calls.get(), 1);
288        assert!(
289            path.exists(),
290            "throttle file should be touched after a check"
291        );
292        remove(&path);
293    }
294
295    #[test]
296    fn update_skips_inner_when_throttled() {
297        let mock = MockRelease::new("mock-throttle-skip");
298        let calls = mock.call_counter();
299        let updater = concrete(mock, DEFAULT_THROTTLE_WINDOW);
300        let path = updater.throttle_path();
301        remove(&path);
302        updater.touch_throttle();
303
304        let status = updater.update().unwrap();
305
306        assert!(matches!(status, Status::UpToDate(_)));
307        assert_eq!(calls.get(), 0);
308        remove(&path);
309    }
310
311    #[test]
312    fn update_extended_runs_inner_when_allowed() {
313        let mock = MockRelease::new("mock-throttle-run-ext");
314        let calls = mock.call_counter();
315        let updater = concrete(mock, DEFAULT_THROTTLE_WINDOW);
316        let path = updater.throttle_path();
317        remove(&path);
318
319        let status = updater.update_extended().unwrap();
320
321        assert!(matches!(status, UpdateStatus::UpToDate));
322        assert_eq!(calls.get(), 1);
323        assert!(path.exists());
324        remove(&path);
325    }
326
327    #[test]
328    fn update_extended_skips_inner_when_throttled() {
329        let mock = MockRelease::new("mock-throttle-skip-ext");
330        let calls = mock.call_counter();
331        let updater = concrete(mock, DEFAULT_THROTTLE_WINDOW);
332        let path = updater.throttle_path();
333        remove(&path);
334        updater.touch_throttle();
335
336        let status = updater.update_extended().unwrap();
337
338        assert!(matches!(status, UpdateStatus::UpToDate));
339        assert_eq!(calls.get(), 0);
340        remove(&path);
341    }
342
343    #[test]
344    fn throttle_window_getter_reports_the_configured_window() {
345        let window = Duration::from_secs(42);
346        let updater = concrete(MockRelease::new("mock-throttle-getter"), window);
347        assert_eq!(updater.throttle_window(), window);
348    }
349
350    #[test]
351    fn build_requires_a_release_update() {
352        assert!(Update::configure().build().is_err());
353    }
354
355    #[test]
356    fn builder_wraps_the_release_update() {
357        let mock = MockRelease::new("mock-throttle-builder");
358        let calls = mock.call_counter();
359        let path = std::env::temp_dir().join("mock-throttle-builder.throttle");
360        remove(&path);
361
362        let updater = Update::configure()
363            .release_update(Box::new(mock))
364            .throttle_window(Duration::from_secs(60))
365            .build()
366            .unwrap();
367
368        updater.update().unwrap();
369
370        assert_eq!(calls.get(), 1);
371        remove(&path);
372    }
373}