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