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