1use std::{
2 cell::{Cell, RefCell},
3 collections::HashMap,
4 net::SocketAddr,
5 rc::{Rc, Weak},
6 time::{Duration, Instant},
7};
8
9use mio::net::TcpStream;
10use rand::{
11 Rng, SeedableRng,
12 rngs::{StdRng, SysRng},
13};
14use sozu_command::{
15 proto::command::{
16 Event, EventKind, HealthCheckConfig, LoadBalancingAlgorithms, LoadBalancingParams,
17 LoadMetric, ShardMode,
18 },
19 secret::ProtectedHeaderSet,
20 state::ClusterId,
21};
22
23use crate::metrics::names;
24use crate::{
25 PeakEWMA,
26 load_balancing::{
27 Candidates, LeastLoaded, LoadBalancingAlgorithm, Maglev, PowerOfTwo, Random, Rendezvous,
28 RoundRobin, ShuffleSharding, count_sibling_step, hrw_score, outranks_at_address,
29 },
30 retry::{self, RetryPolicy},
31 server::{self, push_event},
32};
33
34#[derive(thiserror::Error, Debug)]
35pub enum BackendError {
36 #[error("No backend found for cluster {0}")]
37 NoBackendForCluster(String),
38 #[error("Failed to connect to socket with MIO: {0}")]
39 MioConnection(std::io::Error),
40 #[error("This backend is not in a normal status: status={0:?}")]
41 Status(BackendStatus),
42 #[error("could not connect {cluster_id} to {backend_address:?} ({failures} failures): {error}")]
43 ConnectionFailures {
44 cluster_id: String,
45 backend_address: SocketAddr,
46 failures: usize,
47 error: String,
48 },
49}
50
51#[derive(Debug, PartialEq, Eq, Clone)]
52pub enum BackendStatus {
53 Normal,
54 Closing,
55 Closed,
56}
57
58#[derive(Debug, PartialEq, Eq, Clone, Copy)]
59pub enum HealthStatus {
60 Healthy,
61 Unhealthy,
62}
63
64#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
71pub(crate) enum ClusterAvailability {
72 #[default]
73 Available,
74 AllDown,
75}
76
77#[derive(Debug, Clone, PartialEq)]
78pub struct HealthState {
79 pub status: HealthStatus,
80 pub consecutive_successes: u32,
81 pub consecutive_failures: u32,
82}
83
84impl Default for HealthState {
85 fn default() -> Self {
86 HealthState {
87 status: HealthStatus::Healthy,
88 consecutive_successes: 0,
89 consecutive_failures: 0,
90 }
91 }
92}
93
94impl HealthState {
95 pub fn record_success(&mut self, healthy_threshold: u32) -> bool {
97 let was_unhealthy = self.status == HealthStatus::Unhealthy;
98 let successes_before = self.consecutive_successes;
99 self.consecutive_failures = 0;
100 self.consecutive_successes += 1;
101
102 debug_assert_eq!(
105 self.consecutive_failures, 0,
106 "a success must clear the consecutive-failure streak"
107 );
108 debug_assert_eq!(
109 self.consecutive_successes,
110 successes_before + 1,
111 "a success must advance the success streak by exactly one"
112 );
113
114 if was_unhealthy && self.consecutive_successes >= healthy_threshold {
115 self.status = HealthStatus::Healthy;
116 debug_assert!(
119 self.status == HealthStatus::Healthy,
120 "a reported recovery must leave the status Healthy"
121 );
122 return true;
123 }
124 false
125 }
126
127 pub fn record_failure(&mut self, unhealthy_threshold: u32) -> bool {
129 let was_healthy = self.status == HealthStatus::Healthy;
130 let failures_before = self.consecutive_failures;
131 self.consecutive_successes = 0;
132 self.consecutive_failures += 1;
133
134 debug_assert_eq!(
137 self.consecutive_successes, 0,
138 "a failure must clear the consecutive-success streak"
139 );
140 debug_assert_eq!(
141 self.consecutive_failures,
142 failures_before + 1,
143 "a failure must advance the failure streak by exactly one"
144 );
145
146 if was_healthy && self.consecutive_failures >= unhealthy_threshold {
147 self.status = HealthStatus::Unhealthy;
148 debug_assert!(
149 self.status == HealthStatus::Unhealthy,
150 "a reported drop must leave the status Unhealthy"
151 );
152 return true;
153 }
154 false
155 }
156
157 pub fn is_healthy(&self) -> bool {
158 self.status == HealthStatus::Healthy
159 }
160}
161
162#[derive(Debug, PartialEq, Clone)]
163pub struct Backend {
164 pub sticky_id: Option<String>,
165 pub backend_id: String,
166 pub address: SocketAddr,
167 pub status: BackendStatus,
168 pub retry_policy: retry::RetryPolicyWrapper,
169 pub active_connections: usize,
170 pub active_requests: usize,
171 pub failures: usize,
172 pub load_balancing_parameters: Option<LoadBalancingParams>,
173 pub backup: bool,
174 pub connection_time: PeakEWMA,
175 pub health: HealthState,
176 pub unix_path: Option<String>,
180}
181
182pub fn connect_backend(address: SocketAddr, unix_path: Option<&str>) -> std::io::Result<TcpStream> {
196 let Some(unix_path) = unix_path else {
197 return TcpStream::connect(address);
198 };
199 let unix_address = sozu_command::config::parse_backend_unix_path(unix_path)
200 .map_err(|reason| std::io::Error::new(std::io::ErrorKind::InvalidInput, reason))?;
201 let stream = mio::net::UnixStream::connect_addr(&unix_address)
203 .map_err(|error| std::io::Error::new(error.kind(), format!("{unix_path}: {error}")))?;
204 let stream = TcpStream::from_std(std::net::TcpStream::from(std::os::fd::OwnedFd::from(
205 stream,
206 )));
207 debug_assert_eq!(
209 socket2::SockRef::from(&stream).domain().ok(),
210 Some(socket2::Domain::UNIX),
211 "a unix_path backend dials a UNIX domain socket"
212 );
213 Ok(stream)
214}
215
216impl Backend {
217 pub fn new(
224 backend_id: &str,
225 address: SocketAddr,
226 sticky_id: Option<String>,
227 load_balancing_parameters: Option<LoadBalancingParams>,
228 backup: Option<bool>,
229 ) -> Backend {
230 Self::new_at(
231 backend_id,
232 address,
233 sticky_id,
234 load_balancing_parameters,
235 backup,
236 Instant::now(),
237 )
238 }
239
240 pub fn new_at(
242 backend_id: &str,
243 address: SocketAddr,
244 sticky_id: Option<String>,
245 load_balancing_parameters: Option<LoadBalancingParams>,
246 backup: Option<bool>,
247 now: Instant,
248 ) -> Backend {
249 let desired_policy = retry::ExponentialBackoffPolicy::new(6);
250 Backend {
251 sticky_id,
252 backend_id: backend_id.to_owned(),
253 address,
254 status: BackendStatus::Normal,
255 retry_policy: desired_policy.into(),
256 active_connections: 0,
257 active_requests: 0,
258 failures: 0,
259 load_balancing_parameters,
260 backup: backup.unwrap_or(false),
261 connection_time: PeakEWMA::new(now),
262 health: HealthState::default(),
263 unix_path: None,
264 }
265 }
266
267 pub fn with_unix_path(mut self, unix_path: Option<String>) -> Backend {
270 self.unix_path = unix_path;
271 self
272 }
273
274 pub fn connect_socket(&self) -> std::io::Result<TcpStream> {
276 connect_backend(self.address, self.unix_path.as_deref())
277 }
278
279 pub fn set_closing(&mut self) {
280 self.status = BackendStatus::Closing;
281 }
282
283 pub fn retry_policy(&mut self) -> &mut retry::RetryPolicyWrapper {
284 &mut self.retry_policy
285 }
286
287 pub fn can_open(&self, now: Instant) -> bool {
289 if !self.health.is_healthy() {
290 return false;
291 }
292 if let Some(action) = self.retry_policy.can_try(now) {
293 self.status == BackendStatus::Normal && action == retry::RetryAction::OKAY
294 } else {
295 false
296 }
297 }
298
299 pub fn is_available(&self) -> bool {
310 self.health.is_healthy()
311 && self.status == BackendStatus::Normal
312 && !self.retry_policy.is_down()
313 }
314
315 pub fn inc_connections(&mut self) -> Option<usize> {
316 let before = self.active_connections;
317 if self.status == BackendStatus::Normal {
318 self.active_connections += 1;
319 debug_assert_eq!(
322 self.active_connections,
323 before + 1,
324 "inc_connections must add exactly one active connection"
325 );
326 Some(self.active_connections)
327 } else {
328 debug_assert_eq!(
331 self.active_connections, before,
332 "inc_connections must not touch the count for a non-Normal backend"
333 );
334 None
335 }
336 }
337
338 pub fn dec_connections(&mut self) -> Option<usize> {
340 let before = self.active_connections;
341 match self.status {
342 BackendStatus::Normal => {
343 if self.active_connections > 0 {
344 self.active_connections -= 1;
345 }
346 debug_assert!(
349 self.active_connections <= before,
350 "dec_connections must never increase the active-connection count"
351 );
352 debug_assert_eq!(
353 self.active_connections,
354 before.saturating_sub(1),
355 "dec_connections must drop by exactly one (saturating at zero)"
356 );
357 Some(self.active_connections)
358 }
359 BackendStatus::Closed => {
360 debug_assert_eq!(
362 self.active_connections, before,
363 "dec_connections on a Closed backend must not mutate the count"
364 );
365 None
366 }
367 BackendStatus::Closing => {
368 if self.active_connections > 0 {
369 self.active_connections -= 1;
370 }
371 debug_assert_eq!(
372 self.active_connections,
373 before.saturating_sub(1),
374 "dec_connections must drop by exactly one (saturating at zero)"
375 );
376 if self.active_connections == 0 {
377 self.status = BackendStatus::Closed;
378 debug_assert_eq!(
381 self.status,
382 BackendStatus::Closed,
383 "a fully drained Closing backend must become Closed"
384 );
385 None
386 } else {
387 Some(self.active_connections)
388 }
389 }
390 }
391 }
392
393 pub fn reserve_connection(&mut self) -> Result<(), BackendError> {
401 match self.inc_connections() {
402 Some(_) => Ok(()),
403 None => Err(BackendError::Status(self.status.to_owned())),
404 }
405 }
406
407 pub fn release_failed_dial<R: Rng + ?Sized>(&mut self, now: Instant, rng: &mut R) {
414 let (connections_before, failures_before) = (self.active_connections, self.failures);
415 self.dec_connections();
416 self.failures += 1;
417 self.retry_policy.fail(now, rng);
418 debug_assert_eq!(
419 self.active_connections,
420 connections_before.saturating_sub(1),
421 "a failed dial must release exactly its one reservation (saturating at 0)"
422 );
423 debug_assert_eq!(
424 self.failures,
425 failures_before + 1,
426 "a failed dial must advance the failure counter by exactly one"
427 );
428 }
429
430 pub fn set_connection_time(&mut self, dur: Duration, now: Instant) {
432 self.connection_time.observe(dur.as_nanos() as f64, now);
433 }
434
435 pub fn peak_ewma_connection(&mut self, now: Instant) -> f64 {
437 self.connection_time.get(self.active_connections, now)
438 }
439
440 pub fn try_connect<R: Rng + ?Sized>(
443 &mut self,
444 now: Instant,
445 rng: &mut R,
446 ) -> Result<mio::net::TcpStream, BackendError> {
447 if self.status != BackendStatus::Normal {
448 return Err(BackendError::Status(self.status.to_owned()));
449 }
450 debug_assert_eq!(
453 self.status,
454 BackendStatus::Normal,
455 "try_connect only attempts a connection on a Normal backend"
456 );
457 let failures_before = self.failures;
458 let connections_before = self.active_connections;
459
460 match self.connect_socket() {
461 Ok(tcp_stream) => {
462 self.inc_connections();
464 debug_assert_eq!(
467 self.active_connections,
468 connections_before + 1,
469 "a successful connect must register exactly one active connection"
470 );
471 debug_assert_eq!(
472 self.failures, failures_before,
473 "a successful connect must not bump the failure counter"
474 );
475 Ok(tcp_stream)
476 }
477 Err(io_error) => {
478 self.retry_policy.fail(now, rng);
479 self.failures += 1;
480 debug_assert_eq!(
484 self.failures,
485 failures_before + 1,
486 "a failed connect must advance the failure counter by exactly one"
487 );
488 debug_assert_eq!(
489 self.active_connections, connections_before,
490 "a failed connect must not register an active connection"
491 );
492 Err(BackendError::MioConnection(io_error))
497 }
498 }
499 }
500}
501
502impl std::ops::Drop for Backend {
506 fn drop(&mut self) {
507 server::push_event(Event {
508 kind: EventKind::RemovedBackendHasNoConnections as i32,
509 backend_id: Some(self.backend_id.to_owned()),
510 address: Some(self.address.into()),
511 cluster_id: None,
512 metric_detail: None,
513 });
514 }
515}
516
517#[derive(Debug)]
518pub struct BackendMap {
519 pub backends: HashMap<ClusterId, BackendList>,
520 pub health_check_configs: HashMap<ClusterId, HealthCheckConfig>,
521 pub protected_health_check_headers: HashMap<ClusterId, ProtectedHeaderSet>,
526 pub cluster_http2: HashMap<ClusterId, bool>,
533 rng: StdRng,
542}
543
544fn os_rng() -> StdRng {
547 StdRng::try_from_rng(&mut SysRng).expect("failed to seed random number generator from system")
548}
549
550impl Default for BackendMap {
551 fn default() -> Self {
552 Self::new()
553 }
554}
555
556impl BackendMap {
557 pub fn new() -> BackendMap {
559 Self::with_rng(os_rng())
560 }
561
562 pub fn with_seed(seed: u64) -> BackendMap {
565 Self::with_rng(StdRng::seed_from_u64(seed))
566 }
567
568 fn with_rng(rng: StdRng) -> BackendMap {
569 BackendMap {
570 backends: HashMap::new(),
571 health_check_configs: HashMap::new(),
572 protected_health_check_headers: HashMap::new(),
573 cluster_http2: HashMap::new(),
574 rng,
575 }
576 }
577
578 pub fn rng(&mut self) -> &mut StdRng {
581 &mut self.rng
582 }
583
584 pub(crate) fn record_cluster_availability(&self, cluster_id: &str) {
599 let Some(list) = self.backends.get(cluster_id) else {
600 return;
601 };
602
603 let (available, total) = list.evaluate_availability();
604 debug_assert!(
607 available <= total,
608 "available backends ({available}) cannot exceed total ({total})"
609 );
610 debug_assert_eq!(
611 total,
612 list.backends.len(),
613 "total must equal the number of registered backends"
614 );
615 gauge!(
616 names::cluster::AVAILABLE_BACKENDS,
617 available,
618 Some(cluster_id),
619 None
620 );
621 gauge!(
622 names::cluster::TOTAL_BACKENDS,
623 total,
624 Some(cluster_id),
625 None
626 );
627
628 let new_state = if total > 0 && available == 0 {
629 ClusterAvailability::AllDown
630 } else {
631 ClusterAvailability::Available
632 };
633 debug_assert!(
636 !(total == 0 && new_state == ClusterAvailability::AllDown),
637 "an empty cluster must never be reported AllDown"
638 );
639 debug_assert!(
640 !(available > 0 && new_state == ClusterAvailability::AllDown),
641 "a cluster with an available backend must not be AllDown"
642 );
643
644 let prev = list.availability.replace(new_state);
645 debug_assert_eq!(
647 list.availability.get(),
648 new_state,
649 "the availability cell must latch the newly computed state"
650 );
651 if prev == new_state {
652 return;
653 }
654 match (prev, new_state) {
655 (ClusterAvailability::Available, ClusterAvailability::AllDown) => {
656 error!("cluster {}: all {} backends are down", cluster_id, total);
657 incr!(
658 names::cluster::NO_AVAILABLE_BACKENDS,
659 Some(cluster_id),
660 None
661 );
662 push_event(Event {
663 kind: EventKind::NoAvailableBackends as i32,
664 cluster_id: Some(cluster_id.to_owned()),
665 backend_id: None,
666 address: None,
667 metric_detail: None,
668 });
669 }
670 (ClusterAvailability::AllDown, ClusterAvailability::Available) => {
671 info!(
672 "cluster {}: backends recovered ({}/{} available)",
673 cluster_id, available, total
674 );
675 incr!(names::cluster::AVAILABLE_RECOVERED, Some(cluster_id), None);
676 push_event(Event {
677 kind: EventKind::ClusterRecovered as i32,
678 cluster_id: Some(cluster_id.to_owned()),
679 backend_id: None,
680 address: None,
681 metric_detail: None,
682 });
683 }
684 _ => {}
685 }
686 }
687
688 pub fn remove_cluster(&mut self, cluster_id: &str) {
698 if let Some(list) = self.backends.remove(cluster_id) {
699 for backend in &list.backends {
700 backend.borrow_mut().set_closing();
701 }
702 }
703 self.health_check_configs.remove(cluster_id);
704 self.protected_health_check_headers.remove(cluster_id);
705 self.cluster_http2.remove(cluster_id);
706 debug_assert!(
707 !self.backends.contains_key(cluster_id)
708 && !self.health_check_configs.contains_key(cluster_id)
709 && !self.cluster_http2.contains_key(cluster_id),
710 "remove_cluster must leave nothing keyed by the cluster"
711 );
712 }
713
714 pub fn set_cluster_http2(&mut self, cluster_id: &str, http2: bool) {
719 if http2 {
720 self.cluster_http2.insert(cluster_id.into(), true);
721 } else {
722 self.cluster_http2.remove(cluster_id);
723 }
724 }
725
726 pub fn set_protected_health_check(
730 &mut self,
731 cluster_id: &str,
732 config: HealthCheckConfig,
733 headers: ProtectedHeaderSet,
734 ) {
735 self.health_check_configs.insert(cluster_id.into(), config);
736 self.protected_health_check_headers
737 .insert(cluster_id.into(), headers);
738 }
739
740 pub fn restore_health_check(
744 &mut self,
745 cluster_id: &str,
746 config: Option<HealthCheckConfig>,
747 headers: Option<ProtectedHeaderSet>,
748 ) {
749 match config {
750 Some(config) => {
751 self.health_check_configs.insert(cluster_id.into(), config);
752 }
753 None => {
754 self.health_check_configs.remove(cluster_id);
755 }
756 }
757 match headers {
758 Some(headers) => {
759 self.protected_health_check_headers
760 .insert(cluster_id.into(), headers);
761 }
762 None => {
763 self.protected_health_check_headers.remove(cluster_id);
764 }
765 }
766 }
767
768 pub fn set_health_check_config(&mut self, cluster_id: &str, config: Option<HealthCheckConfig>) {
769 let keeps_values = config.as_ref().is_some_and(|new| {
776 !new.headers.is_empty()
777 && self
778 .protected_health_check_headers
779 .get(cluster_id)
780 .is_some_and(|held| held.satisfies(&new.headers))
781 });
782 if !keeps_values {
783 self.protected_health_check_headers.remove(cluster_id);
784 }
785 match config {
786 Some(c) => {
787 self.health_check_configs.insert(cluster_id.into(), c);
788 }
789 None => {
790 self.health_check_configs.remove(cluster_id);
791 if let Some(backend_list) = self.backends.get(cluster_id) {
799 for backend in &backend_list.backends {
800 backend.borrow_mut().health = HealthState::default();
801 }
802 }
803 self.record_cluster_availability(cluster_id);
807 }
808 }
809 }
810
811 pub fn import_configuration_state(
812 &mut self,
813 backends: &HashMap<ClusterId, Vec<sozu_command::response::Backend>>,
814 ) {
815 let mut cluster_ids: Vec<&ClusterId> = backends.keys().collect();
821 cluster_ids.sort_unstable();
822 for cluster_id in cluster_ids {
823 let list =
824 BackendList::import_configuration_state(&backends[cluster_id], self.rng.next_u64());
825 self.backends.insert(cluster_id.clone(), list);
826 }
827 for cluster_id in backends.keys() {
834 self.record_cluster_availability(cluster_id);
835 }
836 }
837
838 pub fn add_backend(&mut self, cluster_id: &str, backend: Backend) {
839 let address = backend.address;
840 self.get_or_create_backend_list_for_cluster(cluster_id)
841 .add_backend(backend);
842 debug_assert!(
845 self.backends
846 .get(cluster_id)
847 .is_some_and(|list| list.has_backend(&address)),
848 "add_backend must leave the backend present in its cluster"
849 );
850 self.record_cluster_availability(cluster_id);
855 }
856
857 pub fn remove_backend(
862 &mut self,
863 cluster_id: &str,
864 backend_id: &str,
865 backend_address: &SocketAddr,
866 ) -> bool {
867 let removed = if let Some(backends) = self.backends.get_mut(cluster_id) {
868 let removed = backends.remove_backend(backend_id, backend_address);
869 if !removed {
870 warn!(
871 "No backend matches id {} at address {:?} in cluster {}: nothing removed",
872 backend_id, backend_address, cluster_id
873 );
874 }
875 removed
876 } else {
877 error!(
878 "Backend was already removed: cluster id {}, backend id {}, address {:?}",
879 cluster_id, backend_id, backend_address
880 );
881 return false;
882 };
883 debug_assert!(
884 self.backends.get(cluster_id).is_none_or(|list| list
885 .find_backend_by_identity(backend_id, backend_address)
886 .is_none()),
887 "remove_backend must evict the backend it names"
888 );
889 self.record_cluster_availability(cluster_id);
893 removed
894 }
895
896 pub fn close_backend_connection(&mut self, cluster_id: &str, addr: &SocketAddr) {
898 if let Some(cluster_backends) = self.backends.get_mut(cluster_id)
899 && let Some(ref mut backend) = cluster_backends.find_backend(addr)
900 {
901 backend.borrow_mut().dec_connections();
902 }
903 }
904
905 pub fn has_backend_id(&self, cluster_id: &str, backend_id: &str) -> bool {
909 self.backends.get(cluster_id).is_some_and(|backends| {
910 backends
911 .backends
912 .iter()
913 .any(|backend| backend.borrow().backend_id == backend_id)
914 })
915 }
916
917 pub fn has_backend(&self, cluster_id: &str, backend: &Backend) -> bool {
918 self.backends
919 .get(cluster_id)
920 .map(|backends| backends.has_backend(&backend.address))
921 .unwrap_or(false)
922 }
923
924 pub fn backend_from_cluster_id(
937 &mut self,
938 cluster_id: &str,
939 key: Option<u64>,
940 now: Instant,
941 ) -> Result<(Rc<RefCell<Backend>>, TcpStream), BackendError> {
942 let next_backend = self.select_backend(cluster_id, key, now, &[])?;
943
944 let tcp_stream = {
945 let mut borrowed_backend = next_backend.borrow_mut();
946
947 debug!(
948 "Connecting {} -> {:?}",
949 cluster_id,
950 (
951 borrowed_backend.address,
952 borrowed_backend.active_connections,
953 borrowed_backend.failures
954 )
955 );
956
957 borrowed_backend
958 .try_connect(now, &mut self.rng)
959 .map_err(|backend_error| BackendError::ConnectionFailures {
960 cluster_id: cluster_id.to_owned(),
961 backend_address: borrowed_backend.address,
962 failures: borrowed_backend.failures,
963 error: backend_error.to_string(),
964 })?
965 };
966
967 self.record_cluster_availability(cluster_id);
974
975 debug_assert!(
978 self.backends.get(cluster_id).is_some_and(|list| {
979 let picked = next_backend.borrow().address;
980 list.has_backend(&picked)
981 }),
982 "the selected backend must belong to the cluster's live set"
983 );
984
985 Ok((next_backend.clone(), tcp_stream))
986 }
987
988 pub fn reserve_backend(
1005 &mut self,
1006 cluster_id: &str,
1007 key: Option<u64>,
1008 now: Instant,
1009 ) -> Result<Rc<RefCell<Backend>>, BackendError> {
1010 self.reserve_backend_excluding(cluster_id, key, now, &[])
1011 }
1012
1013 pub fn reserve_backend_excluding(
1018 &mut self,
1019 cluster_id: &str,
1020 key: Option<u64>,
1021 now: Instant,
1022 exclude: &[SocketAddr],
1023 ) -> Result<Rc<RefCell<Backend>>, BackendError> {
1024 let backend = self.select_backend(cluster_id, key, now, exclude)?;
1025 backend.borrow_mut().reserve_connection()?;
1026 self.record_cluster_availability(cluster_id);
1030 Ok(backend)
1031 }
1032
1033 pub fn reserve_sticky_backend(
1040 &mut self,
1041 cluster_id: &str,
1042 sticky_session: &str,
1043 key: Option<u64>,
1044 now: Instant,
1045 ) -> Result<Rc<RefCell<Backend>>, BackendError> {
1046 self.reserve_sticky_backend_excluding(cluster_id, sticky_session, key, now, &[])
1047 }
1048
1049 pub fn reserve_sticky_backend_excluding(
1054 &mut self,
1055 cluster_id: &str,
1056 sticky_session: &str,
1057 key: Option<u64>,
1058 now: Instant,
1059 exclude: &[SocketAddr],
1060 ) -> Result<Rc<RefCell<Backend>>, BackendError> {
1061 let sticky = self
1062 .backends
1063 .get_mut(cluster_id)
1064 .and_then(|cluster_backends| cluster_backends.find_sticky(sticky_session, now))
1065 .filter(|backend| !exclude.contains(&backend.borrow().address))
1066 .cloned();
1067 match sticky {
1068 Some(backend) => {
1069 backend.borrow_mut().reserve_connection()?;
1070 Ok(backend)
1071 }
1072 None => {
1073 debug!(
1074 "Couldn't find a backend corresponding to sticky_session {} for cluster {}",
1075 sticky_session, cluster_id
1076 );
1077 self.reserve_backend_excluding(cluster_id, key, now, exclude)
1078 }
1079 }
1080 }
1081
1082 fn select_backend(
1085 &mut self,
1086 cluster_id: &str,
1087 key: Option<u64>,
1088 now: Instant,
1089 exclude: &[SocketAddr],
1090 ) -> Result<Rc<RefCell<Backend>>, BackendError> {
1091 let cluster_backends = self
1092 .backends
1093 .get_mut(cluster_id)
1094 .ok_or(BackendError::NoBackendForCluster(cluster_id.to_owned()))?;
1095
1096 if cluster_backends.backends.is_empty() {
1097 let _ = cluster_backends;
1102 self.record_cluster_availability(cluster_id);
1103 return Err(BackendError::NoBackendForCluster(cluster_id.to_owned()));
1104 }
1105 debug_assert!(
1107 !cluster_backends.backends.is_empty(),
1108 "selection runs only on a non-empty backend list"
1109 );
1110
1111 let (picked, outcome) = cluster_backends.select_with_key_excluding(key, now, exclude);
1112 record_shard_outcome(cluster_id, outcome, picked.is_some());
1113 match picked {
1114 Some(backend) => Ok(backend),
1115 None => {
1116 let _ = cluster_backends;
1122 self.record_cluster_availability(cluster_id);
1123 Err(BackendError::NoBackendForCluster(cluster_id.to_owned()))
1124 }
1125 }
1126 }
1127
1128 pub fn backend_from_cluster_id_with_key(
1142 &mut self,
1143 cluster_id: &str,
1144 key: Option<u64>,
1145 now: Instant,
1146 ) -> Result<(String, SocketAddr), BackendError> {
1147 let cluster_backends = self
1148 .backends
1149 .get_mut(cluster_id)
1150 .ok_or(BackendError::NoBackendForCluster(cluster_id.to_owned()))?;
1151
1152 if cluster_backends.backends.is_empty() {
1153 let _ = cluster_backends;
1154 self.record_cluster_availability(cluster_id);
1155 return Err(BackendError::NoBackendForCluster(cluster_id.to_owned()));
1156 }
1157 debug_assert!(
1158 !cluster_backends.backends.is_empty(),
1159 "keyed selection runs only on a non-empty backend list"
1160 );
1161
1162 let (picked, outcome) = cluster_backends.select_with_key(key, now);
1163 record_shard_outcome(cluster_id, outcome, picked.is_some());
1164 let next_backend = match picked {
1165 Some(nb) => nb,
1166 None => {
1167 let _ = cluster_backends;
1168 self.record_cluster_availability(cluster_id);
1169 return Err(BackendError::NoBackendForCluster(cluster_id.to_owned()));
1170 }
1171 };
1172
1173 let (backend_id, address) = {
1174 let borrowed = next_backend.borrow();
1175 (borrowed.backend_id.to_owned(), borrowed.address)
1176 };
1177 debug_assert!(
1180 cluster_backends.has_backend(&address),
1181 "keyed selection must return a backend in the cluster's live set"
1182 );
1183 Ok((backend_id, address))
1184 }
1185
1186 pub fn set_load_balancing_policy_for_cluster(
1187 &mut self,
1188 cluster_id: &str,
1189 lb_algo: LoadBalancingAlgorithms,
1190 metric: Option<LoadMetric>,
1191 ) {
1192 let seed = self.rng.next_u64();
1195 let cluster_backends = self.get_or_create_backend_list_for_cluster(cluster_id);
1196 cluster_backends.set_load_balancing_policy(lb_algo, metric, seed);
1197 }
1198
1199 pub fn set_shuffle_sharding_for_cluster(
1203 &mut self,
1204 cluster_id: &str,
1205 shuffle_sharding: Option<ShuffleSharding>,
1206 ) {
1207 self.get_or_create_backend_list_for_cluster(cluster_id)
1208 .set_shuffle_sharding(shuffle_sharding);
1209 }
1210
1211 pub fn get_or_create_backend_list_for_cluster(&mut self, cluster_id: &str) -> &mut BackendList {
1212 let rng = &mut self.rng;
1213 self.backends
1214 .entry(cluster_id.into())
1215 .or_insert_with(|| BackendList::with_seed(rng.next_u64()))
1216 }
1217}
1218
1219impl crate::protocol::udp::BackendSource for BackendMap {
1228 fn select(
1229 &mut self,
1230 cluster: &str,
1231 key: Option<u64>,
1232 now: Instant,
1233 ) -> Option<(crate::protocol::udp::BackendId, SocketAddr)> {
1234 self.backend_from_cluster_id_with_key(cluster, key, now)
1235 .ok()
1236 }
1237}
1238
1239#[derive(Debug)]
1240pub struct BackendList {
1241 pub(crate) backends: Vec<Rc<RefCell<Backend>>>,
1248 pub next_id: u32,
1249 pub(crate) load_balancing: Box<dyn LoadBalancingAlgorithm>,
1255 fail_open_warned: bool,
1263 pub(crate) availability: Cell<ClusterAvailability>,
1268 candidates: Vec<usize>,
1274 shuffle_sharding: Option<ShuffleSharding>,
1277 shard_scores: Vec<(f64, usize)>,
1281 shard: Vec<usize>,
1283 shares_address: bool,
1287 sibling_ring: Vec<usize>,
1294 shard_representative: Vec<bool>,
1299 shard_addresses: usize,
1302 collapsed: Vec<usize>,
1305}
1306
1307#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1310pub(crate) enum ShardOutcome {
1311 Unsharded,
1314 InShard,
1316 SpilledOver,
1319 Exhausted,
1322}
1323
1324impl Default for BackendList {
1325 fn default() -> Self {
1326 Self::new()
1327 }
1328}
1329
1330impl BackendList {
1331 pub fn new() -> BackendList {
1333 Self::with_seed(os_rng().next_u64())
1334 }
1335
1336 pub fn with_seed(seed: u64) -> BackendList {
1341 BackendList {
1342 backends: Vec::new(),
1343 next_id: 0,
1344 load_balancing: Box::new(Random::with_seed(seed)),
1345 fail_open_warned: false,
1346 availability: Cell::new(ClusterAvailability::Available),
1347 candidates: Vec::new(),
1348 shuffle_sharding: None,
1349 shard_scores: Vec::new(),
1350 shard: Vec::new(),
1351 shares_address: false,
1352 sibling_ring: Vec::new(),
1353 shard_representative: Vec::new(),
1354 shard_addresses: 0,
1355 collapsed: Vec::new(),
1356 }
1357 }
1358
1359 pub fn backends(&self) -> &[Rc<RefCell<Backend>>] {
1388 &self.backends
1389 }
1390
1391 pub(crate) fn evaluate_availability(&self) -> (usize, usize) {
1395 let total = self.backends.len();
1396 let available = self
1397 .backends
1398 .iter()
1399 .filter(|b| b.borrow().is_available())
1400 .count();
1401 debug_assert!(
1404 available <= total,
1405 "available ({available}) cannot exceed total ({total})"
1406 );
1407 debug_assert_eq!(total, self.backends.len(), "total must equal backend count");
1408 (available, total)
1409 }
1410
1411 #[cfg(debug_assertions)]
1416 fn check_invariants(&self) {
1417 debug_assert!(
1421 self.next_id as usize >= self.backends.len(),
1422 "next_id ({}) must be >= live backend count ({})",
1423 self.next_id,
1424 self.backends.len()
1425 );
1426 let n = self.backends.len();
1431 debug_assert_eq!(self.sibling_ring.len(), n, "one ring slot per backend");
1432 debug_assert_eq!(self.shard_representative.len(), n);
1433 let mut targeted = vec![false; n];
1435 for (index, &next) in self.sibling_ring.iter().enumerate() {
1436 debug_assert!(!targeted[next], "the sibling ring is a permutation");
1437 targeted[next] = true;
1438 debug_assert_eq!(
1439 self.backends[next].borrow().address,
1440 self.backends[index].borrow().address,
1441 "a sibling ring links only ids of one address"
1442 );
1443 }
1444 let mut leaders: HashMap<SocketAddr, usize> = HashMap::new();
1447 for index in 0..n {
1448 let mut leader = index;
1449 let mut next = self.sibling_ring[index];
1450 while next != index {
1451 leader = leader.min(next);
1452 next = self.sibling_ring[next];
1453 }
1454 let address = self.backends[index].borrow().address;
1455 debug_assert_eq!(
1456 *leaders.entry(address).or_insert(leader),
1457 leader,
1458 "an address has one sibling ring"
1459 );
1460 let backend = self.backends[index].borrow();
1461 let mut represents = !backend.backup;
1462 let mut sibling = self.sibling_ring[index];
1463 while sibling != index {
1464 let other = self.backends[sibling].borrow();
1465 represents &=
1466 other.backup || !outranks_at_address(&other, sibling, &backend, index);
1467 sibling = self.sibling_ring[sibling];
1468 }
1469 debug_assert_eq!(
1470 self.shard_representative[index], represents,
1471 "shard_representative must match the configuration"
1472 );
1473 }
1474 debug_assert_eq!(
1475 self.shares_address,
1476 leaders.len() < n,
1477 "shares_address must match the backend list"
1478 );
1479 debug_assert_eq!(
1480 self.shard_addresses,
1481 self.shard_representative.iter().filter(|r| **r).count()
1482 );
1483 for (i, a) in self.backends.iter().enumerate() {
1488 let a = a.borrow();
1489 for b in self.backends.iter().skip(i + 1) {
1490 let b = b.borrow();
1491 debug_assert!(
1492 a.address != b.address || a.backend_id != b.backend_id,
1493 "duplicate (address, backend_id) in the live set: {:?} / {}",
1494 a.address,
1495 a.backend_id
1496 );
1497 }
1498 }
1499 }
1500
1501 pub fn import_configuration_state(
1503 backend_vec: &[sozu_command_lib::response::Backend],
1504 seed: u64,
1505 ) -> BackendList {
1506 let mut list = BackendList::with_seed(seed);
1507 for backend in backend_vec {
1508 let backend = Backend::new(
1509 &backend.backend_id,
1510 backend.address,
1511 backend.sticky_id.clone(),
1512 backend.load_balancing_parameters,
1513 backend.backup,
1514 )
1515 .with_unix_path(backend.unix_path.clone());
1516 list.add_backend(backend);
1517 }
1518
1519 list
1520 }
1521
1522 pub fn add_backend(&mut self, backend: Backend) {
1523 let address = backend.address;
1524 let len_before = self.backends.len();
1525 let next_id_before = self.next_id;
1526 let existed = self.backends.iter().any(|b| {
1527 b.borrow().address == backend.address && b.borrow().backend_id == backend.backend_id
1528 });
1529 match self.backends.iter_mut().find(|b| {
1530 b.borrow().address == backend.address && b.borrow().backend_id == backend.backend_id
1531 }) {
1532 None => {
1533 let backend = Rc::new(RefCell::new(backend));
1534 self.backends.push(backend);
1535 self.next_id += 1;
1536 }
1537 Some(old_backend) => {
1540 let mut b = old_backend.borrow_mut();
1541 b.sticky_id.clone_from(&backend.sticky_id);
1542 b.load_balancing_parameters
1543 .clone_from(&backend.load_balancing_parameters);
1544 b.backup = backend.backup;
1545 b.unix_path.clone_from(&backend.unix_path);
1546 }
1547 }
1548 debug_assert_eq!(
1551 self.backends.len(),
1552 len_before + (!existed) as usize,
1553 "add_backend grows the list by one only on a genuine insert"
1554 );
1555 debug_assert_eq!(
1556 self.next_id,
1557 next_id_before + (!existed) as u32,
1558 "next_id advances by one only on a genuine insert"
1559 );
1560 debug_assert!(
1561 self.has_backend(&address),
1562 "add_backend must leave the backend present in the list"
1563 );
1564 self.load_balancing.rebuild(&self.backends);
1570 self.candidates.clear();
1573 self.candidates.reserve(self.backends.len());
1574 self.shard_scores.clear();
1575 self.shard_scores.reserve(self.backends.len());
1576 self.shard.clear();
1577 self.shard.reserve(self.backends.len());
1578 self.collapsed.clear();
1579 self.collapsed.reserve(self.backends.len());
1580 self.refresh_siblings();
1581 #[cfg(debug_assertions)]
1582 self.check_invariants();
1583 }
1584
1585 fn refresh_siblings(&mut self) {
1592 let mut groups: HashMap<SocketAddr, Vec<usize>> = HashMap::new();
1593 for (index, backend) in self.backends.iter().enumerate() {
1594 groups
1595 .entry(backend.borrow().address)
1596 .or_default()
1597 .push(index);
1598 }
1599 self.sibling_ring.clear();
1600 self.sibling_ring.extend(0..self.backends.len());
1601 self.shares_address = false;
1602 for group in groups.values() {
1603 if group.len() > 1 {
1604 self.shares_address = true;
1605 for (rank, &index) in group.iter().enumerate() {
1606 self.sibling_ring[index] = group[(rank + 1) % group.len()];
1607 }
1608 }
1609 }
1610 self.shard_representative.clear();
1611 for (index, backend) in self.backends.iter().enumerate() {
1612 let backend = backend.borrow();
1613 let mut represents = !backend.backup;
1614 let mut sibling = self.sibling_ring[index];
1615 while represents && sibling != index {
1616 let other = self.backends[sibling].borrow();
1617 represents = other.backup || !outranks_at_address(&other, sibling, &backend, index);
1618 sibling = self.sibling_ring[sibling];
1619 }
1620 self.shard_representative.push(represents);
1621 }
1622 self.shard_addresses = self.shard_representative.iter().filter(|r| **r).count();
1623 }
1624
1625 fn collapse_shared_addresses(&mut self) -> usize {
1646 if !self.shares_address || self.sibling_ring.len() != self.backends.len() {
1647 return self.candidates.len();
1648 }
1649 let backends = &self.backends;
1650 let candidates = &self.candidates;
1651 let ring = &self.sibling_ring;
1652 self.collapsed.clear();
1653 for &position in candidates {
1654 let mut outranked = false;
1655 let mut sibling = ring[position];
1656 if sibling != position {
1657 let backend = backends[position].borrow();
1658 while !outranked && sibling != position {
1659 count_sibling_step();
1660 outranked = candidates.binary_search(&sibling).is_ok()
1661 && outranks_at_address(
1662 &backends[sibling].borrow(),
1663 sibling,
1664 &backend,
1665 position,
1666 );
1667 sibling = ring[sibling];
1668 }
1669 }
1670 if !outranked {
1671 self.collapsed.push(position);
1672 }
1673 }
1674 std::mem::swap(&mut self.candidates, &mut self.collapsed);
1675 debug_assert!(
1678 self.candidates.len() <= self.collapsed.len(),
1679 "collapsing never adds a candidate"
1680 );
1681 debug_assert!(
1682 self.collapsed.is_empty() || !self.candidates.is_empty(),
1683 "every address keeps one candidate"
1684 );
1685 self.candidates.len()
1686 }
1687
1688 pub fn set_shuffle_sharding(&mut self, shuffle_sharding: Option<ShuffleSharding>) {
1691 self.shuffle_sharding = shuffle_sharding;
1692 }
1693
1694 fn compute_shard(&mut self, key: Option<u64>) -> bool {
1714 let (Some(sharding), Some(key)) = (self.shuffle_sharding, key) else {
1715 return false;
1716 };
1717 let backends = &self.backends;
1724 let shares_address = self.shares_address;
1725 let synced = self.shard_representative.len() == backends.len()
1726 && self.sibling_ring.len() == backends.len();
1727 let ring: &[usize] = if synced { &self.sibling_ring } else { &[] };
1728 let primaries = if synced {
1731 self.shard_addresses
1732 } else {
1733 backends
1734 .iter()
1735 .filter(|backend| !backend.borrow().backup)
1736 .count()
1737 };
1738 let Some(k) = sharding.shard_size(primaries) else {
1739 return false;
1740 };
1741 self.shard_scores.clear();
1742 for (index, backend) in backends.iter().enumerate() {
1743 let represents = match self.shard_representative.get(index) {
1744 Some(&represents) if synced => represents,
1745 _ => !backend.borrow().backup,
1746 };
1747 if represents {
1748 self.shard_scores
1749 .push((hrw_score(key, &backend.borrow()), index));
1750 }
1751 }
1752 debug_assert_eq!(
1753 self.shard_scores.len(),
1754 primaries,
1755 "one score per primary address"
1756 );
1757 let by_rank = |a: &(f64, usize), b: &(f64, usize)| b.0.total_cmp(&a.0).then(a.1.cmp(&b.1));
1758 if k < self.shard_scores.len() {
1759 self.shard_scores.select_nth_unstable_by(k - 1, by_rank);
1760 }
1761 self.shard.clear();
1762 for &(_, representative) in &self.shard_scores[..k] {
1763 self.shard.push(representative);
1764 let next = |position: usize| ring.get(position).copied().unwrap_or(representative);
1765 let mut sibling = next(representative);
1766 while sibling != representative {
1767 count_sibling_step();
1768 if !backends[sibling].borrow().backup {
1769 self.shard.push(sibling);
1770 }
1771 sibling = next(sibling);
1772 }
1773 }
1774 self.shard.sort_unstable();
1775 debug_assert!(
1776 self.shard.len() >= k,
1777 "a shard holds every primary id of its k addresses"
1778 );
1779 debug_assert!(
1780 shares_address || self.shard.len() == k,
1781 "without a shared address a shard holds exactly k backends"
1782 );
1783 debug_assert!(
1784 self.shard
1785 .iter()
1786 .all(|&index| !self.backends[index].borrow().backup),
1787 "a shard holds primary backends only"
1788 );
1789 true
1790 }
1791
1792 fn retain_shard_candidates(&mut self) -> usize {
1795 let shard = &self.shard;
1796 self.candidates
1797 .retain(|index| shard.binary_search(index).is_ok());
1798 self.candidates.len()
1799 }
1800
1801 pub fn remove_backend(&mut self, backend_id: &str, backend_address: &SocketAddr) -> bool {
1806 let len_before = self.backends.len();
1807 let mut removed = false;
1808 self.backends.retain(|backend| {
1809 let mut b = backend.borrow_mut();
1810 if b.backend_id == backend_id && &b.address == backend_address {
1811 removed = true;
1812 b.set_closing();
1815 false
1816 } else {
1817 true
1818 }
1819 });
1820 debug_assert_eq!(
1823 self.backends.len(),
1824 len_before - removed as usize,
1825 "remove_backend must drop exactly the backend it reports"
1826 );
1827 debug_assert!(
1828 self.find_backend_by_identity(backend_id, backend_address)
1829 .is_none(),
1830 "remove_backend must evict the backend it names"
1831 );
1832 if removed {
1836 self.load_balancing.rebuild(&self.backends);
1837 self.refresh_siblings();
1838 }
1839 #[cfg(debug_assertions)]
1840 self.check_invariants();
1841 removed
1842 }
1843
1844 pub fn find_backend_by_identity(
1846 &self,
1847 backend_id: &str,
1848 backend_address: &SocketAddr,
1849 ) -> Option<&Rc<RefCell<Backend>>> {
1850 self.backends.iter().find(|backend| {
1851 let b = backend.borrow();
1852 b.backend_id == backend_id && b.address == *backend_address
1853 })
1854 }
1855
1856 pub fn find_incarnation(
1862 &self,
1863 incarnation: &Weak<RefCell<Backend>>,
1864 ) -> Option<&Rc<RefCell<Backend>>> {
1865 self.backends
1866 .iter()
1867 .find(|backend| std::ptr::eq(Rc::as_ptr(backend), incarnation.as_ptr()))
1868 }
1869
1870 pub fn has_backend(&self, backend_address: &SocketAddr) -> bool {
1871 self.backends
1872 .iter()
1873 .any(|backend| backend.borrow().address == *backend_address)
1874 }
1875
1876 pub fn find_backend(
1877 &mut self,
1878 backend_address: &SocketAddr,
1879 ) -> Option<&mut Rc<RefCell<Backend>>> {
1880 self.backends
1881 .iter_mut()
1882 .find(|backend| backend.borrow().address == *backend_address)
1883 }
1884
1885 pub fn find_sticky(
1888 &mut self,
1889 sticky_session: &str,
1890 now: Instant,
1891 ) -> Option<&mut Rc<RefCell<Backend>>> {
1892 self.backends
1893 .iter_mut()
1894 .find(|b| b.borrow().sticky_id.as_deref() == Some(sticky_session))
1895 .filter(|b| b.borrow().can_open(now))
1896 }
1897
1898 pub fn available_backends(&mut self, backup: bool, now: Instant) -> Vec<Rc<RefCell<Backend>>> {
1902 self.backends
1903 .iter()
1904 .filter(|backend| is_tier_candidate(&backend.borrow(), backup, now))
1905 .map(Clone::clone)
1906 .collect()
1907 }
1908
1909 fn collect_candidates(&mut self, keep: impl Fn(&Backend) -> bool) -> usize {
1913 self.candidates.clear();
1914 for (index, backend) in self.backends.iter().enumerate() {
1915 if keep(&backend.borrow()) {
1916 self.candidates.push(index);
1917 }
1918 }
1919 debug_assert!(
1920 self.candidates.len() <= self.backends.len(),
1921 "candidate set cannot be larger than the full backend list"
1922 );
1923 debug_assert!(
1924 self.candidates.windows(2).all(|pair| pair[0] < pair[1]),
1925 "candidate positions must follow list order"
1926 );
1927 self.candidates.len()
1928 }
1929
1930 pub fn next_available_backend(&mut self, now: Instant) -> Option<Rc<RefCell<Backend>>> {
1932 self.next_available_backend_with_key(None, now)
1933 }
1934
1935 pub fn next_available_backend_with_key(
1949 &mut self,
1950 key: Option<u64>,
1951 now: Instant,
1952 ) -> Option<Rc<RefCell<Backend>>> {
1953 self.select_with_key(key, now).0
1954 }
1955
1956 pub(crate) fn select_with_key(
1969 &mut self,
1970 key: Option<u64>,
1971 now: Instant,
1972 ) -> (Option<Rc<RefCell<Backend>>>, ShardOutcome) {
1973 self.select_tiers(key, now, &[])
1974 }
1975
1976 pub(crate) fn select_with_key_excluding(
1994 &mut self,
1995 key: Option<u64>,
1996 now: Instant,
1997 exclude: &[SocketAddr],
1998 ) -> (Option<Rc<RefCell<Backend>>>, ShardOutcome) {
1999 if !exclude.is_empty() {
2000 let selection = self.select_tiers(key, now, exclude);
2001 if selection.0.is_some() {
2002 debug_assert!(
2003 selection
2004 .0
2005 .as_ref()
2006 .is_some_and(|backend| !exclude.contains(&backend.borrow().address)),
2007 "an excluding selection must not return an excluded backend"
2008 );
2009 return selection;
2010 }
2011 }
2012 self.select_tiers(key, now, &[])
2013 }
2014
2015 fn select_tiers(
2020 &mut self,
2021 key: Option<u64>,
2022 now: Instant,
2023 exclude: &[SocketAddr],
2024 ) -> (Option<Rc<RefCell<Backend>>>, ShardOutcome) {
2025 let candidate = |backend: &Backend, backup: bool| {
2026 is_tier_candidate(backend, backup, now) && !exclude.contains(&backend.address)
2027 };
2028 let sharded = self.compute_shard(key);
2029 let strict = sharded
2030 && self
2031 .shuffle_sharding
2032 .is_some_and(|sharding| sharding.mode == ShardMode::Strict);
2033 let mut outcome = if sharded {
2034 ShardOutcome::InShard
2035 } else {
2036 ShardOutcome::Unsharded
2037 };
2038
2039 let mut available = self.collect_candidates(|backend| candidate(backend, false));
2040 if sharded {
2041 let primaries_can_open = available != 0;
2047 let in_shard = self.retain_shard_candidates();
2048 if in_shard != 0 {
2049 available = in_shard;
2050 } else if strict {
2051 if primaries_can_open {
2052 return (None, ShardOutcome::Exhausted);
2055 }
2056 available = 0;
2060 } else {
2061 outcome = ShardOutcome::SpilledOver;
2062 available = self.collect_candidates(|backend| candidate(backend, false));
2063 }
2064 }
2065
2066 if available == 0 && !strict {
2067 available = self.collect_candidates(|backend| candidate(backend, true));
2068 }
2069
2070 if available != 0 {
2071 if self.fail_open_warned {
2073 info!(
2074 "fail-open: cluster recovered, {} backends now healthy",
2075 available
2076 );
2077 self.fail_open_warned = false;
2078 }
2079 self.collapse_shared_addresses();
2080 let picked = self.load_balancing.next_available_backend(
2081 key,
2082 Candidates::with_siblings(
2083 &self.backends,
2084 &self.candidates,
2085 &self.sibling_ring,
2086 now,
2087 ),
2088 );
2089 debug_assert!(
2090 picked.as_ref().is_none_or(|b| {
2091 let addr = b.borrow().address;
2092 self.backends.iter().any(|x| x.borrow().address == addr)
2093 }),
2094 "selection must return a backend present in the live list"
2095 );
2096 debug_assert!(
2097 outcome != ShardOutcome::InShard
2098 || picked.as_ref().is_none_or(|b| {
2099 let addr = b.borrow().address;
2100 self.shard
2101 .iter()
2102 .any(|&index| self.backends[index].borrow().address == addr)
2103 }),
2104 "an in-shard selection must return a member of the shard"
2105 );
2106 let outcome = self.settle_spill_outcome(outcome, picked.as_ref());
2107 return (picked, outcome);
2108 }
2109
2110 if !exclude.is_empty() {
2122 return (None, outcome);
2123 }
2124 let mut available = self.collect_candidates(|backend| {
2125 backend.status == BackendStatus::Normal
2126 && matches!(
2127 backend.retry_policy.can_try(now),
2128 Some(retry::RetryAction::OKAY)
2129 )
2130 });
2131 if strict {
2132 available = self.retain_shard_candidates();
2134 if available == 0 {
2135 return (None, ShardOutcome::Exhausted);
2136 }
2137 }
2138
2139 if available == 0 {
2140 return (None, outcome);
2141 }
2142
2143 if !self.fail_open_warned {
2147 warn!(
2148 "fail-open: all backends unhealthy, routing to {} normal backends with retry-policy OKAY",
2149 available
2150 );
2151 self.fail_open_warned = true;
2152 }
2153 count!(names::backend::FAIL_OPEN, 1);
2154
2155 self.collapse_shared_addresses();
2156 let picked = self.load_balancing.next_available_backend(
2157 key,
2158 Candidates::with_siblings(&self.backends, &self.candidates, &self.sibling_ring, now),
2159 );
2160 let outcome = self.settle_spill_outcome(outcome, picked.as_ref());
2161 (picked, outcome)
2162 }
2163
2164 fn settle_spill_outcome(
2168 &self,
2169 outcome: ShardOutcome,
2170 picked: Option<&Rc<RefCell<Backend>>>,
2171 ) -> ShardOutcome {
2172 match (outcome, picked) {
2173 (ShardOutcome::SpilledOver, Some(picked))
2174 if self
2175 .shard
2176 .iter()
2177 .any(|&index| Rc::ptr_eq(&self.backends[index], picked)) =>
2178 {
2179 ShardOutcome::InShard
2180 }
2181 _ => outcome,
2182 }
2183 }
2184
2185 pub fn set_load_balancing_policy(
2208 &mut self,
2209 load_balancing_policy: LoadBalancingAlgorithms,
2210 metric: Option<LoadMetric>,
2211 seed: u64,
2212 ) {
2213 match load_balancing_policy {
2214 LoadBalancingAlgorithms::RoundRobin => {
2215 let mut round_robin = RoundRobin::new();
2216 round_robin.rebuild(&self.backends);
2219 self.load_balancing = Box::new(round_robin)
2220 }
2221 LoadBalancingAlgorithms::Random => {
2222 self.load_balancing = Box::new(Random::with_seed(seed))
2223 }
2224 LoadBalancingAlgorithms::LeastLoaded => {
2225 self.load_balancing = Box::new(LeastLoaded {
2226 metric: metric.unwrap_or(LoadMetric::Connections),
2227 })
2228 }
2229 LoadBalancingAlgorithms::PowerOfTwo => {
2230 self.load_balancing = Box::new(PowerOfTwo::with_seed(
2231 seed,
2232 metric.unwrap_or(LoadMetric::Connections),
2233 ))
2234 }
2235 LoadBalancingAlgorithms::Hrw => {
2239 let mut rendezvous = Rendezvous::new();
2240 rendezvous.rebuild(&self.backends);
2242 self.load_balancing = Box::new(rendezvous)
2243 }
2244 LoadBalancingAlgorithms::Maglev => {
2245 let mut maglev = Maglev::new();
2246 LoadBalancingAlgorithm::rebuild(&mut maglev, &self.backends);
2250 self.load_balancing = Box::new(maglev);
2251 }
2252 }
2253 }
2254}
2255
2256fn record_shard_outcome(cluster_id: &str, outcome: ShardOutcome, picked: bool) {
2260 match outcome {
2261 ShardOutcome::SpilledOver if picked => {
2262 incr!(names::backend::SHARD_SPILLOVER, Some(cluster_id), None);
2263 }
2264 ShardOutcome::Exhausted => {
2265 debug_assert!(!picked, "a strict exhausted shard selects nothing");
2266 incr!(names::backend::SHARD_EXHAUSTED, Some(cluster_id), None);
2267 }
2268 ShardOutcome::Unsharded | ShardOutcome::InShard | ShardOutcome::SpilledOver => {}
2269 }
2270}
2271
2272fn is_tier_candidate(backend: &Backend, backup: bool, now: Instant) -> bool {
2275 backend.backup == backup && backend.can_open(now)
2276}
2277
2278#[cfg(test)]
2279mod backends_test {
2280
2281 use std::{net::TcpListener, sync::mpsc::*, thread};
2282
2283 use super::*;
2284
2285 fn run_mock_tcp_server(stopper: Receiver<()>) -> SocketAddr {
2291 let mut run = true;
2292 let listener = TcpListener::bind("127.0.0.1:0").unwrap();
2293 let addr = listener.local_addr().unwrap();
2294
2295 thread::spawn(move || {
2296 while run {
2297 for _stream in listener.incoming() {
2298 if let Ok(()) = stopper.try_recv() {
2300 run = false;
2301 }
2302 }
2303 }
2304 });
2305 addr
2306 }
2307
2308 #[test]
2309 fn it_should_retrieve_a_backend_from_cluster_id_when_backends_have_been_recorded() {
2310 let mut backend_map = BackendMap::new();
2311 let cluster_id = "mycluster";
2312
2313 let (sender, receiver) = channel();
2314 let backend_addr = run_mock_tcp_server(receiver);
2315
2316 backend_map.add_backend(
2317 cluster_id,
2318 Backend::new(&format!("{cluster_id}-1"), backend_addr, None, None, None),
2319 );
2320
2321 assert!(
2322 backend_map
2323 .backend_from_cluster_id(cluster_id, None, Instant::now())
2324 .is_ok()
2325 );
2326 sender.send(()).unwrap();
2327 }
2328
2329 #[test]
2333 fn removing_a_backend_or_its_cluster_marks_it_closing() {
2334 let mut backend_map = BackendMap::new();
2335 let first: SocketAddr = "127.0.0.1:9001".parse().unwrap();
2336 let second: SocketAddr = "127.0.0.1:9002".parse().unwrap();
2337 backend_map.add_backend("foo", Backend::new("foo-1", first, None, None, None));
2338 backend_map.add_backend("foo", Backend::new("foo-2", second, None, None, None));
2339 let held: Vec<Rc<RefCell<Backend>>> = backend_map.backends["foo"].backends.clone();
2340
2341 assert!(backend_map.remove_backend("foo", "foo-1", &first));
2342 assert_eq!(held[0].borrow().status, BackendStatus::Closing);
2343 assert_eq!(held[1].borrow().status, BackendStatus::Normal);
2344
2345 backend_map.remove_cluster("foo");
2346 assert_eq!(held[1].borrow().status, BackendStatus::Closing);
2347
2348 backend_map.add_backend("foo", Backend::new("foo-2", second, None, None, None));
2350 let readded = &backend_map.backends["foo"].backends[0];
2351 assert!(!Rc::ptr_eq(readded, &held[1]));
2352 assert_eq!(readded.borrow().status, BackendStatus::Normal);
2353 }
2354
2355 #[test]
2358 fn removing_one_backend_id_keeps_its_same_address_sibling() {
2359 let mut backend_map = BackendMap::new();
2360 let shared: SocketAddr = "127.0.0.1:9003".parse().unwrap();
2361 backend_map.add_backend("foo", Backend::new("foo-a", shared, None, None, None));
2362 backend_map.add_backend("foo", Backend::new("foo-b", shared, None, None, None));
2363 let held: Vec<Rc<RefCell<Backend>>> = backend_map.backends["foo"].backends.clone();
2364
2365 assert!(backend_map.remove_backend("foo", "foo-a", &shared));
2366 assert_eq!(held[0].borrow().status, BackendStatus::Closing);
2367 assert_eq!(held[1].borrow().status, BackendStatus::Normal);
2368 let list = &backend_map.backends["foo"];
2369 assert_eq!(list.backends.len(), 1);
2370 assert!(Rc::ptr_eq(
2371 list.find_backend_by_identity("foo-b", &shared).unwrap(),
2372 &held[1]
2373 ));
2374 assert!(list.find_backend_by_identity("foo-a", &shared).is_none());
2375
2376 assert!(!backend_map.remove_backend("foo", "foo-a", &shared));
2378 assert!(!backend_map.remove_backend("foo", "foo-c", &shared));
2379 assert_eq!(backend_map.backends["foo"].backends.len(), 1);
2380 }
2381
2382 #[test]
2383 fn it_should_not_retrieve_a_backend_from_cluster_id_when_backend_has_not_been_recorded() {
2384 let mut backend_map = BackendMap::new();
2385 let cluster_not_recorded = "not";
2386 backend_map.add_backend(
2387 "foo",
2388 Backend::new("foo-1", "127.0.0.1:9001".parse().unwrap(), None, None, None),
2389 );
2390
2391 assert!(
2392 backend_map
2393 .backend_from_cluster_id(cluster_not_recorded, None, Instant::now())
2394 .is_err()
2395 );
2396 }
2397
2398 #[test]
2399 fn it_should_not_retrieve_a_backend_from_cluster_id_when_backend_list_is_empty() {
2400 let mut backend_map = BackendMap::new();
2401
2402 assert!(
2403 backend_map
2404 .backend_from_cluster_id("dumb", None, Instant::now())
2405 .is_err()
2406 );
2407 }
2408
2409 #[test]
2410 fn it_should_retrieve_a_backend_from_sticky_session_when_the_backend_has_been_recorded() {
2411 let mut backend_map = BackendMap::new();
2412 let cluster_id = "mycluster";
2413 let sticky_session = "server-2";
2414
2415 for (index, port) in [(1u16, 9001u16), (2, 9000), (3, 9002)] {
2416 backend_map.add_backend(
2417 cluster_id,
2418 Backend::new(
2419 &format!("{cluster_id}-{index}"),
2420 SocketAddr::from(([127, 0, 0, 1], port)),
2421 Some(format!("server-{index}")),
2422 None,
2423 None,
2424 ),
2425 );
2426 }
2427
2428 let backend = backend_map
2429 .reserve_sticky_backend(cluster_id, sticky_session, None, Instant::now())
2430 .expect("the backend the sticky session names can take a connection");
2431 let backend = backend.borrow();
2432 assert_eq!(
2433 backend.backend_id, "mycluster-2",
2434 "a sticky session must reach the backend it names"
2435 );
2436 assert_eq!(
2437 backend.active_connections, 1,
2438 "reaching it reserves a connection"
2439 );
2440 }
2441
2442 #[test]
2443 fn it_should_not_retrieve_a_backend_from_sticky_session_when_the_backend_has_not_been_recorded()
2444 {
2445 let mut backend_map = BackendMap::new();
2446 let cluster_id = "mycluster";
2447 let sticky_session = "test";
2448
2449 assert!(
2450 backend_map
2451 .reserve_sticky_backend(cluster_id, sticky_session, None, Instant::now())
2452 .is_err()
2453 );
2454 }
2455
2456 #[test]
2457 fn it_should_not_retrieve_a_backend_from_sticky_session_when_the_backend_list_is_empty() {
2458 let mut backend_map = BackendMap::new();
2459 let mycluster_not_recorded = "mycluster";
2460 let sticky_session = "test";
2461
2462 assert!(
2463 backend_map
2464 .reserve_sticky_backend(
2465 mycluster_not_recorded,
2466 sticky_session,
2467 None,
2468 Instant::now()
2469 )
2470 .is_err()
2471 );
2472 }
2473
2474 #[test]
2475 fn it_should_add_a_backend_when_he_doesnt_already_exist() {
2476 let backend_id = "myback";
2477 let mut backends_list = BackendList::new();
2478 backends_list.add_backend(Backend::new(
2479 backend_id,
2480 "127.0.0.1:80".parse().unwrap(),
2481 None,
2482 None,
2483 None,
2484 ));
2485
2486 assert_eq!(1, backends_list.backends.len());
2487 }
2488
2489 #[test]
2490 fn it_should_not_add_a_backend_when_he_already_exist() {
2491 let backend_id = "myback";
2492 let mut backends_list = BackendList::new();
2493 backends_list.add_backend(Backend::new(
2494 backend_id,
2495 "127.0.0.1:80".parse().unwrap(),
2496 None,
2497 None,
2498 None,
2499 ));
2500
2501 backends_list.add_backend(Backend::new(
2503 backend_id,
2504 "127.0.0.1:80".parse().unwrap(),
2505 None,
2506 None,
2507 None,
2508 ));
2509
2510 assert_eq!(1, backends_list.backends.len());
2511 }
2512
2513 fn unhealthy_backend(id: &str, port: u16) -> Backend {
2516 let mut backend = Backend::new(
2517 id,
2518 format!("127.0.0.1:{port}").parse().unwrap(),
2519 None,
2520 None,
2521 None,
2522 );
2523 backend.health.record_failure(1);
2525 assert!(!backend.health.is_healthy());
2526 backend
2527 }
2528
2529 #[test]
2530 fn fail_open_picks_normal_backend_in_retry_policy_okay() {
2531 let mut list = BackendList::new();
2536 list.add_backend(unhealthy_backend("b1", 9001));
2537 list.add_backend(unhealthy_backend("b2", 9002));
2538
2539 assert!(list.available_backends(false, Instant::now()).is_empty());
2541 assert!(list.available_backends(true, Instant::now()).is_empty());
2542
2543 let picked = list.next_available_backend(Instant::now());
2544 assert!(
2545 picked.is_some(),
2546 "fail-open must pick a Normal+OKAY backend"
2547 );
2548 assert!(list.fail_open_warned, "regime entry must latch the warn!");
2549 }
2550
2551 #[test]
2552 fn fail_open_skips_backend_in_retry_backoff() {
2553 let mut list = BackendList::new();
2559 list.add_backend(unhealthy_backend("b1", 9011));
2560 list.add_backend(unhealthy_backend("b2", 9012));
2561 for backend_rc in &list.backends {
2562 backend_rc
2563 .borrow_mut()
2564 .retry_policy()
2565 .fail(Instant::now(), &mut rand::rng());
2566 assert_eq!(
2567 Some(retry::RetryAction::WAIT),
2568 backend_rc.borrow().retry_policy.can_try(Instant::now()),
2569 "test fixture must place retry policy in WAIT"
2570 );
2571 }
2572
2573 let picked = list.next_available_backend(Instant::now());
2574 assert!(
2575 picked.is_none(),
2576 "fail-open must skip backends whose retry policy is in WAIT"
2577 );
2578 assert!(
2579 !list.fail_open_warned,
2580 "no candidate backends, no regime entry"
2581 );
2582 }
2583
2584 #[test]
2585 fn fail_open_warn_latched() {
2586 let mut list = BackendList::new();
2590 list.add_backend(unhealthy_backend("b1", 9021));
2591 list.add_backend(unhealthy_backend("b2", 9022));
2592
2593 assert!(list.next_available_backend(Instant::now()).is_some());
2594 assert!(list.fail_open_warned, "first fail-open must latch");
2595
2596 assert!(list.next_available_backend(Instant::now()).is_some());
2597 assert!(
2598 list.fail_open_warned,
2599 "subsequent fail-open routing keeps the latch"
2600 );
2601
2602 list.backends[0].borrow_mut().health.status = HealthStatus::Healthy;
2605 let picked = list.next_available_backend(Instant::now());
2606 assert!(
2607 picked.is_some(),
2608 "regular path must select the healed backend"
2609 );
2610 assert!(
2611 !list.fail_open_warned,
2612 "regime exit must clear the latch so the next entry is logged again"
2613 );
2614 }
2615
2616 fn healthy_backend(id: &str, port: u16) -> Backend {
2624 Backend::new(
2625 id,
2626 format!("127.0.0.1:{port}").parse().unwrap(),
2627 None,
2628 None,
2629 None,
2630 )
2631 }
2632
2633 #[test]
2634 fn is_available_requires_health_status_and_retry_policy() {
2635 let mut backend = Backend::new("b", "127.0.0.1:9050".parse().unwrap(), None, None, None);
2637 assert!(backend.is_available(), "fresh backend must be available");
2638
2639 backend.health.record_failure(1);
2641 assert!(!backend.is_available(), "unhealthy must not be available");
2642
2643 backend.health.status = HealthStatus::Healthy;
2649 assert!(backend.is_available());
2650 backend.retry_policy.force_down();
2651 assert!(
2652 backend.retry_policy.is_down(),
2653 "test setup: retry policy budget must be exhausted",
2654 );
2655 assert!(
2656 !backend.is_available(),
2657 "retry-policy backoff must fail the predicate"
2658 );
2659
2660 backend.retry_policy.succeed(Instant::now());
2662 backend.set_closing();
2663 assert!(
2664 !backend.is_available(),
2665 "Closing lifecycle status must fail the predicate"
2666 );
2667 }
2668
2669 #[test]
2670 fn evaluate_availability_empty_list_returns_zero_zero() {
2671 let list = BackendList::new();
2672 assert_eq!((0, 0), list.evaluate_availability());
2673 }
2674
2675 #[test]
2676 fn evaluate_availability_counts_only_healthy_normal_not_in_backoff() {
2677 let mut list = BackendList::new();
2678 list.add_backend(healthy_backend("b-ok-1", 9101));
2679 list.add_backend(healthy_backend("b-ok-2", 9102));
2680 list.add_backend(unhealthy_backend("b-bad", 9103));
2681 let (available, total) = list.evaluate_availability();
2682 assert_eq!(3, total, "every configured backend counts toward total");
2683 assert_eq!(
2684 2, available,
2685 "only the two healthy backends pass the predicate"
2686 );
2687 }
2688
2689 #[test]
2690 fn evaluate_availability_excludes_retry_policy_down() {
2691 let mut list = BackendList::new();
2692 list.add_backend(healthy_backend("b-fresh", 9111));
2693 list.add_backend(healthy_backend("b-fail", 9112));
2694 list.backends[1].borrow_mut().retry_policy.force_down();
2699 let (available, total) = list.evaluate_availability();
2700 assert_eq!(2, total);
2701 assert_eq!(
2702 1, available,
2703 "retry-policy is_down() backend must be excluded even when health.is_healthy()"
2704 );
2705 }
2706
2707 #[test]
2708 fn record_cluster_availability_flips_to_alldown_then_idempotent() {
2709 let mut map = BackendMap::new();
2710 let cluster_id = "c-flap";
2711 map.add_backend(cluster_id, unhealthy_backend("b1", 9201));
2712 let list = map.backends.get(cluster_id).expect("cluster present");
2715 assert_eq!(
2716 ClusterAvailability::AllDown,
2717 list.availability.get(),
2718 "single unhealthy backend must drive the cell to AllDown"
2719 );
2720 map.record_cluster_availability(cluster_id);
2722 let list = map.backends.get(cluster_id).expect("cluster present");
2723 assert_eq!(
2724 ClusterAvailability::AllDown,
2725 list.availability.get(),
2726 "repeat call must keep the cell at AllDown without flipping"
2727 );
2728 }
2729
2730 #[test]
2731 fn record_cluster_availability_recovers_to_available() {
2732 let mut map = BackendMap::new();
2733 let cluster_id = "c-recover";
2734 map.add_backend(cluster_id, unhealthy_backend("b1", 9301));
2735 assert_eq!(
2736 ClusterAvailability::AllDown,
2737 map.backends.get(cluster_id).unwrap().availability.get()
2738 );
2739 map.backends.get_mut(cluster_id).unwrap().backends[0]
2743 .borrow_mut()
2744 .health
2745 .status = HealthStatus::Healthy;
2746 map.record_cluster_availability(cluster_id);
2747 assert_eq!(
2748 ClusterAvailability::Available,
2749 map.backends.get(cluster_id).unwrap().availability.get(),
2750 "healed backend must flip the cell back to Available"
2751 );
2752 }
2753
2754 #[test]
2755 fn record_cluster_availability_empty_cluster_stays_available() {
2756 let mut map = BackendMap::new();
2757 let cluster_id = "c-empty";
2758 map.backends.insert(cluster_id.into(), BackendList::new());
2759 map.record_cluster_availability(cluster_id);
2762 assert_eq!(
2763 ClusterAvailability::Available,
2764 map.backends.get(cluster_id).unwrap().availability.get(),
2765 "empty cluster must keep the cell at the default Available"
2766 );
2767 }
2768
2769 #[test]
2770 fn record_cluster_availability_missing_cluster_is_noop() {
2771 let map = BackendMap::new();
2772 map.record_cluster_availability("c-absent");
2774 assert!(
2775 !map.backends.contains_key("c-absent"),
2776 "helper must not insert a BackendList for an unknown cluster_id"
2777 );
2778 }
2779
2780 #[test]
2781 fn import_configuration_state_latches_cluster_rollup_gauges() {
2782 use crate::metrics::METRICS;
2783 use sozu_command_lib::proto::command::QueryMetricsOptions;
2784 let cluster_id = "c-import-rollup-9701";
2787 let mut map = BackendMap::new();
2788 let mut input = HashMap::new();
2789 input.insert(
2790 cluster_id.into(),
2791 vec![sozu_command_lib::response::Backend {
2792 cluster_id: cluster_id.into(),
2793 backend_id: "b1".to_owned(),
2794 address: "127.0.0.1:9701".parse().unwrap(),
2795 sticky_id: None,
2796 load_balancing_parameters: None,
2797 backup: None,
2798 unix_path: None,
2799 }],
2800 );
2801 map.import_configuration_state(&input);
2802 let response = METRICS
2803 .with(|m| {
2804 m.borrow_mut().query(&QueryMetricsOptions {
2805 metric_names: vec![
2806 names::cluster::AVAILABLE_BACKENDS.to_owned(),
2807 names::cluster::TOTAL_BACKENDS.to_owned(),
2808 ],
2809 cluster_ids: vec![cluster_id.to_owned()],
2810 backend_ids: vec![],
2811 list: false,
2812 no_clusters: false,
2813 workers: false,
2814 })
2815 })
2816 .expect("metrics query succeeds");
2817 let cluster_metrics = match response.content_type {
2818 Some(
2819 sozu_command_lib::proto::command::response_content::ContentType::WorkerMetrics(wm),
2820 ) => wm,
2821 Some(_) => panic!("expected WorkerMetrics, got another response variant"),
2822 None => panic!("expected WorkerMetrics, got no response content"),
2823 };
2824 let cm = cluster_metrics
2825 .clusters
2826 .get(cluster_id)
2827 .expect("imported cluster must have a ClusterMetrics entry");
2828 assert!(
2832 cm.cluster.contains_key(names::cluster::AVAILABLE_BACKENDS),
2833 "cluster.available_backends gauge must be latched at import time"
2834 );
2835 assert!(
2836 cm.cluster.contains_key(names::cluster::TOTAL_BACKENDS),
2837 "cluster.total_backends gauge must be latched at import time"
2838 );
2839 }
2840
2841 #[test]
2842 fn set_health_check_config_none_re_emits_rollup_after_reset() {
2843 let mut map = BackendMap::new();
2844 let cluster_id = "c-hc-reset";
2845 map.add_backend(cluster_id, unhealthy_backend("b1", 9801));
2848 assert_eq!(
2849 ClusterAvailability::AllDown,
2850 map.backends.get(cluster_id).unwrap().availability.get(),
2851 "test setup: unhealthy backend must register the cell at AllDown"
2852 );
2853 map.set_health_check_config(cluster_id, None);
2857 assert_eq!(
2858 ClusterAvailability::Available,
2859 map.backends.get(cluster_id).unwrap().availability.get(),
2860 "set_health_check_config(None) must re-emit the rollup after \
2861 resetting backend health, otherwise dashboards stay stuck at AllDown"
2862 );
2863 }
2864
2865 fn held_with_config(
2870 key: &sozu_command::secret::FingerprintKey,
2871 values: &[(&str, &[u8])],
2872 ) -> (HealthCheckConfig, ProtectedHeaderSet) {
2873 let headers: Vec<_> = values
2874 .iter()
2875 .map(|(name, value)| {
2876 let mut carrier = sozu_command::proto::command::HealthCheckHeaderValue {
2877 name: (*name).to_owned(),
2878 value: value.to_vec(),
2879 };
2880 let secret = sozu_command::secret::SecretValue::take_from_ingress(&mut carrier)
2881 .expect("a valid value");
2882 sozu_command::secret::ProtectedHeader::new(key, (*name).to_owned(), secret)
2883 })
2884 .collect();
2885 let config = HealthCheckConfig {
2886 uri: "/".to_owned(),
2887 interval: 10,
2888 timeout: 5,
2889 healthy_threshold: 3,
2890 unhealthy_threshold: 3,
2891 headers: headers
2892 .iter()
2893 .map(|h| sozu_command::proto::command::HealthCheckHeader {
2894 name: h.name().to_owned(),
2895 fingerprint: h.fingerprint().to_string(),
2896 })
2897 .collect(),
2898 ..Default::default()
2899 };
2900 (config, ProtectedHeaderSet::new(headers))
2901 }
2902
2903 fn test_key(byte: u8, epoch: u8) -> sozu_command::secret::FingerprintKey {
2904 sozu_command::secret::FingerprintKey::from_parts(
2905 sozu_command::secret::Zeroizing::new([byte; 32]),
2906 [epoch; 16],
2907 )
2908 }
2909
2910 fn map_holding(config: &HealthCheckConfig, held: &ProtectedHeaderSet) -> BackendMap {
2911 let mut map = BackendMap::new();
2912 map.set_protected_health_check("app", config.to_owned(), held.to_owned());
2913 map
2914 }
2915
2916 #[test]
2917 fn a_config_update_with_identical_header_metadata_keeps_the_values() {
2918 let key = test_key(1, 2);
2919 let (config, held) =
2920 held_with_config(&key, &[("Authorization", b"s1"), ("X-Token", b"s2")]);
2921 let mut map = map_holding(&config, &held);
2922
2923 let mut updated = config.to_owned();
2925 updated.unhealthy_threshold = 7;
2926 updated.interval = 20;
2927 map.set_health_check_config("app", Some(updated.to_owned()));
2928
2929 assert_eq!(map.health_check_configs["app"], updated);
2930 assert!(
2931 map.protected_health_check_headers["app"].ptr_eq(&held),
2932 "the very same set stays, bytes untouched"
2933 );
2934 }
2935
2936 #[test]
2937 fn a_config_update_that_changes_header_metadata_drops_the_values() {
2938 let key = test_key(1, 2);
2939 let (config, held) =
2940 held_with_config(&key, &[("Authorization", b"s1"), ("X-Token", b"s2")]);
2941
2942 let mut cases: Vec<(&str, Vec<sozu_command::proto::command::HealthCheckHeader>)> =
2943 Vec::new();
2944 cases.push(("clear", Vec::new()));
2946 let mut swapped = config.headers.to_owned();
2948 swapped.reverse();
2949 cases.push(("order", swapped));
2950 let (other, _) = held_with_config(&key, &[("Authorization", b"other"), ("X-Token", b"s2")]);
2952 cases.push(("fingerprint", other.headers));
2953 let (foreign, _) = held_with_config(
2955 &test_key(1, 9),
2956 &[("Authorization", b"s1"), ("X-Token", b"s2")],
2957 );
2958 cases.push(("foreign-epoch", foreign.headers));
2959 let mut malformed = config.headers.to_owned();
2961 malformed[0].fingerprint = "not-a-fingerprint".to_owned();
2962 cases.push(("malformed", malformed));
2963 let mut empty = config.headers.to_owned();
2964 empty[1].fingerprint = String::new();
2965 cases.push(("empty-fingerprint", empty));
2966 cases.push(("count", config.headers[..1].to_vec()));
2967
2968 for (label, headers) in cases {
2969 let mut map = map_holding(&config, &held);
2970 let mut updated = config.to_owned();
2971 updated.headers = headers;
2972 map.set_health_check_config("app", Some(updated.to_owned()));
2973 assert_eq!(map.health_check_configs["app"], updated, "{label}");
2974 assert!(
2975 !map.protected_health_check_headers.contains_key("app"),
2976 "{label}: the values are dropped, never kept by names alone"
2977 );
2978 }
2979 }
2980
2981 #[test]
2982 fn a_config_update_without_a_held_set_stays_missing() {
2983 let key = test_key(1, 2);
2984 let (config, _) = held_with_config(&key, &[("Authorization", b"s1")]);
2985 let mut map = BackendMap::new();
2986 map.set_health_check_config("app", Some(config.to_owned()));
2987 assert_eq!(map.health_check_configs["app"], config);
2988 assert!(!map.protected_health_check_headers.contains_key("app"));
2989 }
2990
2991 #[test]
2992 fn removing_the_health_check_drops_the_values() {
2993 let key = test_key(1, 2);
2994 let (config, held) = held_with_config(&key, &[("Authorization", b"s1")]);
2995 let mut map = map_holding(&config, &held);
2996 map.set_health_check_config("app", None);
2997 assert!(!map.health_check_configs.contains_key("app"));
2998 assert!(!map.protected_health_check_headers.contains_key("app"));
2999 }
3000
3001 fn draws(
3006 map: &mut BackendMap,
3007 cluster_id: &str,
3008 policy: LoadBalancingAlgorithms,
3009 now: Instant,
3010 draws: usize,
3011 ) -> Vec<SocketAddr> {
3012 map.set_load_balancing_policy_for_cluster(
3013 cluster_id,
3014 policy,
3015 Some(LoadMetric::Connections),
3016 );
3017 for index in 0..4u16 {
3018 map.add_backend(
3019 cluster_id,
3020 Backend::new_at(
3021 &format!("{cluster_id}-{index}"),
3022 SocketAddr::from(([127, 0, 0, 1], 9100 + index)),
3023 None,
3024 None,
3025 None,
3026 now,
3027 ),
3028 );
3029 }
3030 let list = map
3031 .backends
3032 .get_mut(cluster_id)
3033 .expect("the policy call created the list");
3034 (0..draws)
3035 .map(|_| {
3036 list.next_available_backend(now)
3037 .expect("four fresh backends can all open")
3038 .borrow()
3039 .address
3040 })
3041 .collect()
3042 }
3043
3044 #[test]
3050 fn one_map_seed_yields_one_selection_sequence() {
3051 let now = Instant::now();
3052 for policy in [
3053 LoadBalancingAlgorithms::Random,
3054 LoadBalancingAlgorithms::PowerOfTwo,
3055 ] {
3056 let first = draws(&mut BackendMap::with_seed(7), "c", policy, now, 48);
3057 assert_eq!(
3058 first,
3059 draws(&mut BackendMap::with_seed(7), "c", policy, now, 48),
3060 "{policy:?}: one seed must yield one sequence of picks"
3061 );
3062 assert_ne!(
3063 first,
3064 draws(&mut BackendMap::with_seed(8), "c", policy, now, 48),
3065 "{policy:?}: the picks must come from the map's seed, so another seed differs"
3066 );
3067 }
3068 }
3069
3070 fn picks_of(
3072 map: &mut BackendMap,
3073 cluster_id: &str,
3074 now: Instant,
3075 draws: usize,
3076 ) -> Vec<SocketAddr> {
3077 let list = map
3078 .backends
3079 .get_mut(cluster_id)
3080 .expect("the cluster was created");
3081 (0..draws)
3082 .map(|_| {
3083 list.next_available_backend(now)
3084 .expect("fresh backends can all open")
3085 .borrow()
3086 .address
3087 })
3088 .collect()
3089 }
3090
3091 #[test]
3097 fn a_cluster_created_by_adding_a_backend_is_seeded_by_the_map() {
3098 let now = Instant::now();
3099 let sequence = || {
3100 let mut map = BackendMap::with_seed(7);
3101 for index in 0..4u16 {
3102 map.add_backend(
3103 "c",
3104 Backend::new_at(
3105 &format!("c-{index}"),
3106 SocketAddr::from(([127, 0, 0, 1], 9300 + index)),
3107 None,
3108 None,
3109 None,
3110 now,
3111 ),
3112 );
3113 }
3114 picks_of(&mut map, "c", now, 48)
3115 };
3116 assert_eq!(
3117 sequence(),
3118 sequence(),
3119 "a list created by add_backend must draw its seed from the map"
3120 );
3121 }
3122
3123 #[test]
3131 fn importing_a_state_seeds_its_clusters_in_a_fixed_order() {
3132 const CLUSTERS: u16 = 16;
3133 let now = Instant::now();
3134 let state = |ids: &mut dyn Iterator<Item = u16>| {
3135 let mut state: HashMap<ClusterId, Vec<sozu_command_lib::response::Backend>> =
3136 HashMap::new();
3137 for id in ids {
3138 let cluster_id: ClusterId = format!("import-{id:02}").into();
3139 let backends = (0..4u16)
3140 .map(|index| sozu_command_lib::response::Backend {
3141 cluster_id: cluster_id.clone(),
3142 backend_id: format!("{cluster_id}-{index}"),
3143 address: SocketAddr::from(([127, 0, 0, 1], 9400 + 4 * id + index)),
3144 sticky_id: None,
3145 load_balancing_parameters: None,
3146 backup: None,
3147 unix_path: None,
3148 })
3149 .collect();
3150 state.insert(cluster_id, backends);
3151 }
3152 state
3153 };
3154 let forward = state(&mut (0..CLUSTERS));
3155 let backward = state(&mut (0..CLUSTERS).rev());
3156
3157 let mut first = BackendMap::with_seed(7);
3158 first.import_configuration_state(&forward);
3159 let mut second = BackendMap::with_seed(7);
3160 second.import_configuration_state(&backward);
3161 for id in 0..CLUSTERS {
3162 let cluster_id = format!("import-{id:02}");
3163 assert_eq!(
3164 picks_of(&mut first, &cluster_id, now, 32),
3165 picks_of(&mut second, &cluster_id, now, 32),
3166 "{cluster_id}: one map seed must give each imported cluster one seed"
3167 );
3168 }
3169 }
3170
3171 #[test]
3183 fn random_policies_seeded_by_the_os_are_not_correlated() {
3184 let now = Instant::now();
3185 let policy = LoadBalancingAlgorithms::Random;
3186 let mut map = BackendMap::new();
3187 let first = draws(&mut map, "a", policy, now, 48);
3188 assert_ne!(
3189 first,
3190 draws(&mut BackendMap::new(), "a", policy, now, 48),
3191 "two BackendMap::new() must NOT seed Random from a shared/correlated keystream"
3192 );
3193 assert_ne!(
3194 first,
3195 draws(&mut map, "b", policy, now, 48),
3196 "two clusters of one map must NOT draw from a shared keystream"
3197 );
3198 }
3199
3200 #[test]
3206 fn power_of_two_policies_seeded_by_the_os_are_not_correlated() {
3207 let now = Instant::now();
3208 let policy = LoadBalancingAlgorithms::PowerOfTwo;
3209 let mut map = BackendMap::new();
3210 let first = draws(&mut map, "a", policy, now, 48);
3211 assert_ne!(
3212 first,
3213 draws(&mut BackendMap::new(), "a", policy, now, 48),
3214 "two BackendMap::new() must NOT seed PowerOfTwo from a shared/correlated keystream"
3215 );
3216 assert_ne!(
3217 first,
3218 draws(&mut map, "b", policy, now, 48),
3219 "two clusters of one map must NOT draw from a shared keystream"
3220 );
3221 }
3222
3223 #[test]
3231 fn the_connection_time_average_decays_on_the_injected_clock() {
3232 let t0 = Instant::now();
3233 let mut backend = Backend::new_at(
3234 "ewma",
3235 SocketAddr::from(([127, 0, 0, 1], 9200)),
3236 None,
3237 None,
3238 None,
3239 t0,
3240 );
3241 let default_rtt = backend.connection_time.rtt;
3242 let cost = backend.peak_ewma_connection(t0 + Duration::from_secs(1));
3243 assert_eq!(
3244 cost,
3245 default_rtt * (-1f64).exp(),
3246 "one second after creation, the average must have decayed by e^-1"
3247 );
3248
3249 let replay = |backend: &mut Backend| {
3251 backend.set_connection_time(Duration::from_millis(80), t0 + Duration::from_millis(10));
3252 backend.set_connection_time(Duration::from_millis(20), t0 + Duration::from_millis(900));
3253 backend.peak_ewma_connection(t0 + Duration::from_millis(2_500))
3254 };
3255 let mut twin = Backend::new_at(
3256 "ewma",
3257 SocketAddr::from(([127, 0, 0, 1], 9200)),
3258 None,
3259 None,
3260 None,
3261 t0,
3262 );
3263 let mut again = Backend::new_at(
3264 "ewma",
3265 SocketAddr::from(([127, 0, 0, 1], 9200)),
3266 None,
3267 None,
3268 None,
3269 t0,
3270 );
3271 assert_eq!(
3272 replay(&mut twin).to_bits(),
3273 replay(&mut again).to_bits(),
3274 "one sequence of instants must yield one average, bit for bit"
3275 );
3276 }
3277
3278 fn sharded_list(
3283 n: u16,
3284 policy: LoadBalancingAlgorithms,
3285 sharding: Option<ShuffleSharding>,
3286 ) -> BackendList {
3287 let now = Instant::now();
3288 let mut list = BackendList::with_seed(524);
3289 list.set_load_balancing_policy(policy, Some(LoadMetric::Connections), 524);
3290 for index in 0..n {
3291 list.add_backend(Backend::new_at(
3292 &format!("shard-{index}"),
3293 SocketAddr::from(([127, 0, 0, 1], 20_000 + index)),
3294 None,
3295 None,
3296 None,
3297 now,
3298 ));
3299 }
3300 list.set_shuffle_sharding(sharding);
3301 list
3302 }
3303
3304 fn sharding(percent: u32, min_backends: u32, mode: ShardMode) -> Option<ShuffleSharding> {
3305 Some(ShuffleSharding {
3306 percent,
3307 min_backends,
3308 mode,
3309 })
3310 }
3311
3312 fn shard_of(list: &mut BackendList, key: u64) -> Vec<SocketAddr> {
3315 assert!(list.compute_shard(Some(key)), "the list must shard");
3316 let mut ranked: Vec<(f64, SocketAddr)> = list
3317 .shard
3318 .iter()
3319 .map(|&index| {
3320 let backend = list.backends[index].borrow();
3321 (hrw_score(key, &backend), backend.address)
3322 })
3323 .collect();
3324 ranked.sort_by(|a, b| b.0.total_cmp(&a.0));
3325 ranked.into_iter().map(|(_, address)| address).collect()
3326 }
3327
3328 #[test]
3332 fn a_shard_is_the_top_of_the_hrw_ranking() {
3333 let unweighted_list = || {
3334 sharded_list(
3335 12,
3336 LoadBalancingAlgorithms::RoundRobin,
3337 sharding(25, 8, ShardMode::Fallback),
3338 )
3339 };
3340 let mut unweighted = unweighted_list();
3341 let mut weighted = unweighted_list();
3344 for (index, backend) in weighted.backends.iter().enumerate() {
3345 backend.borrow_mut().load_balancing_parameters = Some(LoadBalancingParams {
3346 weight: 10 + 90 * (index as i32 % 4),
3347 });
3348 }
3349 let mut weights_moved_a_shard = false;
3350 for list in [&mut unweighted, &mut weighted] {
3351 for key in 0..64u64 {
3352 let shard = shard_of(list, key);
3353 assert_eq!(shard.len(), 3, "25% of 12 backends is a shard of 3");
3354 assert_eq!(shard, shard_of(list, key), "a shard is deterministic");
3355 let mut everyone: Vec<(f64, SocketAddr)> = list
3356 .backends
3357 .iter()
3358 .map(|backend| {
3359 let backend = backend.borrow();
3360 (hrw_score(key, &backend), backend.address)
3361 })
3362 .collect();
3363 everyone.sort_by(|a, b| b.0.total_cmp(&a.0));
3364 let top: Vec<SocketAddr> = everyone.iter().take(3).map(|(_, a)| *a).collect();
3365 assert_eq!(shard, top, "key {key}: the shard is the HRW top 3");
3366 }
3367 }
3368 for key in 0..64u64 {
3369 weights_moved_a_shard |= shard_of(&mut unweighted, key) != shard_of(&mut weighted, key);
3370 }
3371 assert!(
3372 weights_moved_a_shard,
3373 "the weighted ranking must differ from the unweighted one for some key"
3374 );
3375 }
3376
3377 #[test]
3386 fn a_shard_changes_only_at_its_tail_when_backends_change() {
3387 let now = Instant::now();
3388 let extra = |port: u16| {
3389 Backend::new_at(
3390 &format!("extra-{port}"),
3391 SocketAddr::from(([127, 0, 0, 1], port)),
3392 None,
3393 None,
3394 None,
3395 now,
3396 )
3397 };
3398 for key in 0..128u64 {
3399 let mut list = sharded_list(
3402 10,
3403 LoadBalancingAlgorithms::RoundRobin,
3404 sharding(10, 8, ShardMode::Fallback),
3405 );
3406 let before = shard_of(&mut list, key);
3407 list.add_backend(extra(21_000));
3408 let after = shard_of(&mut list, key);
3409 let lost: Vec<_> = before.iter().filter(|a| !after.contains(a)).collect();
3410 assert!(lost.len() <= 1, "key {key}: at most one member displaced");
3411 if let Some(lost) = lost.first() {
3412 assert_eq!(
3413 Some(*lost),
3414 before.last(),
3415 "key {key}: only the lowest-ranked member may be displaced"
3416 );
3417 }
3418
3419 let mut list = sharded_list(
3421 10,
3422 LoadBalancingAlgorithms::RoundRobin,
3423 sharding(20, 8, ShardMode::Fallback),
3424 );
3425 let before = shard_of(&mut list, key);
3426 list.add_backend(extra(21_001));
3427 let after = shard_of(&mut list, key);
3428 assert_eq!((before.len(), after.len()), (2, 3));
3429 assert!(
3430 before.iter().all(|a| after.contains(a)),
3431 "key {key}: a growing shard keeps its members"
3432 );
3433
3434 let mut list = sharded_list(
3437 10,
3438 LoadBalancingAlgorithms::RoundRobin,
3439 sharding(20, 8, ShardMode::Fallback),
3440 );
3441 let before = shard_of(&mut list, key);
3442 let outsider = list
3443 .backends
3444 .iter()
3445 .map(|b| b.borrow().address)
3446 .find(|a| !before.contains(a))
3447 .expect("a shard of 2 among 10 leaves outsiders");
3448 let id_of = |address: &SocketAddr| format!("shard-{}", address.port() - 20_000);
3450 assert!(list.remove_backend(&id_of(&outsider), &outsider));
3451 assert_eq!(shard_of(&mut list, key), before, "key {key}");
3452 assert!(list.remove_backend(&id_of(&before[0]), &before[0]));
3453 assert!(
3454 shard_of(&mut list, key).contains(&before[1]),
3455 "key {key}: removing one member keeps the other"
3456 );
3457 }
3458 }
3459
3460 #[test]
3463 fn a_sharded_selection_stays_in_the_shard() {
3464 let now = Instant::now();
3465 for policy in [
3466 LoadBalancingAlgorithms::RoundRobin,
3467 LoadBalancingAlgorithms::LeastLoaded,
3468 LoadBalancingAlgorithms::Random,
3469 LoadBalancingAlgorithms::Hrw,
3470 ] {
3471 let mut list = sharded_list(8, policy, sharding(25, 8, ShardMode::Fallback));
3472 for key in 0..32u64 {
3473 let shard = shard_of(&mut list, key);
3474 for _ in 0..8 {
3475 let (picked, outcome) = list.select_with_key(Some(key), now);
3476 let picked = picked.expect("a healthy shard selects").borrow().address;
3477 assert_eq!(outcome, ShardOutcome::InShard);
3478 assert!(
3479 shard.contains(&picked),
3480 "{policy:?} key {key}: {picked} is outside {shard:?}"
3481 );
3482 }
3483 }
3484 assert_eq!(list.select_with_key(None, now).1, ShardOutcome::Unsharded);
3486 let mut small = sharded_list(7, policy, sharding(25, 8, ShardMode::Fallback));
3487 assert_eq!(
3488 small.select_with_key(Some(1), now).1,
3489 ShardOutcome::Unsharded
3490 );
3491 }
3492 }
3493
3494 fn fail_shard(list: &mut BackendList, key: u64) -> Vec<SocketAddr> {
3497 let shard = shard_of(list, key);
3498 let now = Instant::now();
3499 for backend in &list.backends {
3500 let mut backend = backend.borrow_mut();
3501 if shard.contains(&backend.address) {
3502 backend.retry_policy.fail(now, &mut rand::rng());
3503 }
3504 }
3505 shard
3506 }
3507
3508 #[test]
3516 fn an_exhausted_shard_refuses_in_strict_and_spills_in_fallback() {
3517 let now = Instant::now();
3518 for key in 0..32u64 {
3519 let mut strict = sharded_list(
3520 8,
3521 LoadBalancingAlgorithms::RoundRobin,
3522 sharding(25, 8, ShardMode::Strict),
3523 );
3524 fail_shard(&mut strict, key);
3525 let (picked, outcome) = strict.select_with_key(Some(key), now);
3526 assert!(picked.is_none(), "key {key}: strict never leaves the shard");
3527 assert_eq!(outcome, ShardOutcome::Exhausted);
3528
3529 let mut fallback = sharded_list(
3530 8,
3531 LoadBalancingAlgorithms::RoundRobin,
3532 sharding(25, 8, ShardMode::Fallback),
3533 );
3534 let shard = fail_shard(&mut fallback, key);
3535 let (picked, outcome) = fallback.select_with_key(Some(key), now);
3536 let picked = picked.expect("fallback spills over").borrow().address;
3537 assert_eq!(outcome, ShardOutcome::SpilledOver);
3538 assert!(
3539 !shard.contains(&picked),
3540 "key {key}: the spill leaves the shard"
3541 );
3542 }
3543 }
3544
3545 #[test]
3549 fn a_strict_shard_fails_open_inside_itself() {
3550 let now = Instant::now();
3551 for key in 0..16u64 {
3552 let mut list = sharded_list(
3553 8,
3554 LoadBalancingAlgorithms::RoundRobin,
3555 sharding(25, 8, ShardMode::Strict),
3556 );
3557 let shard = shard_of(&mut list, key);
3558 for backend in &list.backends {
3559 backend.borrow_mut().health.status = HealthStatus::Unhealthy;
3560 }
3561 let (picked, outcome) = list.select_with_key(Some(key), now);
3562 let picked = picked.expect("fail-open still routes").borrow().address;
3563 assert!(
3564 shard.contains(&picked),
3565 "key {key}: fail-open stays in the shard"
3566 );
3567 assert_eq!(outcome, ShardOutcome::InShard);
3568 }
3569 }
3570
3571 fn unhealthy_shard(list: &mut BackendList, key: u64) -> Vec<SocketAddr> {
3575 let shard = shard_of(list, key);
3576 for backend in &list.backends {
3577 let mut backend = backend.borrow_mut();
3578 if shard.contains(&backend.address) {
3579 backend.health.status = HealthStatus::Unhealthy;
3580 }
3581 }
3582 shard
3583 }
3584
3585 #[test]
3593 fn a_strict_shard_failing_health_checks_refuses_while_others_are_healthy() {
3594 let now = Instant::now();
3595 for key in 0..32u64 {
3596 let mut list = sharded_list(
3597 8,
3598 LoadBalancingAlgorithms::RoundRobin,
3599 sharding(25, 8, ShardMode::Strict),
3600 );
3601 unhealthy_shard(&mut list, key);
3602 let (picked, outcome) = list.select_with_key(Some(key), now);
3603 assert!(
3604 picked.is_none(),
3605 "key {key}: strict must refuse, not fail open"
3606 );
3607 assert_eq!(outcome, ShardOutcome::Exhausted);
3608 }
3609 }
3610
3611 #[test]
3620 fn a_healthy_backup_does_not_turn_a_strict_fail_open_into_a_refusal() {
3621 let now = Instant::now();
3622 for key in 0..32u64 {
3623 let mut list = sharded_list(
3624 8,
3625 LoadBalancingAlgorithms::RoundRobin,
3626 sharding(25, 8, ShardMode::Strict),
3627 );
3628 let shard = shard_of(&mut list, key);
3629 for backend in &list.backends {
3630 backend.borrow_mut().health.status = HealthStatus::Unhealthy;
3631 }
3632 list.add_backend(Backend::new_at(
3633 "healthy-backup",
3634 SocketAddr::from(([127, 0, 0, 1], 23_000)),
3635 None,
3636 None,
3637 Some(true),
3638 now,
3639 ));
3640 let (picked, outcome) = list.select_with_key(Some(key), now);
3641 let picked = picked
3642 .expect("a strict shard fails open inside itself")
3643 .borrow()
3644 .address;
3645 assert!(
3646 shard.contains(&picked),
3647 "key {key}: never the backup, never outside"
3648 );
3649 assert_eq!(outcome, ShardOutcome::InShard);
3650 }
3651 }
3652
3653 #[test]
3661 fn a_fallback_fail_open_pick_inside_the_shard_is_not_a_spill_over() {
3662 let now = Instant::now();
3663 let key = 3;
3664 let mut list = sharded_list(
3665 8,
3666 LoadBalancingAlgorithms::RoundRobin,
3667 sharding(25, 8, ShardMode::Fallback),
3668 );
3669 let shard = shard_of(&mut list, key);
3670 for backend in &list.backends {
3671 backend.borrow_mut().health.status = HealthStatus::Unhealthy;
3672 }
3673 let (mut in_shard, mut outside) = (0, 0);
3674 for _ in 0..8 {
3676 let (picked, outcome) = list.select_with_key(Some(key), now);
3677 let picked = picked.expect("fail-open routes").borrow().address;
3678 if shard.contains(&picked) {
3679 assert_eq!(outcome, ShardOutcome::InShard, "{picked} is a shard member");
3680 in_shard += 1;
3681 } else {
3682 assert_eq!(outcome, ShardOutcome::SpilledOver, "{picked} is outside");
3683 outside += 1;
3684 }
3685 }
3686 assert_eq!(
3687 (in_shard, outside),
3688 (2, 6),
3689 "round-robin visits every backend once"
3690 );
3691 }
3692
3693 #[test]
3696 fn a_sticky_backend_outside_the_shard_wins() {
3697 let now = Instant::now();
3698 let key = 7;
3699 let mut map = BackendMap::with_seed(524);
3700 map.set_load_balancing_policy_for_cluster(
3701 "sharded",
3702 LoadBalancingAlgorithms::RoundRobin,
3703 None,
3704 );
3705 for index in 0..8u16 {
3706 map.add_backend(
3707 "sharded",
3708 Backend::new_at(
3709 &format!("shard-{index}"),
3710 SocketAddr::from(([127, 0, 0, 1], 22_000 + index)),
3711 Some(format!("sticky-{index}")),
3712 None,
3713 None,
3714 now,
3715 ),
3716 );
3717 }
3718 map.set_shuffle_sharding_for_cluster("sharded", sharding(25, 8, ShardMode::Strict));
3719 let list = map.backends.get_mut("sharded").expect("the cluster exists");
3720 let shard = shard_of(list, key);
3721 let outsider = list
3722 .backends
3723 .iter()
3724 .find(|b| !shard.contains(&b.borrow().address))
3725 .map(|b| {
3726 b.borrow()
3727 .sticky_id
3728 .clone()
3729 .expect("every backend has a sticky id")
3730 })
3731 .expect("a shard of 2 among 8 leaves outsiders");
3732 let reserved = map
3733 .reserve_sticky_backend("sharded", &outsider, Some(key), now)
3734 .expect("the sticky backend is live");
3735 assert_eq!(
3736 reserved.borrow().sticky_id.as_deref(),
3737 Some(outsider.as_str())
3738 );
3739 assert!(!shard.contains(&reserved.borrow().address));
3740 }
3741
3742 #[test]
3748 fn a_sharded_selection_allocates_nothing() {
3749 use std::hint::black_box;
3750
3751 use crate::test_allocations::allocations;
3752
3753 let now = Instant::now();
3754 let mut list = sharded_list(
3755 16,
3756 LoadBalancingAlgorithms::LeastLoaded,
3757 sharding(25, 8, ShardMode::Fallback),
3758 );
3759 drop(list.select_with_key(Some(0), now));
3761 let mut allocated = 0;
3762 for key in 0..256u64 {
3763 let before = allocations();
3764 let picked = black_box(&mut list).select_with_key(black_box(Some(key)), now);
3765 allocated += allocations() - before;
3766 assert_eq!(picked.1, ShardOutcome::InShard);
3767 }
3768 assert_eq!(
3769 allocated, 0,
3770 "256 sharded selections made {allocated} allocations"
3771 );
3772 }
3773}
3774
3775#[cfg(test)]
3778mod exclusion_tests {
3779 use super::*;
3780
3781 const POLICIES: [LoadBalancingAlgorithms; 6] = [
3782 LoadBalancingAlgorithms::RoundRobin,
3783 LoadBalancingAlgorithms::Random,
3784 LoadBalancingAlgorithms::LeastLoaded,
3785 LoadBalancingAlgorithms::PowerOfTwo,
3786 LoadBalancingAlgorithms::Hrw,
3787 LoadBalancingAlgorithms::Maglev,
3788 ];
3789
3790 fn address(port: u16) -> SocketAddr {
3791 SocketAddr::from(([127, 0, 0, 1], port))
3792 }
3793
3794 fn list(policy: LoadBalancingAlgorithms, count: u16) -> BackendList {
3796 let mut list = BackendList::with_seed(7);
3797 for index in 0..count {
3798 list.add_backend(Backend::new(
3799 &format!("b{index}"),
3800 address(9100 + index),
3801 None,
3802 None,
3803 None,
3804 ));
3805 }
3806 list.set_load_balancing_policy(policy, None, 11);
3807 list
3808 }
3809
3810 #[test]
3811 fn a_selection_never_returns_an_excluded_backend_under_any_policy() {
3812 let now = Instant::now();
3813 let exclude = [address(9100), address(9101), address(9102)];
3814 for policy in POLICIES {
3815 let mut list = list(policy, 4);
3816 for key in 0..32 {
3817 let (picked, _) = list.select_with_key_excluding(Some(key), now, &exclude);
3819 let picked = picked
3820 .expect("one backend is not excluded")
3821 .borrow()
3822 .address;
3823 assert_eq!(
3824 picked,
3825 address(9103),
3826 "{policy:?} picked {picked}, which the request already tried"
3827 );
3828 }
3829 }
3830 }
3831
3832 #[test]
3833 fn an_empty_exclusion_selects_exactly_as_before() {
3834 let now = Instant::now();
3835 for policy in POLICIES {
3836 let mut excluding = list(policy, 4);
3837 let mut plain = list(policy, 4);
3838 for key in 0..16 {
3839 let left = excluding
3840 .select_with_key_excluding(Some(key), now, &[])
3841 .0
3842 .map(|backend| backend.borrow().address);
3843 let right = plain
3844 .select_with_key(Some(key), now)
3845 .0
3846 .map(|backend| backend.borrow().address);
3847 assert_eq!(left, right, "{policy:?} diverged on key {key}");
3848 }
3849 }
3850 }
3851
3852 #[test]
3853 fn excluding_every_backend_falls_back_to_the_plain_selection() {
3854 let now = Instant::now();
3855 let exclude = [address(9100), address(9101)];
3856 for policy in POLICIES {
3857 let mut list = list(policy, 2);
3858 let (picked, _) = list.select_with_key_excluding(Some(3), now, &exclude);
3859 assert!(
3860 picked.is_some(),
3861 "{policy:?}: a request that tried every backend retries one of them"
3862 );
3863 }
3864 }
3865
3866 #[test]
3867 fn an_excluded_sticky_backend_is_not_followed() {
3868 let now = Instant::now();
3869 let mut map = BackendMap::with_seed(5);
3870 for index in 0..2u16 {
3871 map.add_backend(
3872 "cluster",
3873 Backend::new(
3874 &format!("b{index}"),
3875 address(9200 + index),
3876 Some(format!("sticky{index}")),
3877 None,
3878 None,
3879 ),
3880 );
3881 }
3882 let followed = map
3883 .reserve_sticky_backend_excluding("cluster", "sticky0", None, now, &[])
3884 .expect("a backend is selectable");
3885 assert_eq!(followed.borrow().address, address(9200));
3886 let skipped = map
3887 .reserve_sticky_backend_excluding("cluster", "sticky0", None, now, &[address(9200)])
3888 .expect("a backend is selectable");
3889 assert_eq!(
3890 skipped.borrow().address,
3891 address(9201),
3892 "the cookie names a backend the request already failed to reach"
3893 );
3894 }
3895}
3896
3897#[cfg(test)]
3901mod same_address_tests {
3902 use super::*;
3903 use crate::load_balancing::affinity_key_from_value;
3904
3905 const POLICIES: [LoadBalancingAlgorithms; 6] = [
3906 LoadBalancingAlgorithms::RoundRobin,
3907 LoadBalancingAlgorithms::Random,
3908 LoadBalancingAlgorithms::LeastLoaded,
3909 LoadBalancingAlgorithms::PowerOfTwo,
3910 LoadBalancingAlgorithms::Hrw,
3911 LoadBalancingAlgorithms::Maglev,
3912 ];
3913
3914 const PICKS: u32 = 20_000;
3916
3917 const SHARED: SocketAddr = SocketAddr::new(
3918 std::net::IpAddr::V4(std::net::Ipv4Addr::new(127, 0, 0, 1)),
3919 9300,
3920 );
3921 const OTHER: SocketAddr = SocketAddr::new(
3922 std::net::IpAddr::V4(std::net::Ipv4Addr::new(127, 0, 0, 1)),
3923 9301,
3924 );
3925
3926 fn weighted(id: &str, address: SocketAddr, weight: Option<i32>) -> Backend {
3927 Backend::new(
3928 id,
3929 address,
3930 None,
3931 weight.map(|weight| LoadBalancingParams { weight }),
3932 None,
3933 )
3934 }
3935
3936 fn shared_list(policy: LoadBalancingAlgorithms, a: Backend, b: Backend) -> BackendList {
3938 let mut list = BackendList::with_seed(1839);
3939 list.add_backend(a);
3940 list.add_backend(b);
3941 list.add_backend(weighted("c", OTHER, None));
3942 list.set_load_balancing_policy(policy, Some(LoadMetric::Connections), 1839);
3943 list
3944 }
3945
3946 fn equal_list(policy: LoadBalancingAlgorithms) -> BackendList {
3947 shared_list(
3948 policy,
3949 weighted("a", SHARED, None),
3950 weighted("b", SHARED, None),
3951 )
3952 }
3953
3954 fn picks(list: &mut BackendList) -> Vec<String> {
3958 let now = Instant::now();
3959 (0..PICKS)
3960 .map(|index| {
3961 let key = affinity_key_from_value(&index.to_le_bytes());
3962 let picked = list
3963 .next_available_backend_with_key(Some(key), now)
3964 .expect("three healthy backends");
3965 let mut picked = picked.borrow_mut();
3966 picked.active_connections += 1;
3967 picked.backend_id.to_owned()
3968 })
3969 .collect()
3970 }
3971
3972 fn share_of(ids: &[String], wanted: &[&str]) -> f64 {
3973 ids.iter()
3974 .filter(|id| wanted.contains(&id.as_str()))
3975 .count() as f64
3976 / ids.len() as f64
3977 }
3978
3979 #[test]
3987 fn a_shared_address_takes_the_share_of_one_backend_under_every_policy() {
3988 let skewed: Vec<String> = POLICIES
3989 .into_iter()
3990 .filter_map(|policy| {
3991 let mut list = equal_list(policy);
3992 let shared = share_of(&picks(&mut list), &["a", "b"]);
3993 (!(0.47..=0.53).contains(&shared)).then(|| format!("{policy:?}: {shared:.3}"))
3994 })
3995 .collect();
3996 assert!(
3997 skewed.is_empty(),
3998 "the shared address did not take 1/2 of the selections under {skewed:?}"
3999 );
4000 }
4001
4002 #[test]
4006 fn the_first_id_at_a_shared_address_represents_it() {
4007 for policy in POLICIES {
4008 let mut list = equal_list(policy);
4009 let ids = picks(&mut list);
4010 assert!(
4011 ids.iter().all(|id| id != "b"),
4012 "{policy:?} picked the second id of the shared address"
4013 );
4014 assert!(ids.iter().any(|id| id == "a"), "{policy:?} never picked a");
4015 }
4016 }
4017
4018 #[test]
4021 fn a_down_id_yields_its_address_to_an_eligible_sibling() {
4022 for policy in POLICIES {
4023 let mut down = weighted("a", SHARED, None);
4024 down.health.record_failure(1);
4025 assert!(!down.health.is_healthy());
4026 let mut list = shared_list(policy, down, weighted("b", SHARED, None));
4027 let ids = picks(&mut list);
4028 assert!(
4029 ids.iter().all(|id| id != "a"),
4030 "{policy:?} picked the unhealthy id"
4031 );
4032 let shared = share_of(&ids, &["b"]);
4033 assert!(
4034 (0.47..=0.53).contains(&shared),
4035 "{policy:?} gave the shared address {shared:.3} of the selections, not 1/2"
4036 );
4037 }
4038 }
4039
4040 #[test]
4043 fn the_heaviest_id_at_a_shared_address_represents_it() {
4044 for policy in POLICIES {
4045 let mut list = shared_list(
4046 policy,
4047 weighted("a", SHARED, Some(50)),
4048 weighted("b", SHARED, Some(200)),
4049 );
4050 let ids = picks(&mut list);
4051 assert!(
4052 ids.iter().all(|id| id != "a"),
4053 "{policy:?} picked the lighter id of the shared address"
4054 );
4055 }
4056 }
4057
4058 #[test]
4063 fn a_shard_counts_a_shared_address_once_and_holds_all_its_ids() {
4064 let mut list = BackendList::with_seed(524);
4065 list.set_load_balancing_policy(LoadBalancingAlgorithms::Hrw, None, 524);
4066 for index in 0..8u16 {
4067 list.add_backend(Backend::new(
4068 &format!("shard-{index}"),
4069 SocketAddr::from(([127, 0, 0, 1], 20_000 + index)),
4070 None,
4071 None,
4072 None,
4073 ));
4074 }
4075 let shared = SocketAddr::from(([127, 0, 0, 1], 20_000));
4076 list.add_backend(Backend::new("shard-0-bis", shared, None, None, None));
4077 list.set_shuffle_sharding(Some(ShuffleSharding {
4078 percent: 25,
4079 min_backends: 4,
4080 mode: ShardMode::Fallback,
4081 }));
4082
4083 let mut holding_shared = 0;
4084 for key in 0..256u64 {
4085 assert!(list.compute_shard(Some(key)), "the list must shard");
4086 let mut addresses: Vec<SocketAddr> = list
4087 .shard
4088 .iter()
4089 .map(|&index| list.backends[index].borrow().address)
4090 .collect();
4091 let ids = addresses.len();
4092 addresses.sort_unstable();
4093 addresses.dedup();
4094 assert_eq!(addresses.len(), 2, "key {key}: a shard holds two addresses");
4095 if addresses.contains(&shared) {
4096 holding_shared += 1;
4097 assert_eq!(ids, 3, "key {key}: the shared address brings both ids");
4098 } else {
4099 assert_eq!(ids, 2, "key {key}: one id per unshared address");
4100 }
4101 }
4102 assert!(
4103 holding_shared > 0,
4104 "some shard must hold the shared address"
4105 );
4106 }
4107
4108 #[test]
4113 fn a_load_aware_policy_reads_the_load_of_the_whole_address() {
4114 let now = Instant::now();
4115 for policy in [
4116 LoadBalancingAlgorithms::LeastLoaded,
4117 LoadBalancingAlgorithms::PowerOfTwo,
4118 ] {
4119 let mut list = equal_list(policy);
4120 list.backends[1].borrow_mut().active_connections = 5;
4121 list.backends[2].borrow_mut().active_connections = 3;
4122 for key in 0..64 {
4123 let picked = list
4124 .next_available_backend_with_key(Some(key), now)
4125 .expect("three healthy backends");
4126 assert_eq!(
4127 picked.borrow().backend_id,
4128 "c",
4129 "{policy:?} compared the representative's own load, not its address's"
4130 );
4131 }
4132 }
4133 }
4134
4135 #[test]
4144 fn one_shared_pair_costs_a_selection_constant_work() {
4145 use crate::load_balancing::SIBLING_STEPS;
4146
4147 let now = Instant::now();
4148 let mut list = BackendList::with_seed(1856);
4151 for index in 0..1000u16 {
4152 let address = SocketAddr::from(([10, (index >> 8) as u8, index as u8, 1], 8080));
4153 list.add_backend(Backend::new(
4154 &format!("b{index}"),
4155 address,
4156 None,
4157 None,
4158 None,
4159 ));
4160 }
4161 let shared = SocketAddr::from(([10, 0, 0, 1], 8080));
4162 list.add_backend(Backend::new("dup", shared, None, None, None));
4163 for sharded in [false, true] {
4164 list.set_shuffle_sharding(sharded.then_some(ShuffleSharding {
4165 percent: 25,
4166 min_backends: 8,
4167 mode: ShardMode::Fallback,
4168 }));
4169 for policy in POLICIES {
4170 list.set_load_balancing_policy(policy, Some(LoadMetric::Connections), 1856);
4171 let mut worst = 0;
4172 for key in 0..64u64 {
4173 SIBLING_STEPS.with(|steps| steps.set(0));
4174 assert!(list.select_with_key(Some(key), now).0.is_some());
4175 worst = worst.max(SIBLING_STEPS.with(|steps| steps.get()));
4176 }
4177 assert!(
4178 worst <= 8,
4179 "{policy:?} (sharded: {sharded}) spent {worst} steps on one shared pair"
4180 );
4181 }
4182 }
4183 }
4184
4185 #[test]
4190 fn a_list_edited_directly_still_selects_without_panicking() {
4191 let now = Instant::now();
4192 let sharding = Some(ShuffleSharding {
4193 percent: 50,
4194 min_backends: 2,
4195 mode: ShardMode::Fallback,
4196 });
4197 for policy in POLICIES {
4198 for sharded in [false, true] {
4199 let mut grown = equal_list(policy);
4202 grown.set_shuffle_sharding(sharded.then_some(sharding).flatten());
4203 for (id, port) in [("pushed-sibling", 9300), ("pushed", 9302)] {
4204 let address = SocketAddr::from(([127, 0, 0, 1], port));
4205 grown
4206 .backends
4207 .push(Rc::new(RefCell::new(weighted(id, address, None))));
4208 }
4209 let mut shrunk = equal_list(policy);
4211 shrunk.set_shuffle_sharding(sharded.then_some(sharding).flatten());
4212 shrunk.backends.pop();
4213 for list in [&mut grown, &mut shrunk] {
4214 for key in 0..64 {
4215 assert!(
4216 list.select_with_key(Some(key), now).0.is_some(),
4217 "{policy:?} (sharded: {sharded}) selected nothing"
4218 );
4219 }
4220 }
4221 }
4222 }
4223 }
4224
4225 #[test]
4228 fn a_selection_over_a_shared_address_allocates_nothing() {
4229 use std::hint::black_box;
4230
4231 use crate::test_allocations::allocations;
4232
4233 let now = Instant::now();
4234 for policy in POLICIES {
4235 let mut list = equal_list(policy);
4236 drop(list.select_with_key(Some(0), now));
4237 let mut allocated = 0;
4238 for key in 0..256u64 {
4239 let before = allocations();
4240 let picked = black_box(&mut list).select_with_key(black_box(Some(key)), now);
4241 allocated += allocations() - before;
4242 assert!(picked.0.is_some());
4243 }
4244 assert_eq!(
4245 allocated, 0,
4246 "{policy:?}: 256 selections made {allocated} allocations"
4247 );
4248 }
4249 }
4250}
4251
4252#[cfg(test)]
4253mod unix_path_tests {
4254 use std::{
4255 io::{Read, Write},
4256 os::unix::net::UnixListener,
4257 path::PathBuf,
4258 time::Instant,
4259 };
4260
4261 use super::*;
4262
4263 fn socket_path(tag: &str) -> PathBuf {
4265 let path = std::env::temp_dir().join(format!("sozu-lib-{tag}-{}.sock", std::process::id()));
4266 let _ = std::fs::remove_file(&path);
4267 path
4268 }
4269
4270 fn assert_round_trip(stream: TcpStream, listener: &UnixListener) {
4273 let (mut accepted, _) = listener.accept().expect("the dial reaches the listener");
4274 let mut stream = std::net::TcpStream::from(std::os::fd::OwnedFd::from(stream));
4275 stream.set_nonblocking(false).unwrap();
4276 stream.write_all(b"ping").unwrap();
4277 let mut received = [0u8; 4];
4278 accepted.read_exact(&mut received).unwrap();
4279 assert_eq!(&received, b"ping");
4280 accepted.write_all(b"pong").unwrap();
4281 stream.read_exact(&mut received).unwrap();
4282 assert_eq!(&received, b"pong");
4283 }
4284
4285 #[test]
4286 fn a_unix_path_backend_dials_its_socket_not_its_address() {
4287 let path = socket_path("dial");
4288 let listener = UnixListener::bind(&path).unwrap();
4289 let unused: SocketAddr = "127.0.0.1:1".parse().unwrap();
4291 let stream = connect_backend(unused, path.to_str()).expect("dial the unix socket");
4292 assert_eq!(
4293 socket2::SockRef::from(&stream).domain().unwrap(),
4294 socket2::Domain::UNIX
4295 );
4296 assert_round_trip(stream, &listener);
4297 std::fs::remove_file(&path).unwrap();
4298 }
4299
4300 #[cfg(target_os = "linux")]
4301 #[test]
4302 fn an_abstract_unix_path_backend_dials_the_abstract_namespace() {
4303 use std::os::linux::net::SocketAddrExt;
4304 let name = format!("sozu-lib-abstract-{}", std::process::id());
4305 let address = std::os::unix::net::SocketAddr::from_abstract_name(name.as_bytes()).unwrap();
4306 let listener = UnixListener::bind_addr(&address).unwrap();
4307 let unused: SocketAddr = "127.0.0.1:1".parse().unwrap();
4308 let stream = connect_backend(unused, Some(&format!("@{name}"))).unwrap();
4309 assert_round_trip(stream, &listener);
4310 }
4311
4312 #[test]
4313 fn a_missing_unix_socket_is_a_failed_connect() {
4314 let mut backend = Backend::new("b", "127.0.0.1:1".parse().unwrap(), None, None, None)
4315 .with_unix_path(Some(socket_path("missing").to_str().unwrap().to_owned()));
4316 let mut rng = StdRng::seed_from_u64(0);
4317 let error = backend
4318 .try_connect(Instant::now(), &mut rng)
4319 .expect_err("no socket file means no connection");
4320 assert!(matches!(error, BackendError::MioConnection(_)));
4321 assert_eq!(backend.failures, 1);
4322 assert_eq!(backend.active_connections, 0);
4323 }
4324
4325 #[test]
4326 fn a_unix_backend_connects_through_try_connect() {
4327 let path = socket_path("try-connect");
4328 let listener = UnixListener::bind(&path).unwrap();
4329 let mut backend = Backend::new("b", "127.0.0.1:1".parse().unwrap(), None, None, None)
4330 .with_unix_path(Some(path.to_str().unwrap().to_owned()));
4331 let mut rng = StdRng::seed_from_u64(0);
4332 let stream = backend.try_connect(Instant::now(), &mut rng).unwrap();
4333 assert_eq!(backend.active_connections, 1);
4334 assert_eq!(backend.failures, 0);
4335 assert_round_trip(stream, &listener);
4336 std::fs::remove_file(&path).unwrap();
4337 }
4338
4339 #[test]
4340 fn re_adding_a_backend_updates_its_unix_path() {
4341 let address: SocketAddr = "127.0.0.1:1".parse().unwrap();
4342 let mut list = BackendList::new();
4343 list.add_backend(Backend::new("b", address, None, None, None));
4344 list.add_backend(
4345 Backend::new("b", address, None, None, None)
4346 .with_unix_path(Some("/run/b.sock".to_owned())),
4347 );
4348 assert_eq!(list.backends.len(), 1);
4349 assert_eq!(
4350 list.backends[0].borrow().unix_path.as_deref(),
4351 Some("/run/b.sock")
4352 );
4353 }
4354}