foyer_common/
countdown.rs

1// Copyright 2025 foyer Project Authors
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::sync::atomic::{AtomicBool, AtomicIsize, Ordering};
16
17/// A concurrent count down util.
18#[derive(Debug)]
19pub struct Countdown {
20    finish: AtomicBool,
21    counter: AtomicIsize,
22}
23
24impl Countdown {
25    /// Countdown `counter` times.
26    ///
27    /// # Safety
28    ///
29    /// Panics if `counter` exceeds [`isize::MAX`].
30    pub fn new(counter: usize) -> Self {
31        Self {
32            finish: AtomicBool::new(false),
33            counter: AtomicIsize::new(isize::try_from(counter).expect("`counter` must NOT exceed `isize::MAX`.")),
34        }
35    }
36
37    /// Returns `false` for the first `counter` times, then always returns `true`.
38    pub fn countdown(&self) -> bool {
39        if self.finish.load(Ordering::Relaxed) {
40            return true;
41        }
42        self.counter.fetch_sub(1, Ordering::Relaxed) <= 0
43    }
44
45    /// Reset [`Countdown`] with `counter`.
46    pub fn reset(&self, counter: usize) {
47        self.finish.store(false, Ordering::Relaxed);
48        self.counter.store(
49            isize::try_from(counter).expect("`counter` must NOT exceed `isize::MAX`."),
50            Ordering::Relaxed,
51        );
52    }
53}
54
55#[cfg(test)]
56mod tests {
57    use std::time::Duration;
58
59    use futures_util::future::join_all;
60
61    use super::*;
62
63    async fn case(counter: usize, concurrency: usize) {
64        let cd = Countdown::new(counter);
65        let res = join_all((0..concurrency).map(|_| async {
66            tokio::time::sleep(Duration::from_millis(10)).await;
67            cd.countdown()
68        }))
69        .await;
70        assert_eq!(counter, res.into_iter().filter(|b| !b).count());
71    }
72
73    #[tokio::test]
74    async fn test_countdown() {
75        for counter in [1, 4, 8, 16] {
76            for concurrency in [16, 32, 64, 128] {
77                case(counter, concurrency).await;
78            }
79        }
80    }
81}
foyer_common/countdown.rs

foyer_common/
countdown.rs
