1use 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
35type SourceCounts = HashMap<ClusterId, HashMap<IpAddr, usize>>;
39
40fn subnet_of(cfg: &ClusterConfig, ip: IpAddr) -> IpAddr {
42 subnet_key(&ip, cfg.subnet_ipv4_prefix, cfg.subnet_ipv6_prefix)
43}
44
45fn 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
54fn 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
63fn 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
82pub trait FlowKeyExtractor {
88 fn flow_key(&self, src: SocketAddr, payload: &[u8], cfg: &ClusterConfig) -> Option<FlowKey>;
92}
93
94#[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 if payload.is_empty() {
104 return None;
105 }
106 Some(FlowKey::from_src(src))
107 }
108}
109
110pub struct UdpManager<E: FlowKeyExtractor = SourceTupleExtractor> {
113 table: HashMap<FlowKey, FlowId>,
115 flows: Slab<UdpFlow>,
117 max_flows: usize,
119 flows_per_ip: SourceCounts,
124 flows_per_subnet: SourceCounts,
128 max_rx_datagram_size: usize,
131 cluster: ClusterConfig,
133 hash_seed: u64,
136 extractor: E,
138 draining: bool,
140
141 outputs: VecDeque<Output>,
143 armed_deadline: Option<Instant>,
148
149 #[cfg(debug_assertions)]
156 max_flows_high_water: usize,
157}
158
159impl UdpManager<SourceTupleExtractor> {
160 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 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 pub fn flow_count(&self) -> usize {
208 self.flows.len()
209 }
210
211 pub fn max_flows(&self) -> usize {
213 self.max_flows
214 }
215
216 pub fn cluster_config(&self) -> &ClusterConfig {
218 &self.cluster
219 }
220
221 pub fn is_draining(&self) -> bool {
223 self.draining
224 }
225
226 pub fn flow(&self, flow: FlowId) -> Option<&UdpFlow> {
228 self.flows.get(flow)
229 }
230
231 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 self.debug_assert_invariants();
249 }
250
251 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 pub fn close_all(&mut self, now: Instant) {
282 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 #[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 if payload.len() > self.max_rx_datagram_size {
314 self.drop_datagram(DropReason::Truncated);
315 return;
316 }
317 if self.cluster.cluster.is_empty() {
319 self.drop_datagram(DropReason::NoBackend);
320 return;
321 }
322 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 if let Some(&flow_id) = self.table.get(&key) {
333 self.forward_on_existing_flow(flow_id, payload, now);
334 return;
335 }
336
337 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 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 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 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 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 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 self.outputs.push_back(Output::OpenUpstream {
446 flow: flow_id,
447 backend: backend_addr,
448 });
449
450 self.forward_admission_datagram(flow_id, payload, backend_addr, now);
454 }
455
456 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 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 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 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 #[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 pub fn handle_timeout(&mut self, now: Instant) {
603 self.armed_deadline = None;
608
609 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 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 #[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 #[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 pub fn poll_timeout(&self) -> Option<Instant> {
663 self.armed_deadline
664 }
665
666 pub fn poll_output(&mut self) -> Option<Output> {
668 self.outputs.pop_front()
669 }
670
671 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 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 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 #[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 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 fn drop_datagram(&mut self, reason: DropReason) {
762 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 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 #[cfg(debug_assertions)]
802 fn check_invariants(&self) {
803 use std::collections::HashSet;
804
805 debug_assert_eq!(
807 self.flow_count(),
808 self.flows.len(),
809 "flow_count() must equal flows.len()"
810 );
811
812 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 debug_assert!(
828 seen_ids.len() <= self.flows.len(),
829 "more distinct table targets than live flows"
830 );
831
832 let mut min_live_deadline: Option<Instant> = None;
834 let mut live_count = 0usize;
835 for (id, flow) in self.flows.iter() {
836 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 debug_assert!(
855 flow.backend_addr.is_some(),
856 "live FlowId {id} has no backend address"
857 );
858
859 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 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 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 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 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 #[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 struct TestBackends {
968 backends: Vec<(String, SocketAddr)>,
969 next: usize,
970 }
971
972 impl TestBackends {
973 fn one() -> Self {
976 Self {
977 backends: vec![("b1".to_owned(), backend_addr(1))],
978 next: 0,
979 }
980 }
981
982 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 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 #[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 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 #[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 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 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 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 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 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 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 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 #[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 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 #[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 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 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 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 mgr.handle_timeout(now + Duration::from_secs(10));
1720 assert_eq!(mgr.flow_count(), 1, "refreshed flow survives stale expiry");
1721 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 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 let mut newcfg = cluster("dns");
1783 newcfg.responses = 1;
1784 mgr.handle_input(ManagerInput::Config(ConfigEvent::SetCluster(newcfg)), now);
1785 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 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 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 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 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 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 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 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 mgr.close_all(now);
1971 assert!(drain(&mut mgr).is_empty());
1972 assert_eq!(mgr.flow_count(), 0);
1973
1974 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 mgr.abort_flow(flow, now, CloseReason::Aborted);
2013 assert!(drain(&mut mgr).is_empty());
2014 assert_eq!(mgr.flow_count(), 0);
2015 }
2016
2017 #[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 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 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 use quickcheck::{Arbitrary, Gen, quickcheck};
2102 #[derive(Clone, Debug)]
2103 enum Step {
2104 Client(u8),
2106 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 #[test]
2127 fn every_datagram_of_an_opening_burst_is_forwarded_and_counted_once() {
2128 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 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 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 if mgr.flow_count() > MAX_FLOWS {
2245 return false;
2246 }
2247 }
2248
2249 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 #[test]
2267 fn prop_deadline_revalidation_defeats_stale_close() {
2268 fn prop(refresh_offset: u8) -> bool {
2269 let timeout = 20u64;
2270 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 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 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 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}