common_batcher/
flush_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/// Shared execution budget. One permit represents one complete flush, not one RPC.
20#[derive(Clone)]
21pub struct FlushLimiter {
22    semaphore: Arc<Semaphore>,
23}
24
25impl FlushLimiter {
26    /// Returns `None` if the limit is zero or exceeds Tokio's supported semaphore capacity.
27    pub fn try_new(max_concurrent_flushes: usize) -> Option<Self> {
28        let permits = max_concurrent_flushes;
29        ((1..=Semaphore::MAX_PERMITS).contains(&permits)).then(|| Self {
30            semaphore: Arc::new(Semaphore::new(permits)),
31        })
32    }
33
34    /// The caller must retain the returned permit until its flush finishes.
35    pub async fn acquire(&self) -> Result<OwnedSemaphorePermit, AcquireError> {
36        self.semaphore.clone().acquire_owned().await
37    }
38}
39
40#[cfg(test)]
41mod tests {
42    use std::future::Future;
43
44    use super::*;
45
46    #[test]
47    fn test_invalid_capacity() {
48        for capacity in [0, Semaphore::MAX_PERMITS + 1, usize::MAX] {
49            assert!(FlushLimiter::try_new(capacity).is_none());
50        }
51        for capacity in [1, Semaphore::MAX_PERMITS] {
52            assert!(FlushLimiter::try_new(capacity).is_some());
53        }
54    }
55
56    #[tokio::test]
57    async fn test_clones_share_flush_budget() {
58        let limiter = FlushLimiter::try_new(1).unwrap();
59        let other = limiter.clone();
60        let permit = limiter.acquire().await.unwrap();
61        let waiting = other.acquire();
62        tokio::pin!(waiting);
63        // Poll admission without a sleep or an assumption about scheduler timing.
64        assert!(
65            std::future::poll_fn(|cx| {
66                std::task::Poll::Ready(waiting.as_mut().poll(cx).is_pending())
67            })
68            .await
69        );
70        drop(permit);
71        let _next = waiting.await.unwrap();
72    }
73}
common_batcher/flush_limiter.rs

common_batcher/
flush_limiter.rs
