common_batcher/
request_limiter.rs

1// Copyright 2023 Greptime Team
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::Arc;
16
17use tokio::sync::{AcquireError, OwnedSemaphorePermit, Semaphore};
18
19/// Limits unfinished original requests independently of their number of submissions.
20///
21/// Acquire once for each original request, then clone its permit into all of that
22/// request's submissions. Capacity is released when the last owner drops it.
23#[derive(Clone)]
24pub struct RequestLimiter {
25    semaphore: Arc<Semaphore>,
26}
27
28impl RequestLimiter {
29    /// Returns `None` if the limit is zero or exceeds Tokio's supported semaphore capacity.
30    pub fn try_new(max_inflight_requests: usize) -> Option<Self> {
31        let permits = max_inflight_requests;
32        ((1..=Semaphore::MAX_PERMITS).contains(&permits)).then(|| Self {
33            semaphore: Arc::new(Semaphore::new(permits)),
34        })
35    }
36
37    /// Waits for one request slot and returns a permit shared by its submissions.
38    ///
39    /// Cancelling a pending acquisition does not consume a slot. Once acquired,
40    /// callers must retain the permit until every submission has finished, even
41    /// if the original caller stops waiting for its response.
42    pub async fn acquire(&self) -> Result<Arc<OwnedSemaphorePermit>, AcquireError> {
43        self.semaphore.clone().acquire_owned().await.map(Arc::new)
44    }
45}
46
47#[cfg(test)]
48mod tests {
49    use std::future::{Future, poll_fn};
50    use std::pin::Pin;
51    use std::task::Poll;
52
53    use tokio::sync::Semaphore;
54
55    use crate::request_limiter::RequestLimiter;
56
57    async fn is_pending<F: Future>(mut future: Pin<&mut F>) -> bool {
58        poll_fn(|cx| Poll::Ready(future.as_mut().poll(cx).is_pending())).await
59    }
60
61    #[test]
62    fn test_capacity_boundaries() {
63        for capacity in [1, Semaphore::MAX_PERMITS] {
64            assert!(RequestLimiter::try_new(capacity).is_some());
65        }
66        for capacity in [0, Semaphore::MAX_PERMITS + 1, usize::MAX] {
67            assert!(RequestLimiter::try_new(capacity).is_none());
68        }
69    }
70
71    #[tokio::test]
72    async fn test_last_submission_releases_request_slot() {
73        let limiter = RequestLimiter::try_new(1).unwrap();
74        let other = limiter.clone();
75        let request = limiter.acquire().await.unwrap();
76        let first_submission = request.clone();
77        let second_submission = request.clone();
78        let mut waiting = Box::pin(other.acquire());
79        assert!(is_pending(waiting.as_mut()).await);
80        // The caller and one submission finish, but another submission remains.
81        drop(request);
82        drop(first_submission);
83        assert!(is_pending(waiting.as_mut()).await);
84        drop(second_submission);
85        let next = waiting.await.unwrap();
86        assert_eq!(next.num_permits(), 1);
87    }
88
89    #[tokio::test]
90    async fn test_cancelled_acquisition_does_not_leak_capacity() {
91        let limiter = RequestLimiter::try_new(1).unwrap();
92        let permit = limiter.acquire().await.unwrap();
93        let mut cancelled = Box::pin(limiter.acquire());
94        assert!(is_pending(cancelled.as_mut()).await);
95        drop(cancelled);
96        let mut next = Box::pin(limiter.acquire());
97        assert!(is_pending(next.as_mut()).await);
98        drop(permit);
99        let next = next.await.unwrap();
100        drop(next);
101        let _available = limiter.acquire().await.unwrap();
102    }
103
104    #[tokio::test]
105    async fn test_independent_requests_consume_separate_slots() {
106        let limiter = RequestLimiter::try_new(2).unwrap();
107        let first = limiter.acquire().await.unwrap();
108        let second = limiter.acquire().await.unwrap();
109        let mut third = Box::pin(limiter.acquire());
110        assert!(is_pending(third.as_mut()).await);
111        drop(first);
112        let _third = third.await.unwrap();
113        drop(second);
114    }
115}
common_batcher/request_limiter.rs

common_batcher/
request_limiter.rs
