1use alloc::collections::BTreeMap;
2use alloc::vec::Vec;
3
4use crate::packet::RawPacket;
5
6use super::types::InterfaceId;
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9pub enum OverflowPolicy {
10 DropNewest,
11 DropOldest,
12 DropWorst,
13}
14
15#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
16pub struct AnnounceVerifyKey {
17 pub destination_hash: [u8; 16],
18 pub random_blob: [u8; 10],
19 pub received_from: [u8; 16],
20}
21
22#[derive(Debug, Clone)]
23pub struct PendingAnnounce {
24 pub original_raw: Vec<u8>,
25 pub packet: RawPacket,
26 pub interface: InterfaceId,
27 pub received_from: [u8; 16],
28 pub queued_at: f64,
29 pub best_hops: u8,
30 pub emission_ts: u64,
31 pub random_blob: [u8; 10],
32}
33
34#[derive(Debug, Clone)]
35pub enum QueueEntry {
36 Pending(PendingAnnounce),
37 InFlight(PendingAnnounce),
38}
39
40#[derive(Debug, Clone)]
41pub struct AnnounceVerifyQueue {
42 pending: BTreeMap<AnnounceVerifyKey, QueueEntry>,
43 max_entries: usize,
44 max_bytes: usize,
45 max_stale_secs: f64,
46 overflow_policy: OverflowPolicy,
47 queued_bytes: usize,
48}
49
50impl AnnounceVerifyQueue {
51 pub fn new(max_entries: usize) -> Self {
52 Self::with_limits(max_entries, 256 * 1024, 30.0, OverflowPolicy::DropWorst)
53 }
54
55 pub fn with_limits(
56 max_entries: usize,
57 max_bytes: usize,
58 max_stale_secs: f64,
59 overflow_policy: OverflowPolicy,
60 ) -> Self {
61 Self {
62 pending: BTreeMap::new(),
63 max_entries: max_entries.max(1),
64 max_bytes: max_bytes.max(1),
65 max_stale_secs: max_stale_secs.max(0.001),
66 overflow_policy,
67 queued_bytes: 0,
68 }
69 }
70
71 pub fn enqueue(&mut self, key: AnnounceVerifyKey, entry: PendingAnnounce) -> bool {
72 if let Some(existing) = self.pending.get_mut(&key) {
73 return match existing {
74 QueueEntry::Pending(current) | QueueEntry::InFlight(current) => {
75 if entry.best_hops < current.best_hops {
76 let current_bytes = pending_bytes(current);
77 let replacement_bytes = pending_bytes(&entry);
78 self.queued_bytes = self
79 .queued_bytes
80 .saturating_sub(current_bytes)
81 .saturating_add(replacement_bytes);
82 *current = entry;
83 true
84 } else {
85 false
86 }
87 }
88 };
89 }
90
91 let entry_bytes = pending_bytes(&entry);
92 if entry_bytes > self.max_bytes {
93 return false;
94 }
95
96 while self.pending.len() >= self.max_entries
97 || self.queued_bytes.saturating_add(entry_bytes) > self.max_bytes
98 {
99 let Some(evict_key) = self.select_eviction_candidate(&entry) else {
100 return false;
101 };
102 self.remove_entry(&evict_key);
103 }
104
105 self.queued_bytes = self.queued_bytes.saturating_add(entry_bytes);
106 self.pending.insert(key, QueueEntry::Pending(entry));
107 true
108 }
109
110 pub fn take_pending(&mut self, now: f64) -> Vec<(AnnounceVerifyKey, PendingAnnounce)> {
111 let stale_before = now - self.max_stale_secs;
112 let stale_keys: Vec<_> = self
113 .pending
114 .iter()
115 .filter_map(|(key, entry)| match entry {
116 QueueEntry::Pending(current) | QueueEntry::InFlight(current)
117 if current.queued_at < stale_before =>
118 {
119 Some(*key)
120 }
121 _ => None,
122 })
123 .collect();
124 for key in stale_keys {
125 self.remove_entry(&key);
126 }
127
128 let keys: Vec<_> = self
129 .pending
130 .iter()
131 .filter_map(|(key, entry)| match entry {
132 QueueEntry::Pending(_) => Some(*key),
133 QueueEntry::InFlight(_) => None,
134 })
135 .collect();
136
137 let mut drained = Vec::with_capacity(keys.len());
138 for key in keys {
139 if let Some(entry) = self.pending.get_mut(&key) {
140 if let QueueEntry::Pending(current) = entry {
141 let cloned = current.clone();
142 *entry = QueueEntry::InFlight(cloned.clone());
143 drained.push((key, cloned));
144 }
145 }
146 }
147
148 drained
149 }
150
151 pub fn complete_success(&mut self, key: &AnnounceVerifyKey) -> Option<PendingAnnounce> {
152 match self.remove_entry(key) {
153 Some(QueueEntry::InFlight(entry)) => Some(entry),
154 Some(QueueEntry::Pending(entry)) => Some(entry),
155 None => None,
156 }
157 }
158
159 pub fn complete_failure(&mut self, key: &AnnounceVerifyKey) -> bool {
160 self.remove_entry(key).is_some()
161 }
162
163 pub fn len(&self) -> usize {
164 self.pending.len()
165 }
166
167 pub fn is_empty(&self) -> bool {
168 self.pending.is_empty()
169 }
170
171 pub fn queued_bytes(&self) -> usize {
172 self.queued_bytes
173 }
174
175 fn select_eviction_candidate(
176 &self,
177 incoming_entry: &PendingAnnounce,
178 ) -> Option<AnnounceVerifyKey> {
179 match self.overflow_policy {
180 OverflowPolicy::DropNewest => None,
181 OverflowPolicy::DropOldest => self
182 .pending
183 .iter()
184 .min_by(|a, b| {
185 queued_at_of(a.1)
186 .partial_cmp(&queued_at_of(b.1))
187 .unwrap_or(core::cmp::Ordering::Equal)
188 })
189 .map(|(key, _)| *key),
190 OverflowPolicy::DropWorst => {
191 let candidate = self
192 .pending
193 .iter()
194 .map(|(existing_key, existing_entry)| {
195 (*existing_key, pending_of(existing_entry))
196 })
197 .max_by(|a, b| {
198 a.1.best_hops.cmp(&b.1.best_hops).then_with(|| {
199 a.1.queued_at
200 .partial_cmp(&b.1.queued_at)
201 .unwrap_or(core::cmp::Ordering::Equal)
202 })
203 })?;
204 if incoming_entry.best_hops >= candidate.1.best_hops {
205 None
206 } else {
207 Some(candidate.0)
208 }
209 }
210 }
211 }
212
213 fn remove_entry(&mut self, key: &AnnounceVerifyKey) -> Option<QueueEntry> {
214 let removed = self.pending.remove(key)?;
215 self.queued_bytes = self
216 .queued_bytes
217 .saturating_sub(pending_bytes(pending_of(&removed)));
218 Some(removed)
219 }
220}
221
222fn pending_of(entry: &QueueEntry) -> &PendingAnnounce {
223 match entry {
224 QueueEntry::Pending(current) | QueueEntry::InFlight(current) => current,
225 }
226}
227
228fn queued_at_of(entry: &QueueEntry) -> f64 {
229 pending_of(entry).queued_at
230}
231
232fn pending_bytes(entry: &PendingAnnounce) -> usize {
233 entry.original_raw.len()
234 + entry.packet.data.len()
235 + entry.packet.transport_id.as_ref().map_or(0, |id| id.len())
236}
237
238#[cfg(test)]
239mod tests {
240 use super::*;
241 use crate::constants;
242 use crate::packet::{PacketFlags, RawPacket};
243
244 fn make_packet(dest: [u8; 16], hops: u8, fill: u8) -> RawPacket {
245 RawPacket::pack(
246 PacketFlags {
247 header_type: constants::HEADER_1,
248 context_flag: constants::FLAG_UNSET,
249 transport_type: constants::TRANSPORT_BROADCAST,
250 destination_type: constants::DESTINATION_SINGLE,
251 packet_type: constants::PACKET_TYPE_ANNOUNCE,
252 },
253 hops,
254 &dest,
255 None,
256 constants::CONTEXT_NONE,
257 &[fill; 8],
258 )
259 .unwrap()
260 }
261
262 fn make_pending(
263 dest: [u8; 16],
264 random_blob: [u8; 10],
265 received_from: [u8; 16],
266 hops: u8,
267 ) -> (AnnounceVerifyKey, PendingAnnounce) {
268 (
269 AnnounceVerifyKey {
270 destination_hash: dest,
271 random_blob,
272 received_from,
273 },
274 PendingAnnounce {
275 original_raw: vec![hops],
276 packet: make_packet(dest, hops, hops),
277 interface: InterfaceId(1),
278 received_from,
279 queued_at: 10.0,
280 best_hops: hops,
281 emission_ts: 42,
282 random_blob,
283 },
284 )
285 }
286
287 #[test]
288 fn enqueue_replaces_lower_hops_and_preserves_distinct_paths() {
289 let mut queue = AnnounceVerifyQueue::new(8);
290 let dest = [1; 16];
291 let random = [2; 10];
292 let rx_a = [3; 16];
293 let rx_b = [4; 16];
294
295 let (key_a, entry_a) = make_pending(dest, random, rx_a, 5);
296 assert!(queue.enqueue(key_a, entry_a));
297
298 let (_, better_a) = make_pending(dest, random, rx_a, 3);
299 assert!(queue.enqueue(key_a, better_a));
300 assert_eq!(queue.len(), 1);
301
302 let (key_b, entry_b) = make_pending(dest, random, rx_b, 4);
303 assert!(queue.enqueue(key_b, entry_b));
304 assert_eq!(queue.len(), 2);
305
306 let taken = queue.take_pending(10.0);
307 assert_eq!(taken.len(), 2);
308 assert!(taken
309 .iter()
310 .any(|(key, entry)| *key == key_a && entry.best_hops == 3));
311 assert!(taken
312 .iter()
313 .any(|(key, entry)| *key == key_b && entry.best_hops == 4));
314 }
315
316 #[test]
317 fn enqueue_updates_inflight_and_cleans_stale_entries() {
318 let mut queue = AnnounceVerifyQueue::new(2);
319 let dest = [8; 16];
320 let random = [9; 10];
321 let recv = [10; 16];
322
323 let (key, entry) = make_pending(dest, random, recv, 6);
324 assert!(queue.enqueue(key, entry));
325 let _ = queue.take_pending(20.0);
326
327 let (_, better) = make_pending(dest, random, recv, 2);
328 assert!(queue.enqueue(key, better));
329 let completed = queue.complete_success(&key).unwrap();
330 assert_eq!(completed.best_hops, 2);
331
332 let (stale_key, mut stale) = make_pending([11; 16], [12; 10], [13; 16], 7);
333 stale.queued_at = 1.0;
334 assert!(queue.enqueue(stale_key, stale));
335 assert!(queue.take_pending(40.0).is_empty());
336 assert_eq!(queue.len(), 0);
337 }
338
339 #[test]
340 fn enqueue_evicts_worst_entry_when_full() {
341 let mut queue = AnnounceVerifyQueue::with_limits(2, 1024, 30.0, OverflowPolicy::DropWorst);
342 let (k1, e1) = make_pending([1; 16], [1; 10], [1; 16], 8);
343 let (k2, e2) = make_pending([2; 16], [2; 10], [2; 16], 5);
344 let (k3, e3) = make_pending([3; 16], [3; 10], [3; 16], 4);
345 let (_, e4) = make_pending([4; 16], [4; 10], [4; 16], 9);
346
347 assert!(queue.enqueue(k1, e1));
348 assert!(queue.enqueue(k2, e2));
349 assert!(queue.enqueue(k3, e3));
350 assert_eq!(queue.len(), 2);
351 assert!(!queue.enqueue(
352 AnnounceVerifyKey {
353 destination_hash: [4; 16],
354 random_blob: [4; 10],
355 received_from: [4; 16],
356 },
357 e4
358 ));
359
360 let taken = queue.take_pending(10.0);
361 assert_eq!(taken.len(), 2);
362 assert!(taken.iter().all(|(key, _)| *key != k1));
363 }
364
365 #[test]
366 fn drop_newest_policy_rejects_when_full() {
367 let mut queue = AnnounceVerifyQueue::with_limits(1, 1024, 30.0, OverflowPolicy::DropNewest);
368 let (k1, e1) = make_pending([1; 16], [1; 10], [1; 16], 4);
369 let (k2, e2) = make_pending([2; 16], [2; 10], [2; 16], 1);
370 assert!(queue.enqueue(k1, e1));
371 assert!(!queue.enqueue(k2, e2));
372 let taken = queue.take_pending(10.0);
373 assert_eq!(taken.len(), 1);
374 assert_eq!(taken[0].0, k1);
375 }
376
377 #[test]
378 fn drop_oldest_policy_evicts_oldest_for_byte_cap() {
379 let mut queue = AnnounceVerifyQueue::with_limits(4, 24, 30.0, OverflowPolicy::DropOldest);
380 let (k1, mut e1) = make_pending([1; 16], [1; 10], [1; 16], 4);
381 let (k2, mut e2) = make_pending([2; 16], [2; 10], [2; 16], 3);
382 e1.original_raw = vec![1; 12];
383 e2.original_raw = vec![2; 12];
384 e1.queued_at = 1.0;
385 e2.queued_at = 2.0;
386 assert!(queue.enqueue(k1, e1));
387 assert!(queue.enqueue(k2, e2));
388 assert_eq!(queue.len(), 1);
389 let taken = queue.take_pending(10.0);
390 assert_eq!(taken.len(), 1);
391 assert_eq!(taken[0].0, k2);
392 }
393}