batch_aint_one/
limits.rs

1use std::fmt::{self, Display};
2
3use bon::bon;
4
5/// A policy controlling limits on batch sizes and concurrency.
6///
7/// New items will be rejected when both the limits have been reached.
8///
9/// `max_key_concurrency * max_batch_size` is the number of items that can be processed concurrently.
10///
11/// `max_batch_queue_size * max_batch_size` is the number of items that can be queued.
12#[derive(Debug, Clone, Copy)]
13#[non_exhaustive]
14pub struct Limits {
15    pub(crate) max_batch_size: usize,
16    pub(crate) max_key_concurrency: usize,
17    pub(crate) max_batch_queue_size: usize,
18}
19
20#[bon]
21impl Limits {
22    #[allow(missing_docs)]
23    #[builder]
24    pub fn new(
25        /// Limits the maximum size of a batch.
26        #[builder(default = 100)]
27        max_batch_size: usize,
28        /// Limits the maximum number of batches that can be processed concurrently for a key,
29        /// including resource acquisition.
30        #[builder(default = 10)]
31        max_key_concurrency: usize,
32        /// Limits the maximum number of batches that can be queued concurrently for a key.
33        max_batch_queue_size: Option<usize>,
34    ) -> Self {
35        Self {
36            max_batch_size,
37            max_key_concurrency,
38            max_batch_queue_size: max_batch_queue_size.unwrap_or(max_key_concurrency * 2),
39        }
40    }
41
42    fn max_items_processing_per_key(&self) -> usize {
43        self.max_batch_size * self.max_key_concurrency
44    }
45
46    fn max_items_queued_per_key(&self) -> usize {
47        self.max_batch_size * self.max_batch_queue_size
48    }
49
50    /// The maximum number of items that can be in the system for a given key.
51    pub(crate) fn max_items_in_system_per_key(&self) -> usize {
52        self.max_items_processing_per_key() + self.max_items_queued_per_key()
53    }
54}
55
56impl Default for Limits {
57    fn default() -> Self {
58        let max_batch_size = 100;
59        let max_key_concurrency = 10;
60        let max_batch_queue_size = max_key_concurrency;
61        Self {
62            max_batch_size,
63            max_key_concurrency,
64            max_batch_queue_size,
65        }
66    }
67}
68
69impl Display for Limits {
70    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
71        write!(
72            f,
73            "batch_size: {}, key_concurrency: {}, queue_size: {}",
74            self.max_batch_size, self.max_key_concurrency, self.max_batch_queue_size
75        )
76    }
77}
78
79#[cfg(test)]
80mod tests {
81    use super::*;
82
83    #[test]
84    fn limits_builder_methods() {
85        let limits = Limits::builder()
86            .max_batch_size(50)
87            .max_key_concurrency(5)
88            .build();
89
90        assert_eq!(limits.max_batch_size, 50);
91        assert_eq!(limits.max_key_concurrency, 5);
92    }
93}