tower_async/limit/policy/
concurrent.rs

1//! A policy that limits the number of concurrent requests.
2//!
3//! See [`ConcurrentPolicy`].
4//!
5//! # Examples
6//!
7//! ```
8//! use tower_async::{
9//!     limit::{Limit, policy::ConcurrentPolicy},
10//!     Service, ServiceExt, service_fn,
11//! };
12//! # use std::convert::Infallible;
13//!
14//! # #[tokio::main]
15//! # async fn main() {
16//!
17//! let service = service_fn(|_| async {
18//!     Ok::<_, Infallible>(())
19//! });
20//! let mut service = Limit::new(service, ConcurrentPolicy::new(2));
21//!
22//! let response = service.oneshot(()).await;
23//! assert!(response.is_ok());
24//! # }
25//! ```
26
27use std::{
28    convert::Infallible,
29    sync::{Arc, Mutex},
30};
31
32use crate::util::backoff::Backoff;
33
34use super::{Policy, PolicyOutput};
35
36/// A policy that limits the number of concurrent requests.
37#[derive(Debug)]
38pub struct ConcurrentPolicy<B> {
39    max: usize,
40    current: Arc<Mutex<usize>>,
41    backoff: B,
42}
43
44impl<B> Clone for ConcurrentPolicy<B>
45where
46    B: Clone,
47{
48    fn clone(&self) -> Self {
49        ConcurrentPolicy {
50            max: self.max,
51            current: self.current.clone(),
52            backoff: self.backoff.clone(),
53        }
54    }
55}
56
57impl ConcurrentPolicy<()> {
58    /// Create a new concurrent policy,
59    /// which aborts the request if the limit is reached.
60    pub fn new(max: usize) -> Self {
61        ConcurrentPolicy {
62            max,
63            current: Arc::new(Mutex::new(0)),
64            backoff: (),
65        }
66    }
67}
68
69impl<B> ConcurrentPolicy<B> {
70    /// Create a new concurrent policy,
71    /// which backs off if the limit is reached,
72    /// using the given backoff policy.
73    pub fn with_backoff(max: usize, backoff: B) -> Self {
74        ConcurrentPolicy {
75            max,
76            current: Arc::new(Mutex::new(0)),
77            backoff,
78        }
79    }
80}
81
82/// The guard that releases the concurrent request limit.
83#[derive(Debug)]
84pub struct ConcurrentGuard {
85    current: Arc<Mutex<usize>>,
86}
87
88impl Drop for ConcurrentGuard {
89    fn drop(&mut self) {
90        let mut current = self.current.lock().unwrap();
91        *current -= 1;
92    }
93}
94
95impl<B, Request> Policy<Request> for ConcurrentPolicy<B>
96where
97    B: Backoff,
98{
99    type Guard = ConcurrentGuard;
100    type Error = Infallible;
101
102    async fn check(&self, _: &mut Request) -> PolicyOutput<Self::Guard, Self::Error> {
103        {
104            let mut current = self.current.lock().unwrap();
105            if *current < self.max {
106                *current += 1;
107                return PolicyOutput::Ready(ConcurrentGuard {
108                    current: self.current.clone(),
109                });
110            }
111        }
112
113        self.backoff.next_backoff().await;
114        PolicyOutput::Retry
115    }
116}
117
118/// The error that indicates the request is aborted,
119/// because the concurrent request limit is reached.
120#[derive(Debug)]
121pub struct LimitReached;
122
123impl std::fmt::Display for LimitReached {
124    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
125        f.write_str("LimitReached")
126    }
127}
128
129impl std::error::Error for LimitReached {}
130
131impl<Request> Policy<Request> for ConcurrentPolicy<()> {
132    type Guard = ConcurrentGuard;
133    type Error = LimitReached;
134
135    async fn check(&self, _: &mut Request) -> PolicyOutput<Self::Guard, Self::Error> {
136        let mut current = self.current.lock().unwrap();
137        if *current < self.max {
138            *current += 1;
139            PolicyOutput::Ready(ConcurrentGuard {
140                current: self.current.clone(),
141            })
142        } else {
143            PolicyOutput::Abort(LimitReached)
144        }
145    }
146}
147
148#[cfg(test)]
149mod tests {
150    use super::*;
151
152    fn assert_ready<G, E>(output: PolicyOutput<G, E>) -> G {
153        match output {
154            PolicyOutput::Ready(guard) => guard,
155            _ => panic!("unexpected output, expected ready"),
156        }
157    }
158
159    fn assert_abort<G, E>(output: PolicyOutput<G, E>) {
160        match output {
161            PolicyOutput::Abort(_) => (),
162            _ => panic!("unexpected output, expected abort"),
163        }
164    }
165
166    #[tokio::test]
167    async fn concurrent_policy() {
168        let policy = ConcurrentPolicy::new(2);
169
170        let guard_1 = assert_ready(policy.check(&mut ()).await);
171        let guard_2 = assert_ready(policy.check(&mut ()).await);
172
173        assert_abort(policy.check(&mut ()).await);
174
175        drop(guard_1);
176        let _guard_3 = assert_ready(policy.check(&mut ()).await);
177
178        assert_abort(policy.check(&mut ()).await);
179
180        drop(guard_2);
181        assert_ready(policy.check(&mut ()).await);
182    }
183}
tower_async/limit/policy/concurrent.rs

tower_async/limit/policy/
concurrent.rs
