Skip to main content

sozu_lib/protocol/udp/
manager.rs

1//! The sans-io [`UdpManager`]: admission, flow table, LB request, timers.
2//!
3//! `UdpManager` owns the flow table (`HashMap<FlowKey, FlowId>` over a
4//! `slab::Slab<UdpFlow>`), admission ("allocate nothing for unknown / over-cap
5//! / invalid datagrams"), the flow-table cap and shedding, the per-source flow
6//! limits (`max_flows_per_ip` / `max_flows_per_subnet`), pluggable flow-key
7//! extraction ([`FlowKeyExtractor`]), backend selection for new flows (from
8//! the [`BackendSource`] view the embedder supplies), and the
9//! timer scheduling: a **single armed manager-wide deadline**, with each
10//! expiry **revalidating every flow's idle deadline** so a stale expiry can
11//! never close a refreshed flow.
12//!
13//! Pure: every entry point that depends on time takes `now: Instant`; the hash
14//! seed is injected at construction. The shell drives the manager with
15//! [`ManagerInput`] and drains [`Output`] via [`UdpManager::poll_output`].
16
17use std::{
18    collections::{HashMap, VecDeque},
19    hash::{Hash, Hasher},
20    net::{IpAddr, SocketAddr},
21    time::Instant,
22};
23
24use slab::Slab;
25use sozu_command::state::ClusterId;
26
27use crate::protocol::udp::{
28    BackendSource, ClusterConfig, ConfigEvent, DropReason, FlowId, FlowKey, ManagerInput,
29    MetricEvent, Output, Transmit,
30    flow::{CloseReason, FlowPhase, UdpFlow},
31    proxy_protocol::prepend_dgram_header,
32};
33use crate::server::subnet_key;
34
35/// Live flow count per cluster, then per source key (an IP or a masked
36/// subnet). Nested so a lookup borrows the cluster id instead of cloning it,
37/// as `SessionManager` does for the TCP/HTTP connection counters.
38type SourceCounts = HashMap<ClusterId, HashMap<IpAddr, usize>>;
39
40/// The subnet key of `ip` under the prefixes captured in `cfg`.
41fn subnet_of(cfg: &ClusterConfig, ip: IpAddr) -> IpAddr {
42    subnet_key(&ip, cfg.subnet_ipv4_prefix, cfg.subnet_ipv6_prefix)
43}
44
45/// The live flow count recorded for `(cluster, key)`; `0` when absent.
46fn source_count(counts: &SourceCounts, cluster: &str, key: &IpAddr) -> usize {
47    counts
48        .get(cluster)
49        .and_then(|by_key| by_key.get(key))
50        .copied()
51        .unwrap_or(0)
52}
53
54/// Count one more flow for `(cluster, key)`.
55fn source_incr(counts: &mut SourceCounts, cluster: &ClusterId, key: IpAddr) {
56    *counts
57        .entry(cluster.clone())
58        .or_default()
59        .entry(key)
60        .or_insert(0) += 1;
61}
62
63/// Count one flow less for `(cluster, key)`, removing entries that reach zero
64/// so the maps never outgrow the live flow set.
65fn source_decr(counts: &mut SourceCounts, cluster: &str, key: &IpAddr) {
66    let Some(by_key) = counts.get_mut(cluster) else {
67        debug_assert!(false, "a closing flow's cluster must have a source count");
68        return;
69    };
70    match by_key.get_mut(key) {
71        Some(count) if *count > 1 => *count -= 1,
72        Some(_) => {
73            by_key.remove(key);
74        }
75        None => debug_assert!(false, "a closing flow's source must have a count"),
76    }
77    if by_key.is_empty() {
78        counts.remove(cluster);
79    }
80}
81
82/// Extracts a [`FlowKey`] from an admitted client datagram. The default
83/// [`SourceTupleExtractor`] keys on the real client source address (source IP
84/// and port, whatever the cluster's affinity key). The trait is the
85/// only seam for alternative keying (e.g. a QUIC-CID extractor, a non-goal);
86/// the 4-tuple impl is the only one in scope.
87pub trait FlowKeyExtractor {
88    /// Compute the flow key for a datagram from `src`. Returns `None` to reject
89    /// the datagram (the manager then emits `Drop(Invalid)` and allocates
90    /// nothing). `cfg` is the listener's active cluster config.
91    fn flow_key(&self, src: SocketAddr, payload: &[u8], cfg: &ClusterConfig) -> Option<FlowKey>;
92}
93
94/// The in-scope flow-key extractor: keys on the real client source address,
95/// IP and port. The cluster's `affinity_with_port` knob is not read here: it
96/// shapes the backend-selection hash only ([`UdpManager`]'s `affinity_hash`).
97#[derive(Clone, Copy, Debug, Default)]
98pub struct SourceTupleExtractor;
99
100impl FlowKeyExtractor for SourceTupleExtractor {
101    fn flow_key(&self, src: SocketAddr, payload: &[u8], _cfg: &ClusterConfig) -> Option<FlowKey> {
102        // "Silence is a virtue": an empty datagram is not a valid flow trigger.
103        if payload.is_empty() {
104            return None;
105        }
106        Some(FlowKey::from_src(src))
107    }
108}
109
110/// The pure UDP flow manager. Generic over the flow-key extractor so the seam
111/// stays type-checked; the shell instantiates it with [`SourceTupleExtractor`].
112pub struct UdpManager<E: FlowKeyExtractor = SourceTupleExtractor> {
113    /// `FlowKey -> FlowId` lookup for tracked-flow reuse (two-tier selection).
114    table: HashMap<FlowKey, FlowId>,
115    /// Slab of admitted flows; the slot index is the [`FlowId`].
116    flows: Slab<UdpFlow>,
117    /// Flow-table cap. New flows beyond this are shed (drop + metric).
118    max_flows: usize,
119    /// Live flows per `(cluster, source IP)`, for the cluster's
120    /// `max_flows_per_ip`. Every admitted flow is counted, whatever the
121    /// limit, so a limit raised from `0` at runtime applies to the flows
122    /// already open. Keyed on the cluster each flow captured at admission.
123    flows_per_ip: SourceCounts,
124    /// Live flows per `(cluster, source subnet)`, for `max_flows_per_subnet`.
125    /// The subnet key uses the prefixes the flow captured at admission, so a
126    /// close always decrements the key its admission incremented.
127    flows_per_subnet: SourceCounts,
128    /// Maximum accepted rx datagram size; larger datagrams are dropped as
129    /// truncated.
130    max_rx_datagram_size: usize,
131    /// Active cluster routing + per-cluster knobs for *new* flows.
132    cluster: ClusterConfig,
133    /// Per-worker hash seed, injected once at construction; persisted across
134    /// reconfig for stable affinity.
135    hash_seed: u64,
136    /// Pluggable flow-key extractor.
137    extractor: E,
138    /// Draining: admit no new flows; let existing ones reach teardown.
139    draining: bool,
140
141    /// FIFO of outputs the shell drains via [`poll_output`].
142    outputs: VecDeque<Output>,
143    /// The single armed manager-wide deadline currently reflected to the shell
144    /// via the last `ArmTimer`. `None` means no timer is armed — which
145    /// [`handle_timeout`](Self::handle_timeout) sets on entry, because the
146    /// expiry that called it consumed the shell's wheel entry.
147    armed_deadline: Option<Instant>,
148
149    /// High-water mark of every `max_flows` cap ever set (construction +
150    /// `SetMaxFlows`). The live cap can be shrunk below `flows.len()` by a
151    /// `SetMaxFlows`, so `flows.len() <= max_flows` is NOT an invariant; but a
152    /// flow can only ever have been admitted under *some* cap that was in force
153    /// at admission, so `flows.len() <= max_flows_high_water` always holds.
154    /// Debug-only — only read by [`check_invariants`](Self::check_invariants).
155    #[cfg(debug_assertions)]
156    max_flows_high_water: usize,
157}
158
159impl UdpManager<SourceTupleExtractor> {
160    /// Construct a manager with the default client source-address extractor.
161    pub fn new(
162        cluster: ClusterConfig,
163        max_flows: usize,
164        max_rx_datagram_size: usize,
165        hash_seed: u64,
166    ) -> Self {
167        Self::with_extractor(
168            cluster,
169            max_flows,
170            max_rx_datagram_size,
171            hash_seed,
172            SourceTupleExtractor,
173        )
174    }
175}
176
177impl<E: FlowKeyExtractor> UdpManager<E> {
178    /// Construct a manager with a custom flow-key extractor.
179    pub fn with_extractor(
180        cluster: ClusterConfig,
181        max_flows: usize,
182        max_rx_datagram_size: usize,
183        hash_seed: u64,
184        extractor: E,
185    ) -> Self {
186        UdpManager {
187            table: HashMap::new(),
188            flows: Slab::new(),
189            max_flows,
190            flows_per_ip: HashMap::new(),
191            flows_per_subnet: HashMap::new(),
192            max_rx_datagram_size,
193            cluster,
194            hash_seed,
195            extractor,
196            draining: false,
197            outputs: VecDeque::new(),
198            armed_deadline: None,
199            #[cfg(debug_assertions)]
200            max_flows_high_water: max_flows,
201        }
202    }
203
204    // ---- introspection (used by log macros / tests / the shell) ------------
205
206    /// Number of currently admitted flows. Mirrors `udp.active_flows`.
207    pub fn flow_count(&self) -> usize {
208        self.flows.len()
209    }
210
211    /// The configured flow-table cap.
212    pub fn max_flows(&self) -> usize {
213        self.max_flows
214    }
215
216    /// The cluster routing + per-cluster knobs new flows are admitted under.
217    pub fn cluster_config(&self) -> &ClusterConfig {
218        &self.cluster
219    }
220
221    /// Whether the listener is draining.
222    pub fn is_draining(&self) -> bool {
223        self.draining
224    }
225
226    /// Borrow a flow by id (for the shell's access log on close).
227    pub fn flow(&self, flow: FlowId) -> Option<&UdpFlow> {
228        self.flows.get(flow)
229    }
230
231    // ---- input -------------------------------------------------------------
232
233    /// Feed one input into the manager. Pure: `now` is injected.
234    pub fn handle_input(&mut self, input: ManagerInput<'_>, now: Instant) {
235        match input {
236            ManagerInput::ClientDatagram {
237                src,
238                payload,
239                backends,
240            } => self.on_client_datagram(src, payload, backends, now),
241            ManagerInput::BackendDatagram { flow, payload } => {
242                self.on_backend_datagram(flow, payload, now)
243            }
244            ManagerInput::Config(event) => self.on_config(event, now),
245        }
246        // Post-condition: the public method ran to completion under the caller's
247        // lock, so every structural invariant must hold again.
248        self.debug_assert_invariants();
249    }
250
251    /// Tear down a single flow on demand, emitting the same outputs a normal
252    /// idle close produces — `Output::Metric(MetricEvent::FlowEvicted)` then
253    /// `Output::CloseFlow(flow)` — so the shell draining `poll_output()`
254    /// decrements `udp.active_flows` and frees the upstream socket exactly once.
255    ///
256    /// The shell calls this when it cannot establish a flow it just admitted:
257    /// the upstream `connect()` failed (EMFILE / connection refused) or no
258    /// backend resolved. Without it the flow would sit `AwaitingBackend` /
259    /// `Established` pinning a `max_flows` slot for the full idle timeout.
260    ///
261    /// Idempotent: a missing flow or one already `Closing` is a no-op — no
262    /// double-evict, no gauge underflow. Works in either `AwaitingBackend` or
263    /// `Established`. `now` is accepted for signature symmetry with the other
264    /// time-driven entry points (the teardown itself is time-independent).
265    pub fn abort_flow(&mut self, flow: FlowId, _now: Instant, reason: CloseReason) {
266        self.close_flow(flow, reason);
267        self.debug_assert_invariants();
268    }
269
270    /// Tear down EVERY live flow, emitting — for each — the same outputs a
271    /// normal idle close produces (`Output::Metric(MetricEvent::FlowEvicted)`
272    /// then `Output::CloseFlow(flow)`), so the shell draining `poll_output()`
273    /// decrements `udp.active_flows` and frees each upstream socket exactly once
274    /// per flow.
275    ///
276    /// The shell calls this on listener remove / deactivate / soft-stop before
277    /// dropping the manager, so the active-flows gauge does not leak. Draining
278    /// the flow table + slab to zero; a flow already `Closing` is skipped
279    /// (idempotent, no double-evict, no underflow). After this returns,
280    /// `flow_count() == 0` and no timer is armed.
281    pub fn close_all(&mut self, now: Instant) {
282        // Snapshot the live ids first: `close_flow` mutates the slab, so we must
283        // not iterate it while closing. `Closing` slots are already gone from
284        // the slab (`close_flow` removes them), so every id here is live.
285        let live: Vec<FlowId> = self.flows.iter().map(|(id, _)| id).collect();
286        for flow_id in live {
287            self.abort_flow(flow_id, now, CloseReason::Drain);
288        }
289        // Post: the table and slab are drained to zero and no timer is armed.
290        #[cfg(debug_assertions)]
291        {
292            debug_assert_eq!(self.flow_count(), 0, "close_all must drain every flow");
293            debug_assert!(
294                self.table.is_empty(),
295                "close_all must clear every table entry"
296            );
297            debug_assert!(
298                self.armed_deadline.is_none(),
299                "close_all must leave no armed timer"
300            );
301        }
302        self.debug_assert_invariants();
303    }
304
305    fn on_client_datagram(
306        &mut self,
307        src: SocketAddr,
308        payload: &[u8],
309        backends: &mut dyn BackendSource,
310        now: Instant,
311    ) {
312        // Over-size check first — never allocate for a truncated datagram.
313        if payload.len() > self.max_rx_datagram_size {
314            self.drop_datagram(DropReason::Truncated);
315            return;
316        }
317        // No backend cluster configured → nothing to route to.
318        if self.cluster.cluster.is_empty() {
319            self.drop_datagram(DropReason::NoBackend);
320            return;
321        }
322        // Extract the flow key; rejection allocates nothing.
323        let key = match self.extractor.flow_key(src, payload, &self.cluster) {
324            Some(key) => key,
325            None => {
326                self.drop_datagram(DropReason::Invalid);
327                return;
328            }
329        };
330
331        // Tracked flow → reuse its backend (two-tier selection).
332        if let Some(&flow_id) = self.table.get(&key) {
333            self.forward_on_existing_flow(flow_id, payload, now);
334            return;
335        }
336
337        // New flow. Draining or at cap → shed, allocate nothing.
338        // Drop / shed paths ("silence is a virtue"): a reject must allocate
339        // nothing, so `flows.len()` is unchanged across it. Snapshot the count
340        // to pair-assert that on every early return below. Not debug-gated: the
341        // `debug_assert_eq!`s that read it compile in release too (execution only
342        // is gated); the read is dead there and the optimizer drops the binding.
343        let flows_before_admit = self.flows.len();
344        if self.draining {
345            self.drop_datagram(DropReason::Shed);
346            debug_assert_eq!(
347                self.flows.len(),
348                flows_before_admit,
349                "drain shed must allocate no flow"
350            );
351            return;
352        }
353        if self.source_at_limit(src.ip()) {
354            self.outputs
355                .push_back(Output::Metric(MetricEvent::FlowShedSourceLimit));
356            self.drop_datagram(DropReason::Shed);
357            debug_assert_eq!(
358                self.flows.len(),
359                flows_before_admit,
360                "per-source shed must allocate no flow"
361            );
362            return;
363        }
364        if self.flows.len() >= self.max_flows {
365            self.outputs
366                .push_back(Output::Metric(MetricEvent::FlowShed));
367            self.drop_datagram(DropReason::Shed);
368            debug_assert_eq!(
369                self.flows.len(),
370                flows_before_admit,
371                "cap shed must allocate no flow"
372            );
373            return;
374        }
375
376        // Admit. Pre-conditions (positive + negative space): we are on the admit
377        // path, so there is room under the live cap AND the key is NOT already
378        // tracked (a tracked key would have been served above without a new
379        // slot — a double-insert would orphan the previous flow and leak a slab
380        // slot).
381        debug_assert!(
382            self.flows.len() < self.max_flows,
383            "admit path entered while at/over the cap"
384        );
385        debug_assert!(
386            !self.table.contains_key(&key),
387            "admit path entered for a key already in the table (would orphan a flow)"
388        );
389
390        // Select the backend HERE, in the core, from the view the embedder
391        // handed in with this datagram (#1340, Question 6). The affinity hash
392        // is computed from a provisional flow so HRW / Maglev keep a client
393        // pinned; it reads only `client` and `config`, neither of which the
394        // backend can change.
395        let key_hash = self.affinity_hash(&UdpFlow::new(
396            src,
397            self.cluster.clone(),
398            String::new(),
399            src,
400            now,
401        ));
402        let Some((backend_id, backend_addr)) =
403            backends.select(&self.cluster.cluster, Some(key_hash), now)
404        else {
405            // The cluster has nothing that can serve. Drop without allocating
406            // a slab slot: there is no flow to park and nothing to abort.
407            self.drop_datagram(DropReason::NoBackend);
408            debug_assert_eq!(
409                self.flows.len(),
410                flows_before_admit,
411                "a failed selection must allocate no flow"
412            );
413            return;
414        };
415
416        // Admit: one slab slot. The flow is `Established` from birth — its
417        // backend was chosen a line ago — so there is no window in which it
418        // exists without one, and no datagram to buffer against that window.
419        let flow = UdpFlow::new(src, self.cluster.clone(), backend_id, backend_addr, now);
420        let flow_id = self.flows.insert(flow);
421        self.table.insert(key, flow_id);
422        let subnet = subnet_of(&self.cluster, src.ip());
423        source_incr(&mut self.flows_per_ip, &self.cluster.cluster, src.ip());
424        source_incr(&mut self.flows_per_subnet, &self.cluster.cluster, subnet);
425
426        // Post (admission): exactly one slot was added and the key now maps to
427        // it. Pairs the pre-conditions above (grew by exactly 1, not 0 or 2).
428        debug_assert_eq!(
429            self.flows.len(),
430            flows_before_admit + 1,
431            "admission must add exactly one flow"
432        );
433        debug_assert_eq!(
434            self.table.get(&key),
435            Some(&flow_id),
436            "admission must map the key to the new flow"
437        );
438
439        self.outputs
440            .push_back(Output::Metric(MetricEvent::FlowCreated));
441        // Ask the shell to open the connected upstream socket and register
442        // `upstream_token -> flow` for NAT return demux. This is the returned
443        // request the embedder fulfils — the shape `Output::SelectBackend` used
444        // to have, except the decision it carries is now the core's.
445        self.outputs.push_back(Output::OpenUpstream {
446            flow: flow_id,
447            backend: backend_addr,
448        });
449
450        // Forward the admission datagram immediately. It used to be buffered
451        // until the shell resolved a backend; there is no such window now, so
452        // it is a forward like any other and is counted as one.
453        self.forward_admission_datagram(flow_id, payload, backend_addr, now);
454    }
455
456    /// Send the datagram that caused a flow's admission, counting it as the
457    /// forward it now is.
458    ///
459    /// Before selection moved into the core this datagram sat in a one-slot
460    /// newest-wins buffer until `BackendResolved` arrived, and a second
461    /// datagram in that window replaced it. The window is gone, so the buffer
462    /// is too: every datagram of an opening burst is forwarded.
463    fn forward_admission_datagram(
464        &mut self,
465        flow_id: FlowId,
466        payload: &[u8],
467        backend: SocketAddr,
468        now: Instant,
469    ) {
470        let Some(flow) = self.flows.get_mut(flow_id) else {
471            debug_assert!(false, "the flow was inserted one statement ago");
472            return;
473        };
474        let mut out = payload.to_vec();
475        flow.on_client_datagram(now);
476        if flow.take_proxy_protocol() {
477            prepend_dgram_header(&mut out, flow.client, backend);
478        }
479        let teardown = flow.teardown_reason();
480        self.outputs
481            .push_back(Output::Metric(MetricEvent::DatagramIn(payload.len())));
482        self.outputs.push_back(Output::SendToBackend(Transmit {
483            dst: backend,
484            segment_size: None,
485            payload: out,
486        }));
487        if let Some(reason) = teardown {
488            self.close_flow(flow_id, reason);
489            return;
490        }
491        // `on_client_datagram` already refreshed the deadline + generation.
492        self.reschedule();
493    }
494
495    fn forward_on_existing_flow(&mut self, flow_id: FlowId, payload: &[u8], now: Instant) {
496        let Some(flow) = self.flows.get_mut(flow_id) else {
497            // Table/slab desync should be impossible, but never panic on it.
498            self.drop_datagram(DropReason::UnknownFlow);
499            return;
500        };
501
502        match flow.phase {
503            FlowPhase::Established => {
504                flow.on_client_datagram(now);
505                let backend = flow
506                    .backend_addr
507                    .expect("Established flow always has a backend address");
508                let mut out = payload.to_vec();
509                if flow.take_proxy_protocol() {
510                    prepend_dgram_header(&mut out, flow.client, backend);
511                }
512                let teardown = flow.teardown_reason();
513                self.outputs
514                    .push_back(Output::Metric(MetricEvent::DatagramIn(payload.len())));
515                self.outputs.push_back(Output::SendToBackend(Transmit {
516                    dst: backend,
517                    segment_size: None,
518                    payload: out,
519                }));
520                if let Some(reason) = teardown {
521                    self.close_flow(flow_id, reason);
522                } else {
523                    self.reschedule();
524                }
525            }
526            FlowPhase::Closing => {
527                // Racing datagram against a teardown already decided: drop it.
528                self.drop_datagram(DropReason::Shed);
529            }
530        }
531    }
532
533    fn on_backend_datagram(&mut self, flow_id: FlowId, payload: &[u8], now: Instant) {
534        if payload.len() > self.max_rx_datagram_size {
535            self.drop_datagram(DropReason::Truncated);
536            return;
537        }
538        let Some(flow) = self.flows.get_mut(flow_id) else {
539            self.drop_datagram(DropReason::UnknownFlow);
540            return;
541        };
542        if flow.phase != FlowPhase::Established {
543            self.drop_datagram(DropReason::UnknownFlow);
544            return;
545        }
546        flow.on_backend_datagram(now);
547        let client = flow.client;
548        let teardown = flow.teardown_reason();
549        self.outputs
550            .push_back(Output::Metric(MetricEvent::DatagramOut(payload.len())));
551        self.outputs.push_back(Output::SendToClient(Transmit {
552            dst: client,
553            segment_size: None,
554            payload: payload.to_vec(),
555        }));
556        if let Some(reason) = teardown {
557            self.close_flow(flow_id, reason);
558        } else {
559            self.reschedule();
560        }
561    }
562
563    fn on_config(&mut self, event: ConfigEvent, _now: Instant) {
564        match event {
565            ConfigEvent::SetCluster(cfg) => self.cluster = cfg,
566            ConfigEvent::SetMaxFlows(n) => {
567                self.max_flows = n;
568                // Track the high-water mark: a `SetMaxFlows` may shrink the live
569                // cap below `flows.len()` (documented), so the only durable cap
570                // bound is the largest cap ever in force.
571                #[cfg(debug_assertions)]
572                {
573                    self.max_flows_high_water = self.max_flows_high_water.max(n);
574                }
575            }
576            ConfigEvent::SetMaxRxDatagramSize(n) => self.max_rx_datagram_size = n,
577            ConfigEvent::Drain => self.draining = true,
578        }
579    }
580
581    // ---- timers ------------------------------------------------------------
582
583    /// Fire all flows whose idle deadline has elapsed at `now`. Deadline
584    /// revalidation: a flow is selected, and re-checked before its close, by
585    /// `idle_deadline <= now` alone. A flow that saw traffic since the expiry
586    /// was armed has a later `idle_deadline`, so the stale expiry leaves it
587    /// open; `UdpFlow::timer_gen` is not consulted.
588    ///
589    /// Called ONLY from a wheel expiry: the shell's single timer entry has just
590    /// been delivered and consumed. `now` is therefore the wheel's tick date,
591    /// not the deadline — `crate::timer` rounds a delay to the nearest tick, so
592    /// with a 100 ms tick the entry arrives up to 50 ms EARLY on the tick grid
593    /// (up to 99 ms off it; see `duration_to_tick`) and no flow need be due at
594    /// all. Either way the shell now holds nothing, so
595    /// `armed_deadline` is cleared on entry and `reschedule`
596    /// re-emits `ArmTimer` even when the minimum deadline has not moved.
597    /// Without that, an early expiry is a LOST WAKEUP: nothing is closed,
598    /// nothing is re-armed, and the flow is never reaped until some other
599    /// flow's deadline happens to change the minimum. This is the
600    /// consume-then-reschedule rule, and its cost is one extra wheel wakeup per
601    /// early-fired expiry.
602    pub fn handle_timeout(&mut self, now: Instant) {
603        // The wheel entry that brought us here is gone. Record that BEFORE any
604        // `reschedule` — including the ones `close_flow` runs below — so a
605        // recomputed deadline equal to the old one is still emitted as a fresh
606        // `ArmTimer` instead of being memoized away.
607        self.armed_deadline = None;
608
609        // Collect due flow ids first to avoid borrowing the slab while mutating.
610        let due: Vec<FlowId> = self
611            .flows
612            .iter()
613            .filter(|(_, flow)| flow.idle_deadline <= now)
614            .map(|(id, _)| id)
615            .collect();
616        for flow_id in due {
617            // Re-check under current state (a flow may have been closed by an
618            // earlier iteration's teardown, though here ids are disjoint).
619            if let Some(flow) = self.flows.get(flow_id)
620                && flow.idle_deadline <= now
621                && flow.phase != FlowPhase::Closing
622            {
623                self.close_flow(flow_id, CloseReason::Idle);
624            }
625        }
626        self.reschedule();
627
628        // Strict-advance guard: after firing every flow due at `now`, the next
629        // armed deadline (if any) MUST be strictly greater than `now`. A
630        // deadline `<= now` would make the shell immediately re-fire and spin —
631        // the canonical sans-io busy-loop bug. Deadline revalidation above plus
632        // consume-then-reschedule are what uphold it.
633        #[cfg(debug_assertions)]
634        if let Some(next) = self.armed_deadline {
635            debug_assert!(
636                next > now,
637                "UdpManager::poll_timeout must strictly advance past a firing: \
638                 armed {next:?} <= fired_at {now:?} (busy-loop)"
639            );
640        }
641
642        // Post: every flow due at `now` was reaped — no live flow may retain a
643        // deadline `<= now`. Pair with the strict-advance guard above: that one
644        // proves the next *armed* deadline advanced, this one proves no *flow*
645        // was left behind due (a leak the armed-deadline check alone misses,
646        // since min() over an empty set is None regardless of stragglers).
647        #[cfg(debug_assertions)]
648        for (id, flow) in self.flows.iter() {
649            debug_assert!(
650                flow.idle_deadline > now,
651                "FlowId {id} still due after handle_timeout: deadline {:?} <= now {now:?}",
652                flow.idle_deadline,
653            );
654        }
655        self.debug_assert_invariants();
656    }
657
658    /// The next manager-wide deadline, or `None` if no flow is armed. After a
659    /// [`Self::handle_timeout`] at deadline `d`, the value returned here is guaranteed
660    /// `> d` (or `None`) — the strict-advance invariant `handle_timeout` asserts
661    /// in debug builds, which is what stops the shell busy-looping.
662    pub fn poll_timeout(&self) -> Option<Instant> {
663        self.armed_deadline
664    }
665
666    /// Drain the next queued output, or `None` when the queue is empty.
667    pub fn poll_output(&mut self) -> Option<Output> {
668        self.outputs.pop_front()
669    }
670
671    // ---- internals ---------------------------------------------------------
672
673    /// Recompute the earliest flow deadline and emit `ArmTimer` only when it
674    /// differs from `armed_deadline`, so the shell re-arms its wheel exactly
675    /// once per real change.
676    ///
677    /// The memoization is against what the SHELL currently holds, not against
678    /// the previous minimum: [`handle_timeout`](Self::handle_timeout) clears
679    /// `armed_deadline` on entry because the expiry consumed that wheel entry,
680    /// so an unchanged minimum is re-emitted there rather than swallowed. A
681    /// memoization against the minimum alone would be a lost wakeup on every
682    /// expiry the wheel delivered early.
683    fn reschedule(&mut self) {
684        let next = self
685            .flows
686            .iter()
687            .filter(|(_, f)| f.phase != FlowPhase::Closing)
688            .map(|(_, f)| f.idle_deadline)
689            .min();
690        if next != self.armed_deadline {
691            self.armed_deadline = next;
692            if let Some(deadline) = next {
693                self.outputs.push_back(Output::ArmTimer(deadline));
694            }
695        }
696    }
697
698    /// Tear down a flow: remove it from the table + slab, emit `CloseFlow` and
699    /// the eviction metric. Idempotent — a missing flow is a no-op (never an
700    /// underflow). Re-arms the manager timer.
701    fn close_flow(&mut self, flow_id: FlowId, _reason: CloseReason) {
702        let Some(flow) = self.flows.get_mut(flow_id) else {
703            return;
704        };
705        if flow.phase == FlowPhase::Closing {
706            return;
707        }
708        flow.set_phase(FlowPhase::Closing);
709        let key = FlowKey::from_src(flow.client);
710        // Remove the table entry only if it still points at this flow; a
711        // recreated flow under the same key must not be unmapped.
712        if self.table.get(&key) == Some(&flow_id) {
713            self.table.remove(&key);
714        }
715        let ip = flow.client.ip();
716        let subnet = subnet_of(&flow.config, ip);
717        let cluster = flow.config.cluster.clone();
718        source_decr(&mut self.flows_per_ip, &cluster, &ip);
719        source_decr(&mut self.flows_per_subnet, &cluster, &subnet);
720        self.flows.remove(flow_id);
721        self.outputs
722            .push_back(Output::Metric(MetricEvent::FlowEvicted));
723        self.outputs.push_back(Output::CloseFlow(flow_id));
724        self.reschedule();
725
726        // Post: the slot is gone from the slab AND no table key still maps to it
727        // (a stale table key would dangle — caught by check_invariants (1), but
728        // assert it here too so the local failure points at close_flow). Pair:
729        // positive = removed from slab; negative = no key still references it.
730        #[cfg(debug_assertions)]
731        {
732            debug_assert!(
733                !self.flows.contains(flow_id),
734                "close_flow left FlowId {flow_id} in the slab"
735            );
736            debug_assert!(
737                self.table.values().all(|&id| id != flow_id),
738                "close_flow left a table entry mapping to the removed FlowId {flow_id}"
739            );
740        }
741    }
742
743    /// Whether a new flow from `ip` would exceed the active cluster's
744    /// per-source limits: `max_flows_per_ip` on the address itself, or
745    /// `max_flows_per_subnet` on its masked subnet. `0` disables a limit, and
746    /// the subnet key is only derived when that limit is set.
747    fn source_at_limit(&self, ip: IpAddr) -> bool {
748        let cfg = &self.cluster;
749        if cfg.max_flows_per_ip > 0
750            && source_count(&self.flows_per_ip, &cfg.cluster, &ip) as u64 >= cfg.max_flows_per_ip
751        {
752            return true;
753        }
754        cfg.max_flows_per_subnet > 0
755            && source_count(&self.flows_per_subnet, &cfg.cluster, &subnet_of(cfg, ip)) as u64
756                >= cfg.max_flows_per_subnet
757    }
758
759    /// Emit a drop with its by-reason metric. Allocates nothing per the
760    /// "silence is a virtue" posture.
761    fn drop_datagram(&mut self, reason: DropReason) {
762        // "Silence is a virtue": a drop allocates no flow and frees none — the
763        // slab is untouched across the reject. Snapshot + pair-assert so a future
764        // edit that accidentally mutates the slab on a drop path is caught loudly.
765        // Not debug-gated: the `debug_assert_eq!` reads it in release too (only
766        // execution is gated); dead there, dropped by the optimizer.
767        let flows_before_drop = self.flows.len();
768        self.outputs
769            .push_back(Output::Metric(MetricEvent::DatagramDropped(reason)));
770        self.outputs.push_back(Output::Drop(reason));
771        debug_assert_eq!(
772            self.flows.len(),
773            flows_before_drop,
774            "a drop must allocate nothing and free nothing (flows.len() unchanged)"
775        );
776    }
777
778    /// Affinity hash for a flow: `hash(seed, affinity_key)`. The shell feeds
779    /// this into HRW (`max hash`) / Maglev (`key % M`) / RR (ignores it).
780    fn affinity_hash(&self, flow: &UdpFlow) -> u64 {
781        let mut hasher = std::collections::hash_map::DefaultHasher::new();
782        self.hash_seed.hash(&mut hasher);
783        if flow.config.affinity_with_port {
784            flow.client.hash(&mut hasher);
785        } else {
786            flow.client.ip().hash(&mut hasher);
787        }
788        hasher.finish()
789    }
790
791    /// TigerStyle invariant sweep (TigerBeetle / FoundationDB style). A single
792    /// full check of every structural invariant the manager must preserve,
793    /// asserted at the END of every public mutating method via
794    /// [`debug_assert_invariants`](Self::debug_assert_invariants). Compiled out
795    /// entirely in release (`#[cfg(debug_assertions)]`); on in every test / e2e /
796    /// fuzz / dev build. It must NEVER change runtime behavior — it only reads.
797    ///
798    /// The right place for the sweep is a *post-condition*: these public methods
799    /// run to completion under the caller's lock, so the table/slab/timer state
800    /// is fully reconciled by the time the method returns.
801    #[cfg(debug_assertions)]
802    fn check_invariants(&self) {
803        use std::collections::HashSet;
804
805        // (3) flow_count() is exactly the slab population.
806        debug_assert_eq!(
807            self.flow_count(),
808            self.flows.len(),
809            "flow_count() must equal flows.len()"
810        );
811
812        // (1) Table -> slab consistency + (2) table injectivity: every FlowId in
813        // the table points at a live slab slot, and no two keys share a FlowId.
814        let mut seen_ids: HashSet<FlowId> = HashSet::with_capacity(self.table.len());
815        for (key, &id) in self.table.iter() {
816            debug_assert!(
817                self.flows.contains(id),
818                "table key {key:?} maps to FlowId {id} absent from the slab (dangling key)"
819            );
820            debug_assert!(
821                seen_ids.insert(id),
822                "table injectivity violated: FlowId {id} is the target of two distinct FlowKeys"
823            );
824        }
825        // Pair (negative space): a live flow that is reachable from the table is
826        // mapped exactly once — there is never a live flow with two table keys.
827        debug_assert!(
828            seen_ids.len() <= self.flows.len(),
829            "more distinct table targets than live flows"
830        );
831
832        // Per-flow invariants over the slab.
833        let mut min_live_deadline: Option<Instant> = None;
834        let mut live_count = 0usize;
835        for (id, flow) in self.flows.iter() {
836            // (4) No slab flow is Closing. close_flow sets Closing then removes
837            // the slot in the same call, so a Closing flow must never persist.
838            // Pair: positive space — every live flow is Awaiting or Established.
839            debug_assert_ne!(
840                flow.phase,
841                FlowPhase::Closing,
842                "FlowId {id} persists in the slab while Closing (close_flow must remove it)"
843            );
844            debug_assert_eq!(
845                flow.phase,
846                FlowPhase::Established,
847                "FlowId {id} has an unexpected live phase"
848            );
849
850            // (5) A live flow always carries the backend it was admitted on.
851            // The negative half of this pair — a phase that carried no address
852            // — went with `AwaitingBackend`: selection now happens before the
853            // flow exists, so there is no such state to assert about.
854            debug_assert!(
855                flow.backend_addr.is_some(),
856                "live FlowId {id} has no backend address"
857            );
858
859            // (7) Counters within caps OR a teardown is due. A flow whose cap is
860            // exhausted must report a teardown reason; an exhausted flow that
861            // reports None would be an immortal flow (the cap silently lost).
862            if flow.requests_exhausted() || flow.responses_exhausted() {
863                debug_assert!(
864                    flow.teardown_reason().is_some(),
865                    "FlowId {id} exhausted a cap (req {}/{}, resp {}/{}) but reports no teardown",
866                    flow.requests_seen,
867                    flow.config.requests,
868                    flow.responses_seen,
869                    flow.config.responses,
870                );
871            } else {
872                // Pair (negative): a flow within both caps must NOT report a
873                // cap-driven teardown (idle teardown is handled by the timer,
874                // not teardown_reason()).
875                debug_assert!(
876                    flow.teardown_reason().is_none(),
877                    "FlowId {id} reports a teardown while within both caps"
878                );
879            }
880
881            live_count += 1;
882            min_live_deadline = Some(match min_live_deadline {
883                Some(d) => d.min(flow.idle_deadline),
884                None => flow.idle_deadline,
885            });
886        }
887
888        // (8) Per-source counters are exact: recounting the slab gives both
889        // maps back, key for key, with no zero entry left behind.
890        let mut per_ip = SourceCounts::new();
891        let mut per_subnet = SourceCounts::new();
892        for (_, flow) in self.flows.iter() {
893            let ip = flow.client.ip();
894            source_incr(&mut per_ip, &flow.config.cluster, ip);
895            source_incr(
896                &mut per_subnet,
897                &flow.config.cluster,
898                subnet_of(&flow.config, ip),
899            );
900        }
901        debug_assert_eq!(
902            self.flows_per_ip, per_ip,
903            "per-IP flow counts must match the live flows"
904        );
905        debug_assert_eq!(
906            self.flows_per_subnet, per_subnet,
907            "per-subnet flow counts must match the live flows"
908        );
909
910        // The high-water cap bounds the live population at all times (the live
911        // cap itself can be shrunk below flows.len() by SetMaxFlows, so we assert
912        // against the largest cap ever in force, not max_flows).
913        debug_assert!(
914            self.flows.len() <= self.max_flows_high_water,
915            "live flows {} exceed the high-water cap {}",
916            self.flows.len(),
917            self.max_flows_high_water,
918        );
919
920        // (6) Timer coherence: armed_deadline is Some IFF at least one (non-
921        // Closing) live flow exists, and when Some it equals the minimum
922        // idle_deadline over live flows. Closing flows are never in the slab
923        // (invariant 4), so "live" == "slab" here.
924        debug_assert_eq!(
925            self.armed_deadline.is_some(),
926            live_count > 0,
927            "timer coherence: armed_deadline.is_some() ({}) must match having live flows ({})",
928            self.armed_deadline.is_some(),
929            live_count > 0,
930        );
931        if let Some(armed) = self.armed_deadline {
932            debug_assert_eq!(
933                Some(armed),
934                min_live_deadline,
935                "timer coherence: armed deadline {armed:?} must equal the minimum live idle deadline {min_live_deadline:?}"
936            );
937        }
938    }
939
940    /// Run the full [`check_invariants`](Self::check_invariants) sweep, but only
941    /// in debug builds. A thin wrapper so call sites read as one line and so the
942    /// whole sweep — and any read it performs — is dead-stripped in release.
943    #[inline]
944    fn debug_assert_invariants(&self) {
945        #[cfg(debug_assertions)]
946        self.check_invariants();
947    }
948}
949
950#[cfg(test)]
951mod tests {
952    use std::{
953        net::{IpAddr, Ipv4Addr},
954        time::Duration,
955    };
956
957    use super::*;
958
959    /// The backend set these tests select from.
960    ///
961    /// Selection moved into the core (#1340, Question 6), so a test that
962    /// admits a flow has to supply one. Round-robin over whatever it holds;
963    /// `none()` is the cluster-has-nothing case that makes admission fail,
964    /// which is a live path because
965    /// `BackendMap::backend_from_cluster_id_with_key` has three
966    /// `NoBackendForCluster` returns.
967    struct TestBackends {
968        backends: Vec<(String, SocketAddr)>,
969        next: usize,
970    }
971
972    impl TestBackends {
973        /// One backend, which is what every test that is not about selection
974        /// wants.
975        fn one() -> Self {
976            Self {
977                backends: vec![("b1".to_owned(), backend_addr(1))],
978                next: 0,
979            }
980        }
981
982        /// A cluster with nothing that can serve.
983        fn none() -> Self {
984            Self {
985                backends: Vec::new(),
986                next: 0,
987            }
988        }
989    }
990
991    impl BackendSource for TestBackends {
992        fn select(
993            &mut self,
994            _cluster: &str,
995            _key: Option<u64>,
996            _now: Instant,
997        ) -> Option<(String, SocketAddr)> {
998            if self.backends.is_empty() {
999                return None;
1000            }
1001            let picked = self.backends[self.next % self.backends.len()].clone();
1002            self.next += 1;
1003            Some(picked)
1004        }
1005    }
1006
1007    fn backend_addr(n: u8) -> SocketAddr {
1008        SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 0, 0, n)), 5353)
1009    }
1010
1011    fn cluster(name: &str) -> ClusterConfig {
1012        ClusterConfig {
1013            cluster: name.into(),
1014            front_timeout: Duration::from_secs(30),
1015            back_timeout: Duration::from_secs(30),
1016            ..Default::default()
1017        }
1018    }
1019
1020    fn client(n: u8, port: u16) -> SocketAddr {
1021        SocketAddr::new(IpAddr::V4(Ipv4Addr::new(10, 0, 0, n)), port)
1022    }
1023
1024    /// Drain all outputs into a Vec for assertions.
1025    fn drain(mgr: &mut UdpManager) -> Vec<Output> {
1026        let mut out = Vec::new();
1027        while let Some(o) = mgr.poll_output() {
1028            out.push(o);
1029        }
1030        out
1031    }
1032
1033    #[test]
1034    fn unknown_datagram_allocates_nothing_when_no_cluster() {
1035        let mut mgr = UdpManager::new(ClusterConfig::default(), 16, 65535, 0xABCD);
1036        let now = Instant::now();
1037        mgr.handle_input(
1038            ManagerInput::ClientDatagram {
1039                src: client(1, 1000),
1040                payload: b"hi",
1041                backends: &mut TestBackends::one(),
1042            },
1043            now,
1044        );
1045        assert_eq!(mgr.flow_count(), 0);
1046        let outs = drain(&mut mgr);
1047        assert!(
1048            outs.iter()
1049                .any(|o| matches!(o, Output::Drop(DropReason::NoBackend)))
1050        );
1051        assert!(
1052            !outs
1053                .iter()
1054                .any(|o| matches!(o, Output::OpenUpstream { .. }))
1055        );
1056    }
1057
1058    #[test]
1059    fn empty_datagram_is_invalid() {
1060        let mut mgr = UdpManager::new(cluster("dns"), 16, 65535, 1);
1061        mgr.handle_input(
1062            ManagerInput::ClientDatagram {
1063                src: client(1, 1000),
1064                payload: b"",
1065                backends: &mut TestBackends::one(),
1066            },
1067            Instant::now(),
1068        );
1069        assert_eq!(mgr.flow_count(), 0);
1070        let outs = drain(&mut mgr);
1071        assert!(
1072            outs.iter()
1073                .any(|o| matches!(o, Output::Drop(DropReason::Invalid)))
1074        );
1075    }
1076
1077    #[test]
1078    fn truncated_datagram_dropped_before_admission() {
1079        let mut mgr = UdpManager::new(cluster("dns"), 16, 4, 1);
1080        mgr.handle_input(
1081            ManagerInput::ClientDatagram {
1082                src: client(1, 1000),
1083                payload: b"toolong",
1084                backends: &mut TestBackends::one(),
1085            },
1086            Instant::now(),
1087        );
1088        assert_eq!(mgr.flow_count(), 0);
1089        let outs = drain(&mut mgr);
1090        assert!(
1091            outs.iter()
1092                .any(|o| matches!(o, Output::Drop(DropReason::Truncated)))
1093        );
1094    }
1095
1096    /// A new flow opens an upstream and forwards its admission datagram, in
1097    /// one call.
1098    ///
1099    /// Retargeted from `new_flow_requests_backend_then_forwards_buffered_datagram`.
1100    /// That name described a two-step: the core asked the shell to choose a
1101    /// backend, buffered the datagram, and flushed it when the shell replied.
1102    /// Selection is the core's now, so there is no window and nothing to
1103    /// buffer — the datagram is forwarded on the way in.
1104    ///
1105    /// Its "no SendToBackend yet, backend not resolved" assertion is the one
1106    /// casualty: there is no *yet*. Its partner — that the datagram does reach
1107    /// the backend, intact and at the right address — survives and is below.
1108    #[test]
1109    fn a_new_flow_opens_an_upstream_and_forwards_its_admission_datagram() {
1110        let mut mgr = UdpManager::new(cluster("dns"), 16, 65535, 7);
1111        let now = Instant::now();
1112        let src = client(1, 1000);
1113        mgr.handle_input(
1114            ManagerInput::ClientDatagram {
1115                src,
1116                payload: b"query",
1117                backends: &mut TestBackends::one(),
1118            },
1119            now,
1120        );
1121        assert_eq!(mgr.flow_count(), 1);
1122        let outs = drain(&mut mgr);
1123        let opened = outs
1124            .iter()
1125            .find_map(|o| match o {
1126                Output::OpenUpstream { flow, backend } => Some((*flow, *backend)),
1127                _ => None,
1128            })
1129            .expect("admission must ask the shell to open an upstream");
1130        assert_eq!(
1131            opened.1,
1132            backend_addr(1),
1133            "the upstream must be opened to the backend the core selected"
1134        );
1135        let sent = outs
1136            .iter()
1137            .find_map(|o| match o {
1138                Output::SendToBackend(t) => Some(t.clone()),
1139                _ => None,
1140            })
1141            .expect("the admission datagram must be forwarded in the same call");
1142        assert_eq!(sent.dst, backend_addr(1));
1143        assert_eq!(sent.payload, b"query");
1144    }
1145
1146    /// A backend source that records the affinity key of every selection, so a
1147    /// test can check which client key reached the load balancer.
1148    struct KeyRecordingBackends {
1149        keys: Vec<Option<u64>>,
1150    }
1151
1152    impl BackendSource for KeyRecordingBackends {
1153        fn select(
1154            &mut self,
1155            _cluster: &str,
1156            key: Option<u64>,
1157            _now: Instant,
1158        ) -> Option<(String, SocketAddr)> {
1159            self.keys.push(key);
1160            Some(("b1".to_owned(), backend_addr(1)))
1161        }
1162    }
1163
1164    /// Two client sockets on one source IP are two flows even under the
1165    /// default `SOURCE_IP` affinity (#1732). The affinity key only feeds
1166    /// backend selection: both flows present the same key to the load
1167    /// balancer, so a hash LB pins them to one backend, but each flow keeps its
1168    /// own upstream socket and each backend reply goes back to the client that
1169    /// flow belongs to, not to whichever socket of that IP spoke first.
1170    #[test]
1171    fn same_ip_clients_get_distinct_flows_and_their_own_replies() {
1172        let cfg = cluster("udp");
1173        assert!(!cfg.affinity_with_port, "the default affinity is SOURCE_IP");
1174        let mut mgr = UdpManager::new(cfg, 16, 65535, 7);
1175        let now = Instant::now();
1176        let mut backends = KeyRecordingBackends { keys: Vec::new() };
1177        let a = client(1, 50376);
1178        let b = client(1, 50644);
1179
1180        let mut flows = Vec::new();
1181        for src in [a, b] {
1182            mgr.handle_input(
1183                ManagerInput::ClientDatagram {
1184                    src,
1185                    payload: b"hi",
1186                    backends: &mut backends,
1187                },
1188                now,
1189            );
1190            let opened = drain(&mut mgr)
1191                .into_iter()
1192                .find_map(|o| match o {
1193                    Output::OpenUpstream { flow, .. } => Some(flow),
1194                    _ => None,
1195                })
1196                .expect("each client socket must open its own upstream");
1197            flows.push(opened);
1198        }
1199        assert_eq!(mgr.flow_count(), 2, "one flow per client source address");
1200        assert_ne!(flows[0], flows[1]);
1201        assert_eq!(backends.keys.len(), 2);
1202        assert_eq!(
1203            backends.keys[0], backends.keys[1],
1204            "SOURCE_IP affinity must present one key for every port of an IP"
1205        );
1206
1207        for (flow, src) in [(flows[1], b), (flows[0], a)] {
1208            mgr.handle_input(
1209                ManagerInput::BackendDatagram {
1210                    flow,
1211                    payload: b"reply",
1212                },
1213                now,
1214            );
1215            let dst = drain(&mut mgr)
1216                .into_iter()
1217                .find_map(|o| match o {
1218                    Output::SendToClient(t) => Some(t.dst),
1219                    _ => None,
1220                })
1221                .expect("a backend reply must be returned");
1222            assert_eq!(dst, src, "the reply must go to the flow's own client");
1223        }
1224    }
1225
1226    #[test]
1227    fn tracked_flow_reuses_backend() {
1228        let mut mgr = UdpManager::new(cluster("dns"), 16, 65535, 7);
1229        let now = Instant::now();
1230        let src = client(1, 1000);
1231        mgr.handle_input(
1232            ManagerInput::ClientDatagram {
1233                src,
1234                payload: b"q1",
1235                backends: &mut TestBackends::one(),
1236            },
1237            now,
1238        );
1239        let _select_flow = drain(&mut mgr)
1240            .into_iter()
1241            .find_map(|o| match o {
1242                Output::OpenUpstream { flow, .. } => Some(flow),
1243                _ => None,
1244            })
1245            .unwrap();
1246        drain(&mut mgr);
1247
1248        // Second datagram from same source: no new selection, direct send.
1249        mgr.handle_input(
1250            ManagerInput::ClientDatagram {
1251                src,
1252                payload: b"q2",
1253                backends: &mut TestBackends::one(),
1254            },
1255            now,
1256        );
1257        assert_eq!(mgr.flow_count(), 1);
1258        let outs = drain(&mut mgr);
1259        assert!(
1260            !outs
1261                .iter()
1262                .any(|o| matches!(o, Output::OpenUpstream { .. }))
1263        );
1264        assert!(
1265            outs.iter()
1266                .any(|o| matches!(o, Output::SendToBackend(t) if t.payload == b"q2"))
1267        );
1268    }
1269
1270    #[test]
1271    fn responses_one_closes_flow_after_single_reply() {
1272        let mut cfg = cluster("dns");
1273        cfg.responses = 1;
1274        let mut mgr = UdpManager::new(cfg, 16, 65535, 7);
1275        let now = Instant::now();
1276        let src = client(1, 1000);
1277        mgr.handle_input(
1278            ManagerInput::ClientDatagram {
1279                src,
1280                payload: b"q",
1281                backends: &mut TestBackends::one(),
1282            },
1283            now,
1284        );
1285        let flow = drain(&mut mgr)
1286            .into_iter()
1287            .find_map(|o| match o {
1288                Output::OpenUpstream { flow, .. } => Some(flow),
1289                _ => None,
1290            })
1291            .unwrap();
1292        drain(&mut mgr);
1293
1294        // Single backend reply closes the flow.
1295        mgr.handle_input(
1296            ManagerInput::BackendDatagram {
1297                flow,
1298                payload: b"answer",
1299            },
1300            now,
1301        );
1302        let outs = drain(&mut mgr);
1303        assert!(
1304            outs.iter()
1305                .any(|o| matches!(o, Output::SendToClient(t) if t.payload == b"answer"))
1306        );
1307        assert!(
1308            outs.iter()
1309                .any(|o| matches!(o, Output::CloseFlow(f) if *f == flow))
1310        );
1311        assert_eq!(mgr.flow_count(), 0);
1312    }
1313
1314    #[test]
1315    fn requests_cap_closes_flow() {
1316        let mut cfg = cluster("syslog");
1317        cfg.requests = 2;
1318        let mut mgr = UdpManager::new(cfg, 16, 65535, 7);
1319        let now = Instant::now();
1320        let src = client(1, 1000);
1321        mgr.handle_input(
1322            ManagerInput::ClientDatagram {
1323                src,
1324                payload: b"1",
1325                backends: &mut TestBackends::one(),
1326            },
1327            now,
1328        );
1329        let _flow = drain(&mut mgr)
1330            .into_iter()
1331            .find_map(|o| match o {
1332                Output::OpenUpstream { flow, .. } => Some(flow),
1333                _ => None,
1334            })
1335            .unwrap();
1336        drain(&mut mgr);
1337        // requests_seen is now 1 (the admission datagram). Second hits the cap.
1338        mgr.handle_input(
1339            ManagerInput::ClientDatagram {
1340                src,
1341                payload: b"2",
1342                backends: &mut TestBackends::one(),
1343            },
1344            now,
1345        );
1346        let outs = drain(&mut mgr);
1347        assert!(outs.iter().any(|o| matches!(o, Output::CloseFlow(_))));
1348        assert_eq!(mgr.flow_count(), 0);
1349    }
1350
1351    #[test]
1352    fn flow_table_full_sheds_new_flow() {
1353        let mut mgr = UdpManager::new(cluster("dns"), 1, 65535, 7);
1354        let now = Instant::now();
1355        mgr.handle_input(
1356            ManagerInput::ClientDatagram {
1357                src: client(1, 1000),
1358                payload: b"a",
1359                backends: &mut TestBackends::one(),
1360            },
1361            now,
1362        );
1363        assert_eq!(mgr.flow_count(), 1);
1364        drain(&mut mgr);
1365        // Second distinct source over cap → shed.
1366        mgr.handle_input(
1367            ManagerInput::ClientDatagram {
1368                src: client(2, 1000),
1369                payload: b"b",
1370                backends: &mut TestBackends::one(),
1371            },
1372            now,
1373        );
1374        assert_eq!(mgr.flow_count(), 1);
1375        let outs = drain(&mut mgr);
1376        assert!(
1377            outs.iter()
1378                .any(|o| matches!(o, Output::Metric(MetricEvent::FlowShed)))
1379        );
1380        assert!(
1381            outs.iter()
1382                .any(|o| matches!(o, Output::Drop(DropReason::Shed)))
1383        );
1384    }
1385
1386    /// Feed one client datagram and return the outputs it produced.
1387    fn send_from(mgr: &mut UdpManager, src: SocketAddr, now: Instant) -> Vec<Output> {
1388        mgr.handle_input(
1389            ManagerInput::ClientDatagram {
1390                src,
1391                payload: b"q",
1392                backends: &mut TestBackends::one(),
1393            },
1394            now,
1395        );
1396        drain(mgr)
1397    }
1398
1399    /// Whether `outs` is a per-source shed: the source-limit metric and a
1400    /// `Shed` drop, and no upstream opened.
1401    fn source_shed(outs: &[Output]) -> bool {
1402        outs.iter()
1403            .any(|o| matches!(o, Output::Metric(MetricEvent::FlowShedSourceLimit)))
1404            && outs
1405                .iter()
1406                .any(|o| matches!(o, Output::Drop(DropReason::Shed)))
1407            && !outs
1408                .iter()
1409                .any(|o| matches!(o, Output::OpenUpstream { .. }))
1410    }
1411
1412    fn opened(outs: &[Output]) -> bool {
1413        outs.iter()
1414            .any(|o| matches!(o, Output::OpenUpstream { .. }))
1415    }
1416
1417    #[test]
1418    fn per_ip_limit_sheds_new_ports_of_one_address_only() {
1419        let cfg = ClusterConfig {
1420            max_flows_per_ip: 2,
1421            ..cluster("dns")
1422        };
1423        let mut mgr = UdpManager::new(cfg, 16, 65535, 7);
1424        let now = Instant::now();
1425
1426        assert!(opened(&send_from(&mut mgr, client(1, 1000), now)));
1427        assert!(opened(&send_from(&mut mgr, client(1, 1001), now)));
1428        let outs = send_from(&mut mgr, client(1, 1002), now);
1429        assert!(
1430            source_shed(&outs),
1431            "a third port of one address must be shed: {outs:?}"
1432        );
1433        assert!(
1434            !outs
1435                .iter()
1436                .any(|o| matches!(o, Output::Metric(MetricEvent::FlowShed))),
1437            "a per-source shed is not a table-full shed"
1438        );
1439        assert_eq!(mgr.flow_count(), 2);
1440
1441        assert!(
1442            opened(&send_from(&mut mgr, client(2, 1000), now)),
1443            "another address must still be admitted"
1444        );
1445        let outs = send_from(&mut mgr, client(1, 1000), now);
1446        assert!(
1447            outs.iter().any(|o| matches!(o, Output::SendToBackend(_))),
1448            "an existing flow of the capped address keeps forwarding"
1449        );
1450        assert_eq!(mgr.flow_count(), 3);
1451    }
1452
1453    #[test]
1454    fn closing_a_flow_frees_its_source_slot() {
1455        let cfg = ClusterConfig {
1456            max_flows_per_ip: 1,
1457            front_timeout: Duration::from_secs(1),
1458            back_timeout: Duration::from_secs(1),
1459            ..cluster("dns")
1460        };
1461        let mut mgr = UdpManager::new(cfg, 16, 65535, 7);
1462        let now = Instant::now();
1463        assert!(opened(&send_from(&mut mgr, client(1, 1000), now)));
1464        assert!(source_shed(&send_from(&mut mgr, client(1, 1001), now)));
1465
1466        let later = now + Duration::from_secs(2);
1467        mgr.handle_timeout(later);
1468        assert_eq!(mgr.flow_count(), 0);
1469        drain(&mut mgr);
1470        assert!(
1471            opened(&send_from(&mut mgr, client(1, 1001), later)),
1472            "the closed flow's slot must be free again"
1473        );
1474    }
1475
1476    #[test]
1477    fn per_subnet_limit_groups_addresses_by_prefix() {
1478        let cfg = ClusterConfig {
1479            max_flows_per_subnet: 2,
1480            subnet_ipv4_prefix: 24,
1481            ..cluster("dns")
1482        };
1483        let mut mgr = UdpManager::new(cfg, 16, 65535, 7);
1484        let now = Instant::now();
1485        let at = |a, b, c, d| SocketAddr::new(IpAddr::V4(Ipv4Addr::new(a, b, c, d)), 1000);
1486
1487        assert!(opened(&send_from(&mut mgr, at(192, 0, 2, 1), now)));
1488        assert!(opened(&send_from(&mut mgr, at(192, 0, 2, 2), now)));
1489        assert!(
1490            source_shed(&send_from(&mut mgr, at(192, 0, 2, 3), now)),
1491            "a third address of one /24 must be shed"
1492        );
1493        assert!(
1494            opened(&send_from(&mut mgr, at(192, 0, 3, 1), now)),
1495            "an address of another /24 must still be admitted"
1496        );
1497    }
1498
1499    #[test]
1500    fn a_limit_set_at_runtime_counts_flows_already_open() {
1501        let mut mgr = UdpManager::new(cluster("dns"), 16, 65535, 7);
1502        let now = Instant::now();
1503        assert!(opened(&send_from(&mut mgr, client(1, 1000), now)));
1504        assert!(opened(&send_from(&mut mgr, client(1, 1001), now)));
1505        mgr.handle_input(
1506            ManagerInput::Config(ConfigEvent::SetCluster(ClusterConfig {
1507                max_flows_per_ip: 2,
1508                ..cluster("dns")
1509            })),
1510            now,
1511        );
1512        assert!(source_shed(&send_from(&mut mgr, client(1, 1002), now)));
1513    }
1514
1515    #[test]
1516    fn idle_timeout_closes_flow() {
1517        let mut cfg = cluster("dns");
1518        cfg.front_timeout = Duration::from_secs(10);
1519        let mut mgr = UdpManager::new(cfg, 16, 65535, 7);
1520        let now = Instant::now();
1521        mgr.handle_input(
1522            ManagerInput::ClientDatagram {
1523                src: client(1, 1000),
1524                payload: b"q",
1525                backends: &mut TestBackends::one(),
1526            },
1527            now,
1528        );
1529        drain(&mut mgr);
1530        let deadline = mgr.poll_timeout().expect("timer armed");
1531        assert!(deadline >= now + Duration::from_secs(10));
1532        // Fire after the deadline.
1533        mgr.handle_timeout(now + Duration::from_secs(11));
1534        let outs = drain(&mut mgr);
1535        assert!(outs.iter().any(|o| matches!(o, Output::CloseFlow(_))));
1536        assert_eq!(mgr.flow_count(), 0);
1537        assert!(mgr.poll_timeout().is_none());
1538    }
1539
1540    /// An expiry that finds NOTHING due must still re-arm.
1541    ///
1542    /// The shell's wheel (`crate::timer`) rounds a delay to the nearest tick, so
1543    /// with a 100 ms tick it delivers an entry up to 50 ms EARLY on the grid
1544    /// (`test_timeout_fires_up_to_half_a_tick_early`, `lib/src/timer.rs`) and
1545    /// up to 99 ms off it (see `duration_to_tick`). The
1546    /// shell then calls `handle_timeout` at a `now` that has not reached any
1547    /// flow's deadline: no flow is due, nothing closes, the minimum deadline is
1548    /// unchanged — yet the wheel entry has been CONSUMED. If `reschedule` keeps
1549    /// memoizing on "the deadline changed" it emits no `ArmTimer`, the shell
1550    /// never re-arms, and the flow is never reaped. `handle_timeout` is by
1551    /// contract only ever called from a wheel expiry, so on entry the shell has
1552    /// no armed timer and `armed_deadline` must say so.
1553    ///
1554    /// To SEE THIS RED: remove `self.armed_deadline = None;` from the top of
1555    /// [`UdpManager::handle_timeout`]. `reschedule` then finds `next ==
1556    /// self.armed_deadline`, emits nothing, and the ArmTimer assertion fails
1557    /// with `left: []  right: [Instant { .. }]`.
1558    #[test]
1559    fn early_expiry_that_finds_nothing_due_still_rearms() {
1560        let mut cfg = cluster("dns");
1561        cfg.front_timeout = Duration::from_secs(10);
1562        let mut mgr = UdpManager::new(cfg, 16, 65535, 7);
1563        let now = Instant::now();
1564        mgr.handle_input(
1565            ManagerInput::ClientDatagram {
1566                src: client(1, 1000),
1567                payload: b"q",
1568                backends: &mut TestBackends::one(),
1569            },
1570            now,
1571        );
1572        drain(&mut mgr);
1573        let deadline = mgr.poll_timeout().expect("timer armed on admission");
1574
1575        // The wheel fires before the deadline. One second early here rather than
1576        // the wheel's real earliness — up to 50 ms on the tick grid, up to 99 ms
1577        // off it (see `duration_to_tick`) — so the test does not encode the
1578        // tick size.
1579        let early = deadline - Duration::from_secs(1);
1580        mgr.handle_timeout(early);
1581        let outs = drain(&mut mgr);
1582
1583        assert_eq!(
1584            mgr.flow_count(),
1585            1,
1586            "nothing was due at the early expiry: the flow must survive"
1587        );
1588        assert!(
1589            !outs.iter().any(|o| matches!(
1590                o,
1591                Output::CloseFlow(_) | Output::Metric(MetricEvent::FlowEvicted)
1592            )),
1593            "nothing was due: no flow may be closed, got {outs:?}"
1594        );
1595
1596        let armed: Vec<Instant> = outs
1597            .iter()
1598            .filter_map(|o| match o {
1599                Output::ArmTimer(d) => Some(*d),
1600                _ => None,
1601            })
1602            .collect();
1603        assert_eq!(
1604            armed,
1605            vec![deadline],
1606            "an expiry that found nothing due must re-arm the consumed wheel \
1607             entry with the (unchanged) deadline"
1608        );
1609        assert_eq!(
1610            mgr.poll_timeout(),
1611            Some(deadline),
1612            "the flow's deadline itself must not move"
1613        );
1614    }
1615
1616    /// The re-arm is not a one-shot: a second early expiry must re-arm again.
1617    /// The shell holds exactly one wheel entry for the whole listener, so every
1618    /// consumed entry owes exactly one `ArmTimer` for as long as a flow is live.
1619    ///
1620    /// To SEE THIS RED: the same mutation as
1621    /// `early_expiry_that_finds_nothing_due_still_rearms` — remove
1622    /// `self.armed_deadline = None;` from the top of
1623    /// [`UdpManager::handle_timeout`]. Both expiries then emit nothing and this
1624    /// fails with `left: 0  right: 1` on the first `assert_eq!`.
1625    #[test]
1626    fn repeated_early_expiries_each_rearm() {
1627        let mut cfg = cluster("dns");
1628        cfg.front_timeout = Duration::from_secs(10);
1629        let mut mgr = UdpManager::new(cfg, 16, 65535, 7);
1630        let now = Instant::now();
1631        mgr.handle_input(
1632            ManagerInput::ClientDatagram {
1633                src: client(1, 1000),
1634                payload: b"q",
1635                backends: &mut TestBackends::one(),
1636            },
1637            now,
1638        );
1639        drain(&mut mgr);
1640        let deadline = mgr.poll_timeout().expect("timer armed on admission");
1641
1642        let count_arms = |outs: &[Output]| {
1643            outs.iter()
1644                .filter(|o| matches!(o, Output::ArmTimer(_)))
1645                .count()
1646        };
1647
1648        mgr.handle_timeout(deadline - Duration::from_secs(2));
1649        let first = drain(&mut mgr);
1650        assert_eq!(count_arms(&first), 1, "first early expiry must re-arm");
1651
1652        mgr.handle_timeout(deadline - Duration::from_secs(1));
1653        let second = drain(&mut mgr);
1654        assert_eq!(
1655            count_arms(&second),
1656            1,
1657            "second early expiry must re-arm too"
1658        );
1659
1660        // The flow survived both and is still reaped at its real deadline.
1661        assert_eq!(mgr.flow_count(), 1);
1662        mgr.handle_timeout(deadline);
1663        let outs = drain(&mut mgr);
1664        assert!(outs.iter().any(|o| matches!(o, Output::CloseFlow(_))));
1665        assert_eq!(mgr.flow_count(), 0);
1666        assert_eq!(
1667            count_arms(&outs),
1668            0,
1669            "the last flow is gone: nothing left to arm"
1670        );
1671        assert!(mgr.poll_timeout().is_none());
1672    }
1673
1674    #[test]
1675    fn idle_race_resolved_by_deadline_revalidation() {
1676        // A datagram refreshes the deadline; the stale expiry must NOT close.
1677        let mut cfg = cluster("dns");
1678        cfg.front_timeout = Duration::from_secs(10);
1679        let mut mgr = UdpManager::new(cfg, 16, 65535, 7);
1680        let now = Instant::now();
1681        let src = client(1, 1000);
1682        mgr.handle_input(
1683            ManagerInput::ClientDatagram {
1684                src,
1685                payload: b"q",
1686                backends: &mut TestBackends::one(),
1687            },
1688            now,
1689        );
1690        let flow = drain(&mut mgr)
1691            .into_iter()
1692            .find_map(|o| match o {
1693                Output::OpenUpstream { flow, .. } => Some(flow),
1694                _ => None,
1695            })
1696            .unwrap();
1697        drain(&mut mgr);
1698        let gen0 = mgr.flow(flow).unwrap().timer_gen;
1699        // Datagram at t=5 refreshes deadline to t=15 and bumps generation.
1700        let t5 = now + Duration::from_secs(5);
1701        mgr.handle_input(
1702            ManagerInput::ClientDatagram {
1703                src,
1704                payload: b"q2",
1705                backends: &mut TestBackends::one(),
1706            },
1707            t5,
1708        );
1709        drain(&mut mgr);
1710        let gen1 = mgr.flow(flow).unwrap().timer_gen;
1711        assert_ne!(gen0, gen1, "touch advances the timer_gen refresh counter");
1712        assert_eq!(
1713            mgr.flow(flow).unwrap().idle_deadline,
1714            t5 + Duration::from_secs(10),
1715            "touch pushes the idle deadline, which handle_timeout revalidates"
1716        );
1717        // Stale expiry at the original t=10 deadline must NOT close (deadline is
1718        // now t=15).
1719        mgr.handle_timeout(now + Duration::from_secs(10));
1720        assert_eq!(mgr.flow_count(), 1, "refreshed flow survives stale expiry");
1721        // The real deadline at t=15 closes it.
1722        mgr.handle_timeout(now + Duration::from_secs(16));
1723        assert_eq!(mgr.flow_count(), 0);
1724    }
1725
1726    #[test]
1727    fn drain_sheds_new_flows_but_keeps_existing() {
1728        let mut mgr = UdpManager::new(cluster("dns"), 16, 65535, 7);
1729        let now = Instant::now();
1730        let src = client(1, 1000);
1731        mgr.handle_input(
1732            ManagerInput::ClientDatagram {
1733                src,
1734                payload: b"q",
1735                backends: &mut TestBackends::one(),
1736            },
1737            now,
1738        );
1739        drain(&mut mgr);
1740        mgr.handle_input(ManagerInput::Config(ConfigEvent::Drain), now);
1741        // New flow shed.
1742        mgr.handle_input(
1743            ManagerInput::ClientDatagram {
1744                src: client(2, 1000),
1745                payload: b"q",
1746                backends: &mut TestBackends::one(),
1747            },
1748            now,
1749        );
1750        assert_eq!(mgr.flow_count(), 1, "existing flow kept, new flow shed");
1751        let outs = drain(&mut mgr);
1752        assert!(
1753            outs.iter()
1754                .any(|o| matches!(o, Output::Drop(DropReason::Shed)))
1755        );
1756    }
1757
1758    #[test]
1759    fn reconfig_midflow_preserves_existing_flow_contract() {
1760        let mut cfg = cluster("dns");
1761        cfg.responses = 0;
1762        let mut mgr = UdpManager::new(cfg, 16, 65535, 7);
1763        let now = Instant::now();
1764        let src = client(1, 1000);
1765        mgr.handle_input(
1766            ManagerInput::ClientDatagram {
1767                src,
1768                payload: b"q",
1769                backends: &mut TestBackends::one(),
1770            },
1771            now,
1772        );
1773        let flow = drain(&mut mgr)
1774            .into_iter()
1775            .find_map(|o| match o {
1776                Output::OpenUpstream { flow, .. } => Some(flow),
1777                _ => None,
1778            })
1779            .unwrap();
1780        drain(&mut mgr);
1781        // Reconfigure to responses=1; the live flow keeps responses=0 (captured).
1782        let mut newcfg = cluster("dns");
1783        newcfg.responses = 1;
1784        mgr.handle_input(ManagerInput::Config(ConfigEvent::SetCluster(newcfg)), now);
1785        // A reply does NOT close the existing flow (it captured responses=0).
1786        mgr.handle_input(
1787            ManagerInput::BackendDatagram {
1788                flow,
1789                payload: b"reply",
1790            },
1791            now,
1792        );
1793        assert_eq!(mgr.flow_count(), 1);
1794    }
1795
1796    #[test]
1797    fn reaper_drains_active_flows_to_zero() {
1798        let mut cfg = cluster("dns");
1799        cfg.front_timeout = Duration::from_secs(5);
1800        let mut mgr = UdpManager::new(cfg, 64, 65535, 7);
1801        let now = Instant::now();
1802        for n in 1..=10u8 {
1803            mgr.handle_input(
1804                ManagerInput::ClientDatagram {
1805                    src: client(n, 1000),
1806                    payload: b"q",
1807                    backends: &mut TestBackends::one(),
1808                },
1809                now,
1810            );
1811        }
1812        assert_eq!(mgr.flow_count(), 10);
1813        drain(&mut mgr);
1814        // Reaper after all idle deadlines pass.
1815        mgr.handle_timeout(now + Duration::from_secs(6));
1816        assert_eq!(mgr.flow_count(), 0, "reaper drains every flow, no leak");
1817        let outs = drain(&mut mgr);
1818        let closes = outs
1819            .iter()
1820            .filter(|o| matches!(o, Output::CloseFlow(_)))
1821            .count();
1822        assert_eq!(closes, 10);
1823        assert!(mgr.poll_timeout().is_none(), "no armed timer after drain");
1824    }
1825
1826    #[test]
1827    fn proxy_protocol_first_datagram_only() {
1828        let mut cfg = cluster("dns");
1829        cfg.send_proxy_protocol = true;
1830        cfg.proxy_protocol_every_datagram = false;
1831        let mut mgr = UdpManager::new(cfg, 16, 65535, 7);
1832        let now = Instant::now();
1833        let src = client(1, 1000);
1834        mgr.handle_input(
1835            ManagerInput::ClientDatagram {
1836                src,
1837                payload: b"q1",
1838                backends: &mut TestBackends::one(),
1839            },
1840            now,
1841        );
1842        // One drain: the upstream request and the admission datagram are
1843        // emitted by the same call now, not across a resolve round-trip.
1844        let outs = drain(&mut mgr);
1845        let _flow = outs
1846            .iter()
1847            .find_map(|o| match o {
1848                Output::OpenUpstream { flow, .. } => Some(*flow),
1849                _ => None,
1850            })
1851            .expect("admission must ask the shell to open an upstream");
1852        // First datagram to the backend carries the PPv2 prefix.
1853        let first = outs
1854            .into_iter()
1855            .find_map(|o| match o {
1856                Output::SendToBackend(t) => Some(t.payload),
1857                _ => None,
1858            })
1859            .expect("the admission datagram must be forwarded");
1860        assert!(first.len() > 2, "PPv2 header prepended to first datagram");
1861        assert_eq!(&first[..4], &[0x0D, 0x0A, 0x0D, 0x0A]);
1862        assert_eq!(first[12], 0x21);
1863        assert_eq!(first[13], 0x12);
1864        assert_eq!(&first[first.len() - 2..], b"q1");
1865
1866        // Second datagram: NO prefix.
1867        mgr.handle_input(
1868            ManagerInput::ClientDatagram {
1869                src,
1870                payload: b"q2",
1871                backends: &mut TestBackends::one(),
1872            },
1873            now,
1874        );
1875        let second = drain(&mut mgr)
1876            .into_iter()
1877            .find_map(|o| match o {
1878                Output::SendToBackend(t) => Some(t.payload),
1879                _ => None,
1880            })
1881            .unwrap();
1882        assert_eq!(second, b"q2", "no PPv2 prefix on subsequent datagrams");
1883    }
1884
1885    /// Helper: admit a flow from `src`, resolve it to `backend()`, and return
1886    /// its FlowId. Drains the manager between steps. Leaves the flow
1887    /// `Established`.
1888    /// Admit a flow and return its id.
1889    ///
1890    /// One call now. It used to be two — admit, then reply to
1891    /// `Output::SelectBackend` with `ManagerInput::BackendResolved` — because
1892    /// the shell chose the backend. The core chooses it, so admission
1893    /// establishes in the same call and the flow id comes from the
1894    /// `OpenUpstream` the manager asks for.
1895    fn establish(mgr: &mut UdpManager, src: SocketAddr, now: Instant) -> FlowId {
1896        mgr.handle_input(
1897            ManagerInput::ClientDatagram {
1898                src,
1899                payload: b"q",
1900                backends: &mut TestBackends::one(),
1901            },
1902            now,
1903        );
1904        drain(mgr)
1905            .into_iter()
1906            .find_map(|o| match o {
1907                Output::OpenUpstream { flow, .. } => Some(flow),
1908                _ => None,
1909            })
1910            .expect("admission must ask the shell to open an upstream")
1911    }
1912
1913    #[test]
1914    fn close_all_evicts_every_flow_exactly_once() {
1915        let mut cfg = cluster("dns");
1916        cfg.front_timeout = Duration::from_secs(30);
1917        let mut mgr = UdpManager::new(cfg, 64, 65535, 7);
1918        let now = Instant::now();
1919        // N flows: a mix of Established and AwaitingBackend.
1920        const N: usize = 6;
1921        let mut flow_ids = Vec::new();
1922        for n in 1..=4u8 {
1923            flow_ids.push(establish(&mut mgr, client(n, 1000), now));
1924        }
1925        // Two flows left AwaitingBackend (admitted, not resolved).
1926        for n in 5..=6u8 {
1927            mgr.handle_input(
1928                ManagerInput::ClientDatagram {
1929                    src: client(n, 1000),
1930                    payload: b"q",
1931                    backends: &mut TestBackends::one(),
1932                },
1933                now,
1934            );
1935            let flow = drain(&mut mgr)
1936                .into_iter()
1937                .find_map(|o| match o {
1938                    Output::OpenUpstream { flow, .. } => Some(flow),
1939                    _ => None,
1940                })
1941                .unwrap();
1942            flow_ids.push(flow);
1943        }
1944        assert_eq!(mgr.flow_count(), N);
1945
1946        mgr.close_all(now);
1947        let outs = drain(&mut mgr);
1948        let evicted = outs
1949            .iter()
1950            .filter(|o| matches!(o, Output::Metric(MetricEvent::FlowEvicted)))
1951            .count();
1952        let closed = outs
1953            .iter()
1954            .filter(|o| matches!(o, Output::CloseFlow(_)))
1955            .count();
1956        assert_eq!(evicted, N, "one FlowEvicted per live flow");
1957        assert_eq!(closed, N, "one CloseFlow per live flow");
1958        assert_eq!(mgr.flow_count(), 0, "flow table + slab drained to zero");
1959        assert!(
1960            mgr.poll_timeout().is_none(),
1961            "no armed timer after close_all"
1962        );
1963    }
1964
1965    #[test]
1966    fn close_all_is_idempotent_and_empty_safe() {
1967        let mut mgr = UdpManager::new(cluster("dns"), 16, 65535, 7);
1968        let now = Instant::now();
1969        // Empty manager: close_all is a no-op, no outputs.
1970        mgr.close_all(now);
1971        assert!(drain(&mut mgr).is_empty());
1972        assert_eq!(mgr.flow_count(), 0);
1973
1974        // Now one flow; close it twice. The second pass emits nothing (no
1975        // double-evict / underflow).
1976        establish(&mut mgr, client(1, 1000), now);
1977        assert_eq!(mgr.flow_count(), 1);
1978        mgr.close_all(now);
1979        let first = drain(&mut mgr);
1980        assert_eq!(
1981            first
1982                .iter()
1983                .filter(|o| matches!(o, Output::CloseFlow(_)))
1984                .count(),
1985            1
1986        );
1987        mgr.close_all(now);
1988        assert!(drain(&mut mgr).is_empty(), "second close_all emits nothing");
1989        assert_eq!(mgr.flow_count(), 0);
1990    }
1991
1992    #[test]
1993    fn abort_flow_closes_established_flow() {
1994        let mut mgr = UdpManager::new(cluster("dns"), 16, 65535, 7);
1995        let now = Instant::now();
1996        let flow = establish(&mut mgr, client(1, 1000), now);
1997        assert_eq!(mgr.flow_count(), 1);
1998
1999        mgr.abort_flow(flow, now, CloseReason::Aborted);
2000        let outs = drain(&mut mgr);
2001        assert!(
2002            outs.iter()
2003                .any(|o| matches!(o, Output::Metric(MetricEvent::FlowEvicted)))
2004        );
2005        assert!(
2006            outs.iter()
2007                .any(|o| matches!(o, Output::CloseFlow(f) if *f == flow))
2008        );
2009        assert_eq!(mgr.flow_count(), 0);
2010
2011        // Idempotent: aborting again emits nothing, no underflow.
2012        mgr.abort_flow(flow, now, CloseReason::Aborted);
2013        assert!(drain(&mut mgr).is_empty());
2014        assert_eq!(mgr.flow_count(), 0);
2015    }
2016
2017    /// A flow that cannot get a backend must free its slot immediately.
2018    ///
2019    /// Retargeted from `abort_flow_closes_awaiting_backend_flow`. Its subject
2020    /// was the `AwaitingBackend` phase, which selection-in-core removed — but
2021    /// the *property* it pinned survives, because selection can still fail:
2022    /// `BackendMap::backend_from_cluster_id_with_key` has three
2023    /// `NoBackendForCluster` returns. The failure just happens inside the
2024    /// admitting call now instead of across two.
2025    ///
2026    /// What is asserted is unchanged in substance: nothing is left pinning a
2027    /// `max_flows` slot, and no timer is armed for a flow that never ran.
2028    #[test]
2029    fn a_flow_that_cannot_get_a_backend_frees_its_slot_immediately() {
2030        let mut mgr = UdpManager::new(cluster("dns"), 16, 65535, 7);
2031        let now = Instant::now();
2032        mgr.handle_input(
2033            ManagerInput::ClientDatagram {
2034                src: client(1, 1000),
2035                payload: b"q",
2036                // The cluster has nothing that can serve.
2037                backends: &mut TestBackends::none(),
2038            },
2039            now,
2040        );
2041
2042        let outs = drain(&mut mgr);
2043        assert!(
2044            outs.iter()
2045                .any(|o| matches!(o, Output::Drop(DropReason::NoBackend))),
2046            "a failed selection must surface as a NoBackend drop: {outs:?}"
2047        );
2048        assert!(
2049            !outs
2050                .iter()
2051                .any(|o| matches!(o, Output::OpenUpstream { .. })),
2052            "no upstream may be opened for a flow that got no backend"
2053        );
2054        assert_eq!(
2055            mgr.flow_count(),
2056            0,
2057            "a failed selection must leave no flow pinning a max_flows slot"
2058        );
2059        assert!(
2060            mgr.poll_timeout().is_none(),
2061            "no timer may be armed for a flow that never ran"
2062        );
2063
2064        // The key is reusable: a later datagram from the same source admits a
2065        // fresh flow rather than colliding with the refused one.
2066        mgr.handle_input(
2067            ManagerInput::ClientDatagram {
2068                src: client(1, 1000),
2069                payload: b"q2",
2070                backends: &mut TestBackends::one(),
2071            },
2072            now,
2073        );
2074        assert_eq!(mgr.flow_count(), 1);
2075    }
2076
2077    #[test]
2078    fn abort_flow_unknown_id_is_noop() {
2079        let mut mgr = UdpManager::new(cluster("dns"), 16, 65535, 7);
2080        let now = Instant::now();
2081        mgr.abort_flow(999, now, CloseReason::Aborted);
2082        assert!(drain(&mut mgr).is_empty());
2083        assert_eq!(mgr.flow_count(), 0);
2084    }
2085
2086    /// Every datagram of an opening burst reaches the backend, and each is
2087    /// counted exactly once.
2088    ///
2089    /// Retargeted from the test that pinned the await-window buffering.
2090    /// Its subject was the await window: datagrams arriving while the flow was
2091    /// `AwaitingBackend` went into a one-slot newest-wins buffer, all but the
2092    /// last were discarded, and only the survivor counted toward `requests`.
2093    /// Selection happens inside the admitting call now, so there is no window,
2094    /// no buffer, and nothing discarded — which is the **behaviour change**
2095    /// this changeset carries into the UDP datapath.
2096    ///
2097    /// Two properties, both live: the burst is delivered in full, and
2098    /// `requests` counts real forwards exactly once each so the cap still
2099    /// closes the flow on the right datagram.
2100    // ---- quickcheck property tests (zero sockets, injected Instants) -------
2101    use quickcheck::{Arbitrary, Gen, quickcheck};
2102    #[derive(Clone, Debug)]
2103    enum Step {
2104        /// Client datagram from source `id % 8`.
2105        Client(u8),
2106        /// Advance the injected clock by `secs` seconds and fire the reaper.
2107        Tick(u8),
2108    }
2109
2110    impl Arbitrary for Step {
2111        fn arbitrary(g: &mut Gen) -> Self {
2112            if bool::arbitrary(g) {
2113                Step::Client(u8::arbitrary(g))
2114            } else {
2115                Step::Tick(u8::arbitrary(g) % 40)
2116            }
2117        }
2118    }
2119
2120    /// Property: across any interleaving of datagrams and clock ticks, the flow
2121    /// count never exceeds `max_flows`, and a final long tick reaps every flow
2122    /// back to zero — no leak, no gauge underflow (`CloseFlow` count ==
2123    /// `FlowCreated` count). The busy-loop strict-advance invariant is asserted
2124    /// inside `handle_timeout` itself (debug builds), so every `Tick` exercises
2125    /// it for free.
2126    #[test]
2127    fn every_datagram_of_an_opening_burst_is_forwarded_and_counted_once() {
2128        // Part one: an unlimited flow delivers the whole burst. Before this
2129        // change only `b"3"` would have reached the backend.
2130        let mut mgr = UdpManager::new(cluster("syslog"), 16, 65535, 7);
2131        let now = Instant::now();
2132        let src = client(1, 1000);
2133        for p in [b"1".as_slice(), b"2", b"3"] {
2134            mgr.handle_input(
2135                ManagerInput::ClientDatagram {
2136                    src,
2137                    payload: p,
2138                    backends: &mut TestBackends::one(),
2139                },
2140                now,
2141            );
2142        }
2143        let sent: Vec<Vec<u8>> = drain(&mut mgr)
2144            .into_iter()
2145            .filter_map(|o| match o {
2146                Output::SendToBackend(t) => Some(t.payload),
2147                _ => None,
2148            })
2149            .collect();
2150        assert_eq!(
2151            sent,
2152            vec![b"1".to_vec(), b"2".to_vec(), b"3".to_vec()],
2153            "every datagram of the opening burst must reach the backend, in order"
2154        );
2155
2156        // Part two: `requests` counts each real forward once, so a cap of two
2157        // closes the flow on the second datagram and not before.
2158        let mut cfg = cluster("syslog");
2159        cfg.requests = 2;
2160        let mut mgr = UdpManager::new(cfg, 16, 65535, 7);
2161        let src = client(2, 1000);
2162        mgr.handle_input(
2163            ManagerInput::ClientDatagram {
2164                src,
2165                payload: b"a",
2166                backends: &mut TestBackends::one(),
2167            },
2168            now,
2169        );
2170        let flow = drain(&mut mgr)
2171            .into_iter()
2172            .find_map(|o| match o {
2173                Output::OpenUpstream { flow, .. } => Some(flow),
2174                _ => None,
2175            })
2176            .expect("admission must ask the shell to open an upstream");
2177        assert_eq!(
2178            mgr.flow(flow).unwrap().requests_seen,
2179            1,
2180            "the admission datagram is a forward and is counted exactly once"
2181        );
2182        assert_eq!(mgr.flow_count(), 1, "one forward must not reach requests=2");
2183
2184        mgr.handle_input(
2185            ManagerInput::ClientDatagram {
2186                src,
2187                payload: b"b",
2188                backends: &mut TestBackends::one(),
2189            },
2190            now,
2191        );
2192        let outs = drain(&mut mgr);
2193        assert!(
2194            outs.iter()
2195                .any(|o| matches!(o, Output::SendToBackend(t) if t.payload == b"b")),
2196            "the second datagram must still be forwarded before the cap closes the flow"
2197        );
2198        assert!(
2199            outs.iter().any(|o| matches!(o, Output::CloseFlow(_))),
2200            "the second real forward reaches requests=2 and closes the flow"
2201        );
2202        assert_eq!(mgr.flow_count(), 0);
2203    }
2204
2205    #[test]
2206    fn prop_flow_invariants() {
2207        fn prop(steps: Vec<Step>) -> bool {
2208            const MAX_FLOWS: usize = 4;
2209            let mut cfg = cluster("dns");
2210            cfg.front_timeout = Duration::from_secs(10);
2211            let mut mgr = UdpManager::new(cfg, MAX_FLOWS, 65535, 0x5EED);
2212            let base = Instant::now();
2213            let mut now = base;
2214            let mut created = 0usize;
2215            let mut closed = 0usize;
2216
2217            for step in steps {
2218                match step {
2219                    Step::Client(id) => {
2220                        let src = client(id % 8, 9000 + (id % 8) as u16);
2221                        mgr.handle_input(
2222                            ManagerInput::ClientDatagram {
2223                                src,
2224                                payload: b"q",
2225                                backends: &mut TestBackends::one(),
2226                            },
2227                            now,
2228                        );
2229                    }
2230                    Step::Tick(secs) => {
2231                        now += Duration::from_secs(secs as u64);
2232                        mgr.handle_timeout(now);
2233                    }
2234                }
2235                // Drain outputs, tallying create/close.
2236                while let Some(out) = mgr.poll_output() {
2237                    match out {
2238                        Output::Metric(MetricEvent::FlowCreated) => created += 1,
2239                        Output::CloseFlow(_) => closed += 1,
2240                        _ => {}
2241                    }
2242                }
2243                // The cap is never exceeded.
2244                if mgr.flow_count() > MAX_FLOWS {
2245                    return false;
2246                }
2247            }
2248
2249            // Final long tick must reap everything: drains to zero, no timer.
2250            now += Duration::from_secs(60);
2251            mgr.handle_timeout(now);
2252            while let Some(out) = mgr.poll_output() {
2253                if let Output::CloseFlow(_) = out {
2254                    closed += 1;
2255                }
2256            }
2257            mgr.flow_count() == 0 && mgr.poll_timeout().is_none() && created == closed
2258        }
2259        quickcheck(prop as fn(Vec<Step>) -> bool);
2260    }
2261
2262    /// Property: an idle-timeout race is always resolved by deadline revalidation —
2263    /// for any refresh strictly inside the timeout window, a stale expiry at the
2264    /// original deadline never closes a refreshed flow, and the refreshed
2265    /// deadline eventually does.
2266    #[test]
2267    fn prop_deadline_revalidation_defeats_stale_close() {
2268        fn prop(refresh_offset: u8) -> bool {
2269            let timeout = 20u64;
2270            // Refresh at 1..=timeout-1 seconds (strictly inside the window).
2271            let offset = 1 + (refresh_offset as u64 % (timeout - 1));
2272            let mut cfg = cluster("dns");
2273            cfg.front_timeout = Duration::from_secs(timeout);
2274            let mut mgr = UdpManager::new(cfg, 8, 65535, 1);
2275            let now = Instant::now();
2276            let src = client(1, 1000);
2277            mgr.handle_input(
2278                ManagerInput::ClientDatagram {
2279                    src,
2280                    payload: b"q",
2281                    backends: &mut TestBackends::one(),
2282                },
2283                now,
2284            );
2285            while mgr.poll_output().is_some() {}
2286
2287            // Refresh inside the window: pushes the deadline forward.
2288            let refreshed_at = now + Duration::from_secs(offset);
2289            mgr.handle_input(
2290                ManagerInput::ClientDatagram {
2291                    src,
2292                    payload: b"q2",
2293                    backends: &mut TestBackends::one(),
2294                },
2295                refreshed_at,
2296            );
2297            while mgr.poll_output().is_some() {}
2298
2299            // Stale expiry at the ORIGINAL deadline must not close the flow.
2300            mgr.handle_timeout(now + Duration::from_secs(timeout));
2301            while mgr.poll_output().is_some() {}
2302            if mgr.flow_count() != 1 {
2303                return false;
2304            }
2305            // The refreshed deadline eventually closes it.
2306            mgr.handle_timeout(refreshed_at + Duration::from_secs(timeout + 1));
2307            while mgr.poll_output().is_some() {}
2308            mgr.flow_count() == 0
2309        }
2310        quickcheck(prop as fn(u8) -> bool);
2311    }
2312}