Skip to main content

batch_aint_one/
limits.rs

1use std::fmt::{self, Display};
2
3use bon::bon;
4
5/// Limits on batch sizes, queueing and concurrency, applied per key.
6///
7/// New items are rejected when the batch queue for their key is full: `max_batch_queue_size`
8/// batches are queued, and the newest batch already contains `max_batch_size` items.
9///
10/// `max_key_concurrency * max_batch_size` is the maximum number of items that can be processed
11/// concurrently for a key, and `max_batch_queue_size * max_batch_size` is the maximum number of
12/// items that can be queued for a key.
13#[derive(Debug, Clone, Copy)]
14#[non_exhaustive]
15pub struct Limits {
16    pub(crate) max_batch_size: usize,
17    pub(crate) max_key_concurrency: usize,
18    pub(crate) max_batch_queue_size: usize,
19}
20
21#[bon]
22impl Limits {
23    /// Create new limits.
24    ///
25    /// # Panics
26    ///
27    /// Panics if any of the limits are zero.
28    #[builder]
29    pub fn new(
30        /// Limits the maximum size of a batch.
31        ///
32        /// Defaults to 100.
33        #[builder(default = 100)]
34        max_batch_size: usize,
35        /// Limits the maximum number of batches that can be processed concurrently for a key,
36        /// including resource acquisition.
37        ///
38        /// Defaults to 10.
39        #[builder(default = 10)]
40        max_key_concurrency: usize,
41        /// Limits the maximum number of batches that can be queued concurrently for a key.
42        ///
43        /// Defaults to `max_key_concurrency * 2`.
44        max_batch_queue_size: Option<usize>,
45    ) -> Self {
46        assert!(
47            max_batch_size > 0,
48            "max_batch_size must be greater than zero"
49        );
50        assert!(
51            max_key_concurrency > 0,
52            "max_key_concurrency must be greater than zero"
53        );
54        let max_batch_queue_size = max_batch_queue_size.unwrap_or(max_key_concurrency * 2);
55        assert!(
56            max_batch_queue_size > 0,
57            "max_batch_queue_size must be greater than zero"
58        );
59
60        Self {
61            max_batch_size,
62            max_key_concurrency,
63            max_batch_queue_size,
64        }
65    }
66
67    fn max_items_processing_per_key(&self) -> usize {
68        self.max_batch_size * self.max_key_concurrency
69    }
70
71    fn max_items_queued_per_key(&self) -> usize {
72        self.max_batch_size * self.max_batch_queue_size
73    }
74
75    /// The maximum number of items that can be in the system for a given key.
76    pub(crate) fn max_items_in_system_per_key(&self) -> usize {
77        self.max_items_processing_per_key() + self.max_items_queued_per_key()
78    }
79}
80
81impl Default for Limits {
82    fn default() -> Self {
83        Self::builder().build()
84    }
85}
86
87impl Display for Limits {
88    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
89        write!(
90            f,
91            "batch_size: {}, key_concurrency: {}, queue_size: {}",
92            self.max_batch_size, self.max_key_concurrency, self.max_batch_queue_size
93        )
94    }
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100
101    #[test]
102    fn default_matches_builder_defaults() {
103        let default = Limits::default();
104        let built = Limits::builder().build();
105
106        assert_eq!(default.max_batch_size, built.max_batch_size);
107        assert_eq!(default.max_key_concurrency, built.max_key_concurrency);
108        assert_eq!(default.max_batch_queue_size, built.max_batch_queue_size);
109    }
110
111    #[test]
112    #[should_panic(expected = "max_batch_size must be greater than zero")]
113    fn rejects_zero_max_batch_size() {
114        Limits::builder().max_batch_size(0).build();
115    }
116
117    #[test]
118    #[should_panic(expected = "max_key_concurrency must be greater than zero")]
119    fn rejects_zero_max_key_concurrency() {
120        Limits::builder().max_key_concurrency(0).build();
121    }
122
123    #[test]
124    #[should_panic(expected = "max_batch_queue_size must be greater than zero")]
125    fn rejects_zero_max_batch_queue_size() {
126        Limits::builder().max_batch_queue_size(0).build();
127    }
128
129    #[test]
130    fn limits_builder_methods() {
131        let limits = Limits::builder()
132            .max_batch_size(50)
133            .max_key_concurrency(5)
134            .build();
135
136        assert_eq!(limits.max_batch_size, 50);
137        assert_eq!(limits.max_key_concurrency, 5);
138    }
139}