shardline_server/
transfer_limiter.rs

1use std::{num::NonZeroUsize, sync::Arc};
2
3use tokio::sync::{OwnedSemaphorePermit, Semaphore};
4
5use crate::ServerError;
6
7/// Weighted transfer concurrency limiter based on chunk-equivalent cost.
8#[derive(Debug, Clone)]
9pub struct TransferLimiter {
10    chunk_size_bytes: NonZeroUsize,
11    max_in_flight_chunks: NonZeroUsize,
12    semaphore: Arc<Semaphore>,
13}
14
15impl TransferLimiter {
16    /// Creates a transfer limiter that budgets concurrent response work in
17    /// chunk-equivalent permits.
18    #[must_use]
19    pub fn new(chunk_size_bytes: NonZeroUsize, max_in_flight_chunks: NonZeroUsize) -> Self {
20        Self {
21            chunk_size_bytes,
22            max_in_flight_chunks,
23            semaphore: Arc::new(Semaphore::new(max_in_flight_chunks.get())),
24        }
25    }
26
27    /// Acquires permits for a transfer with the supplied byte length.
28    ///
29    /// # Errors
30    ///
31    /// Returns [`ServerError`] when the limiter capacity cannot be represented or the
32    /// semaphore has been closed.
33    pub async fn acquire_bytes(
34        &self,
35        total_bytes: u64,
36    ) -> Result<OwnedSemaphorePermit, ServerError> {
37        let permits = permits_for_bytes(
38            total_bytes,
39            self.chunk_size_bytes,
40            self.max_in_flight_chunks,
41        )?;
42        self.semaphore
43            .clone()
44            .acquire_many_owned(permits)
45            .await
46            .map_err(|_error| ServerError::TransferLimiterClosed)
47    }
48}
49
50fn permits_for_bytes(
51    total_bytes: u64,
52    chunk_size_bytes: NonZeroUsize,
53    max_in_flight_chunks: NonZeroUsize,
54) -> Result<u32, ServerError> {
55    let effective_bytes = total_bytes.max(1);
56    let chunk_size = u64::try_from(chunk_size_bytes.get())?;
57    let required_chunks = effective_bytes.div_ceil(chunk_size);
58    let max_chunks = u64::try_from(max_in_flight_chunks.get())?;
59    let bounded_chunks = required_chunks.min(max_chunks);
60    u32::try_from(bounded_chunks).map_err(ServerError::from)
61}
62
63#[cfg(test)]
64mod tests {
65    use std::{num::NonZeroUsize, time::Duration};
66
67    use tokio::time::timeout;
68
69    use super::{TransferLimiter, permits_for_bytes};
70
71    const CHUNK_SIZE: NonZeroUsize = match NonZeroUsize::new(4) {
72        Some(value) => value,
73        None => NonZeroUsize::MIN,
74    };
75
76    const MAX_IN_FLIGHT: NonZeroUsize = match NonZeroUsize::new(3) {
77        Some(value) => value,
78        None => NonZeroUsize::MIN,
79    };
80
81    #[test]
82    fn permits_scale_with_bytes_and_cap_at_capacity() {
83        let one_byte = permits_for_bytes(1, CHUNK_SIZE, MAX_IN_FLIGHT);
84        let exact_chunk = permits_for_bytes(4, CHUNK_SIZE, MAX_IN_FLIGHT);
85        let two_chunks = permits_for_bytes(5, CHUNK_SIZE, MAX_IN_FLIGHT);
86        let capped = permits_for_bytes(20, CHUNK_SIZE, MAX_IN_FLIGHT);
87
88        assert!(one_byte.is_ok());
89        assert!(exact_chunk.is_ok());
90        assert!(two_chunks.is_ok());
91        assert!(capped.is_ok());
92        assert_eq!(one_byte.ok(), Some(1));
93        assert_eq!(exact_chunk.ok(), Some(1));
94        assert_eq!(two_chunks.ok(), Some(2));
95        assert_eq!(capped.ok(), Some(3));
96    }
97
98    #[test]
99    fn zero_byte_transfers_still_reserve_one_permit() {
100        let permits = permits_for_bytes(0, CHUNK_SIZE, MAX_IN_FLIGHT);
101
102        assert!(permits.is_ok());
103        assert_eq!(permits.ok(), Some(1));
104    }
105
106    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
107    async fn limiter_blocks_until_capacity_returns() {
108        let limiter = TransferLimiter::new(CHUNK_SIZE, NonZeroUsize::MIN);
109        let first = limiter.acquire_bytes(4).await;
110        assert!(first.is_ok());
111        let Ok(first) = first else {
112            return;
113        };
114
115        let blocked = timeout(Duration::from_millis(50), limiter.acquire_bytes(4)).await;
116        assert!(blocked.is_err());
117
118        drop(first);
119
120        let released = timeout(Duration::from_secs(1), limiter.acquire_bytes(4)).await;
121        assert!(released.is_ok());
122        let Ok(released) = released else {
123            return;
124        };
125        assert!(released.is_ok());
126    }
127}
shardline_server/transfer_limiter.rs

shardline_server/
transfer_limiter.rs
