Skip to main content

sozu_lib/
backends.rs

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    state::ClusterId,
20};
21
22use crate::metrics::names;
23use crate::{
24    PeakEWMA,
25    load_balancing::{
26        Candidates, LeastLoaded, LoadBalancingAlgorithm, Maglev, PowerOfTwo, Random, Rendezvous,
27        RoundRobin, ShuffleSharding, count_sibling_step, hrw_score, outranks_at_address,
28    },
29    retry::{self, RetryPolicy},
30    server::{self, push_event},
31};
32
33#[derive(thiserror::Error, Debug)]
34pub enum BackendError {
35    #[error("No backend found for cluster {0}")]
36    NoBackendForCluster(String),
37    #[error("Failed to connect to socket with MIO: {0}")]
38    MioConnection(std::io::Error),
39    #[error("This backend is not in a normal status: status={0:?}")]
40    Status(BackendStatus),
41    #[error("could not connect {cluster_id} to {backend_address:?} ({failures} failures): {error}")]
42    ConnectionFailures {
43        cluster_id: String,
44        backend_address: SocketAddr,
45        failures: usize,
46        error: String,
47    },
48}
49
50#[derive(Debug, PartialEq, Eq, Clone)]
51pub enum BackendStatus {
52    Normal,
53    Closing,
54    Closed,
55}
56
57#[derive(Debug, PartialEq, Eq, Clone, Copy)]
58pub enum HealthStatus {
59    Healthy,
60    Unhealthy,
61}
62
63/// Per-cluster availability state, owned by `BackendList`. Flips between
64/// `Available` (≥1 backend can serve traffic) and `AllDown` (every backend
65/// fails the `health.is_healthy() && !retry_policy.is_down()` predicate)
66/// every time `BackendMap::record_cluster_availability` is invoked.
67/// Empty clusters never report `AllDown` to avoid log spam during cluster
68/// bootstrap.
69#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
70pub(crate) enum ClusterAvailability {
71    #[default]
72    Available,
73    AllDown,
74}
75
76#[derive(Debug, Clone, PartialEq)]
77pub struct HealthState {
78    pub status: HealthStatus,
79    pub consecutive_successes: u32,
80    pub consecutive_failures: u32,
81}
82
83impl Default for HealthState {
84    fn default() -> Self {
85        HealthState {
86            status: HealthStatus::Healthy,
87            consecutive_successes: 0,
88            consecutive_failures: 0,
89        }
90    }
91}
92
93impl HealthState {
94    /// Record a successful health check. Returns true if the backend transitioned to healthy.
95    pub fn record_success(&mut self, healthy_threshold: u32) -> bool {
96        let was_unhealthy = self.status == HealthStatus::Unhealthy;
97        let successes_before = self.consecutive_successes;
98        self.consecutive_failures = 0;
99        self.consecutive_successes += 1;
100
101        // A success resets the failure streak and advances the success streak by
102        // exactly one — the two counters are never both non-zero afterwards.
103        debug_assert_eq!(
104            self.consecutive_failures, 0,
105            "a success must clear the consecutive-failure streak"
106        );
107        debug_assert_eq!(
108            self.consecutive_successes,
109            successes_before + 1,
110            "a success must advance the success streak by exactly one"
111        );
112
113        if was_unhealthy && self.consecutive_successes >= healthy_threshold {
114            self.status = HealthStatus::Healthy;
115            // The transition is only reported when crossing the threshold from
116            // Unhealthy; a backend that was already Healthy never "transitions".
117            debug_assert!(
118                self.status == HealthStatus::Healthy,
119                "a reported recovery must leave the status Healthy"
120            );
121            return true;
122        }
123        false
124    }
125
126    /// Record a failed health check. Returns true if the backend transitioned to unhealthy.
127    pub fn record_failure(&mut self, unhealthy_threshold: u32) -> bool {
128        let was_healthy = self.status == HealthStatus::Healthy;
129        let failures_before = self.consecutive_failures;
130        self.consecutive_successes = 0;
131        self.consecutive_failures += 1;
132
133        // A failure resets the success streak and advances the failure streak by
134        // exactly one — symmetric with `record_success`.
135        debug_assert_eq!(
136            self.consecutive_successes, 0,
137            "a failure must clear the consecutive-success streak"
138        );
139        debug_assert_eq!(
140            self.consecutive_failures,
141            failures_before + 1,
142            "a failure must advance the failure streak by exactly one"
143        );
144
145        if was_healthy && self.consecutive_failures >= unhealthy_threshold {
146            self.status = HealthStatus::Unhealthy;
147            debug_assert!(
148                self.status == HealthStatus::Unhealthy,
149                "a reported drop must leave the status Unhealthy"
150            );
151            return true;
152        }
153        false
154    }
155
156    pub fn is_healthy(&self) -> bool {
157        self.status == HealthStatus::Healthy
158    }
159}
160
161#[derive(Debug, PartialEq, Clone)]
162pub struct Backend {
163    pub sticky_id: Option<String>,
164    pub backend_id: String,
165    pub address: SocketAddr,
166    pub status: BackendStatus,
167    pub retry_policy: retry::RetryPolicyWrapper,
168    pub active_connections: usize,
169    pub active_requests: usize,
170    pub failures: usize,
171    pub load_balancing_parameters: Option<LoadBalancingParams>,
172    pub backup: bool,
173    pub connection_time: PeakEWMA,
174    pub health: HealthState,
175}
176
177impl Backend {
178    /// A backend created now, on the wall clock.
179    ///
180    /// The production default: the one place backend construction reads the
181    /// clock, because the connection-time average decays from the instant the
182    /// backend was created. A simulator that must not read the wall clock
183    /// uses [`Self::new_at`].
184    pub fn new(
185        backend_id: &str,
186        address: SocketAddr,
187        sticky_id: Option<String>,
188        load_balancing_parameters: Option<LoadBalancingParams>,
189        backup: Option<bool>,
190    ) -> Backend {
191        Self::new_at(
192            backend_id,
193            address,
194            sticky_id,
195            load_balancing_parameters,
196            backup,
197            Instant::now(),
198        )
199    }
200
201    /// A backend created at `now` (#1684).
202    pub fn new_at(
203        backend_id: &str,
204        address: SocketAddr,
205        sticky_id: Option<String>,
206        load_balancing_parameters: Option<LoadBalancingParams>,
207        backup: Option<bool>,
208        now: Instant,
209    ) -> Backend {
210        let desired_policy = retry::ExponentialBackoffPolicy::new(6);
211        Backend {
212            sticky_id,
213            backend_id: backend_id.to_owned(),
214            address,
215            status: BackendStatus::Normal,
216            retry_policy: desired_policy.into(),
217            active_connections: 0,
218            active_requests: 0,
219            failures: 0,
220            load_balancing_parameters,
221            backup: backup.unwrap_or(false),
222            connection_time: PeakEWMA::new(now),
223            health: HealthState::default(),
224        }
225    }
226
227    pub fn set_closing(&mut self) {
228        self.status = BackendStatus::Closing;
229    }
230
231    pub fn retry_policy(&mut self) -> &mut retry::RetryPolicyWrapper {
232        &mut self.retry_policy
233    }
234
235    /// Whether this backend can take a new connection at `now`.
236    pub fn can_open(&self, now: Instant) -> bool {
237        if !self.health.is_healthy() {
238            return false;
239        }
240        if let Some(action) = self.retry_policy.can_try(now) {
241            self.status == BackendStatus::Normal && action == retry::RetryAction::OKAY
242        } else {
243            false
244        }
245    }
246
247    /// Canonical "available" check used by per-backend metrics and cluster
248    /// availability accounting. Slightly more permissive than `can_open()`:
249    /// a backend currently in an exponential-backoff *wait window* still
250    /// counts as available because the next call after the window ends
251    /// will route to it without operator intervention. The dashboard
252    /// reading must reflect "operationally up, not exhausted" rather than
253    /// "ready to receive *this* request" — flicking the gauge to 0 on
254    /// every transient backoff would drown out genuine `is_down()`
255    /// transitions. Pairs with `BackendList::evaluate_availability`,
256    /// which applies the same predicate cluster-wide.
257    pub fn is_available(&self) -> bool {
258        self.health.is_healthy()
259            && self.status == BackendStatus::Normal
260            && !self.retry_policy.is_down()
261    }
262
263    pub fn inc_connections(&mut self) -> Option<usize> {
264        let before = self.active_connections;
265        if self.status == BackendStatus::Normal {
266            self.active_connections += 1;
267            // A `Normal` backend always increments by exactly one and reports the
268            // post-increment count back to the caller.
269            debug_assert_eq!(
270                self.active_connections,
271                before + 1,
272                "inc_connections must add exactly one active connection"
273            );
274            Some(self.active_connections)
275        } else {
276            // Non-`Normal` backends refuse new connections and leave the gauge
277            // untouched (no silent increment on a Closing/Closed backend).
278            debug_assert_eq!(
279                self.active_connections, before,
280                "inc_connections must not touch the count for a non-Normal backend"
281            );
282            None
283        }
284    }
285
286    /// TODO: normalize with saturating_sub()
287    pub fn dec_connections(&mut self) -> Option<usize> {
288        let before = self.active_connections;
289        match self.status {
290            BackendStatus::Normal => {
291                if self.active_connections > 0 {
292                    self.active_connections -= 1;
293                }
294                // The count drops by one when positive, otherwise saturates at
295                // zero — it must never wrap below zero (usize underflow).
296                debug_assert!(
297                    self.active_connections <= before,
298                    "dec_connections must never increase the active-connection count"
299                );
300                debug_assert_eq!(
301                    self.active_connections,
302                    before.saturating_sub(1),
303                    "dec_connections must drop by exactly one (saturating at zero)"
304                );
305                Some(self.active_connections)
306            }
307            BackendStatus::Closed => {
308                // A Closed backend has already been retired: nothing to decrement.
309                debug_assert_eq!(
310                    self.active_connections, before,
311                    "dec_connections on a Closed backend must not mutate the count"
312                );
313                None
314            }
315            BackendStatus::Closing => {
316                if self.active_connections > 0 {
317                    self.active_connections -= 1;
318                }
319                debug_assert_eq!(
320                    self.active_connections,
321                    before.saturating_sub(1),
322                    "dec_connections must drop by exactly one (saturating at zero)"
323                );
324                if self.active_connections == 0 {
325                    self.status = BackendStatus::Closed;
326                    // Draining a Closing backend to zero retires it: the
327                    // lifecycle advances to Closed and we stop reporting a count.
328                    debug_assert_eq!(
329                        self.status,
330                        BackendStatus::Closed,
331                        "a fully drained Closing backend must become Closed"
332                    );
333                    None
334                } else {
335                    Some(self.active_connections)
336                }
337            }
338        }
339    }
340
341    /// Reserve a connection on this backend for a dial about to start.
342    ///
343    /// `active_connections += 1`, exactly the count a successful
344    /// [`Self::try_connect`] takes, but before the connect: selection takes it
345    /// under the same borrow as the choice (`BackendMap::reserve_backend`).
346    /// Only a `Normal` backend can be chosen, and only a `Normal` backend
347    /// accepts the reservation.
348    pub fn reserve_connection(&mut self) -> Result<(), BackendError> {
349        match self.inc_connections() {
350            Some(_) => Ok(()),
351            None => Err(BackendError::Status(self.status.to_owned())),
352        }
353    }
354
355    /// Release the reservation of a dial whose connect failed at `now`, and
356    /// record the failure: `failures += 1` and the retry policy's backoff,
357    /// whose jitter is drawn from `rng`.
358    ///
359    /// Leaves the backend where a failed [`Self::try_connect`] leaves it: the
360    /// reservation it releases is the count that call never took.
361    pub fn release_failed_dial<R: Rng + ?Sized>(&mut self, now: Instant, rng: &mut R) {
362        let (connections_before, failures_before) = (self.active_connections, self.failures);
363        self.dec_connections();
364        self.failures += 1;
365        self.retry_policy.fail(now, rng);
366        debug_assert_eq!(
367            self.active_connections,
368            connections_before.saturating_sub(1),
369            "a failed dial must release exactly its one reservation (saturating at 0)"
370        );
371        debug_assert_eq!(
372            self.failures,
373            failures_before + 1,
374            "a failed dial must advance the failure counter by exactly one"
375        );
376    }
377
378    /// Record that a connection took `dur` to establish, observed at `now`.
379    pub fn set_connection_time(&mut self, dur: Duration, now: Instant) {
380        self.connection_time.observe(dur.as_nanos() as f64, now);
381    }
382
383    /// The connection-time cost of this backend, as seen at `now`.
384    pub fn peak_ewma_connection(&mut self, now: Instant) -> f64 {
385        self.connection_time.get(self.active_connections, now)
386    }
387
388    /// Open a connection to this backend. A failure at `now` arms the retry
389    /// policy, whose backoff jitter is drawn from `rng`.
390    pub fn try_connect<R: Rng + ?Sized>(
391        &mut self,
392        now: Instant,
393        rng: &mut R,
394    ) -> Result<mio::net::TcpStream, BackendError> {
395        if self.status != BackendStatus::Normal {
396            return Err(BackendError::Status(self.status.to_owned()));
397        }
398        // Reaching the connect attempt implies we passed the status gate; the
399        // failure counter is whatever prior attempts accumulated.
400        debug_assert_eq!(
401            self.status,
402            BackendStatus::Normal,
403            "try_connect only attempts a connection on a Normal backend"
404        );
405        let failures_before = self.failures;
406        let connections_before = self.active_connections;
407
408        match mio::net::TcpStream::connect(self.address) {
409            Ok(tcp_stream) => {
410                //self.retry_policy.succeed();
411                self.inc_connections();
412                // Success registers exactly one new active connection and never
413                // touches the failure counter.
414                debug_assert_eq!(
415                    self.active_connections,
416                    connections_before + 1,
417                    "a successful connect must register exactly one active connection"
418                );
419                debug_assert_eq!(
420                    self.failures, failures_before,
421                    "a successful connect must not bump the failure counter"
422                );
423                Ok(tcp_stream)
424            }
425            Err(io_error) => {
426                self.retry_policy.fail(now, rng);
427                self.failures += 1;
428                // A failed connect arms the retry policy and advances the
429                // failure counter by exactly one, leaving the connection gauge
430                // untouched (no connection was established).
431                debug_assert_eq!(
432                    self.failures,
433                    failures_before + 1,
434                    "a failed connect must advance the failure counter by exactly one"
435                );
436                debug_assert_eq!(
437                    self.active_connections, connections_before,
438                    "a failed connect must not register an active connection"
439                );
440                // TODO: handle EINPROGRESS. It is difficult. It is discussed here:
441                // https://docs.rs/mio/latest/mio/net/struct.TcpStream.html#method.connect
442                // with an example code here:
443                // https://github.com/Thomasdezeeuw/heph/blob/0c4f1ab3eaf08bea1d65776528bfd6114c9f8374/src/net/tcp/stream.rs#L560-L622
444                Err(BackendError::MioConnection(io_error))
445            }
446        }
447    }
448}
449
450// when a backend has been removed from configuration and the last connection to
451// it has stopped, it will be dropped, so we can notify that the backend server
452// can be safely stopped
453impl std::ops::Drop for Backend {
454    fn drop(&mut self) {
455        server::push_event(Event {
456            kind: EventKind::RemovedBackendHasNoConnections as i32,
457            backend_id: Some(self.backend_id.to_owned()),
458            address: Some(self.address.into()),
459            cluster_id: None,
460            metric_detail: None,
461        });
462    }
463}
464
465#[derive(Debug)]
466pub struct BackendMap {
467    pub backends: HashMap<ClusterId, BackendList>,
468    pub health_check_configs: HashMap<ClusterId, HealthCheckConfig>,
469    /// Whether the cluster's backends speak HTTP/2 (cluster.http2 = true).
470    /// Mirrors the same backend-capability hint the mux router reads at
471    /// `protocol/mux/router.rs::Router::plan_connect`. The health checker uses
472    /// it to switch the probe wire format from HTTP/1.1 to h2c so an
473    /// h2c-only backend is not probed with an HTTP/1.1 preface that
474    /// would always fail.
475    pub cluster_http2: HashMap<ClusterId, bool>,
476    /// The one source of randomness selection and backoff draw from: it seeds
477    /// every cluster's load-balancing policy and supplies the jitter of every
478    /// backoff window a failed connection arms.
479    ///
480    /// [`Self::new`] seeds it from the OS once, so each worker draws an
481    /// independent keystream, as the thread-local RNG it replaces did;
482    /// [`Self::with_seed`] makes the whole map reproducible for a simulator
483    /// (#1684). Taken at construction, never on the datapath.
484    rng: StdRng,
485}
486
487/// A generator seeded from the OS, for the production defaults of
488/// [`BackendMap::new`] and [`BackendList::new`].
489fn os_rng() -> StdRng {
490    StdRng::try_from_rng(&mut SysRng).expect("failed to seed random number generator from system")
491}
492
493impl Default for BackendMap {
494    fn default() -> Self {
495        Self::new()
496    }
497}
498
499impl BackendMap {
500    /// A backend map whose randomness is seeded from the OS.
501    pub fn new() -> BackendMap {
502        Self::with_rng(os_rng())
503    }
504
505    /// A backend map whose every policy seed and backoff jitter derives from
506    /// `seed`, for a deterministic simulator (#1684).
507    pub fn with_seed(seed: u64) -> BackendMap {
508        Self::with_rng(StdRng::seed_from_u64(seed))
509    }
510
511    fn with_rng(rng: StdRng) -> BackendMap {
512        BackendMap {
513            backends: HashMap::new(),
514            health_check_configs: HashMap::new(),
515            cluster_http2: HashMap::new(),
516            rng,
517        }
518    }
519
520    /// The map's generator, for a caller that arms a backend's retry policy
521    /// outside [`Self::backend_from_cluster_id`] and needs its jitter.
522    pub fn rng(&mut self) -> &mut StdRng {
523        &mut self.rng
524    }
525
526    /// Re-evaluate the availability of `cluster_id`, publish the
527    /// `cluster.available_backends` / `cluster.total_backends` gauges,
528    /// and emit the transition log + counter + `Event` exactly when the
529    /// per-cluster state flips between `Available` and `AllDown`.
530    ///
531    /// Empty clusters (`total == 0`) never report `AllDown` — avoids log
532    /// spam during cluster bootstrap when backends are still being
533    /// registered. The (0, 0) gauges are still published so dashboards
534    /// see "cluster exists, zero backends configured" as a distinct
535    /// state from "cluster doesn't exist".
536    ///
537    /// Takes `&self` so callers that already hold `&mut BackendMap`
538    /// can drop their `&mut BackendList` borrow before invoking it
539    /// without re-borrowing.
540    pub(crate) fn record_cluster_availability(&self, cluster_id: &str) {
541        let Some(list) = self.backends.get(cluster_id) else {
542            return;
543        };
544
545        let (available, total) = list.evaluate_availability();
546        // A subset count can never exceed the whole, and it must match the
547        // live backend vector length the helper just walked.
548        debug_assert!(
549            available <= total,
550            "available backends ({available}) cannot exceed total ({total})"
551        );
552        debug_assert_eq!(
553            total,
554            list.backends.len(),
555            "total must equal the number of registered backends"
556        );
557        gauge!(
558            names::cluster::AVAILABLE_BACKENDS,
559            available,
560            Some(cluster_id),
561            None
562        );
563        gauge!(
564            names::cluster::TOTAL_BACKENDS,
565            total,
566            Some(cluster_id),
567            None
568        );
569
570        let new_state = if total > 0 && available == 0 {
571            ClusterAvailability::AllDown
572        } else {
573            ClusterAvailability::Available
574        };
575        // Empty clusters never report AllDown (avoids bootstrap log spam); a
576        // cluster with at least one available backend is always Available.
577        debug_assert!(
578            !(total == 0 && new_state == ClusterAvailability::AllDown),
579            "an empty cluster must never be reported AllDown"
580        );
581        debug_assert!(
582            !(available > 0 && new_state == ClusterAvailability::AllDown),
583            "a cluster with an available backend must not be AllDown"
584        );
585
586        let prev = list.availability.replace(new_state);
587        // The cell now holds exactly the freshly computed state.
588        debug_assert_eq!(
589            list.availability.get(),
590            new_state,
591            "the availability cell must latch the newly computed state"
592        );
593        if prev == new_state {
594            return;
595        }
596        match (prev, new_state) {
597            (ClusterAvailability::Available, ClusterAvailability::AllDown) => {
598                error!("cluster {}: all {} backends are down", cluster_id, total);
599                incr!(
600                    names::cluster::NO_AVAILABLE_BACKENDS,
601                    Some(cluster_id),
602                    None
603                );
604                push_event(Event {
605                    kind: EventKind::NoAvailableBackends as i32,
606                    cluster_id: Some(cluster_id.to_owned()),
607                    backend_id: None,
608                    address: None,
609                    metric_detail: None,
610                });
611            }
612            (ClusterAvailability::AllDown, ClusterAvailability::Available) => {
613                info!(
614                    "cluster {}: backends recovered ({}/{} available)",
615                    cluster_id, available, total
616                );
617                incr!(names::cluster::AVAILABLE_RECOVERED, Some(cluster_id), None);
618                push_event(Event {
619                    kind: EventKind::ClusterRecovered as i32,
620                    cluster_id: Some(cluster_id.to_owned()),
621                    backend_id: None,
622                    address: None,
623                    metric_detail: None,
624                });
625            }
626            _ => {}
627        }
628    }
629
630    /// Forget everything the map holds for `cluster_id`: its backend list,
631    /// with the load-balancing policy and shuffle sharding it carries, its
632    /// health-check configuration and its `http2` hint. A session that already
633    /// holds one of its backends keeps it until the session closes;
634    /// [`Self::close_backend_connection`] ignores a cluster that is gone.
635    /// Every backend it drops is marked `Closing`, as `BackendList::remove_backend`
636    /// marks the ones it drops, so a connection pooled on it drains the
637    /// requests it carries and takes no new one, even once a cluster of the
638    /// same id is added again.
639    pub fn remove_cluster(&mut self, cluster_id: &str) {
640        if let Some(list) = self.backends.remove(cluster_id) {
641            for backend in &list.backends {
642                backend.borrow_mut().set_closing();
643            }
644        }
645        self.health_check_configs.remove(cluster_id);
646        self.cluster_http2.remove(cluster_id);
647        debug_assert!(
648            !self.backends.contains_key(cluster_id)
649                && !self.health_check_configs.contains_key(cluster_id)
650                && !self.cluster_http2.contains_key(cluster_id),
651            "remove_cluster must leave nothing keyed by the cluster"
652        );
653    }
654
655    /// Record (or clear) the `cluster.http2` backend-capability hint for
656    /// `cluster_id`. The health checker reads the resulting map at probe
657    /// time so the wire format follows what the mux router will use to
658    /// connect to the same backends.
659    pub fn set_cluster_http2(&mut self, cluster_id: &str, http2: bool) {
660        if http2 {
661            self.cluster_http2.insert(cluster_id.into(), true);
662        } else {
663            self.cluster_http2.remove(cluster_id);
664        }
665    }
666
667    pub fn set_health_check_config(&mut self, cluster_id: &str, config: Option<HealthCheckConfig>) {
668        match config {
669            Some(c) => {
670                self.health_check_configs.insert(cluster_id.into(), c);
671            }
672            None => {
673                self.health_check_configs.remove(cluster_id);
674                // When the operator drops the health check, any
675                // previously-recorded `HealthState::Unhealthy` would
676                // otherwise stick — `next_available_backend` keeps
677                // skipping the backend even though we have stopped
678                // probing it. Reset every backend in the cluster to a
679                // pristine healthy state so the load balancer can
680                // route again.
681                if let Some(backend_list) = self.backends.get(cluster_id) {
682                    for backend in &backend_list.backends {
683                        backend.borrow_mut().health = HealthState::default();
684                    }
685                }
686                // Re-emit the rollup gauges so dashboards reflect the
687                // post-reset availability instead of holding the last
688                // health-check value indefinitely.
689                self.record_cluster_availability(cluster_id);
690            }
691        }
692    }
693
694    pub fn import_configuration_state(
695        &mut self,
696        backends: &HashMap<ClusterId, Vec<sozu_command::response::Backend>>,
697    ) {
698        // Seed the clusters in id order, not in the `HashMap`'s: its
699        // `RandomState` iterates differently in every process, so drawing in
700        // that order would hand a seeded map's seeds to different clusters on
701        // every run (#1684). Sorting the ids allocates once, on this
702        // control-plane replay path only.
703        let mut cluster_ids: Vec<&ClusterId> = backends.keys().collect();
704        cluster_ids.sort_unstable();
705        for cluster_id in cluster_ids {
706            let list =
707                BackendList::import_configuration_state(&backends[cluster_id], self.rng.next_u64());
708            self.backends.insert(cluster_id.clone(), list);
709        }
710        // Replay path inserts every cluster's backend list without
711        // touching the gauge emission sites used by add/remove/health.
712        // Latch `cluster.available_backends` and `.total_backends` here
713        // so a freshly-loaded worker reports correct values on the very
714        // first `QueryMetrics` instead of zero until something else
715        // mutates each cluster.
716        for cluster_id in backends.keys() {
717            self.record_cluster_availability(cluster_id);
718        }
719    }
720
721    pub fn add_backend(&mut self, cluster_id: &str, backend: Backend) {
722        let address = backend.address;
723        self.get_or_create_backend_list_for_cluster(cluster_id)
724            .add_backend(backend);
725        // Adding a backend must leave the cluster present and containing the
726        // just-added address (whether it created the entry or updated in place).
727        debug_assert!(
728            self.backends
729                .get(cluster_id)
730                .is_some_and(|list| list.has_backend(&address)),
731            "add_backend must leave the backend present in its cluster"
732        );
733        // Publish initial gauges and surface the corner case where a fresh
734        // cluster's first backend is already down (e.g. registered with a
735        // pre-existing failed retry policy). For an `Available` initial
736        // backend this is just a (1, 1) gauge emission with no transition.
737        self.record_cluster_availability(cluster_id);
738    }
739
740    // TODO: return <Result, BackendError>, log the error downstream
741    /// Remove the backend identified by `(backend_id, backend_address)` from
742    /// `cluster_id` and return whether it was present. A backend of another
743    /// id at the same address stays (#1821).
744    pub fn remove_backend(
745        &mut self,
746        cluster_id: &str,
747        backend_id: &str,
748        backend_address: &SocketAddr,
749    ) -> bool {
750        let removed = if let Some(backends) = self.backends.get_mut(cluster_id) {
751            let removed = backends.remove_backend(backend_id, backend_address);
752            if !removed {
753                warn!(
754                    "No backend matches id {} at address {:?} in cluster {}: nothing removed",
755                    backend_id, backend_address, cluster_id
756                );
757            }
758            removed
759        } else {
760            error!(
761                "Backend was already removed: cluster id {}, backend id {}, address {:?}",
762                cluster_id, backend_id, backend_address
763            );
764            return false;
765        };
766        debug_assert!(
767            self.backends.get(cluster_id).is_none_or(|list| list
768                .find_backend_by_identity(backend_id, backend_address)
769                .is_none()),
770            "remove_backend must evict the backend it names"
771        );
772        // Re-evaluate so removing the last backend logs an explicit
773        // `AllDown` transition (or, with `total == 0`, drops back to
774        // silent gauges).
775        self.record_cluster_availability(cluster_id);
776        removed
777    }
778
779    // TODO: return <Result, BackendError>, log the error downstream
780    pub fn close_backend_connection(&mut self, cluster_id: &str, addr: &SocketAddr) {
781        if let Some(cluster_backends) = self.backends.get_mut(cluster_id)
782            && let Some(ref mut backend) = cluster_backends.find_backend(addr)
783        {
784            backend.borrow_mut().dec_connections();
785        }
786    }
787
788    /// Whether `cluster_id` still holds a backend carrying `backend_id`, at
789    /// any address. Backend metrics are labelled by id alone, so every entry
790    /// sharing an id feeds one metrics row.
791    pub fn has_backend_id(&self, cluster_id: &str, backend_id: &str) -> bool {
792        self.backends.get(cluster_id).is_some_and(|backends| {
793            backends
794                .backends
795                .iter()
796                .any(|backend| backend.borrow().backend_id == backend_id)
797        })
798    }
799
800    pub fn has_backend(&self, cluster_id: &str, backend: &Backend) -> bool {
801        self.backends
802            .get(cluster_id)
803            .map(|backends| backends.has_backend(&backend.address))
804            .unwrap_or(false)
805    }
806
807    /// Select a backend of `cluster_id` at `now` and connect to it.
808    ///
809    /// The TCP proxy's entry point: it selects and connects in one call, and
810    /// `Backend::try_connect` counts the connection once it is established.
811    /// The HTTP mux reserves with [`Self::reserve_backend`] instead and
812    /// connects on its own side.
813    ///
814    /// `key` is the client's affinity key, which `HRW` and `MAGLEV` pin the
815    /// client on and every other policy ignores: the HTTP, HTTPS and TCP
816    /// datapaths pass the one they derived for a cluster using either
817    /// policy, and `None` otherwise (see
818    /// [`BackendList::next_available_backend_with_key`]).
819    pub fn backend_from_cluster_id(
820        &mut self,
821        cluster_id: &str,
822        key: Option<u64>,
823        now: Instant,
824    ) -> Result<(Rc<RefCell<Backend>>, TcpStream), BackendError> {
825        let next_backend = self.select_backend(cluster_id, key, now, &[])?;
826
827        let tcp_stream = {
828            let mut borrowed_backend = next_backend.borrow_mut();
829
830            debug!(
831                "Connecting {} -> {:?}",
832                cluster_id,
833                (
834                    borrowed_backend.address,
835                    borrowed_backend.active_connections,
836                    borrowed_backend.failures
837                )
838            );
839
840            borrowed_backend
841                .try_connect(now, &mut self.rng)
842                .map_err(|backend_error| BackendError::ConnectionFailures {
843                    cluster_id: cluster_id.to_owned(),
844                    backend_address: borrowed_backend.address,
845                    failures: borrowed_backend.failures,
846                    error: backend_error.to_string(),
847                })?
848        };
849
850        // Connection succeeded: re-evaluate so we capture an
851        // AllDown -> Available recovery transition the moment a request
852        // first hits a healthy backend after an outage. `next_backend` is
853        // not borrowed here (the inner block dropped `borrowed_backend`),
854        // so the helper's `BackendList::evaluate_availability` walk is
855        // free to call `borrow()` on every backend.
856        self.record_cluster_availability(cluster_id);
857
858        // The selected backend is a live member of the cluster it was drawn
859        // from — selection never fabricates or returns a stale backend.
860        debug_assert!(
861            self.backends.get(cluster_id).is_some_and(|list| {
862                let picked = next_backend.borrow().address;
863                list.has_backend(&picked)
864            }),
865            "the selected backend must belong to the cluster's live set"
866        );
867
868        Ok((next_backend.clone(), tcp_stream))
869    }
870
871    /// Select a backend of `cluster_id` at `now` and reserve a connection on
872    /// it, without connecting (#1684).
873    ///
874    /// The reservation is `active_connections += 1`, taken under the same
875    /// borrow as the choice, so the next selection weighs it exactly as it
876    /// weighed the connection `Backend::try_connect` used to count at the
877    /// dial. The caller connects. A `connect(2)` that fails releases the
878    /// reservation through [`Backend::release_failed_dial`], which also
879    /// records the failure. A dial that connects but is abandoned before its
880    /// connection is registered releases it with a plain
881    /// `Backend::dec_connections`, recording no failure (the mux's
882    /// `BackendChange::ConnectionClosed`, #1713). A connection that starts
883    /// keeps it, and its close releases it like any other.
884    ///
885    /// `key` is the client's affinity key, as for
886    /// [`Self::backend_from_cluster_id`].
887    pub fn reserve_backend(
888        &mut self,
889        cluster_id: &str,
890        key: Option<u64>,
891        now: Instant,
892    ) -> Result<Rc<RefCell<Backend>>, BackendError> {
893        self.reserve_backend_excluding(cluster_id, key, now, &[])
894    }
895
896    /// [`Self::reserve_backend`], preferring a backend whose address is not in
897    /// `exclude`: the backends a request already failed to connect to
898    /// (sozu-proxy/sozu#1800). See `BackendList::select_with_key_excluding`
899    /// for what happens when every selectable backend is excluded.
900    pub fn reserve_backend_excluding(
901        &mut self,
902        cluster_id: &str,
903        key: Option<u64>,
904        now: Instant,
905        exclude: &[SocketAddr],
906    ) -> Result<Rc<RefCell<Backend>>, BackendError> {
907        let backend = self.select_backend(cluster_id, key, now, exclude)?;
908        backend.borrow_mut().reserve_connection()?;
909        // Re-evaluate on a successful selection, as the connecting path does
910        // on a successful connect, so an AllDown -> Available recovery is
911        // reported the moment a request reaches a backend again.
912        self.record_cluster_availability(cluster_id);
913        Ok(backend)
914    }
915
916    /// Reserve a connection on the backend of `cluster_id` that
917    /// `sticky_session` names, if it can take one at `now`, falling back to
918    /// [`Self::reserve_backend`] under the client's affinity `key`.
919    ///
920    /// A cookie naming a live backend wins over the key: the key only decides
921    /// where a client without a usable cookie lands.
922    pub fn reserve_sticky_backend(
923        &mut self,
924        cluster_id: &str,
925        sticky_session: &str,
926        key: Option<u64>,
927        now: Instant,
928    ) -> Result<Rc<RefCell<Backend>>, BackendError> {
929        self.reserve_sticky_backend_excluding(cluster_id, sticky_session, key, now, &[])
930    }
931
932    /// [`Self::reserve_sticky_backend`], except that a cookie naming a backend
933    /// in `exclude` is not followed: the request already failed to connect to
934    /// it, so it goes to [`Self::reserve_backend_excluding`] instead
935    /// (sozu-proxy/sozu#1800).
936    pub fn reserve_sticky_backend_excluding(
937        &mut self,
938        cluster_id: &str,
939        sticky_session: &str,
940        key: Option<u64>,
941        now: Instant,
942        exclude: &[SocketAddr],
943    ) -> Result<Rc<RefCell<Backend>>, BackendError> {
944        let sticky = self
945            .backends
946            .get_mut(cluster_id)
947            .and_then(|cluster_backends| cluster_backends.find_sticky(sticky_session, now))
948            .filter(|backend| !exclude.contains(&backend.borrow().address))
949            .cloned();
950        match sticky {
951            Some(backend) => {
952                backend.borrow_mut().reserve_connection()?;
953                Ok(backend)
954            }
955            None => {
956                debug!(
957                    "Couldn't find a backend corresponding to sticky_session {} for cluster {}",
958                    sticky_session, cluster_id
959                );
960                self.reserve_backend_excluding(cluster_id, key, now, exclude)
961            }
962        }
963    }
964
965    /// Choose a backend of `cluster_id` at `now`, publishing the cluster's
966    /// availability when nothing can be chosen.
967    fn select_backend(
968        &mut self,
969        cluster_id: &str,
970        key: Option<u64>,
971        now: Instant,
972        exclude: &[SocketAddr],
973    ) -> Result<Rc<RefCell<Backend>>, BackendError> {
974        let cluster_backends = self
975            .backends
976            .get_mut(cluster_id)
977            .ok_or(BackendError::NoBackendForCluster(cluster_id.to_owned()))?;
978
979        if cluster_backends.backends.is_empty() {
980            // Drop the &mut BackendList borrow before the &self helper call.
981            // `total == 0` falls into the "never report AllDown" branch in
982            // record_cluster_availability, so this just publishes the (0, 0)
983            // gauges.
984            let _ = cluster_backends;
985            self.record_cluster_availability(cluster_id);
986            return Err(BackendError::NoBackendForCluster(cluster_id.to_owned()));
987        }
988        // Past the empty guard there is at least one backend to pick from.
989        debug_assert!(
990            !cluster_backends.backends.is_empty(),
991            "selection runs only on a non-empty backend list"
992        );
993
994        let (picked, outcome) = cluster_backends.select_with_key_excluding(key, now, exclude);
995        record_shard_outcome(cluster_id, outcome, picked.is_some());
996        match picked {
997            Some(backend) => Ok(backend),
998            None => {
999                // Drop the &mut BackendList before the &self helper call.
1000                // The helper observes (available=0, total>0) and emits the
1001                // Available -> AllDown transition (log + counter + Event)
1002                // exactly once per regime entry. Subsequent calls in the
1003                // same AllDown regime are no-ops.
1004                let _ = cluster_backends;
1005                self.record_cluster_availability(cluster_id);
1006                Err(BackendError::NoBackendForCluster(cluster_id.to_owned()))
1007            }
1008        }
1009    }
1010
1011    /// Select a backend for `cluster_id`, optionally pinned by an affinity
1012    /// `key`, and return its `(backend_id, address)` **without** opening any
1013    /// connection.
1014    ///
1015    /// This is the UDP datapath's selection entry point. Unlike
1016    /// [`backend_from_cluster_id`](Self::backend_from_cluster_id) — which is
1017    /// TCP-specific because it calls `Backend::try_connect` and hands back a
1018    /// `TcpStream` — UDP owns its own per-flow connected `UdpSocket` (created in
1019    /// the shell via `socket::udp_connect`), so all the map needs to surface is
1020    /// the chosen endpoint identity. `key` is `Some(flow_hash)` so HRW/Maglev
1021    /// keep a client flow pinned to one backend; `None` behaves like the legacy
1022    /// round-robin selection. Fail-open (all-unhealthy ⇒ LB over the full set)
1023    /// is inherited from [`BackendList::next_available_backend_with_key`].
1024    pub fn backend_from_cluster_id_with_key(
1025        &mut self,
1026        cluster_id: &str,
1027        key: Option<u64>,
1028        now: Instant,
1029    ) -> Result<(String, SocketAddr), BackendError> {
1030        let cluster_backends = self
1031            .backends
1032            .get_mut(cluster_id)
1033            .ok_or(BackendError::NoBackendForCluster(cluster_id.to_owned()))?;
1034
1035        if cluster_backends.backends.is_empty() {
1036            let _ = cluster_backends;
1037            self.record_cluster_availability(cluster_id);
1038            return Err(BackendError::NoBackendForCluster(cluster_id.to_owned()));
1039        }
1040        debug_assert!(
1041            !cluster_backends.backends.is_empty(),
1042            "keyed selection runs only on a non-empty backend list"
1043        );
1044
1045        let (picked, outcome) = cluster_backends.select_with_key(key, now);
1046        record_shard_outcome(cluster_id, outcome, picked.is_some());
1047        let next_backend = match picked {
1048            Some(nb) => nb,
1049            None => {
1050                let _ = cluster_backends;
1051                self.record_cluster_availability(cluster_id);
1052                return Err(BackendError::NoBackendForCluster(cluster_id.to_owned()));
1053            }
1054        };
1055
1056        let (backend_id, address) = {
1057            let borrowed = next_backend.borrow();
1058            (borrowed.backend_id.to_owned(), borrowed.address)
1059        };
1060        // The keyed selection returns a live member of the cluster — the
1061        // surfaced identity (id, address) belongs to a registered backend.
1062        debug_assert!(
1063            cluster_backends.has_backend(&address),
1064            "keyed selection must return a backend in the cluster's live set"
1065        );
1066        Ok((backend_id, address))
1067    }
1068
1069    pub fn set_load_balancing_policy_for_cluster(
1070        &mut self,
1071        cluster_id: &str,
1072        lb_algo: LoadBalancingAlgorithms,
1073        metric: Option<LoadMetric>,
1074    ) {
1075        // The cluster can be created before the backends were registered because of the async config messages.
1076        // So when we set the load balancing policy, we have to create the backend list if if it doesn't exist yet.
1077        let seed = self.rng.next_u64();
1078        let cluster_backends = self.get_or_create_backend_list_for_cluster(cluster_id);
1079        cluster_backends.set_load_balancing_policy(lb_algo, metric, seed);
1080    }
1081
1082    /// Set or clear the shuffle sharding of `cluster_id`, creating its
1083    /// backend list if the cluster's backends have not arrived yet, as
1084    /// [`Self::set_load_balancing_policy_for_cluster`] does.
1085    pub fn set_shuffle_sharding_for_cluster(
1086        &mut self,
1087        cluster_id: &str,
1088        shuffle_sharding: Option<ShuffleSharding>,
1089    ) {
1090        self.get_or_create_backend_list_for_cluster(cluster_id)
1091            .set_shuffle_sharding(shuffle_sharding);
1092    }
1093
1094    pub fn get_or_create_backend_list_for_cluster(&mut self, cluster_id: &str) -> &mut BackendList {
1095        let rng = &mut self.rng;
1096        self.backends
1097            .entry(cluster_id.into())
1098            .or_insert_with(|| BackendList::with_seed(rng.next_u64()))
1099    }
1100}
1101
1102/// The worker's backend registry, seen by the UDP core as a selection view.
1103///
1104/// `UdpManager` selects inside the core now (#1340, Question 6) and reaches
1105/// the registry only through this, for the duration of one admission. It
1106/// delegates to [`BackendMap::backend_from_cluster_id_with_key`], the
1107/// connection-free selection entry point — the UDP datapath owns its own
1108/// per-flow `UdpSocket`, so all the map has to surface is the chosen
1109/// endpoint's identity.
1110impl crate::protocol::udp::BackendSource for BackendMap {
1111    fn select(
1112        &mut self,
1113        cluster: &str,
1114        key: Option<u64>,
1115        now: Instant,
1116    ) -> Option<(crate::protocol::udp::BackendId, SocketAddr)> {
1117        self.backend_from_cluster_id_with_key(cluster, key, now)
1118            .ok()
1119    }
1120}
1121
1122#[derive(Debug)]
1123pub struct BackendList {
1124    /// The cluster's backends, in registration order. Crate-private: the
1125    /// sibling state (`sibling_ring`, `shard_representative`,
1126    /// `shares_address`) and the policy table are derived from it and
1127    /// refreshed only by [`Self::add_backend`] and [`Self::remove_backend`],
1128    /// so a direct push or removal would leave them stale. Read it through
1129    /// [`Self::backends`].
1130    pub(crate) backends: Vec<Rc<RefCell<Backend>>>,
1131    pub next_id: u32,
1132    /// The cluster's selection policy. Crate-private: a Maglev policy keeps
1133    /// a lookup table built from `backends`, and only
1134    /// [`Self::set_load_balancing_policy`] seeds it when the policy is
1135    /// installed, so a direct assignment would leave the table empty or
1136    /// stale.
1137    pub(crate) load_balancing: Box<dyn LoadBalancingAlgorithm>,
1138    /// Latches the fail-open `warn!`. Set to `true` when fail-open routing
1139    /// emits its entry warning so subsequent routing decisions in the same
1140    /// regime stay quiet; reset to `false` when a healthy backend is
1141    /// available again, so the regime-exit transition is logged exactly once.
1142    /// Without this latch the warning fired per request, which under a
1143    /// universal outage (the exact scenario fail-open targets) is also the
1144    /// highest request-rate scenario — log volume would become catastrophic.
1145    fail_open_warned: bool,
1146    /// Per-cluster availability latched by `BackendMap::record_cluster_availability`.
1147    /// `Cell` (not `RefCell`) because the receiver is `&self` and the
1148    /// state is `Copy`. Worker runtime is single-threaded, so a `Cell` is
1149    /// sound — no synchronisation needed.
1150    pub(crate) availability: Cell<ClusterAvailability>,
1151    /// Positions in `backends` of the current selection's candidates, refilled
1152    /// by `collect_candidates` on every selection and lent to the policy as a
1153    /// [`Candidates`] view. Reused rather than rebuilt: `add_backend` reserves
1154    /// room for every backend, so a selection never allocates and clones no
1155    /// `Rc` but the one it returns.
1156    candidates: Vec<usize>,
1157    /// The cluster's shuffle sharding, `None` when it is off
1158    /// (sozu-proxy/sozu#524). Set by [`Self::set_shuffle_sharding`].
1159    shuffle_sharding: Option<ShuffleSharding>,
1160    /// Reused scratch for a keyed selection's shard: the HRW score of each
1161    /// primary backend and its position. Reserved in `add_backend` like
1162    /// `candidates`, so computing a shard never allocates.
1163    shard_scores: Vec<(f64, usize)>,
1164    /// The positions of the current selection's shard members, sorted.
1165    shard: Vec<usize>,
1166    /// Whether two backends of the list share an address. Recomputed by
1167    /// `refresh_siblings`, so a list without a shared address — the common
1168    /// case — skips `collapse_shared_addresses` entirely.
1169    shares_address: bool,
1170    /// Sibling ring: `sibling_ring[i]` is the next position whose backend
1171    /// shares the address of position `i`, cyclically through each address's
1172    /// ids in list order, and `i` itself for an address with one id. Lets a
1173    /// selection reach the siblings of an id without scanning the list.
1174    /// Recomputed by `refresh_siblings` on every mutation of the list or of
1175    /// a backend's configuration (`add_backend`, `remove_backend`).
1176    sibling_ring: Vec<usize>,
1177    /// Whether position `i` ranks for shuffle sharding: a primary backend no
1178    /// primary sibling outranks ([`outranks_at_address`]). A function of the
1179    /// configuration alone (address, weight, `backup`, list position), so it
1180    /// is precomputed by `refresh_siblings`.
1181    shard_representative: Vec<bool>,
1182    /// How many positions are `shard_representative`: the distinct primary
1183    /// addresses a shard size is computed from.
1184    shard_addresses: usize,
1185    /// Reused scratch `collapse_shared_addresses` filters the candidates
1186    /// into, reserved in `add_backend` like `candidates`.
1187    collapsed: Vec<usize>,
1188}
1189
1190/// How shuffle sharding shaped one selection, for the caller that knows the
1191/// cluster to count it under.
1192#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1193pub(crate) enum ShardOutcome {
1194    /// Sharding is off, inactive below `shard_min_backends`, or the request
1195    /// had no key: the whole cluster was eligible.
1196    Unsharded,
1197    /// The selection stayed inside the client's shard.
1198    InShard,
1199    /// No member of the shard could take a connection and `FALLBACK`
1200    /// selected over the rest of the cluster.
1201    SpilledOver,
1202    /// No member of the shard could take a connection and `STRICT` selected
1203    /// nothing.
1204    Exhausted,
1205}
1206
1207impl Default for BackendList {
1208    fn default() -> Self {
1209        Self::new()
1210    }
1211}
1212
1213impl BackendList {
1214    /// An empty list whose default `Random` policy is seeded from the OS.
1215    pub fn new() -> BackendList {
1216        Self::with_seed(os_rng().next_u64())
1217    }
1218
1219    /// An empty list whose default `Random` policy is seeded with `seed`.
1220    ///
1221    /// [`BackendMap`] builds every list it holds this way, from its own
1222    /// generator, so a seeded map seeds every cluster (#1684).
1223    pub fn with_seed(seed: u64) -> BackendList {
1224        BackendList {
1225            backends: Vec::new(),
1226            next_id: 0,
1227            load_balancing: Box::new(Random::with_seed(seed)),
1228            fail_open_warned: false,
1229            availability: Cell::new(ClusterAvailability::Available),
1230            candidates: Vec::new(),
1231            shuffle_sharding: None,
1232            shard_scores: Vec::new(),
1233            shard: Vec::new(),
1234            shares_address: false,
1235            sibling_ring: Vec::new(),
1236            shard_representative: Vec::new(),
1237            shard_addresses: 0,
1238            collapsed: Vec::new(),
1239        }
1240    }
1241
1242    /// The cluster's backends, in registration order.
1243    ///
1244    /// Read-only: the list is mutated only through [`Self::add_backend`] and
1245    /// [`Self::remove_backend`], which refresh the sibling state and the
1246    /// policy table derived from it. Code outside this crate cannot reach
1247    /// the field itself:
1248    ///
1249    /// ```
1250    /// use std::{cell::RefCell, rc::Rc};
1251    ///
1252    /// use sozu_lib::backends::{Backend, BackendList};
1253    ///
1254    /// let mut list = BackendList::with_seed(1);
1255    /// let backend = Backend::new("b1", "127.0.0.1:8080".parse().unwrap(), None, None, None);
1256    /// list.add_backend(backend);
1257    /// let backends: &[Rc<RefCell<Backend>>] = list.backends();
1258    /// assert_eq!(backends.len(), 1);
1259    /// ```
1260    ///
1261    /// ```compile_fail,E0616
1262    /// use std::{cell::RefCell, rc::Rc};
1263    ///
1264    /// use sozu_lib::backends::{Backend, BackendList};
1265    ///
1266    /// let mut list = BackendList::with_seed(1);
1267    /// let backend = Backend::new("b1", "127.0.0.1:8080".parse().unwrap(), None, None, None);
1268    /// list.backends.push(Rc::new(RefCell::new(backend)));
1269    /// ```
1270    pub fn backends(&self) -> &[Rc<RefCell<Backend>>] {
1271        &self.backends
1272    }
1273
1274    /// Count `(available, total)` for this cluster. Delegates to
1275    /// `Backend::is_available` so the per-cluster aggregate and the
1276    /// per-backend `backend.available` gauge stay in lock-step.
1277    pub(crate) fn evaluate_availability(&self) -> (usize, usize) {
1278        let total = self.backends.len();
1279        let available = self
1280            .backends
1281            .iter()
1282            .filter(|b| b.borrow().is_available())
1283            .count();
1284        // The available count is a filtered subset of the total, so it can
1285        // never exceed it, and `total` mirrors the backend vector length.
1286        debug_assert!(
1287            available <= total,
1288            "available ({available}) cannot exceed total ({total})"
1289        );
1290        debug_assert_eq!(total, self.backends.len(), "total must equal backend count");
1291        (available, total)
1292    }
1293
1294    /// Full invariant sweep over the backend list. Called as a `debug_assert!`
1295    /// postcondition by every mutating method so any cross-field corruption
1296    /// (duplicate addresses, a `next_id` that underflowed below the registered
1297    /// count) surfaces immediately under test/fuzz.
1298    #[cfg(debug_assertions)]
1299    fn check_invariants(&self) {
1300        // `next_id` is a monotonically incremented registration counter: it is
1301        // bumped once per *newly inserted* backend (never decremented, even on
1302        // removal), so it is always at least the current live count.
1303        debug_assert!(
1304            self.next_id as usize >= self.backends.len(),
1305            "next_id ({}) must be >= live backend count ({})",
1306            self.next_id,
1307            self.backends.len()
1308        );
1309        // The precomputed sibling state gates and drives every selection: a
1310        // stale value would silently hand a policy one candidate per id
1311        // again, or rank a wrong representative. Check it in `O(n)` plus the
1312        // rings' lengths, so the sweep stays affordable on large lists.
1313        let n = self.backends.len();
1314        debug_assert_eq!(self.sibling_ring.len(), n, "one ring slot per backend");
1315        debug_assert_eq!(self.shard_representative.len(), n);
1316        // A permutation whose links join ids of one address...
1317        let mut targeted = vec![false; n];
1318        for (index, &next) in self.sibling_ring.iter().enumerate() {
1319            debug_assert!(!targeted[next], "the sibling ring is a permutation");
1320            targeted[next] = true;
1321            debug_assert_eq!(
1322                self.backends[next].borrow().address,
1323                self.backends[index].borrow().address,
1324                "a sibling ring links only ids of one address"
1325            );
1326        }
1327        // ...with exactly one cycle per address: every id of an address
1328        // reaches the same cycle leader, the cycle's smallest position.
1329        let mut leaders: HashMap<SocketAddr, usize> = HashMap::new();
1330        for index in 0..n {
1331            let mut leader = index;
1332            let mut next = self.sibling_ring[index];
1333            while next != index {
1334                leader = leader.min(next);
1335                next = self.sibling_ring[next];
1336            }
1337            let address = self.backends[index].borrow().address;
1338            debug_assert_eq!(
1339                *leaders.entry(address).or_insert(leader),
1340                leader,
1341                "an address has one sibling ring"
1342            );
1343            let backend = self.backends[index].borrow();
1344            let mut represents = !backend.backup;
1345            let mut sibling = self.sibling_ring[index];
1346            while sibling != index {
1347                let other = self.backends[sibling].borrow();
1348                represents &=
1349                    other.backup || !outranks_at_address(&other, sibling, &backend, index);
1350                sibling = self.sibling_ring[sibling];
1351            }
1352            debug_assert_eq!(
1353                self.shard_representative[index], represents,
1354                "shard_representative must match the configuration"
1355            );
1356        }
1357        debug_assert_eq!(
1358            self.shares_address,
1359            leaders.len() < n,
1360            "shares_address must match the backend list"
1361        );
1362        debug_assert_eq!(
1363            self.shard_addresses,
1364            self.shard_representative.iter().filter(|r| **r).count()
1365        );
1366        // `(backend_id, address)` is the identity `remove_backend` keys on; two
1367        // live backends may legitimately share an address (A/B variant) but
1368        // must then differ by `backend_id`. The pair is therefore unique
1369        // across the live set.
1370        for (i, a) in self.backends.iter().enumerate() {
1371            let a = a.borrow();
1372            for b in self.backends.iter().skip(i + 1) {
1373                let b = b.borrow();
1374                debug_assert!(
1375                    a.address != b.address || a.backend_id != b.backend_id,
1376                    "duplicate (address, backend_id) in the live set: {:?} / {}",
1377                    a.address,
1378                    a.backend_id
1379                );
1380            }
1381        }
1382    }
1383
1384    /// A list holding `backend_vec`, its default policy seeded with `seed`.
1385    pub fn import_configuration_state(
1386        backend_vec: &[sozu_command_lib::response::Backend],
1387        seed: u64,
1388    ) -> BackendList {
1389        let mut list = BackendList::with_seed(seed);
1390        for backend in backend_vec {
1391            let backend = Backend::new(
1392                &backend.backend_id,
1393                backend.address,
1394                backend.sticky_id.clone(),
1395                backend.load_balancing_parameters,
1396                backend.backup,
1397            );
1398            list.add_backend(backend);
1399        }
1400
1401        list
1402    }
1403
1404    pub fn add_backend(&mut self, backend: Backend) {
1405        let address = backend.address;
1406        let len_before = self.backends.len();
1407        let next_id_before = self.next_id;
1408        let existed = self.backends.iter().any(|b| {
1409            b.borrow().address == backend.address && b.borrow().backend_id == backend.backend_id
1410        });
1411        match self.backends.iter_mut().find(|b| {
1412            b.borrow().address == backend.address && b.borrow().backend_id == backend.backend_id
1413        }) {
1414            None => {
1415                let backend = Rc::new(RefCell::new(backend));
1416                self.backends.push(backend);
1417                self.next_id += 1;
1418            }
1419            // the backend already exists, update the configuration while
1420            // keeping connection retry state
1421            Some(old_backend) => {
1422                let mut b = old_backend.borrow_mut();
1423                b.sticky_id.clone_from(&backend.sticky_id);
1424                b.load_balancing_parameters
1425                    .clone_from(&backend.load_balancing_parameters);
1426                b.backup = backend.backup;
1427            }
1428        }
1429        // Insert grows the list by exactly one and bumps `next_id`; an update
1430        // in place leaves both untouched. The address is present either way.
1431        debug_assert_eq!(
1432            self.backends.len(),
1433            len_before + (!existed) as usize,
1434            "add_backend grows the list by one only on a genuine insert"
1435        );
1436        debug_assert_eq!(
1437            self.next_id,
1438            next_id_before + (!existed) as u32,
1439            "next_id advances by one only on a genuine insert"
1440        );
1441        debug_assert!(
1442            self.has_backend(&address),
1443            "add_backend must leave the backend present in the list"
1444        );
1445        // Refresh table-based policies (Maglev) off the datapath whenever the
1446        // full backend set or a weight changes. This is the ONLY place the
1447        // table is rebuilt on mutation; selection never rebuilds. The default
1448        // `rebuild` is a no-op for the stateless policies.
1449        self.load_balancing.rebuild(&self.backends);
1450        // Size the candidate buffer for the whole list here, on the control
1451        // plane, so no selection ever grows it.
1452        self.candidates.clear();
1453        self.candidates.reserve(self.backends.len());
1454        self.shard_scores.clear();
1455        self.shard_scores.reserve(self.backends.len());
1456        self.shard.clear();
1457        self.shard.reserve(self.backends.len());
1458        self.collapsed.clear();
1459        self.collapsed.reserve(self.backends.len());
1460        self.refresh_siblings();
1461        #[cfg(debug_assertions)]
1462        self.check_invariants();
1463    }
1464
1465    /// Recompute the sibling state — [`Self::shares_address`],
1466    /// [`Self::sibling_ring`], [`Self::shard_representative`] and
1467    /// [`Self::shard_addresses`] — from the backend list and its
1468    /// configuration. Control plane only: it allocates a map of the
1469    /// addresses, linear in the list length, plus, for each shared address
1470    /// of `g` ids, up to `g²` comparisons to rank them.
1471    fn refresh_siblings(&mut self) {
1472        let mut groups: HashMap<SocketAddr, Vec<usize>> = HashMap::new();
1473        for (index, backend) in self.backends.iter().enumerate() {
1474            groups
1475                .entry(backend.borrow().address)
1476                .or_default()
1477                .push(index);
1478        }
1479        self.sibling_ring.clear();
1480        self.sibling_ring.extend(0..self.backends.len());
1481        self.shares_address = false;
1482        for group in groups.values() {
1483            if group.len() > 1 {
1484                self.shares_address = true;
1485                for (rank, &index) in group.iter().enumerate() {
1486                    self.sibling_ring[index] = group[(rank + 1) % group.len()];
1487                }
1488            }
1489        }
1490        self.shard_representative.clear();
1491        for (index, backend) in self.backends.iter().enumerate() {
1492            let backend = backend.borrow();
1493            let mut represents = !backend.backup;
1494            let mut sibling = self.sibling_ring[index];
1495            while represents && sibling != index {
1496                let other = self.backends[sibling].borrow();
1497                represents = other.backup || !outranks_at_address(&other, sibling, &backend, index);
1498                sibling = self.sibling_ring[sibling];
1499            }
1500            self.shard_representative.push(represents);
1501        }
1502        self.shard_addresses = self.shard_representative.iter().filter(|r| **r).count();
1503    }
1504
1505    /// Keep one candidate per address, its representative
1506    /// ([`outranks_at_address`]: the heaviest eligible id, the first in list
1507    /// order among equals), so every policy gives an address the share of
1508    /// one backend however many ids it carries (decided 2026-10-04). Runs on
1509    /// the candidates that passed the tier and shard filters, so a down or
1510    /// backing-off id never hides an eligible sibling. A filter: the
1511    /// candidates stay in list order, as `Candidates::new` requires.
1512    ///
1513    /// Free when no address is shared. Otherwise one pass over the
1514    /// candidates, where only an id with a sibling walks its address's ring
1515    /// (`sibling_ring`) and looks each sibling up among the candidates by
1516    /// binary search: linear in the `c` candidates plus, for each shared
1517    /// address of `g` ids, up to `g²` ring steps (about `g²/2` when the
1518    /// walks stop at the first outranking sibling), each with an
1519    /// `O(log c)` lookup — never a scan of the list per candidate.
1520    /// Allocates nothing (`collapsed` is reserved by `add_backend`).
1521    ///
1522    /// A list whose `backends` a caller grew or shrank directly, bypassing
1523    /// `add_backend`/`remove_backend`, has a ring of another length: the
1524    /// collapse is then skipped rather than indexing a stale ring.
1525    fn collapse_shared_addresses(&mut self) -> usize {
1526        if !self.shares_address || self.sibling_ring.len() != self.backends.len() {
1527            return self.candidates.len();
1528        }
1529        let backends = &self.backends;
1530        let candidates = &self.candidates;
1531        let ring = &self.sibling_ring;
1532        self.collapsed.clear();
1533        for &position in candidates {
1534            let mut outranked = false;
1535            let mut sibling = ring[position];
1536            if sibling != position {
1537                let backend = backends[position].borrow();
1538                while !outranked && sibling != position {
1539                    count_sibling_step();
1540                    outranked = candidates.binary_search(&sibling).is_ok()
1541                        && outranks_at_address(
1542                            &backends[sibling].borrow(),
1543                            sibling,
1544                            &backend,
1545                            position,
1546                        );
1547                    sibling = ring[sibling];
1548                }
1549            }
1550            if !outranked {
1551                self.collapsed.push(position);
1552            }
1553        }
1554        std::mem::swap(&mut self.candidates, &mut self.collapsed);
1555        // `collapsed` now holds the uncollapsed candidates: collapsing only
1556        // drops ids, and keeps one per address, so a non-empty set stays so.
1557        debug_assert!(
1558            self.candidates.len() <= self.collapsed.len(),
1559            "collapsing never adds a candidate"
1560        );
1561        debug_assert!(
1562            self.collapsed.is_empty() || !self.candidates.is_empty(),
1563            "every address keeps one candidate"
1564        );
1565        self.candidates.len()
1566    }
1567
1568    /// Set or clear the cluster's shuffle sharding. Takes effect from the
1569    /// next selection.
1570    pub fn set_shuffle_sharding(&mut self, shuffle_sharding: Option<ShuffleSharding>) {
1571        self.shuffle_sharding = shuffle_sharding;
1572    }
1573
1574    /// Fill `self.shard` with the positions of `key`'s shard, sorted, and
1575    /// say whether sharding applies to this selection at all.
1576    ///
1577    /// The shard is the top `k` of the HRW ranking of `key` over the
1578    /// **configured** primary addresses, healthy or not, with every primary
1579    /// id at those addresses (one address is one share, see
1580    /// `collapse_shared_addresses`): a backend going down
1581    /// must not pull another one into the shard, or a shard could never be
1582    /// exhausted and isolation would leak exactly when it matters. Ties are
1583    /// broken by list position so the shard is a function of the scores.
1584    /// Partitioning with `select_nth_unstable_by` is `O(N)`; each address's
1585    /// representative is precomputed (`shard_representative`), and a shared
1586    /// address brings its other ids in by walking its ring, so a shared
1587    /// address adds work only for its own ids. The buffers are reserved by
1588    /// `add_backend`, so this allocates nothing. If a caller grew or shrank
1589    /// the public `backends` directly, the precomputed state has another
1590    /// length and is ignored: every primary then ranks as its own address,
1591    /// as before shared addresses were collapsed, and nothing is indexed out
1592    /// of bounds.
1593    fn compute_shard(&mut self, key: Option<u64>) -> bool {
1594        let (Some(sharding), Some(key)) = (self.shuffle_sharding, key) else {
1595            return false;
1596        };
1597        // One address is one share: a shard counts and ranks primary
1598        // ADDRESSES, each scored through its representative primary id
1599        // (`outranks_at_address`), and takes in every primary id at the
1600        // addresses it selects, so the sibling of a down representative still
1601        // serves inside the shard. Without a shared address every primary is
1602        // its own representative.
1603        let backends = &self.backends;
1604        let shares_address = self.shares_address;
1605        let synced = self.shard_representative.len() == backends.len()
1606            && self.sibling_ring.len() == backends.len();
1607        let ring: &[usize] = if synced { &self.sibling_ring } else { &[] };
1608        // Size the shard before hashing anything, so a cluster below
1609        // `shard_min_backends` pays no walk and no hash.
1610        let primaries = if synced {
1611            self.shard_addresses
1612        } else {
1613            backends
1614                .iter()
1615                .filter(|backend| !backend.borrow().backup)
1616                .count()
1617        };
1618        let Some(k) = sharding.shard_size(primaries) else {
1619            return false;
1620        };
1621        self.shard_scores.clear();
1622        for (index, backend) in backends.iter().enumerate() {
1623            let represents = match self.shard_representative.get(index) {
1624                Some(&represents) if synced => represents,
1625                _ => !backend.borrow().backup,
1626            };
1627            if represents {
1628                self.shard_scores
1629                    .push((hrw_score(key, &backend.borrow()), index));
1630            }
1631        }
1632        debug_assert_eq!(
1633            self.shard_scores.len(),
1634            primaries,
1635            "one score per primary address"
1636        );
1637        let by_rank = |a: &(f64, usize), b: &(f64, usize)| b.0.total_cmp(&a.0).then(a.1.cmp(&b.1));
1638        if k < self.shard_scores.len() {
1639            self.shard_scores.select_nth_unstable_by(k - 1, by_rank);
1640        }
1641        self.shard.clear();
1642        for &(_, representative) in &self.shard_scores[..k] {
1643            self.shard.push(representative);
1644            let next = |position: usize| ring.get(position).copied().unwrap_or(representative);
1645            let mut sibling = next(representative);
1646            while sibling != representative {
1647                count_sibling_step();
1648                if !backends[sibling].borrow().backup {
1649                    self.shard.push(sibling);
1650                }
1651                sibling = next(sibling);
1652            }
1653        }
1654        self.shard.sort_unstable();
1655        debug_assert!(
1656            self.shard.len() >= k,
1657            "a shard holds every primary id of its k addresses"
1658        );
1659        debug_assert!(
1660            shares_address || self.shard.len() == k,
1661            "without a shared address a shard holds exactly k backends"
1662        );
1663        debug_assert!(
1664            self.shard
1665                .iter()
1666                .all(|&index| !self.backends[index].borrow().backup),
1667            "a shard holds primary backends only"
1668        );
1669        true
1670    }
1671
1672    /// Keep only the candidates that belong to the current shard; the
1673    /// candidates stay in list order, as `Candidates::new` requires.
1674    fn retain_shard_candidates(&mut self) -> usize {
1675        let shard = &self.shard;
1676        self.candidates
1677            .retain(|index| shard.binary_search(index).is_ok());
1678        self.candidates.len()
1679    }
1680
1681    /// Remove the backend identified by `(backend_id, backend_address)`, the
1682    /// identity `AddBackend` and `ConfigState` key backends on, and return
1683    /// whether it was present. Another id at the same address (A/B test,
1684    /// weighted variant) is a distinct backend and stays (#1821).
1685    pub fn remove_backend(&mut self, backend_id: &str, backend_address: &SocketAddr) -> bool {
1686        let len_before = self.backends.len();
1687        let mut removed = false;
1688        self.backends.retain(|backend| {
1689            let mut b = backend.borrow_mut();
1690            if b.backend_id == backend_id && &b.address == backend_address {
1691                removed = true;
1692                // A session may still hold this backend: retire it so none
1693                // of its pooled connections takes a new request.
1694                b.set_closing();
1695                false
1696            } else {
1697                true
1698            }
1699        });
1700        // The list holds each `(backend_id, address)` at most once, so it
1701        // shrinks by exactly one entry when the identity was present.
1702        debug_assert_eq!(
1703            self.backends.len(),
1704            len_before - removed as usize,
1705            "remove_backend must drop exactly the backend it reports"
1706        );
1707        debug_assert!(
1708            self.find_backend_by_identity(backend_id, backend_address)
1709                .is_none(),
1710            "remove_backend must evict the backend it names"
1711        );
1712        // Rebuild table-based policies (Maglev) off the datapath after the set
1713        // shrinks, only when something was actually removed. No-op for the
1714        // stateless policies.
1715        if removed {
1716            self.load_balancing.rebuild(&self.backends);
1717            self.refresh_siblings();
1718        }
1719        #[cfg(debug_assertions)]
1720        self.check_invariants();
1721        removed
1722    }
1723
1724    /// The live backend identified by `(backend_id, backend_address)`.
1725    pub fn find_backend_by_identity(
1726        &self,
1727        backend_id: &str,
1728        backend_address: &SocketAddr,
1729    ) -> Option<&Rc<RefCell<Backend>>> {
1730        self.backends.iter().find(|backend| {
1731            let b = backend.borrow();
1732            b.backend_id == backend_id && b.address == *backend_address
1733        })
1734    }
1735
1736    /// The live entry that is the `incarnation` a health probe captured when
1737    /// it launched, or `None` once that backend left the list: removed, or
1738    /// replaced by a backend re-added under the same id and address, which is
1739    /// a new incarnation. Comparing allocations is sound because the `Weak`
1740    /// keeps its allocation alive, so no later backend can reuse it (#1821).
1741    pub fn find_incarnation(
1742        &self,
1743        incarnation: &Weak<RefCell<Backend>>,
1744    ) -> Option<&Rc<RefCell<Backend>>> {
1745        self.backends
1746            .iter()
1747            .find(|backend| std::ptr::eq(Rc::as_ptr(backend), incarnation.as_ptr()))
1748    }
1749
1750    pub fn has_backend(&self, backend_address: &SocketAddr) -> bool {
1751        self.backends
1752            .iter()
1753            .any(|backend| backend.borrow().address == *backend_address)
1754    }
1755
1756    pub fn find_backend(
1757        &mut self,
1758        backend_address: &SocketAddr,
1759    ) -> Option<&mut Rc<RefCell<Backend>>> {
1760        self.backends
1761            .iter_mut()
1762            .find(|backend| backend.borrow().address == *backend_address)
1763    }
1764
1765    /// The backend `sticky_session` names, if it can take a connection at
1766    /// `now`.
1767    pub fn find_sticky(
1768        &mut self,
1769        sticky_session: &str,
1770        now: Instant,
1771    ) -> Option<&mut Rc<RefCell<Backend>>> {
1772        self.backends
1773            .iter_mut()
1774            .find(|b| b.borrow().sticky_id.as_deref() == Some(sticky_session))
1775            .filter(|b| b.borrow().can_open(now))
1776    }
1777
1778    /// The backends of one tier that can take a connection now, cloned into
1779    /// a fresh `Vec`. Selection does not call this: it walks the same
1780    /// predicate, `is_tier_candidate`, into a reused buffer instead.
1781    pub fn available_backends(&mut self, backup: bool, now: Instant) -> Vec<Rc<RefCell<Backend>>> {
1782        self.backends
1783            .iter()
1784            .filter(|backend| is_tier_candidate(&backend.borrow(), backup, now))
1785            .map(Clone::clone)
1786            .collect()
1787    }
1788
1789    /// Refill `candidates` with the position of every backend `keep` accepts,
1790    /// in list order, and return how many there are. Allocation-free once
1791    /// `add_backend` has reserved the buffer.
1792    fn collect_candidates(&mut self, keep: impl Fn(&Backend) -> bool) -> usize {
1793        self.candidates.clear();
1794        for (index, backend) in self.backends.iter().enumerate() {
1795            if keep(&backend.borrow()) {
1796                self.candidates.push(index);
1797            }
1798        }
1799        debug_assert!(
1800            self.candidates.len() <= self.backends.len(),
1801            "candidate set cannot be larger than the full backend list"
1802        );
1803        debug_assert!(
1804            self.candidates.windows(2).all(|pair| pair[0] < pair[1]),
1805            "candidate positions must follow list order"
1806        );
1807        self.candidates.len()
1808    }
1809
1810    /// Pick the next available backend at `now`.
1811    pub fn next_available_backend(&mut self, now: Instant) -> Option<Rc<RefCell<Backend>>> {
1812        self.next_available_backend_with_key(None, now)
1813    }
1814
1815    /// Pick the next available backend, optionally pinned by an affinity `key`.
1816    ///
1817    /// `key` is only consulted by consistent-hashing policies (HRW/Maglev);
1818    /// every other policy ignores it, so `next_available_backend_with_key(None)`
1819    /// is byte-for-byte the legacy behavior. The UDP datapath calls this with
1820    /// `Some(flow_hash)` to keep a client flow pinned to one backend; HTTP,
1821    /// HTTPS and TCP reach it through [`BackendMap::backend_from_cluster_id`]
1822    /// with the client affinity key of a cluster that uses HRW or Maglev.
1823    ///
1824    /// `now` is the instant the selection happens at: it decides which
1825    /// backends are out of their backoff window and how far each
1826    /// connection-time average has decayed, so the same `now` over the same
1827    /// state picks the same backend.
1828    pub fn next_available_backend_with_key(
1829        &mut self,
1830        key: Option<u64>,
1831        now: Instant,
1832    ) -> Option<Rc<RefCell<Backend>>> {
1833        self.select_with_key(key, now).0
1834    }
1835
1836    /// [`Self::next_available_backend_with_key`], also saying how shuffle
1837    /// sharding shaped the selection, for [`BackendMap`] to count under the
1838    /// cluster's id.
1839    ///
1840    /// Tiers, in order: the primary backends that can open, then the backup
1841    /// backends, then fail-open over every backend whose retry policy allows
1842    /// a try. When the cluster shards and `key` is known, the primary tier is
1843    /// first narrowed to the client's shard. An exhausted shard then either
1844    /// spills over to the unsharded tiers (`FALLBACK`) or ends the selection
1845    /// empty (`STRICT`), except in the fail-open regime — no backend of the
1846    /// whole cluster can open — where `STRICT` fails open inside its shard
1847    /// only.
1848    pub(crate) fn select_with_key(
1849        &mut self,
1850        key: Option<u64>,
1851        now: Instant,
1852    ) -> (Option<Rc<RefCell<Backend>>>, ShardOutcome) {
1853        self.select_tiers(key, now, &[])
1854    }
1855
1856    /// [`Self::select_with_key`] for a request that already failed to connect
1857    /// to the backends at the addresses in `exclude` (sozu-proxy/sozu#1800).
1858    ///
1859    /// Every policy picks from the candidate set this list builds, so leaving
1860    /// those backends out of it is what makes a retry skip them under round
1861    /// robin, random, least loaded, power of two, HRW and Maglev alike.
1862    ///
1863    /// The exclusion only narrows the healthy tiers, primary then backup. When
1864    /// it leaves both empty — every backend that can take a connection was
1865    /// already tried by this request — the selection runs again without it,
1866    /// exactly as [`Self::select_with_key`] would: the request retries a
1867    /// backend it already tried rather than failing, and the fail-open regime
1868    /// still decides when no backend can take a connection at all. A backend
1869    /// that failed is usually out of the healthy tiers already, its retry
1870    /// policy holding it in back-off; the exclusion matters when that back-off
1871    /// expires while the request is still retrying, which a connect timeout as
1872    /// long as the back-off makes routine.
1873    pub(crate) fn select_with_key_excluding(
1874        &mut self,
1875        key: Option<u64>,
1876        now: Instant,
1877        exclude: &[SocketAddr],
1878    ) -> (Option<Rc<RefCell<Backend>>>, ShardOutcome) {
1879        if !exclude.is_empty() {
1880            let selection = self.select_tiers(key, now, exclude);
1881            if selection.0.is_some() {
1882                debug_assert!(
1883                    selection
1884                        .0
1885                        .as_ref()
1886                        .is_some_and(|backend| !exclude.contains(&backend.borrow().address)),
1887                    "an excluding selection must not return an excluded backend"
1888                );
1889                return selection;
1890            }
1891        }
1892        self.select_tiers(key, now, &[])
1893    }
1894
1895    /// The selection itself. A non-empty `exclude` removes those addresses
1896    /// from the healthy tiers and skips the fail-open regime, which
1897    /// [`Self::select_with_key_excluding`] reaches through a second call
1898    /// without exclusion.
1899    fn select_tiers(
1900        &mut self,
1901        key: Option<u64>,
1902        now: Instant,
1903        exclude: &[SocketAddr],
1904    ) -> (Option<Rc<RefCell<Backend>>>, ShardOutcome) {
1905        let candidate = |backend: &Backend, backup: bool| {
1906            is_tier_candidate(backend, backup, now) && !exclude.contains(&backend.address)
1907        };
1908        let sharded = self.compute_shard(key);
1909        let strict = sharded
1910            && self
1911                .shuffle_sharding
1912                .is_some_and(|sharding| sharding.mode == ShardMode::Strict);
1913        let mut outcome = if sharded {
1914            ShardOutcome::InShard
1915        } else {
1916            ShardOutcome::Unsharded
1917        };
1918
1919        let mut available = self.collect_candidates(|backend| candidate(backend, false));
1920        if sharded {
1921            // Primary backends able to take a connection, inside the shard
1922            // or not: under `STRICT` they, and only they, decide between
1923            // refusing and failing open. Backups sit outside every shard, so
1924            // a healthy backup must not turn a strict shard's fail-open into
1925            // a refusal.
1926            let primaries_can_open = available != 0;
1927            let in_shard = self.retain_shard_candidates();
1928            if in_shard != 0 {
1929                available = in_shard;
1930            } else if strict {
1931                if primaries_can_open {
1932                    // Primary backends outside the shard could serve, but
1933                    // strict isolation forbids them.
1934                    return (None, ShardOutcome::Exhausted);
1935                }
1936                // Fail-open regime of the primary tier: fall through with an
1937                // empty primary tier; the backup tier is skipped below for a
1938                // strict shard, and fail-open is narrowed to the shard.
1939                available = 0;
1940            } else {
1941                outcome = ShardOutcome::SpilledOver;
1942                available = self.collect_candidates(|backend| candidate(backend, false));
1943            }
1944        }
1945
1946        if available == 0 && !strict {
1947            available = self.collect_candidates(|backend| candidate(backend, true));
1948        }
1949
1950        if available != 0 {
1951            // Healthy regime: log the fail-open exit transition exactly once.
1952            if self.fail_open_warned {
1953                info!(
1954                    "fail-open: cluster recovered, {} backends now healthy",
1955                    available
1956                );
1957                self.fail_open_warned = false;
1958            }
1959            self.collapse_shared_addresses();
1960            let picked = self.load_balancing.next_available_backend(
1961                key,
1962                Candidates::with_siblings(
1963                    &self.backends,
1964                    &self.candidates,
1965                    &self.sibling_ring,
1966                    now,
1967                ),
1968            );
1969            debug_assert!(
1970                picked.as_ref().is_none_or(|b| {
1971                    let addr = b.borrow().address;
1972                    self.backends.iter().any(|x| x.borrow().address == addr)
1973                }),
1974                "selection must return a backend present in the live list"
1975            );
1976            debug_assert!(
1977                outcome != ShardOutcome::InShard
1978                    || picked.as_ref().is_none_or(|b| {
1979                        let addr = b.borrow().address;
1980                        self.shard
1981                            .iter()
1982                            .any(|&index| self.backends[index].borrow().address == addr)
1983                    }),
1984                "an in-shard selection must return a member of the shard"
1985            );
1986            let outcome = self.settle_spill_outcome(outcome, picked.as_ref());
1987            return (picked, outcome);
1988        }
1989
1990        // Fail-open: when no backend passes the full `can_open()` gate,
1991        // route to backends that are administratively `Normal` AND whose
1992        // retry policy reports `OKAY` (i.e., not currently in
1993        // exponential-backoff). This prevents a shared dependency outage
1994        // (e.g., database) from making the entire cluster unavailable while
1995        // still respecting the per-backend back-off window — hammering a
1996        // backend at line rate during its back-off would defeat the back-off
1997        // itself. Ref: Amazon "Implementing Health Checks".
1998        //
1999        // An excluding selection stops before this regime: it is decided by
2000        // the second, unexcluding call `select_with_key_excluding` makes.
2001        if !exclude.is_empty() {
2002            return (None, outcome);
2003        }
2004        let mut available = self.collect_candidates(|backend| {
2005            backend.status == BackendStatus::Normal
2006                && matches!(
2007                    backend.retry_policy.can_try(now),
2008                    Some(retry::RetryAction::OKAY)
2009                )
2010        });
2011        if strict {
2012            // A strict shard fails open inside itself only.
2013            available = self.retain_shard_candidates();
2014            if available == 0 {
2015                return (None, ShardOutcome::Exhausted);
2016            }
2017        }
2018
2019        if available == 0 {
2020            return (None, outcome);
2021        }
2022
2023        // Latched warning + per-decision counter: the warn! fires once on
2024        // regime entry; the counter is the operator-visible per-request
2025        // signal that does not drown logs under universal outage.
2026        if !self.fail_open_warned {
2027            warn!(
2028                "fail-open: all backends unhealthy, routing to {} normal backends with retry-policy OKAY",
2029                available
2030            );
2031            self.fail_open_warned = true;
2032        }
2033        count!(names::backend::FAIL_OPEN, 1);
2034
2035        self.collapse_shared_addresses();
2036        let picked = self.load_balancing.next_available_backend(
2037            key,
2038            Candidates::with_siblings(&self.backends, &self.candidates, &self.sibling_ring, now),
2039        );
2040        let outcome = self.settle_spill_outcome(outcome, picked.as_ref());
2041        (picked, outcome)
2042    }
2043
2044    /// A `FALLBACK` selection that left its exhausted shard but whose pick is
2045    /// a shard member after all — fail-open over the whole cluster can land
2046    /// back in the shard — stayed in the shard, and is not a spill-over.
2047    fn settle_spill_outcome(
2048        &self,
2049        outcome: ShardOutcome,
2050        picked: Option<&Rc<RefCell<Backend>>>,
2051    ) -> ShardOutcome {
2052        match (outcome, picked) {
2053            (ShardOutcome::SpilledOver, Some(picked))
2054                if self
2055                    .shard
2056                    .iter()
2057                    .any(|&index| Rc::ptr_eq(&self.backends[index], picked)) =>
2058            {
2059                ShardOutcome::InShard
2060            }
2061            _ => outcome,
2062        }
2063    }
2064
2065    /// Replace the cluster's policy. `seed` seeds the policies that draw at
2066    /// random, `Random` and `PowerOfTwo`; the others ignore it.
2067    ///
2068    /// This is the only way to change the policy: it builds the Maglev
2069    /// lookup table from the current backends when it installs `Maglev`.
2070    /// Code outside this crate cannot assign the field directly:
2071    ///
2072    /// ```
2073    /// use sozu_command_lib::proto::command::LoadBalancingAlgorithms;
2074    /// use sozu_lib::backends::{Backend, BackendList};
2075    ///
2076    /// let mut list = BackendList::with_seed(1);
2077    /// list.add_backend(Backend::new("b1", "127.0.0.1:8080".parse().unwrap(), None, None, None));
2078    /// list.set_load_balancing_policy(LoadBalancingAlgorithms::Maglev, None, 1);
2079    /// ```
2080    ///
2081    /// ```compile_fail,E0616
2082    /// use sozu_lib::{backends::BackendList, load_balancing::Maglev};
2083    ///
2084    /// let mut list = BackendList::with_seed(1);
2085    /// list.load_balancing = Box::new(Maglev::new());
2086    /// ```
2087    pub fn set_load_balancing_policy(
2088        &mut self,
2089        load_balancing_policy: LoadBalancingAlgorithms,
2090        metric: Option<LoadMetric>,
2091        seed: u64,
2092    ) {
2093        match load_balancing_policy {
2094            LoadBalancingAlgorithms::RoundRobin => {
2095                self.load_balancing = Box::new(RoundRobin::new())
2096            }
2097            LoadBalancingAlgorithms::Random => {
2098                self.load_balancing = Box::new(Random::with_seed(seed))
2099            }
2100            LoadBalancingAlgorithms::LeastLoaded => {
2101                self.load_balancing = Box::new(LeastLoaded {
2102                    metric: metric.unwrap_or(LoadMetric::Connections),
2103                })
2104            }
2105            LoadBalancingAlgorithms::PowerOfTwo => {
2106                self.load_balancing = Box::new(PowerOfTwo::with_seed(
2107                    seed,
2108                    metric.unwrap_or(LoadMetric::Connections),
2109                ))
2110            }
2111            // Affinity policies. They consult the client key every datapath
2112            // derives (UDP affinity key; HTTP/HTTPS/TCP source IP, header or
2113            // cookie); with `None` they fall back to round-robin.
2114            LoadBalancingAlgorithms::Hrw => self.load_balancing = Box::new(Rendezvous::new()),
2115            LoadBalancingAlgorithms::Maglev => {
2116                let mut maglev = Maglev::new();
2117                // Seed the lookup table from the currently-known backends so
2118                // selection is correct before the first control-plane rebuild.
2119                maglev.rebuild(&self.backends);
2120                self.load_balancing = Box::new(maglev);
2121            }
2122        }
2123    }
2124}
2125
2126/// Count what shuffle sharding did to one selection of `cluster_id`: a
2127/// spill-over that found a backend outside the shard, or a strict refusal.
2128/// A selection that stayed in its shard, or was not sharded, counts nothing.
2129fn record_shard_outcome(cluster_id: &str, outcome: ShardOutcome, picked: bool) {
2130    match outcome {
2131        ShardOutcome::SpilledOver if picked => {
2132            incr!(names::backend::SHARD_SPILLOVER, Some(cluster_id), None);
2133        }
2134        ShardOutcome::Exhausted => {
2135            debug_assert!(!picked, "a strict exhausted shard selects nothing");
2136            incr!(names::backend::SHARD_EXHAUSTED, Some(cluster_id), None);
2137        }
2138        ShardOutcome::Unsharded | ShardOutcome::InShard | ShardOutcome::SpilledOver => {}
2139    }
2140}
2141
2142/// Whether `backend` is a candidate of the primary (`backup == false`) or the
2143/// backup tier: it belongs to that tier and can take a connection at `now`.
2144fn is_tier_candidate(backend: &Backend, backup: bool, now: Instant) -> bool {
2145    backend.backup == backup && backend.can_open(now)
2146}
2147
2148#[cfg(test)]
2149mod backends_test {
2150
2151    use std::{net::TcpListener, sync::mpsc::*, thread};
2152
2153    use super::*;
2154
2155    /// Start a TCP server that accepts connections, and return its address.
2156    ///
2157    /// The kernel picks the port (`:0`): a fixed port collides with any other
2158    /// process bound to it, including a concurrent `cargo test` of this crate,
2159    /// and fails the test with `AddrInUse` before it asserts anything.
2160    fn run_mock_tcp_server(stopper: Receiver<()>) -> SocketAddr {
2161        let mut run = true;
2162        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
2163        let addr = listener.local_addr().unwrap();
2164
2165        thread::spawn(move || {
2166            while run {
2167                for _stream in listener.incoming() {
2168                    // accept connections
2169                    if let Ok(()) = stopper.try_recv() {
2170                        run = false;
2171                    }
2172                }
2173            }
2174        });
2175        addr
2176    }
2177
2178    #[test]
2179    fn it_should_retrieve_a_backend_from_cluster_id_when_backends_have_been_recorded() {
2180        let mut backend_map = BackendMap::new();
2181        let cluster_id = "mycluster";
2182
2183        let (sender, receiver) = channel();
2184        let backend_addr = run_mock_tcp_server(receiver);
2185
2186        backend_map.add_backend(
2187            cluster_id,
2188            Backend::new(&format!("{cluster_id}-1"), backend_addr, None, None, None),
2189        );
2190
2191        assert!(
2192            backend_map
2193                .backend_from_cluster_id(cluster_id, None, Instant::now())
2194                .is_ok()
2195        );
2196        sender.send(()).unwrap();
2197    }
2198
2199    /// A session keeps the `Rc` of a backend it dialled after the map drops
2200    /// it; the mux pool reads that handle's status to refuse reusing its
2201    /// connection, so removal must leave it `Closing`.
2202    #[test]
2203    fn removing_a_backend_or_its_cluster_marks_it_closing() {
2204        let mut backend_map = BackendMap::new();
2205        let first: SocketAddr = "127.0.0.1:9001".parse().unwrap();
2206        let second: SocketAddr = "127.0.0.1:9002".parse().unwrap();
2207        backend_map.add_backend("foo", Backend::new("foo-1", first, None, None, None));
2208        backend_map.add_backend("foo", Backend::new("foo-2", second, None, None, None));
2209        let held: Vec<Rc<RefCell<Backend>>> = backend_map.backends["foo"].backends.clone();
2210
2211        assert!(backend_map.remove_backend("foo", "foo-1", &first));
2212        assert_eq!(held[0].borrow().status, BackendStatus::Closing);
2213        assert_eq!(held[1].borrow().status, BackendStatus::Normal);
2214
2215        backend_map.remove_cluster("foo");
2216        assert_eq!(held[1].borrow().status, BackendStatus::Closing);
2217
2218        // The same address added again is a new backend, not the retired one.
2219        backend_map.add_backend("foo", Backend::new("foo-2", second, None, None, None));
2220        let readded = &backend_map.backends["foo"].backends[0];
2221        assert!(!Rc::ptr_eq(readded, &held[1]));
2222        assert_eq!(readded.borrow().status, BackendStatus::Normal);
2223    }
2224
2225    /// `AddBackend` admits two ids at one address; removing one of them must
2226    /// keep the other routable (#1821).
2227    #[test]
2228    fn removing_one_backend_id_keeps_its_same_address_sibling() {
2229        let mut backend_map = BackendMap::new();
2230        let shared: SocketAddr = "127.0.0.1:9003".parse().unwrap();
2231        backend_map.add_backend("foo", Backend::new("foo-a", shared, None, None, None));
2232        backend_map.add_backend("foo", Backend::new("foo-b", shared, None, None, None));
2233        let held: Vec<Rc<RefCell<Backend>>> = backend_map.backends["foo"].backends.clone();
2234
2235        assert!(backend_map.remove_backend("foo", "foo-a", &shared));
2236        assert_eq!(held[0].borrow().status, BackendStatus::Closing);
2237        assert_eq!(held[1].borrow().status, BackendStatus::Normal);
2238        let list = &backend_map.backends["foo"];
2239        assert_eq!(list.backends.len(), 1);
2240        assert!(Rc::ptr_eq(
2241            list.find_backend_by_identity("foo-b", &shared).unwrap(),
2242            &held[1]
2243        ));
2244        assert!(list.find_backend_by_identity("foo-a", &shared).is_none());
2245
2246        // Removing it again, or an id never added, removes nothing.
2247        assert!(!backend_map.remove_backend("foo", "foo-a", &shared));
2248        assert!(!backend_map.remove_backend("foo", "foo-c", &shared));
2249        assert_eq!(backend_map.backends["foo"].backends.len(), 1);
2250    }
2251
2252    #[test]
2253    fn it_should_not_retrieve_a_backend_from_cluster_id_when_backend_has_not_been_recorded() {
2254        let mut backend_map = BackendMap::new();
2255        let cluster_not_recorded = "not";
2256        backend_map.add_backend(
2257            "foo",
2258            Backend::new("foo-1", "127.0.0.1:9001".parse().unwrap(), None, None, None),
2259        );
2260
2261        assert!(
2262            backend_map
2263                .backend_from_cluster_id(cluster_not_recorded, None, Instant::now())
2264                .is_err()
2265        );
2266    }
2267
2268    #[test]
2269    fn it_should_not_retrieve_a_backend_from_cluster_id_when_backend_list_is_empty() {
2270        let mut backend_map = BackendMap::new();
2271
2272        assert!(
2273            backend_map
2274                .backend_from_cluster_id("dumb", None, Instant::now())
2275                .is_err()
2276        );
2277    }
2278
2279    #[test]
2280    fn it_should_retrieve_a_backend_from_sticky_session_when_the_backend_has_been_recorded() {
2281        let mut backend_map = BackendMap::new();
2282        let cluster_id = "mycluster";
2283        let sticky_session = "server-2";
2284
2285        for (index, port) in [(1u16, 9001u16), (2, 9000), (3, 9002)] {
2286            backend_map.add_backend(
2287                cluster_id,
2288                Backend::new(
2289                    &format!("{cluster_id}-{index}"),
2290                    SocketAddr::from(([127, 0, 0, 1], port)),
2291                    Some(format!("server-{index}")),
2292                    None,
2293                    None,
2294                ),
2295            );
2296        }
2297
2298        let backend = backend_map
2299            .reserve_sticky_backend(cluster_id, sticky_session, None, Instant::now())
2300            .expect("the backend the sticky session names can take a connection");
2301        let backend = backend.borrow();
2302        assert_eq!(
2303            backend.backend_id, "mycluster-2",
2304            "a sticky session must reach the backend it names"
2305        );
2306        assert_eq!(
2307            backend.active_connections, 1,
2308            "reaching it reserves a connection"
2309        );
2310    }
2311
2312    #[test]
2313    fn it_should_not_retrieve_a_backend_from_sticky_session_when_the_backend_has_not_been_recorded()
2314    {
2315        let mut backend_map = BackendMap::new();
2316        let cluster_id = "mycluster";
2317        let sticky_session = "test";
2318
2319        assert!(
2320            backend_map
2321                .reserve_sticky_backend(cluster_id, sticky_session, None, Instant::now())
2322                .is_err()
2323        );
2324    }
2325
2326    #[test]
2327    fn it_should_not_retrieve_a_backend_from_sticky_session_when_the_backend_list_is_empty() {
2328        let mut backend_map = BackendMap::new();
2329        let mycluster_not_recorded = "mycluster";
2330        let sticky_session = "test";
2331
2332        assert!(
2333            backend_map
2334                .reserve_sticky_backend(
2335                    mycluster_not_recorded,
2336                    sticky_session,
2337                    None,
2338                    Instant::now()
2339                )
2340                .is_err()
2341        );
2342    }
2343
2344    #[test]
2345    fn it_should_add_a_backend_when_he_doesnt_already_exist() {
2346        let backend_id = "myback";
2347        let mut backends_list = BackendList::new();
2348        backends_list.add_backend(Backend::new(
2349            backend_id,
2350            "127.0.0.1:80".parse().unwrap(),
2351            None,
2352            None,
2353            None,
2354        ));
2355
2356        assert_eq!(1, backends_list.backends.len());
2357    }
2358
2359    #[test]
2360    fn it_should_not_add_a_backend_when_he_already_exist() {
2361        let backend_id = "myback";
2362        let mut backends_list = BackendList::new();
2363        backends_list.add_backend(Backend::new(
2364            backend_id,
2365            "127.0.0.1:80".parse().unwrap(),
2366            None,
2367            None,
2368            None,
2369        ));
2370
2371        //same backend id
2372        backends_list.add_backend(Backend::new(
2373            backend_id,
2374            "127.0.0.1:80".parse().unwrap(),
2375            None,
2376            None,
2377            None,
2378        ));
2379
2380        assert_eq!(1, backends_list.backends.len());
2381    }
2382
2383    /// Build a backend addressed at 127.0.0.1:port and force it Unhealthy
2384    /// without going through the health-check loop.
2385    fn unhealthy_backend(id: &str, port: u16) -> Backend {
2386        let mut backend = Backend::new(
2387            id,
2388            format!("127.0.0.1:{port}").parse().unwrap(),
2389            None,
2390            None,
2391            None,
2392        );
2393        // Threshold = 1 transitions on the first failure.
2394        backend.health.record_failure(1);
2395        assert!(!backend.health.is_healthy());
2396        backend
2397    }
2398
2399    #[test]
2400    fn fail_open_picks_normal_backend_in_retry_policy_okay() {
2401        // All backends are unhealthy but their retry policy is fresh (OKAY),
2402        // so fail-open must select one. A fresh ExponentialBackoffPolicy
2403        // returns OKAY on `can_try()` until the first `fail()` arms a wait
2404        // window.
2405        let mut list = BackendList::new();
2406        list.add_backend(unhealthy_backend("b1", 9001));
2407        list.add_backend(unhealthy_backend("b2", 9002));
2408
2409        // Sanity: `available_backends` returns nothing (the regular path).
2410        assert!(list.available_backends(false, Instant::now()).is_empty());
2411        assert!(list.available_backends(true, Instant::now()).is_empty());
2412
2413        let picked = list.next_available_backend(Instant::now());
2414        assert!(
2415            picked.is_some(),
2416            "fail-open must pick a Normal+OKAY backend"
2417        );
2418        assert!(list.fail_open_warned, "regime entry must latch the warn!");
2419    }
2420
2421    #[test]
2422    fn fail_open_skips_backend_in_retry_backoff() {
2423        // Same shape as above, but each backend's retry policy is in the
2424        // WAIT window after a recorded failure. Fail-open must NOT pick any
2425        // of them — hammering a backend at line rate during its back-off
2426        // window is exactly what the back-off is protecting against — and
2427        // the regime-entry warn! must NOT latch (no log spam either).
2428        let mut list = BackendList::new();
2429        list.add_backend(unhealthy_backend("b1", 9011));
2430        list.add_backend(unhealthy_backend("b2", 9012));
2431        for backend_rc in &list.backends {
2432            backend_rc
2433                .borrow_mut()
2434                .retry_policy()
2435                .fail(Instant::now(), &mut rand::rng());
2436            assert_eq!(
2437                Some(retry::RetryAction::WAIT),
2438                backend_rc.borrow().retry_policy.can_try(Instant::now()),
2439                "test fixture must place retry policy in WAIT"
2440            );
2441        }
2442
2443        let picked = list.next_available_backend(Instant::now());
2444        assert!(
2445            picked.is_none(),
2446            "fail-open must skip backends whose retry policy is in WAIT"
2447        );
2448        assert!(
2449            !list.fail_open_warned,
2450            "no candidate backends, no regime entry"
2451        );
2452    }
2453
2454    #[test]
2455    fn fail_open_warn_latched() {
2456        // First call enters the regime → latch flips, warn! fires.
2457        // Second call stays in the regime → latch stays, no second warn!.
2458        // Recovering one backend → latch clears on the next routing call.
2459        let mut list = BackendList::new();
2460        list.add_backend(unhealthy_backend("b1", 9021));
2461        list.add_backend(unhealthy_backend("b2", 9022));
2462
2463        assert!(list.next_available_backend(Instant::now()).is_some());
2464        assert!(list.fail_open_warned, "first fail-open must latch");
2465
2466        assert!(list.next_available_backend(Instant::now()).is_some());
2467        assert!(
2468            list.fail_open_warned,
2469            "subsequent fail-open routing keeps the latch"
2470        );
2471
2472        // Heal one backend — the next routing call takes the healthy path
2473        // and must clear the latch (regime exit logged once).
2474        list.backends[0].borrow_mut().health.status = HealthStatus::Healthy;
2475        let picked = list.next_available_backend(Instant::now());
2476        assert!(
2477            picked.is_some(),
2478            "regular path must select the healed backend"
2479        );
2480        assert!(
2481            !list.fail_open_warned,
2482            "regime exit must clear the latch so the next entry is logged again"
2483        );
2484    }
2485
2486    // ── #892: per-cluster availability tracker ──────────────────────────
2487
2488    /// Build a backend that passes the `evaluate_availability` predicate
2489    /// (`status == Normal && health.is_healthy() && !retry_policy.is_down()`).
2490    /// A `Backend::new` returns Normal/Healthy with a fresh
2491    /// ExponentialBackoffPolicy that reports `is_down() == false` until the
2492    /// first `fail()`, so this is just `Backend::new` with a stable address.
2493    fn healthy_backend(id: &str, port: u16) -> Backend {
2494        Backend::new(
2495            id,
2496            format!("127.0.0.1:{port}").parse().unwrap(),
2497            None,
2498            None,
2499            None,
2500        )
2501    }
2502
2503    #[test]
2504    fn is_available_requires_health_status_and_retry_policy() {
2505        // Fresh backend: Healthy + Normal + retry_policy fresh (OKAY).
2506        let mut backend = Backend::new("b", "127.0.0.1:9050".parse().unwrap(), None, None, None);
2507        assert!(backend.is_available(), "fresh backend must be available");
2508
2509        // Unhealthy fails the predicate even with everything else OK.
2510        backend.health.record_failure(1);
2511        assert!(!backend.is_available(), "unhealthy must not be available");
2512
2513        // Restore health, then drive retry policy into the exhausted-budget
2514        // state via the test-only helper. Calling `fail()` in a tight loop
2515        // would early-return on the second invocation because the
2516        // exponential-backoff window has not elapsed yet, so the natural
2517        // path needs real-time sleeps the unit test cannot afford.
2518        backend.health.status = HealthStatus::Healthy;
2519        assert!(backend.is_available());
2520        backend.retry_policy.force_down();
2521        assert!(
2522            backend.retry_policy.is_down(),
2523            "test setup: retry policy budget must be exhausted",
2524        );
2525        assert!(
2526            !backend.is_available(),
2527            "retry-policy backoff must fail the predicate"
2528        );
2529
2530        // Reset retry, switch lifecycle to Closing.
2531        backend.retry_policy.succeed(Instant::now());
2532        backend.set_closing();
2533        assert!(
2534            !backend.is_available(),
2535            "Closing lifecycle status must fail the predicate"
2536        );
2537    }
2538
2539    #[test]
2540    fn evaluate_availability_empty_list_returns_zero_zero() {
2541        let list = BackendList::new();
2542        assert_eq!((0, 0), list.evaluate_availability());
2543    }
2544
2545    #[test]
2546    fn evaluate_availability_counts_only_healthy_normal_not_in_backoff() {
2547        let mut list = BackendList::new();
2548        list.add_backend(healthy_backend("b-ok-1", 9101));
2549        list.add_backend(healthy_backend("b-ok-2", 9102));
2550        list.add_backend(unhealthy_backend("b-bad", 9103));
2551        let (available, total) = list.evaluate_availability();
2552        assert_eq!(3, total, "every configured backend counts toward total");
2553        assert_eq!(
2554            2, available,
2555            "only the two healthy backends pass the predicate"
2556        );
2557    }
2558
2559    #[test]
2560    fn evaluate_availability_excludes_retry_policy_down() {
2561        let mut list = BackendList::new();
2562        list.add_backend(healthy_backend("b-fresh", 9111));
2563        list.add_backend(healthy_backend("b-fail", 9112));
2564        // Drive the second backend's retry policy into is_down() via the
2565        // test-only force_down() helper. A natural exhaustion would
2566        // require waiting through the exponential-backoff windows
2567        // (`fail()` early-returns when called inside one).
2568        list.backends[1].borrow_mut().retry_policy.force_down();
2569        let (available, total) = list.evaluate_availability();
2570        assert_eq!(2, total);
2571        assert_eq!(
2572            1, available,
2573            "retry-policy is_down() backend must be excluded even when health.is_healthy()"
2574        );
2575    }
2576
2577    #[test]
2578    fn record_cluster_availability_flips_to_alldown_then_idempotent() {
2579        let mut map = BackendMap::new();
2580        let cluster_id = "c-flap";
2581        map.add_backend(cluster_id, unhealthy_backend("b1", 9201));
2582        // After add_backend the helper has run; total=1, available=0,
2583        // so the cell must already be AllDown.
2584        let list = map.backends.get(cluster_id).expect("cluster present");
2585        assert_eq!(
2586            ClusterAvailability::AllDown,
2587            list.availability.get(),
2588            "single unhealthy backend must drive the cell to AllDown"
2589        );
2590        // Calling again in the same regime is a no-op (Cell already AllDown).
2591        map.record_cluster_availability(cluster_id);
2592        let list = map.backends.get(cluster_id).expect("cluster present");
2593        assert_eq!(
2594            ClusterAvailability::AllDown,
2595            list.availability.get(),
2596            "repeat call must keep the cell at AllDown without flipping"
2597        );
2598    }
2599
2600    #[test]
2601    fn record_cluster_availability_recovers_to_available() {
2602        let mut map = BackendMap::new();
2603        let cluster_id = "c-recover";
2604        map.add_backend(cluster_id, unhealthy_backend("b1", 9301));
2605        assert_eq!(
2606            ClusterAvailability::AllDown,
2607            map.backends.get(cluster_id).unwrap().availability.get()
2608        );
2609        // Heal the backend in place and re-evaluate. Without going through
2610        // a routing call the helper still fires from add_backend / HC, so
2611        // here we drive it manually.
2612        map.backends.get_mut(cluster_id).unwrap().backends[0]
2613            .borrow_mut()
2614            .health
2615            .status = HealthStatus::Healthy;
2616        map.record_cluster_availability(cluster_id);
2617        assert_eq!(
2618            ClusterAvailability::Available,
2619            map.backends.get(cluster_id).unwrap().availability.get(),
2620            "healed backend must flip the cell back to Available"
2621        );
2622    }
2623
2624    #[test]
2625    fn record_cluster_availability_empty_cluster_stays_available() {
2626        let mut map = BackendMap::new();
2627        let cluster_id = "c-empty";
2628        map.backends.insert(cluster_id.into(), BackendList::new());
2629        // total == 0 path: never report AllDown — avoids log spam during
2630        // cluster bootstrap when backends are still being registered.
2631        map.record_cluster_availability(cluster_id);
2632        assert_eq!(
2633            ClusterAvailability::Available,
2634            map.backends.get(cluster_id).unwrap().availability.get(),
2635            "empty cluster must keep the cell at the default Available"
2636        );
2637    }
2638
2639    #[test]
2640    fn record_cluster_availability_missing_cluster_is_noop() {
2641        let map = BackendMap::new();
2642        // No panic, no insert — just an early return.
2643        map.record_cluster_availability("c-absent");
2644        assert!(
2645            !map.backends.contains_key("c-absent"),
2646            "helper must not insert a BackendList for an unknown cluster_id"
2647        );
2648    }
2649
2650    #[test]
2651    fn import_configuration_state_latches_cluster_rollup_gauges() {
2652        use crate::metrics::METRICS;
2653        use sozu_command_lib::proto::command::QueryMetricsOptions;
2654        // Unique cluster id so the assertion is not perturbed by gauges
2655        // left in the thread-local METRICS aggregator by sibling tests.
2656        let cluster_id = "c-import-rollup-9701";
2657        let mut map = BackendMap::new();
2658        let mut input = HashMap::new();
2659        input.insert(
2660            cluster_id.into(),
2661            vec![sozu_command_lib::response::Backend {
2662                cluster_id: cluster_id.into(),
2663                backend_id: "b1".to_owned(),
2664                address: "127.0.0.1:9701".parse().unwrap(),
2665                sticky_id: None,
2666                load_balancing_parameters: None,
2667                backup: None,
2668            }],
2669        );
2670        map.import_configuration_state(&input);
2671        let response = METRICS
2672            .with(|m| {
2673                m.borrow_mut().query(&QueryMetricsOptions {
2674                    metric_names: vec![
2675                        names::cluster::AVAILABLE_BACKENDS.to_owned(),
2676                        names::cluster::TOTAL_BACKENDS.to_owned(),
2677                    ],
2678                    cluster_ids: vec![cluster_id.to_owned()],
2679                    backend_ids: vec![],
2680                    list: false,
2681                    no_clusters: false,
2682                    workers: false,
2683                })
2684            })
2685            .expect("metrics query succeeds");
2686        let cluster_metrics = match response.content_type {
2687            Some(
2688                sozu_command_lib::proto::command::response_content::ContentType::WorkerMetrics(wm),
2689            ) => wm,
2690            Some(_) => panic!("expected WorkerMetrics, got another response variant"),
2691            None => panic!("expected WorkerMetrics, got no response content"),
2692        };
2693        let cm = cluster_metrics
2694            .clusters
2695            .get(cluster_id)
2696            .expect("imported cluster must have a ClusterMetrics entry");
2697        // Without the import-time `record_cluster_availability` call the
2698        // two rollup gauges would be absent here. The fix guarantees the
2699        // pair lands without waiting for any follow-up backend mutation.
2700        assert!(
2701            cm.cluster.contains_key(names::cluster::AVAILABLE_BACKENDS),
2702            "cluster.available_backends gauge must be latched at import time"
2703        );
2704        assert!(
2705            cm.cluster.contains_key(names::cluster::TOTAL_BACKENDS),
2706            "cluster.total_backends gauge must be latched at import time"
2707        );
2708    }
2709
2710    #[test]
2711    fn set_health_check_config_none_re_emits_rollup_after_reset() {
2712        let mut map = BackendMap::new();
2713        let cluster_id = "c-hc-reset";
2714        // Seed the cluster with an unhealthy backend so `add_backend`
2715        // drives the `availability` cell to AllDown.
2716        map.add_backend(cluster_id, unhealthy_backend("b1", 9801));
2717        assert_eq!(
2718            ClusterAvailability::AllDown,
2719            map.backends.get(cluster_id).unwrap().availability.get(),
2720            "test setup: unhealthy backend must register the cell at AllDown"
2721        );
2722        // Disabling the health check resets backend health to the default
2723        // pristine state AND must re-emit the rollup so the cell reflects
2724        // the post-reset availability instead of the stale AllDown.
2725        map.set_health_check_config(cluster_id, None);
2726        assert_eq!(
2727            ClusterAvailability::Available,
2728            map.backends.get(cluster_id).unwrap().availability.get(),
2729            "set_health_check_config(None) must re-emit the rollup after \
2730             resetting backend health, otherwise dashboards stay stuck at AllDown"
2731        );
2732    }
2733
2734    // ----- #1684: selection reads neither the wall clock nor an ambient RNG -----
2735
2736    /// Four backends of `cluster_id` under `policy`, created at `now` in a map
2737    /// that seeded the policy, and the addresses of `draws` selections at `now`.
2738    fn draws(
2739        map: &mut BackendMap,
2740        cluster_id: &str,
2741        policy: LoadBalancingAlgorithms,
2742        now: Instant,
2743        draws: usize,
2744    ) -> Vec<SocketAddr> {
2745        map.set_load_balancing_policy_for_cluster(
2746            cluster_id,
2747            policy,
2748            Some(LoadMetric::Connections),
2749        );
2750        for index in 0..4u16 {
2751            map.add_backend(
2752                cluster_id,
2753                Backend::new_at(
2754                    &format!("{cluster_id}-{index}"),
2755                    SocketAddr::from(([127, 0, 0, 1], 9100 + index)),
2756                    None,
2757                    None,
2758                    None,
2759                    now,
2760                ),
2761            );
2762        }
2763        let list = map
2764            .backends
2765            .get_mut(cluster_id)
2766            .expect("the policy call created the list");
2767        (0..draws)
2768            .map(|_| {
2769                list.next_available_backend(now)
2770                    .expect("four fresh backends can all open")
2771                    .borrow()
2772                    .address
2773            })
2774            .collect()
2775    }
2776
2777    /// #1684: a map built from a seed seeds every policy it creates from it,
2778    /// so one seed yields one sequence of picks.
2779    ///
2780    /// TO SEE THIS RED: build `BackendMap::with_seed` from `os_rng()` and
2781    /// ignore the seed; two maps of one seed then disagree.
2782    #[test]
2783    fn one_map_seed_yields_one_selection_sequence() {
2784        let now = Instant::now();
2785        for policy in [
2786            LoadBalancingAlgorithms::Random,
2787            LoadBalancingAlgorithms::PowerOfTwo,
2788        ] {
2789            let first = draws(&mut BackendMap::with_seed(7), "c", policy, now, 48);
2790            assert_eq!(
2791                first,
2792                draws(&mut BackendMap::with_seed(7), "c", policy, now, 48),
2793                "{policy:?}: one seed must yield one sequence of picks"
2794            );
2795            assert_ne!(
2796                first,
2797                draws(&mut BackendMap::with_seed(8), "c", policy, now, 48),
2798                "{policy:?}: the picks must come from the map's seed, so another seed differs"
2799            );
2800        }
2801    }
2802
2803    /// The addresses of `draws` selections at `now` from `cluster_id`.
2804    fn picks_of(
2805        map: &mut BackendMap,
2806        cluster_id: &str,
2807        now: Instant,
2808        draws: usize,
2809    ) -> Vec<SocketAddr> {
2810        let list = map
2811            .backends
2812            .get_mut(cluster_id)
2813            .expect("the cluster was created");
2814        (0..draws)
2815            .map(|_| {
2816                list.next_available_backend(now)
2817                    .expect("fresh backends can all open")
2818                    .borrow()
2819                    .address
2820            })
2821            .collect()
2822    }
2823
2824    /// #1684: a cluster the map creates while adding a backend, before any
2825    /// policy is set, is seeded from the map too, not from the OS.
2826    ///
2827    /// TO SEE THIS RED: create the list in `BackendMap::add_backend` with
2828    /// `or_default()`, which seeds its `Random` policy from the OS.
2829    #[test]
2830    fn a_cluster_created_by_adding_a_backend_is_seeded_by_the_map() {
2831        let now = Instant::now();
2832        let sequence = || {
2833            let mut map = BackendMap::with_seed(7);
2834            for index in 0..4u16 {
2835                map.add_backend(
2836                    "c",
2837                    Backend::new_at(
2838                        &format!("c-{index}"),
2839                        SocketAddr::from(([127, 0, 0, 1], 9300 + index)),
2840                        None,
2841                        None,
2842                        None,
2843                        now,
2844                    ),
2845                );
2846            }
2847            picks_of(&mut map, "c", now, 48)
2848        };
2849        assert_eq!(
2850            sequence(),
2851            sequence(),
2852            "a list created by add_backend must draw its seed from the map"
2853        );
2854    }
2855
2856    /// #1684: importing a configuration state seeds each cluster from the
2857    /// map in cluster-id order, so one map seed gives every cluster the same
2858    /// seed whatever order the state's `HashMap` iterates in.
2859    ///
2860    /// TO SEE THIS RED: draw the seeds while iterating `backends` directly in
2861    /// `BackendMap::import_configuration_state`; each `HashMap` has its own
2862    /// `RandomState`, so two imports hand the seeds out in different orders.
2863    #[test]
2864    fn importing_a_state_seeds_its_clusters_in_a_fixed_order() {
2865        const CLUSTERS: u16 = 16;
2866        let now = Instant::now();
2867        let state = |ids: &mut dyn Iterator<Item = u16>| {
2868            let mut state: HashMap<ClusterId, Vec<sozu_command_lib::response::Backend>> =
2869                HashMap::new();
2870            for id in ids {
2871                let cluster_id: ClusterId = format!("import-{id:02}").into();
2872                let backends = (0..4u16)
2873                    .map(|index| sozu_command_lib::response::Backend {
2874                        cluster_id: cluster_id.clone(),
2875                        backend_id: format!("{cluster_id}-{index}"),
2876                        address: SocketAddr::from(([127, 0, 0, 1], 9400 + 4 * id + index)),
2877                        sticky_id: None,
2878                        load_balancing_parameters: None,
2879                        backup: None,
2880                    })
2881                    .collect();
2882                state.insert(cluster_id, backends);
2883            }
2884            state
2885        };
2886        let forward = state(&mut (0..CLUSTERS));
2887        let backward = state(&mut (0..CLUSTERS).rev());
2888
2889        let mut first = BackendMap::with_seed(7);
2890        first.import_configuration_state(&forward);
2891        let mut second = BackendMap::with_seed(7);
2892        second.import_configuration_state(&backward);
2893        for id in 0..CLUSTERS {
2894            let cluster_id = format!("import-{id:02}");
2895            assert_eq!(
2896                picks_of(&mut first, &cluster_id, now, 32),
2897                picks_of(&mut second, &cluster_id, now, 32),
2898                "{cluster_id}: one map seed must give each imported cluster one seed"
2899            );
2900        }
2901    }
2902
2903    /// Regression guard moved from `load_balancing.rs`, where it held
2904    /// `Random::new()`: the production seed must never be a shared constant.
2905    /// If it were, every worker, on every cold start, would draw from a
2906    /// bit-for-bit identical keystream and pick the same backend at every
2907    /// step — the correlated-load event uniform selection exists to prevent,
2908    /// arriving at a synchronised redeploy. The OS seed now enters through
2909    /// `BackendMap::new`, which seeds every cluster's policy, so two maps (two
2910    /// workers) and two clusters of one map must each diverge.
2911    ///
2912    /// 48 draws over four backends collide with probability ~4^-48 if the
2913    /// keystreams are independent, so this cannot flake.
2914    #[test]
2915    fn random_policies_seeded_by_the_os_are_not_correlated() {
2916        let now = Instant::now();
2917        let policy = LoadBalancingAlgorithms::Random;
2918        let mut map = BackendMap::new();
2919        let first = draws(&mut map, "a", policy, now, 48);
2920        assert_ne!(
2921            first,
2922            draws(&mut BackendMap::new(), "a", policy, now, 48),
2923            "two BackendMap::new() must NOT seed Random from a shared/correlated keystream"
2924        );
2925        assert_ne!(
2926            first,
2927            draws(&mut map, "b", policy, now, 48),
2928            "two clusters of one map must NOT draw from a shared keystream"
2929        );
2930    }
2931
2932    /// The same guard for `PowerOfTwo`'s tie-break, moved from
2933    /// `load_balancing.rs`. Every backend starts at zero load, so a tie (and
2934    /// its coin flip) is the common case: a shared seed would resolve every
2935    /// worker's first ties identically after a synchronised redeploy,
2936    /// reintroducing the herding P2C exists to prevent.
2937    #[test]
2938    fn power_of_two_policies_seeded_by_the_os_are_not_correlated() {
2939        let now = Instant::now();
2940        let policy = LoadBalancingAlgorithms::PowerOfTwo;
2941        let mut map = BackendMap::new();
2942        let first = draws(&mut map, "a", policy, now, 48);
2943        assert_ne!(
2944            first,
2945            draws(&mut BackendMap::new(), "a", policy, now, 48),
2946            "two BackendMap::new() must NOT seed PowerOfTwo from a shared/correlated keystream"
2947        );
2948        assert_ne!(
2949            first,
2950            draws(&mut map, "b", policy, now, 48),
2951            "two clusters of one map must NOT draw from a shared keystream"
2952        );
2953    }
2954
2955    /// #1684: the connection-time average decays on the instants it is
2956    /// given. A backend created at `t0` and read at `t0 + 1s` has aged by
2957    /// exactly one decay constant, `e^-1`, whatever the wall clock did.
2958    ///
2959    /// TO SEE THIS RED: take `now` from `Instant::now()` inside
2960    /// `PeakEWMA::observe`; almost no real time passes, so the average does
2961    /// not decay and the cost stays at its 50 ms default.
2962    #[test]
2963    fn the_connection_time_average_decays_on_the_injected_clock() {
2964        let t0 = Instant::now();
2965        let mut backend = Backend::new_at(
2966            "ewma",
2967            SocketAddr::from(([127, 0, 0, 1], 9200)),
2968            None,
2969            None,
2970            None,
2971            t0,
2972        );
2973        let default_rtt = backend.connection_time.rtt;
2974        let cost = backend.peak_ewma_connection(t0 + Duration::from_secs(1));
2975        assert_eq!(
2976            cost,
2977            default_rtt * (-1f64).exp(),
2978            "one second after creation, the average must have decayed by e^-1"
2979        );
2980
2981        // One sequence of instants, one average, bit for bit.
2982        let replay = |backend: &mut Backend| {
2983            backend.set_connection_time(Duration::from_millis(80), t0 + Duration::from_millis(10));
2984            backend.set_connection_time(Duration::from_millis(20), t0 + Duration::from_millis(900));
2985            backend.peak_ewma_connection(t0 + Duration::from_millis(2_500))
2986        };
2987        let mut twin = Backend::new_at(
2988            "ewma",
2989            SocketAddr::from(([127, 0, 0, 1], 9200)),
2990            None,
2991            None,
2992            None,
2993            t0,
2994        );
2995        let mut again = Backend::new_at(
2996            "ewma",
2997            SocketAddr::from(([127, 0, 0, 1], 9200)),
2998            None,
2999            None,
3000            None,
3001            t0,
3002        );
3003        assert_eq!(
3004            replay(&mut twin).to_bits(),
3005            replay(&mut again).to_bits(),
3006            "one sequence of instants must yield one average, bit for bit"
3007        );
3008    }
3009
3010    // ----- #524: shuffle sharding over the HRW ranking -----
3011
3012    /// `n` primary backends on distinct loopback ports, under `policy`, sharded
3013    /// by `sharding`.
3014    fn sharded_list(
3015        n: u16,
3016        policy: LoadBalancingAlgorithms,
3017        sharding: Option<ShuffleSharding>,
3018    ) -> BackendList {
3019        let now = Instant::now();
3020        let mut list = BackendList::with_seed(524);
3021        list.set_load_balancing_policy(policy, Some(LoadMetric::Connections), 524);
3022        for index in 0..n {
3023            list.add_backend(Backend::new_at(
3024                &format!("shard-{index}"),
3025                SocketAddr::from(([127, 0, 0, 1], 20_000 + index)),
3026                None,
3027                None,
3028                None,
3029                now,
3030            ));
3031        }
3032        list.set_shuffle_sharding(sharding);
3033        list
3034    }
3035
3036    fn sharding(percent: u32, min_backends: u32, mode: ShardMode) -> Option<ShuffleSharding> {
3037        Some(ShuffleSharding {
3038            percent,
3039            min_backends,
3040            mode,
3041        })
3042    }
3043
3044    /// The addresses of `key`'s shard in `list`, and its members ranked best
3045    /// first.
3046    fn shard_of(list: &mut BackendList, key: u64) -> Vec<SocketAddr> {
3047        assert!(list.compute_shard(Some(key)), "the list must shard");
3048        let mut ranked: Vec<(f64, SocketAddr)> = list
3049            .shard
3050            .iter()
3051            .map(|&index| {
3052                let backend = list.backends[index].borrow();
3053                (hrw_score(key, &backend), backend.address)
3054            })
3055            .collect();
3056        ranked.sort_by(|a, b| b.0.total_cmp(&a.0));
3057        ranked.into_iter().map(|(_, address)| address).collect()
3058    }
3059
3060    /// #524: a shard is the top `k` of the key's HRW ranking over every
3061    /// primary backend — the same backends whatever their order in the list,
3062    /// on every call.
3063    #[test]
3064    fn a_shard_is_the_top_of_the_hrw_ranking() {
3065        let unweighted_list = || {
3066            sharded_list(
3067                12,
3068                LoadBalancingAlgorithms::RoundRobin,
3069                sharding(25, 8, ShardMode::Fallback),
3070            )
3071        };
3072        let mut unweighted = unweighted_list();
3073        // The same backends with non-uniform weights: the ranking is the
3074        // weighted HRW score, so heavier backends rank higher more often.
3075        let mut weighted = unweighted_list();
3076        for (index, backend) in weighted.backends.iter().enumerate() {
3077            backend.borrow_mut().load_balancing_parameters = Some(LoadBalancingParams {
3078                weight: 10 + 90 * (index as i32 % 4),
3079            });
3080        }
3081        let mut weights_moved_a_shard = false;
3082        for list in [&mut unweighted, &mut weighted] {
3083            for key in 0..64u64 {
3084                let shard = shard_of(list, key);
3085                assert_eq!(shard.len(), 3, "25% of 12 backends is a shard of 3");
3086                assert_eq!(shard, shard_of(list, key), "a shard is deterministic");
3087                let mut everyone: Vec<(f64, SocketAddr)> = list
3088                    .backends
3089                    .iter()
3090                    .map(|backend| {
3091                        let backend = backend.borrow();
3092                        (hrw_score(key, &backend), backend.address)
3093                    })
3094                    .collect();
3095                everyone.sort_by(|a, b| b.0.total_cmp(&a.0));
3096                let top: Vec<SocketAddr> = everyone.iter().take(3).map(|(_, a)| *a).collect();
3097                assert_eq!(shard, top, "key {key}: the shard is the HRW top 3");
3098            }
3099        }
3100        for key in 0..64u64 {
3101            weights_moved_a_shard |= shard_of(&mut unweighted, key) != shard_of(&mut weighted, key);
3102        }
3103        assert!(
3104            weights_moved_a_shard,
3105            "the weighted ranking must differ from the unweighted one for some key"
3106        );
3107    }
3108
3109    /// #524: a change of N moves a client only at the tail of its ranking.
3110    /// Adding a backend with k unchanged replaces at most the lowest-ranked
3111    /// member; adding one that makes k grow keeps every member; removing a
3112    /// non-member changes nothing, and removing a member keeps the others.
3113    ///
3114    /// TO SEE THIS RED: rank by list position instead of HRW score in
3115    /// `BackendList::compute_shard`; adding a backend then reshuffles shards
3116    /// that should not move.
3117    #[test]
3118    fn a_shard_changes_only_at_its_tail_when_backends_change() {
3119        let now = Instant::now();
3120        let extra = |port: u16| {
3121            Backend::new_at(
3122                &format!("extra-{port}"),
3123                SocketAddr::from(([127, 0, 0, 1], port)),
3124                None,
3125                None,
3126                None,
3127                now,
3128            )
3129        };
3130        for key in 0..128u64 {
3131            // 10% of 10 and of 11 both give k = max(2, ceil(1.0 | 1.1)) = 2,
3132            // so k stays 2 across the add.
3133            let mut list = sharded_list(
3134                10,
3135                LoadBalancingAlgorithms::RoundRobin,
3136                sharding(10, 8, ShardMode::Fallback),
3137            );
3138            let before = shard_of(&mut list, key);
3139            list.add_backend(extra(21_000));
3140            let after = shard_of(&mut list, key);
3141            let lost: Vec<_> = before.iter().filter(|a| !after.contains(a)).collect();
3142            assert!(lost.len() <= 1, "key {key}: at most one member displaced");
3143            if let Some(lost) = lost.first() {
3144                assert_eq!(
3145                    Some(*lost),
3146                    before.last(),
3147                    "key {key}: only the lowest-ranked member may be displaced"
3148                );
3149            }
3150
3151            // k grows with N: every member stays.
3152            let mut list = sharded_list(
3153                10,
3154                LoadBalancingAlgorithms::RoundRobin,
3155                sharding(20, 8, ShardMode::Fallback),
3156            );
3157            let before = shard_of(&mut list, key);
3158            list.add_backend(extra(21_001));
3159            let after = shard_of(&mut list, key);
3160            assert_eq!((before.len(), after.len()), (2, 3));
3161            assert!(
3162                before.iter().all(|a| after.contains(a)),
3163                "key {key}: a growing shard keeps its members"
3164            );
3165
3166            // Removing a non-member changes nothing; removing a member keeps
3167            // the others.
3168            let mut list = sharded_list(
3169                10,
3170                LoadBalancingAlgorithms::RoundRobin,
3171                sharding(20, 8, ShardMode::Fallback),
3172            );
3173            let before = shard_of(&mut list, key);
3174            let outsider = list
3175                .backends
3176                .iter()
3177                .map(|b| b.borrow().address)
3178                .find(|a| !before.contains(a))
3179                .expect("a shard of 2 among 10 leaves outsiders");
3180            // `sharded_list` names each backend after its port.
3181            let id_of = |address: &SocketAddr| format!("shard-{}", address.port() - 20_000);
3182            assert!(list.remove_backend(&id_of(&outsider), &outsider));
3183            assert_eq!(shard_of(&mut list, key), before, "key {key}");
3184            assert!(list.remove_backend(&id_of(&before[0]), &before[0]));
3185            assert!(
3186                shard_of(&mut list, key).contains(&before[1]),
3187                "key {key}: removing one member keeps the other"
3188            );
3189        }
3190    }
3191
3192    /// #524: every selection of a keyed client stays in its shard, under a
3193    /// policy that would otherwise visit every backend.
3194    #[test]
3195    fn a_sharded_selection_stays_in_the_shard() {
3196        let now = Instant::now();
3197        for policy in [
3198            LoadBalancingAlgorithms::RoundRobin,
3199            LoadBalancingAlgorithms::LeastLoaded,
3200            LoadBalancingAlgorithms::Random,
3201            LoadBalancingAlgorithms::Hrw,
3202        ] {
3203            let mut list = sharded_list(8, policy, sharding(25, 8, ShardMode::Fallback));
3204            for key in 0..32u64 {
3205                let shard = shard_of(&mut list, key);
3206                for _ in 0..8 {
3207                    let (picked, outcome) = list.select_with_key(Some(key), now);
3208                    let picked = picked.expect("a healthy shard selects").borrow().address;
3209                    assert_eq!(outcome, ShardOutcome::InShard);
3210                    assert!(
3211                        shard.contains(&picked),
3212                        "{policy:?} key {key}: {picked} is outside {shard:?}"
3213                    );
3214                }
3215            }
3216            // No key, or fewer primaries than the threshold: not sharded.
3217            assert_eq!(list.select_with_key(None, now).1, ShardOutcome::Unsharded);
3218            let mut small = sharded_list(7, policy, sharding(25, 8, ShardMode::Fallback));
3219            assert_eq!(
3220                small.select_with_key(Some(1), now).1,
3221                ShardOutcome::Unsharded
3222            );
3223        }
3224    }
3225
3226    /// Put every member of `key`'s shard out of selection, as a failed
3227    /// connection does: its retry policy enters its back-off window.
3228    fn fail_shard(list: &mut BackendList, key: u64) -> Vec<SocketAddr> {
3229        let shard = shard_of(list, key);
3230        let now = Instant::now();
3231        for backend in &list.backends {
3232            let mut backend = backend.borrow_mut();
3233            if shard.contains(&backend.address) {
3234                backend.retry_policy.fail(now, &mut rand::rng());
3235            }
3236        }
3237        shard
3238    }
3239
3240    /// #524: with its whole shard down, `STRICT` selects nothing although the
3241    /// rest of the cluster is healthy, and `FALLBACK` spills over to a backend
3242    /// outside the shard.
3243    ///
3244    /// TO SEE THIS RED: in `BackendList::select_with_key`, spill over in both
3245    /// modes (drop the `strict` early return); the strict selection then
3246    /// finds a backend outside the shard.
3247    #[test]
3248    fn an_exhausted_shard_refuses_in_strict_and_spills_in_fallback() {
3249        let now = Instant::now();
3250        for key in 0..32u64 {
3251            let mut strict = sharded_list(
3252                8,
3253                LoadBalancingAlgorithms::RoundRobin,
3254                sharding(25, 8, ShardMode::Strict),
3255            );
3256            fail_shard(&mut strict, key);
3257            let (picked, outcome) = strict.select_with_key(Some(key), now);
3258            assert!(picked.is_none(), "key {key}: strict never leaves the shard");
3259            assert_eq!(outcome, ShardOutcome::Exhausted);
3260
3261            let mut fallback = sharded_list(
3262                8,
3263                LoadBalancingAlgorithms::RoundRobin,
3264                sharding(25, 8, ShardMode::Fallback),
3265            );
3266            let shard = fail_shard(&mut fallback, key);
3267            let (picked, outcome) = fallback.select_with_key(Some(key), now);
3268            let picked = picked.expect("fallback spills over").borrow().address;
3269            assert_eq!(outcome, ShardOutcome::SpilledOver);
3270            assert!(
3271                !shard.contains(&picked),
3272                "key {key}: the spill leaves the shard"
3273            );
3274        }
3275    }
3276
3277    /// #524: a strict shard fails open inside itself only. When no backend of
3278    /// the whole cluster passes its health check (the fail-open regime), a
3279    /// strict selection still picks a shard member.
3280    #[test]
3281    fn a_strict_shard_fails_open_inside_itself() {
3282        let now = Instant::now();
3283        for key in 0..16u64 {
3284            let mut list = sharded_list(
3285                8,
3286                LoadBalancingAlgorithms::RoundRobin,
3287                sharding(25, 8, ShardMode::Strict),
3288            );
3289            let shard = shard_of(&mut list, key);
3290            for backend in &list.backends {
3291                backend.borrow_mut().health.status = HealthStatus::Unhealthy;
3292            }
3293            let (picked, outcome) = list.select_with_key(Some(key), now);
3294            let picked = picked.expect("fail-open still routes").borrow().address;
3295            assert!(
3296                shard.contains(&picked),
3297                "key {key}: fail-open stays in the shard"
3298            );
3299            assert_eq!(outcome, ShardOutcome::InShard);
3300        }
3301    }
3302
3303    /// Fail every member of `key`'s shard by health check, as a probe does:
3304    /// they stay `Normal` with an `OKAY` retry policy, so fail-open may still
3305    /// pick them.
3306    fn unhealthy_shard(list: &mut BackendList, key: u64) -> Vec<SocketAddr> {
3307        let shard = shard_of(list, key);
3308        for backend in &list.backends {
3309            let mut backend = backend.borrow_mut();
3310            if shard.contains(&backend.address) {
3311                backend.health.status = HealthStatus::Unhealthy;
3312            }
3313        }
3314        shard
3315    }
3316
3317    /// #524 review L3: `STRICT` refuses when its shard fails its health
3318    /// checks while primary backends outside it are healthy — the fail-open
3319    /// regime is not reached, because the cluster is not all down.
3320    ///
3321    /// TO SEE THIS RED: make the `primaries_can_open` guard of
3322    /// `BackendList::select_with_key` false; the strict shard then fails open
3323    /// and picks one of its unhealthy members.
3324    #[test]
3325    fn a_strict_shard_failing_health_checks_refuses_while_others_are_healthy() {
3326        let now = Instant::now();
3327        for key in 0..32u64 {
3328            let mut list = sharded_list(
3329                8,
3330                LoadBalancingAlgorithms::RoundRobin,
3331                sharding(25, 8, ShardMode::Strict),
3332            );
3333            unhealthy_shard(&mut list, key);
3334            let (picked, outcome) = list.select_with_key(Some(key), now);
3335            assert!(
3336                picked.is_none(),
3337                "key {key}: strict must refuse, not fail open"
3338            );
3339            assert_eq!(outcome, ShardOutcome::Exhausted);
3340        }
3341    }
3342
3343    /// #524 review M1: backups sit outside every shard, so a healthy backup
3344    /// must not decide a strict shard's fate. With every primary failing its
3345    /// health checks, a strict shard fails open inside itself whether or not
3346    /// a backup is healthy.
3347    ///
3348    /// TO SEE THIS RED: count the backup tier in `primaries_can_open`
3349    /// (`available != 0 || <any backup can open>`); the healthy backup then
3350    /// turns the fail-open into a refusal.
3351    #[test]
3352    fn a_healthy_backup_does_not_turn_a_strict_fail_open_into_a_refusal() {
3353        let now = Instant::now();
3354        for key in 0..32u64 {
3355            let mut list = sharded_list(
3356                8,
3357                LoadBalancingAlgorithms::RoundRobin,
3358                sharding(25, 8, ShardMode::Strict),
3359            );
3360            let shard = shard_of(&mut list, key);
3361            for backend in &list.backends {
3362                backend.borrow_mut().health.status = HealthStatus::Unhealthy;
3363            }
3364            list.add_backend(Backend::new_at(
3365                "healthy-backup",
3366                SocketAddr::from(([127, 0, 0, 1], 23_000)),
3367                None,
3368                None,
3369                Some(true),
3370                now,
3371            ));
3372            let (picked, outcome) = list.select_with_key(Some(key), now);
3373            let picked = picked
3374                .expect("a strict shard fails open inside itself")
3375                .borrow()
3376                .address;
3377            assert!(
3378                shard.contains(&picked),
3379                "key {key}: never the backup, never outside"
3380            );
3381            assert_eq!(outcome, ShardOutcome::InShard);
3382        }
3383    }
3384
3385    /// #524 review L2: in `FALLBACK`, fail-open over the whole cluster can
3386    /// pick a member of the exhausted shard; that pick stayed in the shard
3387    /// and must not count as a spill-over.
3388    ///
3389    /// TO SEE THIS RED: drop the `settle_spill_outcome` call after the
3390    /// fail-open pick in `BackendList::select_with_key`; an in-shard pick is
3391    /// then reported `SpilledOver`.
3392    #[test]
3393    fn a_fallback_fail_open_pick_inside_the_shard_is_not_a_spill_over() {
3394        let now = Instant::now();
3395        let key = 3;
3396        let mut list = sharded_list(
3397            8,
3398            LoadBalancingAlgorithms::RoundRobin,
3399            sharding(25, 8, ShardMode::Fallback),
3400        );
3401        let shard = shard_of(&mut list, key);
3402        for backend in &list.backends {
3403            backend.borrow_mut().health.status = HealthStatus::Unhealthy;
3404        }
3405        let (mut in_shard, mut outside) = (0, 0);
3406        // Round-robin over the fail-open set visits every backend in 8 picks.
3407        for _ in 0..8 {
3408            let (picked, outcome) = list.select_with_key(Some(key), now);
3409            let picked = picked.expect("fail-open routes").borrow().address;
3410            if shard.contains(&picked) {
3411                assert_eq!(outcome, ShardOutcome::InShard, "{picked} is a shard member");
3412                in_shard += 1;
3413            } else {
3414                assert_eq!(outcome, ShardOutcome::SpilledOver, "{picked} is outside");
3415                outside += 1;
3416            }
3417        }
3418        assert_eq!(
3419            (in_shard, outside),
3420            (2, 6),
3421            "round-robin visits every backend once"
3422        );
3423    }
3424
3425    /// #524: a sticky cookie naming a live backend wins, even outside the
3426    /// client's shard.
3427    #[test]
3428    fn a_sticky_backend_outside_the_shard_wins() {
3429        let now = Instant::now();
3430        let key = 7;
3431        let mut map = BackendMap::with_seed(524);
3432        map.set_load_balancing_policy_for_cluster(
3433            "sharded",
3434            LoadBalancingAlgorithms::RoundRobin,
3435            None,
3436        );
3437        for index in 0..8u16 {
3438            map.add_backend(
3439                "sharded",
3440                Backend::new_at(
3441                    &format!("shard-{index}"),
3442                    SocketAddr::from(([127, 0, 0, 1], 22_000 + index)),
3443                    Some(format!("sticky-{index}")),
3444                    None,
3445                    None,
3446                    now,
3447                ),
3448            );
3449        }
3450        map.set_shuffle_sharding_for_cluster("sharded", sharding(25, 8, ShardMode::Strict));
3451        let list = map.backends.get_mut("sharded").expect("the cluster exists");
3452        let shard = shard_of(list, key);
3453        let outsider = list
3454            .backends
3455            .iter()
3456            .find(|b| !shard.contains(&b.borrow().address))
3457            .map(|b| {
3458                b.borrow()
3459                    .sticky_id
3460                    .clone()
3461                    .expect("every backend has a sticky id")
3462            })
3463            .expect("a shard of 2 among 8 leaves outsiders");
3464        let reserved = map
3465            .reserve_sticky_backend("sharded", &outsider, Some(key), now)
3466            .expect("the sticky backend is live");
3467        assert_eq!(
3468            reserved.borrow().sticky_id.as_deref(),
3469            Some(outsider.as_str())
3470        );
3471        assert!(!shard.contains(&reserved.borrow().address));
3472    }
3473
3474    /// #524: a sharded selection allocates nothing: the scores, the shard and
3475    /// the candidates live in buffers `add_backend` reserved.
3476    ///
3477    /// TO SEE THIS RED: collect the shard into a fresh `Vec` in
3478    /// `BackendList::compute_shard`; each selection then allocates.
3479    #[test]
3480    fn a_sharded_selection_allocates_nothing() {
3481        use std::hint::black_box;
3482
3483        use crate::test_allocations::allocations;
3484
3485        let now = Instant::now();
3486        let mut list = sharded_list(
3487            16,
3488            LoadBalancingAlgorithms::LeastLoaded,
3489            sharding(25, 8, ShardMode::Fallback),
3490        );
3491        // Warm-up: the first selection may settle lazily built policy state.
3492        drop(list.select_with_key(Some(0), now));
3493        let mut allocated = 0;
3494        for key in 0..256u64 {
3495            let before = allocations();
3496            let picked = black_box(&mut list).select_with_key(black_box(Some(key)), now);
3497            allocated += allocations() - before;
3498            assert_eq!(picked.1, ShardOutcome::InShard);
3499        }
3500        assert_eq!(
3501            allocated, 0,
3502            "256 sharded selections made {allocated} allocations"
3503        );
3504    }
3505}
3506
3507/// A retry skips the backends its request already failed to connect to,
3508/// under every load-balancing policy (sozu-proxy/sozu#1800).
3509#[cfg(test)]
3510mod exclusion_tests {
3511    use super::*;
3512
3513    const POLICIES: [LoadBalancingAlgorithms; 6] = [
3514        LoadBalancingAlgorithms::RoundRobin,
3515        LoadBalancingAlgorithms::Random,
3516        LoadBalancingAlgorithms::LeastLoaded,
3517        LoadBalancingAlgorithms::PowerOfTwo,
3518        LoadBalancingAlgorithms::Hrw,
3519        LoadBalancingAlgorithms::Maglev,
3520    ];
3521
3522    fn address(port: u16) -> SocketAddr {
3523        SocketAddr::from(([127, 0, 0, 1], port))
3524    }
3525
3526    /// A list of `count` healthy backends on ports 9100.. under `policy`.
3527    fn list(policy: LoadBalancingAlgorithms, count: u16) -> BackendList {
3528        let mut list = BackendList::with_seed(7);
3529        for index in 0..count {
3530            list.add_backend(Backend::new(
3531                &format!("b{index}"),
3532                address(9100 + index),
3533                None,
3534                None,
3535                None,
3536            ));
3537        }
3538        list.set_load_balancing_policy(policy, None, 11);
3539        list
3540    }
3541
3542    #[test]
3543    fn a_selection_never_returns_an_excluded_backend_under_any_policy() {
3544        let now = Instant::now();
3545        let exclude = [address(9100), address(9101), address(9102)];
3546        for policy in POLICIES {
3547            let mut list = list(policy, 4);
3548            for key in 0..32 {
3549                // HRW and Maglev pin on the key; the others ignore it.
3550                let (picked, _) = list.select_with_key_excluding(Some(key), now, &exclude);
3551                let picked = picked
3552                    .expect("one backend is not excluded")
3553                    .borrow()
3554                    .address;
3555                assert_eq!(
3556                    picked,
3557                    address(9103),
3558                    "{policy:?} picked {picked}, which the request already tried"
3559                );
3560            }
3561        }
3562    }
3563
3564    #[test]
3565    fn an_empty_exclusion_selects_exactly_as_before() {
3566        let now = Instant::now();
3567        for policy in POLICIES {
3568            let mut excluding = list(policy, 4);
3569            let mut plain = list(policy, 4);
3570            for key in 0..16 {
3571                let left = excluding
3572                    .select_with_key_excluding(Some(key), now, &[])
3573                    .0
3574                    .map(|backend| backend.borrow().address);
3575                let right = plain
3576                    .select_with_key(Some(key), now)
3577                    .0
3578                    .map(|backend| backend.borrow().address);
3579                assert_eq!(left, right, "{policy:?} diverged on key {key}");
3580            }
3581        }
3582    }
3583
3584    #[test]
3585    fn excluding_every_backend_falls_back_to_the_plain_selection() {
3586        let now = Instant::now();
3587        let exclude = [address(9100), address(9101)];
3588        for policy in POLICIES {
3589            let mut list = list(policy, 2);
3590            let (picked, _) = list.select_with_key_excluding(Some(3), now, &exclude);
3591            assert!(
3592                picked.is_some(),
3593                "{policy:?}: a request that tried every backend retries one of them"
3594            );
3595        }
3596    }
3597
3598    #[test]
3599    fn an_excluded_sticky_backend_is_not_followed() {
3600        let now = Instant::now();
3601        let mut map = BackendMap::with_seed(5);
3602        for index in 0..2u16 {
3603            map.add_backend(
3604                "cluster",
3605                Backend::new(
3606                    &format!("b{index}"),
3607                    address(9200 + index),
3608                    Some(format!("sticky{index}")),
3609                    None,
3610                    None,
3611                ),
3612            );
3613        }
3614        let followed = map
3615            .reserve_sticky_backend_excluding("cluster", "sticky0", None, now, &[])
3616            .expect("a backend is selectable");
3617        assert_eq!(followed.borrow().address, address(9200));
3618        let skipped = map
3619            .reserve_sticky_backend_excluding("cluster", "sticky0", None, now, &[address(9200)])
3620            .expect("a backend is selectable");
3621        assert_eq!(
3622            skipped.borrow().address,
3623            address(9201),
3624            "the cookie names a backend the request already failed to reach"
3625        );
3626    }
3627}
3628
3629/// Two backend ids at one address share that address's load, under every
3630/// policy: the address is one candidate, whatever the number of ids
3631/// configured on it.
3632#[cfg(test)]
3633mod same_address_tests {
3634    use super::*;
3635    use crate::load_balancing::affinity_key_from_value;
3636
3637    const POLICIES: [LoadBalancingAlgorithms; 6] = [
3638        LoadBalancingAlgorithms::RoundRobin,
3639        LoadBalancingAlgorithms::Random,
3640        LoadBalancingAlgorithms::LeastLoaded,
3641        LoadBalancingAlgorithms::PowerOfTwo,
3642        LoadBalancingAlgorithms::Hrw,
3643        LoadBalancingAlgorithms::Maglev,
3644    ];
3645
3646    /// Picks per measured distribution.
3647    const PICKS: u32 = 20_000;
3648
3649    const SHARED: SocketAddr = SocketAddr::new(
3650        std::net::IpAddr::V4(std::net::Ipv4Addr::new(127, 0, 0, 1)),
3651        9300,
3652    );
3653    const OTHER: SocketAddr = SocketAddr::new(
3654        std::net::IpAddr::V4(std::net::Ipv4Addr::new(127, 0, 0, 1)),
3655        9301,
3656    );
3657
3658    fn weighted(id: &str, address: SocketAddr, weight: Option<i32>) -> Backend {
3659        Backend::new(
3660            id,
3661            address,
3662            None,
3663            weight.map(|weight| LoadBalancingParams { weight }),
3664            None,
3665        )
3666    }
3667
3668    /// `a` and `b` at [`SHARED`], `c` at [`OTHER`], in that list order.
3669    fn shared_list(policy: LoadBalancingAlgorithms, a: Backend, b: Backend) -> BackendList {
3670        let mut list = BackendList::with_seed(1839);
3671        list.add_backend(a);
3672        list.add_backend(b);
3673        list.add_backend(weighted("c", OTHER, None));
3674        list.set_load_balancing_policy(policy, Some(LoadMetric::Connections), 1839);
3675        list
3676    }
3677
3678    fn equal_list(policy: LoadBalancingAlgorithms) -> BackendList {
3679        shared_list(
3680            policy,
3681            weighted("a", SHARED, None),
3682            weighted("b", SHARED, None),
3683        )
3684    }
3685
3686    /// Run [`PICKS`] selections and return the picked backend ids. A
3687    /// load-aware policy sees each pick as one more open connection, and an
3688    /// affinity policy sees a fresh, well-spread client key per pick.
3689    fn picks(list: &mut BackendList) -> Vec<String> {
3690        let now = Instant::now();
3691        (0..PICKS)
3692            .map(|index| {
3693                let key = affinity_key_from_value(&index.to_le_bytes());
3694                let picked = list
3695                    .next_available_backend_with_key(Some(key), now)
3696                    .expect("three healthy backends");
3697                let mut picked = picked.borrow_mut();
3698                picked.active_connections += 1;
3699                picked.backend_id.to_owned()
3700            })
3701            .collect()
3702    }
3703
3704    fn share_of(ids: &[String], wanted: &[&str]) -> f64 {
3705        ids.iter()
3706            .filter(|id| wanted.contains(&id.as_str()))
3707            .count() as f64
3708            / ids.len() as f64
3709    }
3710
3711    /// The decision of 2026-10-04: one address, one share. Two ids at one
3712    /// address and one id at another, equal weights: each address takes
3713    /// half the selections, as HRW already gave it.
3714    ///
3715    /// TO SEE THIS RED: remove the `collapse_shared_addresses` calls from
3716    /// `BackendList::select_tiers`; round robin, random, least loaded and
3717    /// power of two then give the shared address two thirds.
3718    #[test]
3719    fn a_shared_address_takes_the_share_of_one_backend_under_every_policy() {
3720        let skewed: Vec<String> = POLICIES
3721            .into_iter()
3722            .filter_map(|policy| {
3723                let mut list = equal_list(policy);
3724                let shared = share_of(&picks(&mut list), &["a", "b"]);
3725                (!(0.47..=0.53).contains(&shared)).then(|| format!("{policy:?}: {shared:.3}"))
3726            })
3727            .collect();
3728        assert!(
3729            skewed.is_empty(),
3730            "the shared address did not take 1/2 of the selections under {skewed:?}"
3731        );
3732    }
3733
3734    /// The shared address is served by its first id in list order when the
3735    /// weights are equal: the second id is never picked while the first can
3736    /// take a connection.
3737    #[test]
3738    fn the_first_id_at_a_shared_address_represents_it() {
3739        for policy in POLICIES {
3740            let mut list = equal_list(policy);
3741            let ids = picks(&mut list);
3742            assert!(
3743                ids.iter().all(|id| id != "b"),
3744                "{policy:?} picked the second id of the shared address"
3745            );
3746            assert!(ids.iter().any(|id| id == "a"), "{policy:?} never picked a");
3747        }
3748    }
3749
3750    /// A down id never serves its address: its eligible sibling does, and
3751    /// the address keeps the share of one backend.
3752    #[test]
3753    fn a_down_id_yields_its_address_to_an_eligible_sibling() {
3754        for policy in POLICIES {
3755            let mut down = weighted("a", SHARED, None);
3756            down.health.record_failure(1);
3757            assert!(!down.health.is_healthy());
3758            let mut list = shared_list(policy, down, weighted("b", SHARED, None));
3759            let ids = picks(&mut list);
3760            assert!(
3761                ids.iter().all(|id| id != "a"),
3762                "{policy:?} picked the unhealthy id"
3763            );
3764            let shared = share_of(&ids, &["b"]);
3765            assert!(
3766                (0.47..=0.53).contains(&shared),
3767                "{policy:?} gave the shared address {shared:.3} of the selections, not 1/2"
3768            );
3769        }
3770    }
3771
3772    /// Unequal weights at one address: the heaviest id represents it, as the
3773    /// HRW score — monotonic in weight for one address — has always decided.
3774    #[test]
3775    fn the_heaviest_id_at_a_shared_address_represents_it() {
3776        for policy in POLICIES {
3777            let mut list = shared_list(
3778                policy,
3779                weighted("a", SHARED, Some(50)),
3780                weighted("b", SHARED, Some(200)),
3781            );
3782            let ids = picks(&mut list);
3783            assert!(
3784                ids.iter().all(|id| id != "a"),
3785                "{policy:?} picked the lighter id of the shared address"
3786            );
3787        }
3788    }
3789
3790    /// A shard counts addresses, not ids: eight addresses, one carrying two
3791    /// ids, at 25% make shards of two addresses, and a shard that holds the
3792    /// shared address holds both its ids, so the sibling of a down
3793    /// representative still serves inside the shard.
3794    #[test]
3795    fn a_shard_counts_a_shared_address_once_and_holds_all_its_ids() {
3796        let mut list = BackendList::with_seed(524);
3797        list.set_load_balancing_policy(LoadBalancingAlgorithms::Hrw, None, 524);
3798        for index in 0..8u16 {
3799            list.add_backend(Backend::new(
3800                &format!("shard-{index}"),
3801                SocketAddr::from(([127, 0, 0, 1], 20_000 + index)),
3802                None,
3803                None,
3804                None,
3805            ));
3806        }
3807        let shared = SocketAddr::from(([127, 0, 0, 1], 20_000));
3808        list.add_backend(Backend::new("shard-0-bis", shared, None, None, None));
3809        list.set_shuffle_sharding(Some(ShuffleSharding {
3810            percent: 25,
3811            min_backends: 4,
3812            mode: ShardMode::Fallback,
3813        }));
3814
3815        let mut holding_shared = 0;
3816        for key in 0..256u64 {
3817            assert!(list.compute_shard(Some(key)), "the list must shard");
3818            let mut addresses: Vec<SocketAddr> = list
3819                .shard
3820                .iter()
3821                .map(|&index| list.backends[index].borrow().address)
3822                .collect();
3823            let ids = addresses.len();
3824            addresses.sort_unstable();
3825            addresses.dedup();
3826            assert_eq!(addresses.len(), 2, "key {key}: a shard holds two addresses");
3827            if addresses.contains(&shared) {
3828                holding_shared += 1;
3829                assert_eq!(ids, 3, "key {key}: the shared address brings both ids");
3830            } else {
3831                assert_eq!(ids, 2, "key {key}: one id per unshared address");
3832            }
3833        }
3834        assert!(
3835            holding_shared > 0,
3836            "some shard must hold the shared address"
3837        );
3838    }
3839
3840    /// `LEAST_LOADED` and `POWER_OF_TWO` read an address's load as the sum
3841    /// over its ids: the representative `a` is idle, but its sibling `b`
3842    /// carries 5 connections, so the address weighs 5 against `c`'s 3 and
3843    /// `c` is picked.
3844    #[test]
3845    fn a_load_aware_policy_reads_the_load_of_the_whole_address() {
3846        let now = Instant::now();
3847        for policy in [
3848            LoadBalancingAlgorithms::LeastLoaded,
3849            LoadBalancingAlgorithms::PowerOfTwo,
3850        ] {
3851            let mut list = equal_list(policy);
3852            list.backends[1].borrow_mut().active_connections = 5;
3853            list.backends[2].borrow_mut().active_connections = 3;
3854            for key in 0..64 {
3855                let picked = list
3856                    .next_available_backend_with_key(Some(key), now)
3857                    .expect("three healthy backends");
3858                assert_eq!(
3859                    picked.borrow().backend_id,
3860                    "c",
3861                    "{policy:?} compared the representative's own load, not its address's"
3862                );
3863            }
3864        }
3865    }
3866
3867    /// The cost guard: one shared pair among 1000 backends costs a selection
3868    /// a few units of shared-address work (`SIBLING_STEPS`: one per
3869    /// representative comparison or sibling-ring step), not one per
3870    /// candidate. A scan of the candidates per candidate costs a million.
3871    ///
3872    /// TO SEE THIS RED: replace the ring walk in
3873    /// `BackendList::collapse_shared_addresses` with
3874    /// `candidates.iter().any(|&other| outranks_at_address(..))`.
3875    #[test]
3876    fn one_shared_pair_costs_a_selection_constant_work() {
3877        use crate::load_balancing::SIBLING_STEPS;
3878
3879        let now = Instant::now();
3880        // Built once: the debug-only invariant sweep after each insertion is
3881        // quadratic in the list, so a list per case would dominate the run.
3882        let mut list = BackendList::with_seed(1856);
3883        for index in 0..1000u16 {
3884            let address = SocketAddr::from(([10, (index >> 8) as u8, index as u8, 1], 8080));
3885            list.add_backend(Backend::new(
3886                &format!("b{index}"),
3887                address,
3888                None,
3889                None,
3890                None,
3891            ));
3892        }
3893        let shared = SocketAddr::from(([10, 0, 0, 1], 8080));
3894        list.add_backend(Backend::new("dup", shared, None, None, None));
3895        for sharded in [false, true] {
3896            list.set_shuffle_sharding(sharded.then_some(ShuffleSharding {
3897                percent: 25,
3898                min_backends: 8,
3899                mode: ShardMode::Fallback,
3900            }));
3901            for policy in POLICIES {
3902                list.set_load_balancing_policy(policy, Some(LoadMetric::Connections), 1856);
3903                let mut worst = 0;
3904                for key in 0..64u64 {
3905                    SIBLING_STEPS.with(|steps| steps.set(0));
3906                    assert!(list.select_with_key(Some(key), now).0.is_some());
3907                    worst = worst.max(SIBLING_STEPS.with(|steps| steps.get()));
3908                }
3909                assert!(
3910                    worst <= 8,
3911                    "{policy:?} (sharded: {sharded}) spent {worst} steps on one shared pair"
3912                );
3913            }
3914        }
3915    }
3916
3917    /// `backends` is public: a library user may push or pop a backend
3918    /// without `add_backend`/`remove_backend`, leaving the sibling state of
3919    /// another length. Selection must then neither panic nor index a stale
3920    /// ring, under every policy, sharded or not.
3921    #[test]
3922    fn a_list_edited_directly_still_selects_without_panicking() {
3923        let now = Instant::now();
3924        let sharding = Some(ShuffleSharding {
3925            percent: 50,
3926            min_backends: 2,
3927            mode: ShardMode::Fallback,
3928        });
3929        for policy in POLICIES {
3930            for sharded in [false, true] {
3931                // Grown: a pushed sibling of the shared address, and a
3932                // pushed backend at a new address.
3933                let mut grown = equal_list(policy);
3934                grown.set_shuffle_sharding(sharded.then_some(sharding).flatten());
3935                for (id, port) in [("pushed-sibling", 9300), ("pushed", 9302)] {
3936                    let address = SocketAddr::from(([127, 0, 0, 1], port));
3937                    grown
3938                        .backends
3939                        .push(Rc::new(RefCell::new(weighted(id, address, None))));
3940                }
3941                // Shrunk: the last backend popped, the ring still covers it.
3942                let mut shrunk = equal_list(policy);
3943                shrunk.set_shuffle_sharding(sharded.then_some(sharding).flatten());
3944                shrunk.backends.pop();
3945                for list in [&mut grown, &mut shrunk] {
3946                    for key in 0..64 {
3947                        assert!(
3948                            list.select_with_key(Some(key), now).0.is_some(),
3949                            "{policy:?} (sharded: {sharded}) selected nothing"
3950                        );
3951                    }
3952                }
3953            }
3954        }
3955    }
3956
3957    /// Collapsing allocates nothing: the scratch buffer is reserved on the
3958    /// control plane like the candidate buffer.
3959    #[test]
3960    fn a_selection_over_a_shared_address_allocates_nothing() {
3961        use std::hint::black_box;
3962
3963        use crate::test_allocations::allocations;
3964
3965        let now = Instant::now();
3966        for policy in POLICIES {
3967            let mut list = equal_list(policy);
3968            drop(list.select_with_key(Some(0), now));
3969            let mut allocated = 0;
3970            for key in 0..256u64 {
3971                let before = allocations();
3972                let picked = black_box(&mut list).select_with_key(black_box(Some(key)), now);
3973                allocated += allocations() - before;
3974                assert!(picked.0.is_some());
3975            }
3976            assert_eq!(
3977                allocated, 0,
3978                "{policy:?}: 256 selections made {allocated} allocations"
3979            );
3980        }
3981    }
3982}