throttle/
lib.rs

1// throttle
2// Copyright (C) SOFe
3//
4// Licensed under the Apache License, Version 2.0 (the "License");
5// you may not use this file except in compliance with the License.
6// You may obtain a copy of the License at
7//
8//     http://www.apache.org/licenses/LICENSE-2.0
9//
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16use std::collections::VecDeque;
17use std::time::{Duration, Instant};
18
19/// Throttle is a simple utility for rate-limiting operations.
20///
21/// ```
22/// use std::time::Duration;
23/// use throttle::Throttle;
24///
25/// let unit = Duration::from_millis(100); // we use 100ms to have sufficient time elapse
26/// let mut throttle = Throttle::new(unit * 4, 3);
27///
28/// throttle.accept().expect("The throttle is empty");
29/// assert_eq!(throttle.size(), 1);
30/// std::thread::sleep(unit * 2); // we sleep for 2t, still not expired
31///
32/// assert_eq!(throttle.size(), 1);
33/// throttle.accept().expect("The throttle has one more space");
34/// assert_eq!(throttle.size(), 2);
35/// std::thread::sleep(unit); // time is now +3t
36///
37/// assert_eq!(throttle.size(), 2);
38/// throttle.accept().expect("The last accept before throttle is full");
39/// assert_eq!(throttle.size(), 3);
40/// throttle.accept().expect_err("The throttle should be full");
41/// assert_eq!(throttle.size(), 3);
42///
43/// std::thread::sleep(unit * 2); // time is now +5t, and the first accept should have expired
44/// assert_eq!(throttle.size(), 2);
45/// throttle.accept().expect("The first accept should have expired");
46/// assert_eq!(throttle.size(), 3);
47/// throttle.accept().expect_err("The second accept should not have expired yet");
48/// assert_eq!(throttle.size(), 3);
49///
50/// std::thread::sleep(unit * 10); // time is now +10t, and all accepts should have expired
51/// assert_eq!(throttle.size(), 0);
52/// ```
53pub struct Throttle {
54    timeout: Duration,
55    threshold: usize,
56    deque: VecDeque<Instant>,
57}
58
59impl Throttle {
60    /// Creates a new Throttle
61    pub fn new(timeout: Duration, threshold: usize) -> Throttle {
62        Throttle {
63            timeout,
64            threshold,
65            deque: Default::default(),
66        }
67    }
68
69    fn flush(&mut self) {
70        while let Some(first) = self.deque.front() {
71            if first.elapsed() >= self.timeout.clone() {
72                self.deque.pop_front();
73            } else {
74                break;
75            }
76        }
77    }
78
79    /// Returns the number of remaining items in the throttle
80    pub fn size(&mut self) -> usize {
81        self.flush();
82        self.deque.len()
83    }
84
85    /// Attempts to accept an operation and increment the throttle.
86    ///
87    /// On success, Ok is returned and the counter increments.
88    ///
89    /// On failure, Err is returned with an Instant indicating the time that the throttle is
90    /// available again.
91    pub fn accept(&mut self) -> Result<(), Instant> {
92        self.flush();
93        if self.deque.len() >= self.threshold {
94            return Err(self.deque.front().unwrap().clone() + self.timeout.clone());
95        }
96
97        self.deque.push_back(Instant::now());
98        Ok(())
99    }
100}
throttle/lib.rs

throttle/
lib.rs
