awa_model/
queue_fanout.rs

1use crate::job::InsertOpts;
2use crate::queue_storage::shard_for_ordering_key;
3use std::collections::HashSet;
4
5const DEFAULT_PHYSICAL_QUEUE_SUFFIX: &str = "__p";
6
7/// Errors returned when constructing a [`QueueFanout`].
8#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
9pub enum QueueFanoutError {
10    #[error("queue fanout logical queue must not be empty")]
11    EmptyLogicalQueue,
12    #[error("queue fanout width must be > 0")]
13    ZeroWidth,
14    #[error("queue fanout supports at most {max} physical queues; got {got}")]
15    TooManyPhysicalQueues { got: usize, max: usize },
16    #[error("queue fanout physical queue must not be empty")]
17    EmptyPhysicalQueue,
18    #[error("queue fanout physical queue '{queue}' is duplicated")]
19    DuplicatePhysicalQueue { queue: String },
20}
21
22/// A deterministic set of physical queues for one hot logical queue.
23///
24/// Awa still stores and executes jobs from ordinary queue names. This helper
25/// gives producers and workers the same stable list of physical queues, so an
26/// application can fan one logical workload out over several queues without
27/// hand-rolling naming and routing in every process.
28#[derive(Debug, Clone, PartialEq, Eq)]
29pub struct QueueFanout {
30    logical_queue: String,
31    physical_queues: Vec<String>,
32}
33
34impl QueueFanout {
35    /// Build a fanout using Awa's default physical queue naming.
36    ///
37    /// Width `1` maps to the logical queue name itself. Widths above `1`
38    /// produce `{logical_queue}__p0`, `{logical_queue}__p1`, and so on.
39    pub fn new(logical_queue: impl Into<String>, width: usize) -> Result<Self, QueueFanoutError> {
40        let logical_queue = logical_queue.into();
41        validate_logical_queue(&logical_queue)?;
42        validate_physical_count(width)?;
43
44        let physical_queues = if width == 1 {
45            vec![logical_queue.clone()]
46        } else {
47            (0..width)
48                .map(|idx| format!("{logical_queue}{DEFAULT_PHYSICAL_QUEUE_SUFFIX}{idx}"))
49                .collect()
50        };
51
52        Ok(Self {
53            logical_queue,
54            physical_queues,
55        })
56    }
57
58    /// Build a fanout from explicit physical queue names.
59    ///
60    /// Use this when an application already has queue names it wants to keep,
61    /// or when migrating an existing manually-fanned-out deployment to the
62    /// shared helper.
63    pub fn from_physical_queues<I, S>(
64        logical_queue: impl Into<String>,
65        physical_queues: I,
66    ) -> Result<Self, QueueFanoutError>
67    where
68        I: IntoIterator<Item = S>,
69        S: Into<String>,
70    {
71        let logical_queue = logical_queue.into();
72        validate_logical_queue(&logical_queue)?;
73
74        let physical_queues: Vec<String> = physical_queues.into_iter().map(Into::into).collect();
75        validate_physical_count(physical_queues.len())?;
76        validate_physical_queues(&physical_queues)?;
77
78        Ok(Self {
79            logical_queue,
80            physical_queues,
81        })
82    }
83
84    /// Logical queue name used by the application.
85    pub fn logical_queue(&self) -> &str {
86        &self.logical_queue
87    }
88
89    /// Physical queues that must be declared on worker runtimes.
90    pub fn physical_queues(&self) -> &[String] {
91        &self.physical_queues
92    }
93
94    /// Number of physical queues in the fanout.
95    pub fn width(&self) -> usize {
96        self.physical_queues.len()
97    }
98
99    /// Select a physical queue by stable routing key.
100    ///
101    /// The same key always maps to the same physical queue, using the same
102    /// portable hash Awa uses for queue-storage enqueue shards.
103    pub fn queue_for_key(&self, key: impl AsRef<[u8]>) -> &str {
104        let shard = shard_for_ordering_key(key.as_ref(), self.width() as i16) as usize;
105        &self.physical_queues[shard]
106    }
107
108    /// Select a physical queue by caller-supplied sequence number.
109    ///
110    /// This is useful for bulk producers that want round-robin fanout and do
111    /// not need per-key ordering.
112    pub fn queue_for_index(&self, index: usize) -> &str {
113        &self.physical_queues[index % self.width()]
114    }
115
116    /// Return insert options routed by key.
117    ///
118    /// This sets both the physical queue and `ordering_key`, so per-key FIFO is
119    /// preserved even if the selected physical queue later uses multiple
120    /// queue-storage enqueue shards.
121    pub fn route_opts_by_key(&self, mut opts: InsertOpts, key: impl AsRef<[u8]>) -> InsertOpts {
122        let key = key.as_ref();
123        opts.queue = self.queue_for_key(key).to_string();
124        opts.ordering_key = Some(key.to_vec());
125        opts
126    }
127
128    /// Return insert options routed by round-robin index.
129    pub fn route_opts_by_index(&self, mut opts: InsertOpts, index: usize) -> InsertOpts {
130        opts.queue = self.queue_for_index(index).to_string();
131        opts
132    }
133}
134
135impl AsRef<[String]> for QueueFanout {
136    fn as_ref(&self) -> &[String] {
137        self.physical_queues()
138    }
139}
140
141impl<'a> IntoIterator for &'a QueueFanout {
142    type Item = &'a String;
143    type IntoIter = std::slice::Iter<'a, String>;
144
145    fn into_iter(self) -> Self::IntoIter {
146        self.physical_queues.iter()
147    }
148}
149
150fn validate_logical_queue(logical_queue: &str) -> Result<(), QueueFanoutError> {
151    if logical_queue.is_empty() {
152        return Err(QueueFanoutError::EmptyLogicalQueue);
153    }
154    Ok(())
155}
156
157fn validate_physical_count(count: usize) -> Result<(), QueueFanoutError> {
158    if count == 0 {
159        return Err(QueueFanoutError::ZeroWidth);
160    }
161
162    let max = i16::MAX as usize;
163    if count > max {
164        return Err(QueueFanoutError::TooManyPhysicalQueues { got: count, max });
165    }
166
167    Ok(())
168}
169
170fn validate_physical_queues(queues: &[String]) -> Result<(), QueueFanoutError> {
171    let mut seen = HashSet::with_capacity(queues.len());
172    for queue in queues {
173        if queue.is_empty() {
174            return Err(QueueFanoutError::EmptyPhysicalQueue);
175        }
176        if !seen.insert(queue.as_str()) {
177            return Err(QueueFanoutError::DuplicatePhysicalQueue {
178                queue: queue.clone(),
179            });
180        }
181    }
182    Ok(())
183}
184
185#[cfg(test)]
186mod tests {
187    use super::*;
188
189    #[test]
190    fn default_names_preserve_single_queue_shape() {
191        let fanout = QueueFanout::new("email", 1).expect("fanout should build");
192
193        assert_eq!(fanout.logical_queue(), "email");
194        assert_eq!(fanout.physical_queues(), &["email".to_string()]);
195        assert_eq!(fanout.queue_for_index(42), "email");
196    }
197
198    #[test]
199    fn default_names_are_stable_for_multiple_queues() {
200        let fanout = QueueFanout::new("email", 4).expect("fanout should build");
201
202        assert_eq!(
203            fanout.physical_queues(),
204            &[
205                "email__p0".to_string(),
206                "email__p1".to_string(),
207                "email__p2".to_string(),
208                "email__p3".to_string(),
209            ]
210        );
211        assert_eq!(fanout.queue_for_index(0), "email__p0");
212        assert_eq!(fanout.queue_for_index(5), "email__p1");
213    }
214
215    #[test]
216    fn key_routing_sets_queue_and_ordering_key() {
217        let fanout = QueueFanout::new("customer-updates", 4).expect("fanout should build");
218
219        let opts = fanout.route_opts_by_key(InsertOpts::default(), b"customer-42");
220
221        assert_eq!(fanout.queue_for_key(b"customer-42"), "customer-updates__p0");
222        assert_eq!(opts.queue, fanout.queue_for_key(b"customer-42"));
223        assert_eq!(opts.ordering_key.as_deref(), Some(&b"customer-42"[..]));
224    }
225
226    #[test]
227    fn explicit_queues_reject_empty_and_duplicate_names() {
228        let empty = QueueFanout::from_physical_queues("email", ["email-a", ""]);
229        assert!(matches!(empty, Err(QueueFanoutError::EmptyPhysicalQueue)));
230
231        let duplicate = QueueFanout::from_physical_queues("email", ["email-a", "email-a"]);
232        assert!(matches!(
233            duplicate,
234            Err(QueueFanoutError::DuplicatePhysicalQueue { queue }) if queue == "email-a"
235        ));
236    }
237
238    #[test]
239    fn explicit_queues_preserve_caller_order() {
240        let fanout =
241            QueueFanout::from_physical_queues("email", ["email-fast", "email-bulk", "email-slow"])
242                .expect("fanout should build");
243
244        assert_eq!(
245            fanout.physical_queues(),
246            &[
247                "email-fast".to_string(),
248                "email-bulk".to_string(),
249                "email-slow".to_string(),
250            ]
251        );
252        assert_eq!(fanout.queue_for_index(0), "email-fast");
253        assert_eq!(fanout.queue_for_index(4), "email-bulk");
254    }
255
256    #[test]
257    fn width_is_bounded_by_queue_storage_shard_type() {
258        let too_wide = QueueFanout::from_physical_queues(
259            "email",
260            (0..=(i16::MAX as usize)).map(|idx| format!("email-{idx}")),
261        );
262
263        assert!(matches!(
264            too_wide,
265            Err(QueueFanoutError::TooManyPhysicalQueues { got, max })
266                if got == i16::MAX as usize + 1 && max == i16::MAX as usize
267        ));
268    }
269}
awa_model/queue_fanout.rs

awa_model/
queue_fanout.rs
