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 pending_interface(&self, key: &AnnounceVerifyKey) -> Option<InterfaceId> {
164 self.pending
165 .get(key)
166 .map(|entry| pending_of(entry).interface)
167 }
168
169 pub fn pending_identity_hash(&self, key: &AnnounceVerifyKey) -> Option<[u8; 16]> {
170 let packet = &pending_of(self.pending.get(key)?).packet;
171 Some(crate::hash::truncated_hash(packet.data.get(..64)?))
172 }
173
174 pub fn len(&self) -> usize {
175 self.pending.len()
176 }
177
178 pub fn is_empty(&self) -> bool {
179 self.pending.is_empty()
180 }
181
182 pub fn queued_bytes(&self) -> usize {
183 self.queued_bytes
184 }
185
186 pub fn clear(&mut self) {
187 self.pending.clear();
188 self.queued_bytes = 0;
189 }
190
191 fn select_eviction_candidate(
192 &self,
193 incoming_entry: &PendingAnnounce,
194 ) -> Option<AnnounceVerifyKey> {
195 match self.overflow_policy {
196 OverflowPolicy::DropNewest => None,
197 OverflowPolicy::DropOldest => self
198 .pending
199 .iter()
200 .min_by(|a, b| {
201 queued_at_of(a.1)
202 .partial_cmp(&queued_at_of(b.1))
203 .unwrap_or(core::cmp::Ordering::Equal)
204 })
205 .map(|(key, _)| *key),
206 OverflowPolicy::DropWorst => {
207 let candidate = self
208 .pending
209 .iter()
210 .map(|(existing_key, existing_entry)| {
211 (*existing_key, pending_of(existing_entry))
212 })
213 .max_by(|a, b| {
214 a.1.best_hops.cmp(&b.1.best_hops).then_with(|| {
215 a.1.queued_at
216 .partial_cmp(&b.1.queued_at)
217 .unwrap_or(core::cmp::Ordering::Equal)
218 })
219 })?;
220 if incoming_entry.best_hops >= candidate.1.best_hops {
221 None
222 } else {
223 Some(candidate.0)
224 }
225 }
226 }
227 }
228
229 fn remove_entry(&mut self, key: &AnnounceVerifyKey) -> Option<QueueEntry> {
230 let removed = self.pending.remove(key)?;
231 self.queued_bytes = self
232 .queued_bytes
233 .saturating_sub(pending_bytes(pending_of(&removed)));
234 Some(removed)
235 }
236}
237
238fn pending_of(entry: &QueueEntry) -> &PendingAnnounce {
239 match entry {
240 QueueEntry::Pending(current) | QueueEntry::InFlight(current) => current,
241 }
242}
243
244fn queued_at_of(entry: &QueueEntry) -> f64 {
245 pending_of(entry).queued_at
246}
247
248fn pending_bytes(entry: &PendingAnnounce) -> usize {
249 entry.original_raw.len()
250 + entry.packet.data.len()
251 + entry.packet.transport_id.as_ref().map_or(0, |id| id.len())
252}
253
254#[cfg(test)]
255mod tests {
256 use super::*;
257 use crate::constants;
258 use crate::packet::{PacketFlags, RawPacket};
259
260 fn make_packet(dest: [u8; 16], hops: u8, fill: u8) -> RawPacket {
261 RawPacket::pack(
262 PacketFlags {
263 header_type: constants::HEADER_1,
264 context_flag: constants::FLAG_UNSET,
265 transport_type: constants::TRANSPORT_BROADCAST,
266 destination_type: constants::DESTINATION_SINGLE,
267 packet_type: constants::PACKET_TYPE_ANNOUNCE,
268 },
269 hops,
270 &dest,
271 None,
272 constants::CONTEXT_NONE,
273 &[fill; 8],
274 )
275 .unwrap()
276 }
277
278 fn make_pending(
279 dest: [u8; 16],
280 random_blob: [u8; 10],
281 received_from: [u8; 16],
282 hops: u8,
283 ) -> (AnnounceVerifyKey, PendingAnnounce) {
284 (
285 AnnounceVerifyKey {
286 destination_hash: dest,
287 random_blob,
288 received_from,
289 },
290 PendingAnnounce {
291 original_raw: vec![hops],
292 packet: make_packet(dest, hops, hops),
293 interface: InterfaceId(1),
294 received_from,
295 queued_at: 10.0,
296 best_hops: hops,
297 emission_ts: 42,
298 random_blob,
299 },
300 )
301 }
302
303 #[test]
304 fn enqueue_replaces_lower_hops_and_preserves_distinct_paths() {
305 let mut queue = AnnounceVerifyQueue::new(8);
306 let dest = [1; 16];
307 let random = [2; 10];
308 let rx_a = [3; 16];
309 let rx_b = [4; 16];
310
311 let (key_a, entry_a) = make_pending(dest, random, rx_a, 5);
312 assert!(queue.enqueue(key_a, entry_a));
313
314 let (_, better_a) = make_pending(dest, random, rx_a, 3);
315 assert!(queue.enqueue(key_a, better_a));
316 assert_eq!(queue.len(), 1);
317
318 let (key_b, entry_b) = make_pending(dest, random, rx_b, 4);
319 assert!(queue.enqueue(key_b, entry_b));
320 assert_eq!(queue.len(), 2);
321
322 let taken = queue.take_pending(10.0);
323 assert_eq!(taken.len(), 2);
324 assert!(taken
325 .iter()
326 .any(|(key, entry)| *key == key_a && entry.best_hops == 3));
327 assert!(taken
328 .iter()
329 .any(|(key, entry)| *key == key_b && entry.best_hops == 4));
330 }
331
332 #[test]
333 fn enqueue_updates_inflight_and_cleans_stale_entries() {
334 let mut queue = AnnounceVerifyQueue::new(2);
335 let dest = [8; 16];
336 let random = [9; 10];
337 let recv = [10; 16];
338
339 let (key, entry) = make_pending(dest, random, recv, 6);
340 assert!(queue.enqueue(key, entry));
341 let _ = queue.take_pending(20.0);
342
343 let (_, better) = make_pending(dest, random, recv, 2);
344 assert!(queue.enqueue(key, better));
345 let completed = queue.complete_success(&key).unwrap();
346 assert_eq!(completed.best_hops, 2);
347
348 let (stale_key, mut stale) = make_pending([11; 16], [12; 10], [13; 16], 7);
349 stale.queued_at = 1.0;
350 assert!(queue.enqueue(stale_key, stale));
351 assert!(queue.take_pending(40.0).is_empty());
352 assert_eq!(queue.len(), 0);
353 }
354
355 #[test]
356 fn enqueue_evicts_worst_entry_when_full() {
357 let mut queue = AnnounceVerifyQueue::with_limits(2, 1024, 30.0, OverflowPolicy::DropWorst);
358 let (k1, e1) = make_pending([1; 16], [1; 10], [1; 16], 8);
359 let (k2, e2) = make_pending([2; 16], [2; 10], [2; 16], 5);
360 let (k3, e3) = make_pending([3; 16], [3; 10], [3; 16], 4);
361 let (_, e4) = make_pending([4; 16], [4; 10], [4; 16], 9);
362
363 assert!(queue.enqueue(k1, e1));
364 assert!(queue.enqueue(k2, e2));
365 assert!(queue.enqueue(k3, e3));
366 assert_eq!(queue.len(), 2);
367 assert!(!queue.enqueue(
368 AnnounceVerifyKey {
369 destination_hash: [4; 16],
370 random_blob: [4; 10],
371 received_from: [4; 16],
372 },
373 e4
374 ));
375
376 let taken = queue.take_pending(10.0);
377 assert_eq!(taken.len(), 2);
378 assert!(taken.iter().all(|(key, _)| *key != k1));
379 }
380
381 #[test]
382 fn drop_newest_policy_rejects_when_full() {
383 let mut queue = AnnounceVerifyQueue::with_limits(1, 1024, 30.0, OverflowPolicy::DropNewest);
384 let (k1, e1) = make_pending([1; 16], [1; 10], [1; 16], 4);
385 let (k2, e2) = make_pending([2; 16], [2; 10], [2; 16], 1);
386 assert!(queue.enqueue(k1, e1));
387 assert!(!queue.enqueue(k2, e2));
388 let taken = queue.take_pending(10.0);
389 assert_eq!(taken.len(), 1);
390 assert_eq!(taken[0].0, k1);
391 }
392
393 #[test]
394 fn drop_oldest_policy_evicts_oldest_for_byte_cap() {
395 let mut queue = AnnounceVerifyQueue::with_limits(4, 24, 30.0, OverflowPolicy::DropOldest);
396 let (k1, mut e1) = make_pending([1; 16], [1; 10], [1; 16], 4);
397 let (k2, mut e2) = make_pending([2; 16], [2; 10], [2; 16], 3);
398 e1.original_raw = vec![1; 12];
399 e2.original_raw = vec![2; 12];
400 e1.queued_at = 1.0;
401 e2.queued_at = 2.0;
402 assert!(queue.enqueue(k1, e1));
403 assert!(queue.enqueue(k2, e2));
404 assert_eq!(queue.len(), 1);
405 let taken = queue.take_pending(10.0);
406 assert_eq!(taken.len(), 1);
407 assert_eq!(taken[0].0, k2);
408 }
409
410 #[test]
411 fn clear_removes_pending_and_inflight_entries_and_resets_bytes() {
412 let mut queue = AnnounceVerifyQueue::new(4);
413 let (pending_key, pending) = make_pending([1; 16], [1; 10], [1; 16], 4);
414 let (inflight_key, inflight) = make_pending([2; 16], [2; 10], [2; 16], 3);
415 assert!(queue.enqueue(pending_key, pending));
416 assert!(queue.enqueue(inflight_key, inflight));
417 let _ = queue.take_pending(10.0);
418
419 assert_eq!(queue.len(), 2);
420 assert!(queue.queued_bytes() > 0);
421
422 queue.clear();
423
424 assert!(queue.is_empty());
425 assert_eq!(queue.queued_bytes(), 0);
426 assert!(queue.take_pending(10.0).is_empty());
427 }
428}