Skip to main content

sozu_lib/
load_balancing.rs

1use std::{
2    cell::RefCell, cmp::Ordering, fmt::Debug, hash::Hasher, net::SocketAddr, ops::Index, rc::Rc,
3    time::Instant,
4};
5
6use rand::{
7    RngExt, SeedableRng,
8    distr::uniform::{UniformInt, UniformSampler},
9    prelude::IndexedRandom,
10    rngs::StdRng,
11};
12
13use sozu_command::config::DEFAULT_SHARD_MIN_BACKENDS;
14
15use crate::{
16    backends::Backend,
17    sozu_command::proto::command::{LoadMetric, ShardMode},
18};
19
20/// Default weight applied when a backend declares no explicit
21/// `load_balancing_parameters.weight`. Mirrors the value used by the
22/// `Random` algorithm so weighting stays consistent across policies.
23const DEFAULT_WEIGHT: i32 = 100;
24
25/// Fixed seed used by the affinity hashers (HRW / Maglev). It must NOT be a
26/// `RandomState`: the selection has to be reproducible across workers and
27/// across process restarts so that every Sōzu instance routing the same flow
28/// key lands on the same backend. The constant is arbitrary but stable.
29pub const DEFAULT_HASH_SEED: u64 = 0x9E37_79B9_7F4A_7C15;
30
31/// Read a backend's weight, defaulting to [`DEFAULT_WEIGHT`] when unset. Weight
32/// is clamped to at least `1` so a `0`/negative configured weight never zeroes
33/// out a backend's share (which would make weighted hashing degenerate).
34fn backend_weight(backend: &Backend) -> u32 {
35    let weight = backend
36        .load_balancing_parameters
37        .as_ref()
38        .map(|p| p.weight)
39        .unwrap_or(DEFAULT_WEIGHT)
40        .max(1) as u32;
41    // The clamp guarantees every backend carries a strictly positive weight,
42    // otherwise weighted hashing (HRW score, Maglev slot share) would zero out
43    // the backend's share and skew or divide-by-zero the distribution.
44    debug_assert!(weight >= 1, "backend weight must be clamped to at least 1");
45    weight
46}
47
48/// Whether `other`, at list position `other_position`, takes the place of
49/// `backend`, at `position`, as the one candidate of their shared address.
50///
51/// One address is one share: however many backend ids a cluster configures
52/// at one address, the policies see that address as ONE candidate. Its
53/// representative is the heaviest of its ids ([`backend_weight`], clamped),
54/// the earliest in list order among equally heavy ones; the address then
55/// weighs what its representative weighs. That is the rule [`Rendezvous`]
56/// has always applied, and it is spelled here so every other policy and
57/// shuffle sharding apply it too: the HRW score of an address grows with the
58/// weight, and a tie keeps the earliest backend, so HRW picks exactly this
59/// representative among the ids of an address. With equal weights, the
60/// common case, the representative is the first id in list order.
61///
62/// Callers compare candidates that already passed their eligibility filter,
63/// so a down or backing-off id never represents an address an eligible
64/// sibling could serve.
65pub(crate) fn outranks_at_address(
66    other: &Backend,
67    other_position: usize,
68    backend: &Backend,
69    position: usize,
70) -> bool {
71    count_sibling_step();
72    if other.address != backend.address || other_position == position {
73        return false;
74    }
75    let (other_weight, weight) = (backend_weight(other), backend_weight(backend));
76    other_weight > weight || (other_weight == weight && other_position < position)
77}
78
79#[cfg(test)]
80thread_local! {
81    /// Work spent on shared addresses: one per [`outranks_at_address`]
82    /// comparison and one per step along a sibling ring. Lets a test hold
83    /// that a selection pays for the ids that share an address, not for the
84    /// size of the list.
85    pub(crate) static SIBLING_STEPS: std::cell::Cell<u64> = const { std::cell::Cell::new(0) };
86}
87
88/// Count one unit of shared-address work, in tests only (see
89/// `SIBLING_STEPS`); compiles to nothing otherwise.
90#[inline(always)]
91pub(crate) fn count_sibling_step() {
92    #[cfg(test)]
93    SIBLING_STEPS.with(|steps| steps.set(steps.get() + 1));
94}
95
96/// Deterministic, seedable 64-bit hash over the backend's STABLE identifier.
97///
98/// We hash the backend **socket address** (`SocketAddr`) rather than the
99/// `backend_id` string: the address is where traffic actually goes, so it is
100/// the routing-stable part of a backend and survives a reconfiguration that
101/// merely re-emits the same backend. Backend lifecycle identity is the
102/// `(backend_id, address)` pair (`BackendList::remove_backend` removes only
103/// the backend matching both, so two ids may share an address), but placement
104/// deliberately keys on the address alone: HRW/Maglev then stay stable when
105/// the control plane renames a backend without moving it.
106///
107/// Keying on the address also means one address is one share: two ids at one
108/// address hash identically, so under HRW they compete once, and Maglev
109/// claims table slots once per distinct address. Every policy follows the
110/// same rule, since `BackendList` hands it one candidate per address, the
111/// representative [`outranks_at_address`] defines.
112///
113/// Uses `std::hash::SipHasher13` indirectly via a tiny FNV-1a construction with
114/// an injected seed — fully reproducible, never `RandomState`.
115fn hash_backend(seed: u64, key: u64, addr: &SocketAddr) -> u64 {
116    let mut h = FnvHasher::with_seed(seed);
117    h.write_u64(key);
118    match addr {
119        SocketAddr::V4(v4) => {
120            h.write_u8(4);
121            h.write(&v4.ip().octets());
122            h.write_u16(v4.port());
123        }
124        SocketAddr::V6(v6) => {
125            h.write_u8(6);
126            h.write(&v6.ip().octets());
127            h.write_u16(v6.port());
128        }
129    }
130    // FNV-1a alone has weak avalanche on structured input (the low bits barely
131    // mix), which skews rendezvous/Maglev distribution. Apply a splitmix64
132    // finalizer to scatter the bits before the consumers reduce mod M or feed
133    // it into the score function. Deterministic and seed-independent.
134    splitmix64_finalize(h.finish())
135}
136
137/// splitmix64 mixing step — a strong, reversible 64-bit avalanche finalizer.
138/// Used to scatter the FNV output so downstream `% M` / `ln(h)` reductions see
139/// well-distributed bits.
140fn splitmix64_finalize(mut z: u64) -> u64 {
141    z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
142    z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
143    z ^ (z >> 31)
144}
145
146/// Domain tag of an affinity key derived from a client IP address. Keeps an
147/// address and a header value spelling the same bytes on distinct keys.
148const AFFINITY_DOMAIN_IP: u8 = 0x49; // 'I'
149/// Domain tag of an affinity key derived from a request header or cookie value.
150const AFFINITY_DOMAIN_VALUE: u8 = 0x56; // 'V'
151
152/// The affinity key [`Rendezvous`] and [`Maglev`] pin a client with, derived
153/// from its source address.
154///
155/// The address is canonicalised first, so a dual-stack listener that reports
156/// an IPv4 client as `::ffff:a.b.c.d` keys it exactly as an IPv4 listener
157/// would. The hash is [`DEFAULT_HASH_SEED`]-seeded FNV-1a with the splitmix64
158/// finalizer: a pure function of the address, so every worker and every
159/// restart derives the same key for the same client, and it reads no clock,
160/// draws no randomness and allocates nothing.
161pub fn affinity_key_from_ip(ip: std::net::IpAddr) -> u64 {
162    let mut h = FnvHasher::with_seed(DEFAULT_HASH_SEED);
163    h.write_u8(AFFINITY_DOMAIN_IP);
164    match ip.to_canonical() {
165        std::net::IpAddr::V4(v4) => {
166            h.write_u8(4);
167            h.write(&v4.octets());
168        }
169        std::net::IpAddr::V6(v6) => {
170            h.write_u8(6);
171            h.write(&v6.octets());
172        }
173    }
174    splitmix64_finalize(h.finish())
175}
176
177/// The affinity key derived from the bytes of a request header or cookie
178/// value, for a cluster that keys its clients on one.
179///
180/// Same construction as [`affinity_key_from_ip`] under another domain tag, so
181/// a value never collides with an address by spelling its bytes. The bytes
182/// are hashed where they lie in the request buffer: nothing is copied.
183pub fn affinity_key_from_value(value: &[u8]) -> u64 {
184    let mut h = FnvHasher::with_seed(DEFAULT_HASH_SEED);
185    h.write_u8(AFFINITY_DOMAIN_VALUE);
186    h.write(value);
187    splitmix64_finalize(h.finish())
188}
189
190/// A cluster's shuffle-sharding policy (sozu-proxy/sozu#524).
191///
192/// Shuffle sharding (Colm MacCárthaigh, *Workload isolation using shuffle
193/// sharding*, Amazon Builders' Library, 2019) gives each client a shard: a
194/// small subset of the cluster's backends, drawn so that two clients rarely
195/// share all of theirs. A client that overloads or poisons its shard then
196/// takes down only the clients whose shard is entirely inside its own. With
197/// `N` backends in shards of `k` there are `C(N, k)` shards, against `N / k`
198/// with plain partitioning: `C(8, 2) = 28` against 4.
199///
200/// Here a client's shard is the top `k` of the rendezvous (HRW) ranking of
201/// its affinity key over the cluster's primary backends ([`hrw_score`];
202/// David Thaler and Chinya Ravishankar, *Using name-based mappings to
203/// increase hit rates*, IEEE/ACM Transactions on Networking, 1998). That
204/// makes a shard stable: removing a backend outside a client's shard leaves
205/// the shard unchanged, and adding one can only displace the shard's
206/// lowest-ranked member, so a change of `N` moves clients only at the tail of
207/// their ranking. The cluster's load-balancing policy then selects inside the
208/// shard.
209#[derive(Clone, Copy, Debug, PartialEq, Eq)]
210pub struct ShuffleSharding {
211    /// Share of the primary backends in a shard, `1..=100`.
212    pub percent: u32,
213    /// Sharding applies only from this many primary backends.
214    pub min_backends: u32,
215    /// What an exhausted shard does.
216    pub mode: ShardMode,
217}
218
219impl ShuffleSharding {
220    /// The sharding a cluster configures, `None` when `shard_percent` is
221    /// unset. `shard_min_backends` defaults to
222    /// [`DEFAULT_SHARD_MIN_BACKENDS`], `shard_mode` to `FALLBACK`.
223    pub fn from_cluster(
224        shard_percent: Option<u32>,
225        shard_min_backends: Option<u32>,
226        shard_mode: Option<i32>,
227    ) -> Option<Self> {
228        let percent = shard_percent?;
229        Some(Self {
230            percent,
231            min_backends: shard_min_backends.unwrap_or(DEFAULT_SHARD_MIN_BACKENDS),
232            mode: shard_mode
233                .and_then(|mode| ShardMode::try_from(mode).ok())
234                .unwrap_or(ShardMode::Fallback),
235        })
236    }
237
238    /// The shard size `k` for a cluster of `primaries` primary backends:
239    /// `max(2, ceil(percent × primaries / 100))`, capped at `primaries`.
240    /// `None` below `min_backends`, where the cluster is not sharded.
241    pub fn shard_size(&self, primaries: usize) -> Option<usize> {
242        if primaries < self.min_backends as usize || primaries == 0 {
243            return None;
244        }
245        let k = (self.percent as usize * primaries)
246            .div_ceil(100)
247            .max(2)
248            .min(primaries);
249        debug_assert!(
250            (1..=primaries).contains(&k),
251            "a shard holds at least one and at most every primary backend"
252        );
253        debug_assert!(
254            k >= 2 || primaries < 2,
255            "a shard keeps room for a retry whenever the cluster has two backends"
256        );
257        Some(k)
258    }
259}
260
261/// Smallest prime `>= n`. Maglev requires a prime table size `M`: it keeps the
262/// permutation stride coprime with `M` so the population loop visits every
263/// slot, and `M >= 2` so `skip = h2 % (M - 1) + 1` never divides by zero.
264/// Runs once at construction (off the datapath); a trial-division check is more
265/// than fast enough for the small sizes Sōzu uses (default 65537).
266fn next_prime(n: usize) -> usize {
267    let mut candidate = n.max(2);
268    while !is_prime(candidate) {
269        candidate += 1;
270    }
271    // The result is the smallest prime not below `max(n, 2)`: it is prime, it
272    // is at least 2, and it never went below the requested floor.
273    debug_assert!(is_prime(candidate), "next_prime must return a prime");
274    debug_assert!(candidate >= 2, "next_prime must return at least 2");
275    debug_assert!(
276        candidate >= n,
277        "next_prime ({candidate}) must be >= the requested floor ({n})"
278    );
279    candidate
280}
281
282/// Trial-division primality test. Module-level (not nested in `next_prime`) so
283/// the Maglev `rebuild` post-condition can re-assert that the table size stayed
284/// the prime it was constructed with.
285fn is_prime(x: usize) -> bool {
286    if x < 2 {
287        return false;
288    }
289    if x.is_multiple_of(2) {
290        return x == 2;
291    }
292    let mut d = 3;
293    while d * d <= x {
294        if x.is_multiple_of(d) {
295            return false;
296        }
297        d += 2;
298    }
299    true
300}
301
302/// Minimal FNV-1a 64-bit hasher with a seeded offset basis. Deterministic and
303/// dependency-free — used for the affinity algorithms so results are
304/// reproducible across runs (unlike `std::collections::hash_map::RandomState`).
305struct FnvHasher {
306    state: u64,
307}
308
309impl FnvHasher {
310    const PRIME: u64 = 0x0000_0100_0000_01B3;
311    const OFFSET: u64 = 0xcbf2_9ce4_8422_2325;
312
313    fn with_seed(seed: u64) -> Self {
314        Self {
315            state: Self::OFFSET ^ seed,
316        }
317    }
318}
319
320impl Hasher for FnvHasher {
321    fn finish(&self) -> u64 {
322        self.state
323    }
324
325    fn write(&mut self, bytes: &[u8]) {
326        for &b in bytes {
327            self.state ^= u64::from(b);
328            self.state = self.state.wrapping_mul(Self::PRIME);
329        }
330    }
331}
332
333/// The candidates of one selection: a borrowed view of a cluster's backend
334/// list, restricted to the positions a caller retained.
335///
336/// `BackendList` retains one position per address: when several eligible ids
337/// share an address, only their representative (`outranks_at_address`) is
338/// a candidate, so every policy gives that address the share of one backend.
339///
340/// Position `i` of the view is `backends[indices[i]]`, so a policy that
341/// indexes the view (`RoundRobin`'s cursor, `PowerOfTwo`'s two samples,
342/// `Random`'s draw) computes exactly the index it computed over the
343/// `Vec` of cloned candidates this view replaces. Building the view borrows
344/// and copies nothing: the caller owns `indices`, and only the backend a
345/// policy returns has its `Rc` cloned.
346///
347/// The view also carries the instant the selection happens at, which a
348/// policy reading `LoadMetric::ConnectionTime` decays each average to: no
349/// policy reads a clock of its own (#1684).
350#[derive(Clone, Copy, Debug)]
351pub struct Candidates<'a> {
352    backends: &'a [Rc<RefCell<Backend>>],
353    indices: &'a [usize],
354    /// Sibling ring of `backends`: `siblings[i]` is the next position whose
355    /// backend shares the address of position `i`, cyclically, and `i` itself
356    /// for an address with one id. Empty when the caller tracks no siblings.
357    siblings: &'a [usize],
358    now: Instant,
359}
360
361impl<'a> Candidates<'a> {
362    /// View `backends` through `indices` for a selection at `now`. Every
363    /// index must address a slot of `backends`, and the indices must be
364    /// strictly increasing.
365    pub fn new(backends: &'a [Rc<RefCell<Backend>>], indices: &'a [usize], now: Instant) -> Self {
366        debug_assert!(
367            indices.iter().all(|&index| index < backends.len()),
368            "a candidate index must address a slot of the backend list"
369        );
370        // Strictly increasing positions keep the view in list order, which
371        // the first-wins tie-breaks of `LeastLoaded` and `Rendezvous` and the
372        // probe order of `Maglev` depend on, and rule out a duplicate.
373        debug_assert!(
374            indices.windows(2).all(|pair| pair[0] < pair[1]),
375            "candidate indices must be strictly increasing"
376        );
377        Self {
378            backends,
379            indices,
380            siblings: &[],
381            now,
382        }
383    }
384
385    /// [`Candidates::new`] with the sibling ring of `backends` (see the
386    /// field docs), which [`Candidates::siblings`] walks so a load-aware
387    /// policy can read the load of a whole address.
388    pub fn with_siblings(
389        backends: &'a [Rc<RefCell<Backend>>],
390        indices: &'a [usize],
391        siblings: &'a [usize],
392        now: Instant,
393    ) -> Self {
394        // A ring of another length was built for another list (a caller
395        // edited the public backend list directly): ignore it rather than
396        // follow stale positions, as if no address were shared.
397        let siblings = if siblings.len() == backends.len() {
398            siblings
399        } else {
400            &[]
401        };
402        debug_assert!(
403            siblings.iter().all(|&next| next < backends.len()),
404            "a sibling must address a slot of the backend list"
405        );
406        Self {
407            siblings,
408            ..Self::new(backends, indices, now)
409        }
410    }
411
412    /// Every OTHER backend at the address of the candidate at `index`,
413    /// eligible or not: the ids a collapsed address hides behind its
414    /// representative. Empty without a sibling ring, and for an address with
415    /// one id, which costs one read.
416    pub fn siblings(
417        &self,
418        index: usize,
419    ) -> impl Iterator<Item = &'a Rc<RefCell<Backend>>> + use<'a> {
420        let start = self.indices[index];
421        let ring = self.siblings;
422        let backends = self.backends;
423        let next = move |position: usize| ring.get(position).copied().unwrap_or(start);
424        std::iter::successors(Some(next(start)), move |&position| Some(next(position)))
425            .take_while(move |&position| position != start)
426            .map(move |position| {
427                count_sibling_step();
428                &backends[position]
429            })
430    }
431
432    /// The position in the backend list of the candidate at `index`.
433    fn position(&self, index: usize) -> usize {
434        self.indices[index]
435    }
436
437    /// The length of the backend list the view reads.
438    fn list_len(&self) -> usize {
439        self.backends.len()
440    }
441
442    /// The instant this selection happens at.
443    pub fn now(&self) -> Instant {
444        self.now
445    }
446
447    pub fn len(&self) -> usize {
448        self.indices.len()
449    }
450
451    pub fn is_empty(&self) -> bool {
452        self.indices.is_empty()
453    }
454
455    /// The candidate at position `index` of the view.
456    pub fn get(&self, index: usize) -> Option<&'a Rc<RefCell<Backend>>> {
457        self.indices.get(index).map(|&slot| &self.backends[slot])
458    }
459
460    /// The candidates in view order.
461    pub fn iter(&self) -> impl ExactSizeIterator<Item = &'a Rc<RefCell<Backend>>> + use<'a> {
462        let backends = self.backends;
463        self.indices.iter().map(move |&slot| &backends[slot])
464    }
465
466    /// A uniformly chosen candidate, drawn exactly as
467    /// `IndexedRandom::choose` draws over a slice of the same length.
468    fn choose<R: rand::Rng + ?Sized>(&self, rng: &mut R) -> Option<&'a Rc<RefCell<Backend>>> {
469        self.indices.choose(rng).map(|&slot| &self.backends[slot])
470    }
471}
472
473impl Index<usize> for Candidates<'_> {
474    type Output = Rc<RefCell<Backend>>;
475
476    fn index(&self, index: usize) -> &Self::Output {
477        &self.backends[self.indices[index]]
478    }
479}
480
481pub trait LoadBalancingAlgorithm: Debug {
482    /// Select the next backend among `candidates`.
483    ///
484    /// `key` carries an optional affinity hash: the UDP affinity key, or the
485    /// client key of an HTTP, HTTPS or TCP request ([`affinity_key_from_ip`],
486    /// [`affinity_key_from_value`]). The
487    /// stateless/round-robin policies ignore it; the consistent-hashing
488    /// policies ([`Rendezvous`], [`Maglev`]) use it to pin a key to a backend.
489    /// Passing `None` preserves the historical, key-agnostic behavior.
490    ///
491    /// A policy reads the candidates through the borrowed [`Candidates`] view
492    /// and clones the `Rc` of the one backend it returns, nothing else, so a
493    /// selection allocates nothing.
494    fn next_available_backend(
495        &mut self,
496        key: Option<u64>,
497        candidates: Candidates<'_>,
498    ) -> Option<Rc<RefCell<Backend>>>;
499
500    /// Called by the control plane when the live backend set for a cluster
501    /// changes, so table-based policies (Maglev) can recompute their lookup
502    /// table off the datapath. The default is a no-op; [`Maglev`] rebuilds
503    /// its table, and [`RoundRobin`] (with the round-robin fallback of
504    /// [`Rendezvous`] and [`Maglev`]) resets its running weights.
505    fn rebuild(&mut self, _backends: &[Rc<RefCell<Backend>>]) {}
506}
507
508/// Round-robin selection, weighted by `weight` through smooth weighted round
509/// robin (the nginx `ngx_http_upstream_get_peer` scheme).
510///
511/// # Equal weights: the cursor
512///
513/// When every candidate carries the same weight — the default, since an unset
514/// `weight` reads as `DEFAULT_WEIGHT` — the selection is the historical
515/// cursor over the candidate view: position `next_backend % len`, then the
516/// cursor advances. The weighted scheme would visit a fixed set in the same
517/// order, but the cursor is what decides how a selection continues once the
518/// set shrinks or grows (a backend backs off, recovers, or is excluded by a
519/// retry), so keeping it for equal weights keeps every such sequence exactly
520/// as it was.
521///
522/// # Unequal weights: smooth weighted round robin
523///
524/// Each backend of the list carries a running `current` weight. A selection
525/// adds every candidate's weight to its `current`, picks the candidate with
526/// the largest `current` (the first in view order on a tie), and subtracts the
527/// sum of the candidates' weights from the winner's. Over a fixed set, a
528/// backend of weight `w` is picked exactly `w` times in every `Σw` selections,
529/// and the picks interleave rather than bunch: weights `5, 1, 1` yield
530/// `a a b a c a a`, not `a a a a a b c`.
531///
532/// The weight is `backend_weight`'s, clamped to at least `1`, as for
533/// [`Rendezvous`] and [`Maglev`]: a `0` or negative weight leaves a backend in
534/// the rotation at the smallest share, never out of it.
535///
536/// The running weights are indexed by list position. [`Self::rebuild`], which
537/// the backend list calls whenever a backend is added, removed or reweighted,
538/// resets them, off the datapath; a backend that is not a candidate of one
539/// selection keeps its running weight for the next.
540///
541/// Telling the two cases apart reads every candidate's weight on each
542/// selection. That is deliberate: the candidates are the eligible subset of
543/// the list, which changes with health and retries, so a flag cached on the
544/// list could not say whether THIS subset is uniform.
545#[derive(Debug)]
546pub struct RoundRobin {
547    pub next_backend: u32,
548    /// Smooth weighted round robin's running weight of each backend, by list
549    /// position. Sized to the list by [`Self::rebuild`]; a selection over a
550    /// list of another length resizes it first.
551    current_weights: Vec<i64>,
552}
553
554impl LoadBalancingAlgorithm for RoundRobin {
555    fn next_available_backend(
556        &mut self,
557        _key: Option<u64>,
558        backends: Candidates<'_>,
559    ) -> Option<Rc<RefCell<Backend>>> {
560        // Guard against an empty set: `% backends.len()` would panic with a
561        // divide-by-zero. This also covers the `Rendezvous`/`Maglev` policies,
562        // which delegate here on `key == None`.
563        if backends.is_empty() {
564            return None;
565        }
566        debug_assert!(
567            !backends.is_empty(),
568            "round-robin index math runs only on a non-empty set"
569        );
570
571        if !uniform_weights(backends) {
572            return Some(self.smooth_weighted(backends));
573        }
574
575        let index = self.next_backend as usize % backends.len();
576        // The reduced index always addresses a real slot, so the lookup yields
577        // a backend (never the `None` arm of `get`).
578        debug_assert!(index < backends.len(), "round-robin index out of bounds");
579        let res = backends.get(index).cloned();
580        debug_assert!(
581            res.is_some(),
582            "round-robin must select a backend from a non-empty set"
583        );
584
585        self.next_backend = (self.next_backend + 1) % backends.len() as u32;
586        // The cursor stays a valid index into the current set after advancing,
587        // so the next call never starts out of range.
588        debug_assert!(
589            (self.next_backend as usize) < backends.len(),
590            "round-robin cursor must stay within bounds"
591        );
592        res
593    }
594
595    fn rebuild(&mut self, backends: &[Rc<RefCell<Backend>>]) {
596        self.current_weights.clear();
597        self.current_weights.resize(backends.len(), 0);
598    }
599}
600
601/// Whether every candidate carries the same [`backend_weight`].
602fn uniform_weights(backends: Candidates<'_>) -> bool {
603    let mut weights = backends.iter().map(|b| backend_weight(&b.borrow()));
604    match weights.next() {
605        Some(first) => weights.all(|weight| weight == first),
606        None => true,
607    }
608}
609
610impl Default for RoundRobin {
611    fn default() -> Self {
612        Self::new()
613    }
614}
615
616impl RoundRobin {
617    pub fn new() -> Self {
618        Self {
619            next_backend: 0,
620            current_weights: Vec::new(),
621        }
622    }
623
624    /// One step of smooth weighted round robin over a non-empty `backends`.
625    fn smooth_weighted(&mut self, backends: Candidates<'_>) -> Rc<RefCell<Backend>> {
626        debug_assert!(
627            !backends.is_empty(),
628            "smooth weighted round robin runs only on a non-empty set"
629        );
630        if self.current_weights.len() != backends.list_len() {
631            // A list `rebuild` never saw: start every running weight afresh.
632            self.current_weights.clear();
633            self.current_weights.resize(backends.list_len(), 0);
634        }
635
636        let mut total: i64 = 0;
637        let mut best: Option<usize> = None;
638        for index in 0..backends.len() {
639            let position = backends.position(index);
640            let weight = i64::from(backend_weight(&backends[index].borrow()));
641            self.current_weights[position] += weight;
642            total += weight;
643            let current = self.current_weights[position];
644            // Strictly greater: the first candidate in view order wins a tie.
645            if best.is_none_or(|best| current > self.current_weights[backends.position(best)]) {
646                best = Some(index);
647            }
648        }
649        let best = best.expect("a non-empty set yields a winner");
650        let position = backends.position(best);
651        self.current_weights[position] -= total;
652        // The candidates' running weights were raised by `total` in all and
653        // the winner's lowered by it: a step moves no weight in or out, so
654        // the running weights of the whole list still sum to zero, as they
655        // did when reset.
656        debug_assert!(total >= backends.len() as i64, "every weight is at least 1");
657        debug_assert_eq!(
658            self.current_weights.iter().sum::<i64>(),
659            0,
660            "smooth weighted round robin's running weights must sum to zero"
661        );
662        backends[best].clone()
663    }
664}
665
666/// Uniform (optionally weight-biased) random backend selection.
667///
668/// The RNG is seeded once, at construction, and never from an ambient source
669/// on the selection hot path — but, unlike [`Rendezvous`]/[`Maglev`], it is NOT seeded from the shared
670/// [`DEFAULT_HASH_SEED`] constant. `Rendezvous`/`Maglev` feed their seed into
671/// a *pure, stateless* hash function (`hash_backend`): same `(seed, key,
672/// addr)` in, same score out, independent of call history, so sharing the
673/// seed fleet-wide is exactly the point — it is what makes the same affinity
674/// key land on the same backend across workers and restarts.
675/// `Random` feeds its seed into a *stateful, advancing* `StdRng`, where the
676/// n-th pick depends on the whole call history. Seeding that from a shared
677/// compile-time constant would make every worker, on every cold start, with
678/// an identically-ordered backend list, draw from a bit-for-bit identical
679/// keystream — so every fresh worker's first pick over N equal-weight
680/// backends would be the same backend index, fleet-wide, on every restart:
681/// exactly the correlated-load event uniform selection exists to prevent,
682/// arriving at the worst possible moment (a synchronised redeploy, when every
683/// worker's call counter resets together and initial load is otherwise
684/// indistinguishable). So the seed is the caller's, and the production caller
685/// draws it from an OS-seeded generator: `BackendMap::new` seeds its own
686/// generator from the OS once, at construction, and hands every cluster's
687/// policy a seed drawn from it (`crate::backends::BackendMap`, #1684). Each
688/// worker's keystream is therefore independent, as with the thread-local
689/// `rng()` this replaced, while a simulator that builds the map with
690/// `BackendMap::with_seed` gets a fixed, reproducible sequence (mirrors
691/// `quinn-proto`'s `rng_seed` construction shape).
692///
693/// This does NOT make `Random` return a fixed backend: `next_available_backend`
694/// *advances* the internal RNG on every call, so a sequence of calls still
695/// draws a well-distributed spread of backends. "Random" here means
696/// "statistically spread, and — given `with_seed` — reproducible for a fixed
697/// seed and exact call sequence" — not "constant". Two instances built with
698/// `with_seed` on the same seed and driven through the same sequence of
699/// backend sets reproduce the same sequence of picks (what makes this
700/// testable); a single instance called repeatedly keeps drawing fresh values
701/// from its advancing RNG state, exactly like a real RNG would; two lists of
702/// one map, or two maps, draw from independent keystreams.
703#[derive(Debug)]
704pub struct Random {
705    rng: StdRng,
706}
707
708impl Random {
709    /// A policy whose draws are the keystream of `seed`. Never
710    /// `DEFAULT_HASH_SEED` in production — see the struct docs for why a
711    /// shared compile-time constant would correlate every worker's picks.
712    pub fn with_seed(seed: u64) -> Self {
713        Self {
714            rng: StdRng::seed_from_u64(seed),
715        }
716    }
717}
718
719impl Random {
720    /// Weighted draw over `backends`, identical, RNG draw for RNG draw, to
721    /// `rand::distr::weighted::WeightedIndex::<i32>` built over the same
722    /// weights, without the two `Vec`s building one allocates (the weights and
723    /// their cumulative sums).
724    ///
725    /// `WeightedIndex::new` accepts a set that is non-empty, holds no negative
726    /// weight, sums without `i32` overflow, and sums to more than zero; this
727    /// accepts exactly those, and returns `None` otherwise so the caller falls
728    /// back to the uniform draw as it did on `new`'s error. On acceptance it
729    /// draws once from the same `UniformInt::<i32>` over `[0, total)`, then
730    /// returns the first position whose inclusive prefix sum exceeds the draw:
731    /// the `partition_point` `WeightedIndex::sample` runs over its cumulative
732    /// weights, as a scan. `random_weighted_pick_matches_weighted_index`
733    /// holds the equivalence.
734    fn weighted_index(&mut self, backends: Candidates<'_>) -> Option<usize> {
735        if backends.is_empty() {
736            return None;
737        }
738        let mut total: i32 = 0;
739        for backend in backends.iter() {
740            let weight = random_weight(&backend.borrow());
741            if weight < 0 {
742                return None;
743            }
744            total = total.checked_add(weight)?;
745        }
746        if total == 0 {
747            return None;
748        }
749        let drawn = UniformInt::<i32>::new(0, total)
750            .expect("a positive total is a non-empty range")
751            .sample(&mut self.rng);
752        debug_assert!(
753            (0..total).contains(&drawn),
754            "the weighted draw must land inside [0, total)"
755        );
756
757        let last = backends.len() - 1;
758        let mut prefix: i32 = 0;
759        for (index, backend) in backends.iter().take(last).enumerate() {
760            prefix += random_weight(&backend.borrow());
761            if prefix > drawn {
762                return Some(index);
763            }
764        }
765        // Every earlier prefix is at most the draw, and the full sum `total`
766        // exceeds it: the draw falls in the last backend's share.
767        Some(last)
768    }
769}
770
771/// The weight `Random` draws with: the configured weight, or `DEFAULT_WEIGHT`
772/// when none is set. Unlike [`backend_weight`] it is not clamped, so a zero or
773/// negative weight reaches `Random::weighted_index` and is rejected there.
774fn random_weight(backend: &Backend) -> i32 {
775    backend
776        .load_balancing_parameters
777        .as_ref()
778        .map(|p| p.weight)
779        .unwrap_or(DEFAULT_WEIGHT)
780}
781
782impl LoadBalancingAlgorithm for Random {
783    fn next_available_backend(
784        &mut self,
785        _key: Option<u64>,
786        backends: Candidates<'_>,
787    ) -> Option<Rc<RefCell<Backend>>> {
788        let len = backends.len();
789        if let Some(index) = self.weighted_index(backends) {
790            // The weighted draw only returns a position of the view.
791            debug_assert!(index < len, "Random sampled an out-of-range index");
792            backends.get(index).cloned()
793        } else {
794            // The weighted draw declines only an empty set, a negative weight,
795            // an overflowing sum or an all-zero set; the uniform `choose` then
796            // selects iff non-empty.
797            let chosen = backends.choose(&mut self.rng).cloned();
798            debug_assert_eq!(
799                chosen.is_some(),
800                len > 0,
801                "Random fallback selects iff the set is non-empty"
802            );
803            chosen
804        }
805    }
806}
807
808/// The load of the candidate at `index` under a count metric, summed over
809/// every id at its address ([`Candidates::siblings`]): what a server serving
810/// two ids actually carries. The connection-time average is a latency, not a
811/// count, and is read from the representative alone.
812fn address_count(
813    candidates: Candidates<'_>,
814    index: usize,
815    count: impl Fn(&Backend) -> usize,
816) -> usize {
817    let own = count(&candidates[index].borrow());
818    let total = candidates.siblings(index).fold(own, |sum, sibling| {
819        sum.saturating_add(count(&sibling.borrow()))
820    });
821    debug_assert!(
822        total >= own,
823        "an address carries at least its representative's load"
824    );
825    total
826}
827
828/// Order two candidates by load per unit of weight, `load_a / weight_a`
829/// against `load_b / weight_b`, without a division: the two are compared
830/// cross-multiplied, exactly, in `u128`. With equal weights this is the order
831/// of the raw loads, so an unweighted cluster selects exactly as before.
832fn cmp_count_per_weight(load_a: usize, weight_a: u32, load_b: usize, weight_b: u32) -> Ordering {
833    debug_assert!(
834        weight_a >= 1 && weight_b >= 1,
835        "a weight reaches the comparison clamped to at least 1"
836    );
837    let order =
838        (load_a as u128 * u128::from(weight_b)).cmp(&(load_b as u128 * u128::from(weight_a)));
839    debug_assert!(
840        weight_a != weight_b || order == load_a.cmp(&load_b),
841        "equal weights must order the raw loads"
842    );
843    order
844}
845
846/// [`cmp_count_per_weight`] for a load that is a latency (`f64`). Equal weights
847/// compare the raw values, untouched by any rounding, so the comparison is
848/// then exactly the one the policies made before they read `weight`; unequal
849/// weights compare the cross-multiplied values. `None` only when a value is
850/// NaN, as `f64::partial_cmp`.
851fn cmp_cost_per_weight(cost_a: f64, weight_a: u32, cost_b: f64, weight_b: u32) -> Option<Ordering> {
852    debug_assert!(
853        weight_a >= 1 && weight_b >= 1,
854        "a weight reaches the comparison clamped to at least 1"
855    );
856    if weight_a == weight_b {
857        cost_a.partial_cmp(&cost_b)
858    } else {
859        (cost_a * f64::from(weight_b)).partial_cmp(&(cost_b * f64::from(weight_a)))
860    }
861}
862
863/// The least loaded candidate, load measured per unit of weight.
864///
865/// A candidate's load (`address_count` under a count metric, the
866/// representative's peak-EWMA connection time under
867/// [`LoadMetric::ConnectionTime`]) is divided by its `backend_weight`: a
868/// backend of weight `200` carrying `10` connections is as loaded as one of
869/// weight `100` carrying `5`. The weight is clamped to at least `1`, as for
870/// every weighted policy but [`Random`], so a `0` or negative weight counts as
871/// `1`, the smallest share, and never excludes a backend.
872///
873/// The first minimum in view order wins a tie. With equal weights (the
874/// default: an unset weight reads as `DEFAULT_WEIGHT`) the comparison is the
875/// one on raw loads, so the selection is the one made before the weight was
876/// read.
877#[derive(Debug)]
878pub struct LeastLoaded {
879    pub metric: LoadMetric,
880}
881
882impl LoadBalancingAlgorithm for LeastLoaded {
883    fn next_available_backend(
884        &mut self,
885        _key: Option<u64>,
886        backends: Candidates<'_>,
887    ) -> Option<Rc<RefCell<Backend>>> {
888        let was_empty = backends.is_empty();
889        let weight = |index: usize| backend_weight(&backends[index].borrow());
890        // One address is one share: a count metric reads the address's load,
891        // the sum over its ids (`address_count`), and the address weighs what
892        // its representative weighs. `min_by` keeps the first minimum.
893        let count = |metric: fn(&Backend) -> usize| {
894            (0..backends.len())
895                .map(|index| (address_count(backends, index, metric), weight(index), index))
896                .min_by(|&(load_a, weight_a, _), &(load_b, weight_b, _)| {
897                    cmp_count_per_weight(load_a, weight_a, load_b, weight_b)
898                })
899                .and_then(|(_, _, index)| backends.get(index))
900        };
901        let opt_b = match self.metric {
902            LoadMetric::Connections => count(|b| b.active_connections),
903            LoadMetric::Requests => count(|b| b.active_requests),
904            LoadMetric::ConnectionTime => {
905                let mut b: Option<(f64, u32, &Rc<RefCell<Backend>>)> = None;
906                let now = backends.now();
907                for backend in backends.iter() {
908                    let cost2 = backend.borrow_mut().peak_ewma_connection(now);
909                    let weight2 = backend_weight(&backend.borrow());
910
911                    match b.take() {
912                        None => b = Some((cost2, weight2, backend)),
913                        Some((cost1, weight1, back1)) => {
914                            // Keep the earlier unless the later is strictly
915                            // lighter (or the earlier's cost is NaN), as the
916                            // raw `cost1 <= cost2` did.
917                            if matches!(
918                                cmp_cost_per_weight(cost1, weight1, cost2, weight2),
919                                Some(Ordering::Less | Ordering::Equal)
920                            ) {
921                                b = Some((cost1, weight1, back1));
922                            } else {
923                                b = Some((cost2, weight2, backend));
924                            }
925                        }
926                    }
927                }
928
929                b.map(|(_cost, _weight, backend)| backend)
930            }
931        };
932        // Least-loaded over a non-empty set always finds a minimum; an empty
933        // set yields nothing. (`min_by` / the manual fold both have this
934        // shape.)
935        debug_assert_eq!(
936            opt_b.is_some(),
937            !was_empty,
938            "LeastLoaded selects iff the candidate set is non-empty"
939        );
940        opt_b.cloned()
941    }
942}
943
944/// Power-of-two-choices (P2C) load-aware selection: sample two candidates,
945/// keep the lighter one, and coin-flip a tie.
946///
947/// # Cost: two load reads per selection, whatever the cluster size
948///
949/// The two candidates are drawn uniformly at random from the candidate set
950/// and only those two are measured, so a selection costs `O(1)` load reads
951/// against [`LeastLoaded`]'s `O(n)` scan.
952///
953/// That is a statement about load READS and nothing more. It is NOT a claim
954/// that a selection is `O(1)`: every selection, under every policy, first
955/// collects the candidate positions in
956/// [`crate::backends::BackendList::next_available_backend_with_key`], which
957/// walks the cluster's backend list and records the position of each healthy
958/// backend in a buffer the list reuses across selections, before any policy
959/// runs. That walk is in the caller, this policy does not remove it, and no
960/// policy here is sub-linear per request.
961/// So P2C is not the "cheap" alternative to a scan — picking it to shorten a
962/// per-request walk that lives somewhere else buys nothing.
963///
964/// What the two-read bound does buy is real but narrower: under
965/// [`LoadMetric::ConnectionTime`] a load read is a `PeakEWMA::observe` call
966/// that decays the backend's average to the selection's instant, so
967/// two of them per selection instead of `n` is measurable work removed —
968/// under `Connections`/`Requests` a load read is a field read and the saving
969/// is small.
970///
971/// The reason to choose this policy over [`LeastLoaded`] is balance, not
972/// cost: P2C buys most of least-loaded's balance — the classic result is a
973/// maximum load of `O(log log n)` where uniform-random gives `O(log n)` —
974/// without ever computing a global minimum. Sōzu's workers each select from
975/// their own view, and that missing global minimum is what keeps them from
976/// herding: no worker can steer toward "the least loaded backend" because no
977/// worker ever computes one.
978///
979/// `power_of_two_touches_exactly_two_backends` pins the read count, and
980/// `power_of_two_sample_size_is_two_not_the_whole_set` measures the same
981/// property from the selection distribution alone. Both exist because a test
982/// that only checks "a backend came back" passes for an `O(n)` scan too.
983///
984/// # Tie-break: seeded-random, deliberately NOT deterministic-by-id
985///
986/// Unlike [`Rendezvous`]/[`Maglev`], `PowerOfTwo` carries no affinity `key` —
987/// every call is an independent, unlabeled selection event, so there is no
988/// "the same key must land on the same backend across workers/restarts"
989/// requirement to *preserve* here. But answering "must we preserve
990/// agreement?" is not the same question as "does this mechanism *create*
991/// agreement?" — see the next section for why that distinction is the whole
992/// reason its seed is never [`DEFAULT_HASH_SEED`].
993///
994/// A deterministic tie-break (e.g. "lowest backend id wins") was considered
995/// and rejected: ties are common, not rare — every backend starts at zero
996/// load, so a cold start, a post-scale-up rebalance, or any tick where two
997/// backends carry identical load all resolve through this branch. Always
998/// awarding the tie to the same backend (e.g. the lowest id) would bias the
999/// distribution toward that backend on exactly the events P2C exists to
1000/// spread out, reintroducing the herding effect P2C is designed to avoid —
1001/// on every one of Sōzu's single-threaded workers simultaneously, since they
1002/// would all observe the same tie and resolve it the same deterministic way.
1003/// A seeded RNG that *advances* per call keeps ties spread across the
1004/// candidate set over the process lifetime while still satisfying the
1005/// no-ambient-entropy-on-the-hot-path rule.
1006///
1007/// # Seed source: the backend map's OS-seeded generator, NOT `DEFAULT_HASH_SEED`
1008///
1009/// `Rendezvous`/`Maglev` feed their seed into a *pure, stateless* hash
1010/// function: sharing [`DEFAULT_HASH_SEED`] fleet-wide is exactly what gives
1011/// them cross-worker/cross-restart agreement for the same key. `PowerOfTwo`
1012/// feeds its seed into a *stateful, advancing* `StdRng`, where the n-th
1013/// tie-break depends on the whole call history. Seeding that from the same
1014/// shared constant would make every worker, on every cold start, draw from a
1015/// bit-for-bit identical keystream — so the very first tie any two workers
1016/// hit after a synchronised redeploy (and ties are the common case here, see
1017/// above) would resolve identically, fleet-wide, on every restart. That is
1018/// the deterministic-tie-break herding problem above, reintroduced by the
1019/// "fix" instead of prevented by it. The production seed therefore comes from
1020/// `BackendMap`'s generator, which reads the OS once, at construction — off
1021/// the datapath — so each worker's keystream is independent, the same
1022/// property the old thread-local `rng()` had, while `BackendMap::with_seed`
1023/// gives a simulator a fixed, reproducible sequence (#1684).
1024#[derive(Debug)]
1025pub struct PowerOfTwo {
1026    pub metric: LoadMetric,
1027    rng: StdRng,
1028}
1029
1030impl PowerOfTwo {
1031    /// A policy whose tie-breaks are the keystream of `seed`. Never
1032    /// `DEFAULT_HASH_SEED` in production — see the struct docs.
1033    pub fn with_seed(seed: u64, metric: LoadMetric) -> Self {
1034        Self {
1035            metric,
1036            rng: StdRng::seed_from_u64(seed),
1037        }
1038    }
1039
1040    /// Read the load of the candidates at `first` and `second` under the
1041    /// configured metric and order them by load per unit of weight, `first`
1042    /// against `second`. `None` only for a NaN connection time.
1043    ///
1044    /// Each backend is read exactly once, `first` then `second`, whatever the
1045    /// cluster size — that call count is the algorithm, not an implementation
1046    /// detail of it. The `ConnectionTime` arm needs `borrow_mut` because
1047    /// `peak_ewma_connection` decays the EWMA, to the selection's instant
1048    /// `now`, as it reads it.
1049    ///
1050    /// A count metric reads the load of the candidate's whole address, the
1051    /// sum over its ids (`address_count`), and compares the two counts as
1052    /// integers ([`cmp_count_per_weight`]): converting a count to `f64` would
1053    /// round it above 2^53 and could order two different loads as equal.
1054    fn compare(&self, backends: Candidates<'_>, first: usize, second: usize) -> Option<Ordering> {
1055        let weight = |index: usize| backend_weight(&backends[index].borrow());
1056        let counts = |metric: fn(&Backend) -> usize| {
1057            let first_load = address_count(backends, first, metric);
1058            let second_load = address_count(backends, second, metric);
1059            Some(cmp_count_per_weight(
1060                first_load,
1061                weight(first),
1062                second_load,
1063                weight(second),
1064            ))
1065        };
1066        match self.metric {
1067            LoadMetric::Connections => counts(|b| b.active_connections),
1068            LoadMetric::Requests => counts(|b| b.active_requests),
1069            LoadMetric::ConnectionTime => {
1070                let first_load = backends[first]
1071                    .borrow_mut()
1072                    .peak_ewma_connection(backends.now());
1073                let second_load = backends[second]
1074                    .borrow_mut()
1075                    .peak_ewma_connection(backends.now());
1076                cmp_cost_per_weight(first_load, weight(first), second_load, weight(second))
1077            }
1078        }
1079    }
1080}
1081
1082impl LoadBalancingAlgorithm for PowerOfTwo {
1083    fn next_available_backend(
1084        &mut self,
1085        _key: Option<u64>,
1086        backends: Candidates<'_>,
1087    ) -> Option<Rc<RefCell<Backend>>> {
1088        let len = backends.len();
1089        match len {
1090            0 => return None,
1091            // A singleton set has no second candidate to compare against, so
1092            // the sample degenerates to the only backend there is.
1093            1 => return backends.get(0).cloned(),
1094            _ => {}
1095        }
1096
1097        // Sample two DISTINCT backends uniformly at random. The second index
1098        // is drawn from the `len - 1` remaining slots and shifted past the
1099        // first, which is a uniform draw over "every index except `first`"
1100        // with no rejection loop — so the sampling has no unbounded worst
1101        // case, and the whole selection reads exactly two backends whatever
1102        // the cluster size.
1103        let first = self.rng.random_range(0..len);
1104        let mut second = self.rng.random_range(0..len - 1);
1105        if second >= first {
1106            second += 1;
1107        }
1108        debug_assert_ne!(
1109            first, second,
1110            "power-of-two must sample two distinct backends"
1111        );
1112        debug_assert!(
1113            second < len,
1114            "the shifted second index must stay inside the candidate set"
1115        );
1116
1117        // Load per unit of weight: with equal weights, the raw loads.
1118        let order = self.compare(backends, first, second);
1119
1120        // Keep the lighter of the two samples. An exact tie is broken by a
1121        // coin flip rather than by index order: ties are the common case
1122        // (every backend starts at zero load), and always awarding them to
1123        // the lower index would reintroduce the herding P2C exists to avoid
1124        // — see the struct docs.
1125        let chosen = match order {
1126            Some(Ordering::Less) => first,
1127            Some(Ordering::Greater) => second,
1128            _ if self.rng.random_bool(0.5) => first,
1129            _ => second,
1130        };
1131        // Asserted on the already-computed measures on purpose: re-reading a
1132        // backend under `LoadMetric::ConnectionTime` decays its EWMA, and a
1133        // `debug_assert!` that mutates would make debug and release builds
1134        // diverge.
1135        debug_assert!(
1136            (chosen == first && order != Some(Ordering::Greater))
1137                || (chosen == second && order != Some(Ordering::Less)),
1138            "power-of-two must never keep the strictly heavier of its two samples"
1139        );
1140
1141        backends.get(chosen).cloned()
1142    }
1143}
1144
1145/// Weighted Rendezvous (Highest Random Weight) hashing.
1146///
1147/// For an affinity `key`, the chosen backend is the one maximizing a stable,
1148/// per-(key, backend) score. With no key it degrades to plain round-robin;
1149/// every datapath supplies one for a client it can identify (the UDP affinity
1150/// key, and the client key of HTTP, HTTPS and TCP, sozu-proxy/sozu#524), so the
1151/// fallback is reached only for a request with no source address.
1152///
1153/// # Weighting
1154///
1155/// We use the standard continuous weighted-rendezvous score
1156/// (Schindelhauer–Schomaker / "logarithmic method"):
1157///
1158/// ```text
1159///     score(key, backend) = -weight / ln(h)
1160/// ```
1161///
1162/// where `h = hash64(seed, key, backend_addr) / 2^64 ∈ (0, 1)`. Because
1163/// `ln(h) < 0`, the score is positive and grows with `weight`, so heavier
1164/// backends win proportionally more keys while the choice stays deterministic
1165/// and minimally disruptive: removing a non-winning backend cannot change the
1166/// winner, and adding one only steals the keys for which it now scores highest.
1167/// Selection is `O(N)` per key.
1168#[derive(Debug)]
1169pub struct Rendezvous {
1170    /// Reproducible hash seed (NOT a `RandomState`).
1171    seed: u64,
1172    /// Round-robin cursor used for the `key == None` fallback.
1173    round_robin: RoundRobin,
1174}
1175
1176impl Default for Rendezvous {
1177    fn default() -> Self {
1178        Self::new()
1179    }
1180}
1181
1182impl Rendezvous {
1183    pub fn new() -> Self {
1184        Self::with_seed(DEFAULT_HASH_SEED)
1185    }
1186
1187    pub fn with_seed(seed: u64) -> Self {
1188        Self {
1189            seed,
1190            round_robin: RoundRobin::new(),
1191        }
1192    }
1193
1194    /// Continuous weighted-rendezvous score for a (key, backend) pair.
1195    /// Larger is better. See the struct docs for the formula.
1196    fn score(&self, key: u64, backend: &Backend) -> f64 {
1197        hrw_score_with_seed(self.seed, key, backend)
1198    }
1199}
1200
1201/// The weighted rendezvous (HRW) score of `backend` for the affinity `key`,
1202/// under [`DEFAULT_HASH_SEED`]: the score [`Rendezvous`] ranks a cluster's
1203/// backends by. Larger is better. Shuffle sharding
1204/// ([`ShuffleSharding`]) takes a client's shard as the top of this ranking,
1205/// and a test that needs to predict a shard computes it with this function.
1206pub fn hrw_score(key: u64, backend: &Backend) -> f64 {
1207    hrw_score_with_seed(DEFAULT_HASH_SEED, key, backend)
1208}
1209
1210/// [`hrw_score`] under an explicit seed. See [`Rendezvous`] for the formula.
1211fn hrw_score_with_seed(seed: u64, key: u64, backend: &Backend) -> f64 {
1212    let weight = backend_weight(backend) as f64;
1213    // Map the 64-bit hash into the open interval (0, 1). Guard the
1214    // endpoints so `ln` stays finite: 0 would give -inf, 1 would give 0.
1215    let h = hash_backend(seed, key, &backend.address);
1216    // (h + 0.5) / 2^64 keeps the value strictly inside (0, 1).
1217    let unit = (h as f64 + 0.5) / (u64::MAX as f64 + 1.0);
1218    // `unit` must land strictly inside (0, 1) so `ln(unit) < 0` and the
1219    // score stays a positive, finite number — the whole HRW ordering relies
1220    // on a well-defined comparable score for every backend.
1221    debug_assert!(
1222        unit > 0.0 && unit < 1.0,
1223        "HRW unit must lie strictly inside (0, 1), got {unit}"
1224    );
1225    let score = -weight / unit.ln();
1226    debug_assert!(
1227        score.is_finite() && score > 0.0,
1228        "HRW score must be finite and positive, got {score}"
1229    );
1230    score
1231}
1232
1233impl LoadBalancingAlgorithm for Rendezvous {
1234    fn next_available_backend(
1235        &mut self,
1236        key: Option<u64>,
1237        backends: Candidates<'_>,
1238    ) -> Option<Rc<RefCell<Backend>>> {
1239        let Some(key) = key else {
1240            // No affinity key: behave exactly like RoundRobin.
1241            return self.round_robin.next_available_backend(None, backends);
1242        };
1243
1244        if backends.is_empty() {
1245            return None;
1246        }
1247
1248        let mut best: Option<(f64, &Rc<RefCell<Backend>>)> = None;
1249        for backend in backends.iter() {
1250            let score = self.score(key, &backend.borrow());
1251            match best {
1252                Some((best_score, _)) if best_score >= score => {}
1253                _ => best = Some((score, backend)),
1254            }
1255        }
1256        // A non-empty set always yields a winner, and that winner's score is the
1257        // maximum over the whole set (HRW = highest random weight).
1258        debug_assert!(
1259            best.is_some(),
1260            "HRW must select a winner from a non-empty set"
1261        );
1262        #[cfg(debug_assertions)]
1263        if let Some((best_score, _)) = best {
1264            for backend in backends.iter() {
1265                debug_assert!(
1266                    best_score >= self.score(key, &backend.borrow()),
1267                    "HRW winner does not maximize the score over the set"
1268                );
1269            }
1270        }
1271        best.map(|(_, backend)| backend.clone())
1272    }
1273
1274    fn rebuild(&mut self, backends: &[Rc<RefCell<Backend>>]) {
1275        // Only the round-robin fallback keeps state across selections.
1276        self.round_robin.rebuild(backends);
1277    }
1278}
1279
1280/// Maglev consistent hashing (Google, NSDI'16).
1281///
1282/// Builds a precomputed lookup table of `M` (a prime, default 65537) slots from
1283/// the **stable, full** backend set. Each backend derives a `(offset, skip)`
1284/// permutation from two seeded hashes and claims slots in permutation order
1285/// until the table is full; a backend's share of slots is proportional to its
1286/// weight.
1287///
1288/// The table holds one entry per distinct **address**: several backend ids at
1289/// one address claim the slots of one backend, weighted by the heaviest of
1290/// them over the full set (`outranks_at_address`), and the lookup returns
1291/// the one candidate the caller retained at that address.
1292///
1293/// # Rebuild discipline (never on the hot path)
1294///
1295/// The table is rebuilt **only when the full backend set changes**
1296/// ([`Maglev::rebuild`], wired into `BackendList::add_backend` /
1297/// `remove_backend`), never per packet. Building the table is `O(M)` = 65537
1298/// slot writes; doing it per datagram would be a DoS amplifier (one unhealthy
1299/// backend would trigger a full rebuild on every selection) and would also
1300/// destroy Maglev's stability, since the table would track a shifting subset.
1301///
1302/// # Selection (handles a shrunk healthy subset without rebuilding)
1303///
1304/// At selection time the caller passes the **healthy subset** of backends.
1305/// Lookup is near-`O(1)`: compute `slot = key % M`, then probe forward through
1306/// the table (`table[(slot + i) % M]`) and return the first entry whose address
1307/// is present in the healthy subset. The table maps to the full set it was
1308/// built from; membership is checked against the subset, so an unhealthy
1309/// backend is simply skipped — no rebuild, and healthy keys stay pinned. If no
1310/// table entry resolves to a healthy backend, fall back to round-robin over the
1311/// subset. With no key it falls back to round-robin, which, as for
1312/// [`Rendezvous`], only a request with no source address reaches.
1313#[derive(Debug)]
1314pub struct Maglev {
1315    seed: u64,
1316    /// Prime table size `M`.
1317    size: usize,
1318    /// `table[i]` is an index into `backends` (the set captured at the last
1319    /// [`rebuild`]). Empty when there are no backends.
1320    table: Vec<usize>,
1321    /// Backend addresses captured at the last rebuild, in the order used to
1322    /// build `table`. The lookup maps a table entry back to the live backend
1323    /// by address so it survives `Vec` reordering between rebuilds.
1324    backend_addrs: Vec<SocketAddr>,
1325    /// Round-robin cursor used for the `key == None` fallback.
1326    round_robin: RoundRobin,
1327}
1328
1329impl Default for Maglev {
1330    fn default() -> Self {
1331        Self::new()
1332    }
1333}
1334
1335impl Maglev {
1336    /// Default prime table size. 65537 is the smallest prime above 2^16; large
1337    /// enough to keep per-key disruption small on backend churn, small enough
1338    /// to rebuild cheaply off the datapath.
1339    pub const DEFAULT_TABLE_SIZE: usize = 65537;
1340
1341    pub fn new() -> Self {
1342        Self::with_seed(DEFAULT_HASH_SEED)
1343    }
1344
1345    pub fn with_seed(seed: u64) -> Self {
1346        Self::with_seed_and_size(seed, Self::DEFAULT_TABLE_SIZE)
1347    }
1348
1349    pub fn with_seed_and_size(seed: u64, size: usize) -> Self {
1350        Self {
1351            seed,
1352            // The Maglev permutation uses `skip = h2 % (m - 1) + 1`, which
1353            // panics (divide-by-zero) when `m == 1`, and a single-slot table
1354            // is degenerate anyway. Clamp to the next prime `>= max(size, 2)`
1355            // so the table is always usable and the permutation stride is
1356            // coprime with `m` (a prime), guaranteeing the population loop
1357            // visits every slot.
1358            size: next_prime(size.max(2)),
1359            table: Vec::new(),
1360            backend_addrs: Vec::new(),
1361            round_robin: RoundRobin::new(),
1362        }
1363    }
1364
1365    /// Rebuild the lookup table from `backends`. Called on backend-set change,
1366    /// NOT per packet. Honors backend weight via proportional slot share.
1367    pub fn rebuild(&mut self, backends: &[Rc<RefCell<Backend>>]) {
1368        self.rebuild_from(backends.iter());
1369    }
1370
1371    /// [`Maglev::rebuild`] over any ordered backend sequence, so the cold-start
1372    /// build in `next_available_backend` can read its [`Candidates`] view
1373    /// without collecting it.
1374    fn rebuild_from<'a>(
1375        &mut self,
1376        backends: impl ExactSizeIterator<Item = &'a Rc<RefCell<Backend>>>,
1377    ) {
1378        let ids = backends.len();
1379        self.backend_addrs.clear();
1380        self.table.clear();
1381        if ids == 0 || self.size == 0 {
1382            return;
1383        }
1384
1385        let m = self.size;
1386
1387        // Per-address (offset, skip) permutation parameters and weights. One
1388        // address is one share: a second id at an address already captured
1389        // claims no slots of its own, it only raises the address's weight to
1390        // the heaviest of its ids. Linear search, on the control plane.
1391        let mut offsets = Vec::with_capacity(ids);
1392        let mut skips = Vec::with_capacity(ids);
1393        let mut weights: Vec<u64> = Vec::with_capacity(ids);
1394        for backend in backends {
1395            let b = backend.borrow();
1396            let addr = b.address;
1397            let w = backend_weight(&b) as u64;
1398            if let Some(known) = self.backend_addrs.iter().position(|a| *a == addr) {
1399                weights[known] = weights[known].max(w);
1400                continue;
1401            }
1402            self.backend_addrs.push(addr);
1403            // Two independent seeded hashes give the permutation seeds. We mix
1404            // distinct domain separators into the `key` slot of `hash_backend`
1405            // so `offset` and `skip` are uncorrelated.
1406            let h1 = hash_backend(self.seed, 0x6F66_6673_6574, &addr); // "offset"
1407            let h2 = hash_backend(self.seed, 0x736B_6970_5F5F, &addr); // "skip__"
1408            let offset = (h1 % m as u64) as usize;
1409            let skip = (h2 % (m as u64 - 1)) as usize + 1;
1410            // The permutation parameters must address valid slots and keep a
1411            // non-zero stride; `skip ∈ [1, m-1]` is coprime with the prime `m`,
1412            // which is what guarantees the population loop visits every slot.
1413            debug_assert!(offset < m, "Maglev offset must be a valid slot");
1414            debug_assert!(
1415                (1..m).contains(&skip),
1416                "Maglev skip must lie in [1, m-1] to stay coprime with the prime table"
1417            );
1418            offsets.push(offset);
1419            skips.push(skip);
1420            weights.push(w);
1421        }
1422        let n = self.backend_addrs.len();
1423        let total_weight: u64 = weights.iter().sum();
1424        debug_assert!(
1425            (1..=ids).contains(&n),
1426            "Maglev captures between one and every backend's address"
1427        );
1428        // Every backend contributes weight >= 1, so a non-empty set has a
1429        // strictly positive total — the proportional target math divides by it.
1430        debug_assert!(
1431            total_weight > 0,
1432            "Maglev total weight must be positive for a non-empty backend set"
1433        );
1434
1435        // Target slot count per backend, proportional to weight. The sum of
1436        // targets equals `m` (remainder handed to the heaviest/first backends).
1437        let mut targets = vec![0usize; n];
1438        let mut assigned = 0usize;
1439        for (i, &w) in weights.iter().enumerate() {
1440            let t = ((w as u128 * m as u128) / total_weight as u128) as usize;
1441            targets[i] = t;
1442            assigned += t;
1443        }
1444        // Distribute the rounding remainder so the table fills exactly.
1445        let mut i = 0;
1446        while assigned < m {
1447            targets[i % n] += 1;
1448            assigned += 1;
1449            i += 1;
1450        }
1451        // The targets must sum to exactly `m`: this is the termination
1452        // guarantee for the population loop below — it writes one slot per unit
1453        // of target budget and stops at `count == m`, so an under/over sum would
1454        // either leave holes (`usize::MAX` entries) or loop forever.
1455        debug_assert_eq!(
1456            assigned, m,
1457            "Maglev target budget must sum to the table size"
1458        );
1459        debug_assert_eq!(
1460            targets.iter().sum::<usize>(),
1461            m,
1462            "Maglev per-backend targets must sum to the table size"
1463        );
1464
1465        // Standard Maglev population loop, capped per backend by `targets`.
1466        let mut table = vec![usize::MAX; m];
1467        let mut next = vec![0usize; n];
1468        let mut filled = vec![0usize; n];
1469        let mut count = 0usize;
1470        while count < m {
1471            for b in 0..n {
1472                if filled[b] >= targets[b] {
1473                    continue;
1474                }
1475                // Find this backend's next free preferred slot.
1476                let mut c = (offsets[b] + next[b] * skips[b]) % m;
1477                while table[c] != usize::MAX {
1478                    next[b] += 1;
1479                    c = (offsets[b] + next[b] * skips[b]) % m;
1480                }
1481                table[c] = b;
1482                next[b] += 1;
1483                filled[b] += 1;
1484                count += 1;
1485                if count >= m {
1486                    break;
1487                }
1488            }
1489        }
1490        // The loop terminates with every slot claimed exactly once: `count`
1491        // reached `m`, no slot still holds the `usize::MAX` sentinel, and each
1492        // backend filled precisely its target budget.
1493        debug_assert_eq!(count, m, "Maglev population loop must fill exactly m slots");
1494        debug_assert!(
1495            table.iter().all(|&slot| slot != usize::MAX),
1496            "Maglev population loop left an unfilled slot"
1497        );
1498        debug_assert!(
1499            filled == targets,
1500            "Maglev filled counts must match the per-backend targets"
1501        );
1502
1503        self.table = table;
1504
1505        // Post-conditions (TigerStyle). The table is either fully built or empty:
1506        //   * `size` is the prime chosen at construction — rebuild never changes
1507        //     it (the permutation math depends on a coprime stride over a prime).
1508        //   * a populated table is exactly `size` slots, every entry a valid
1509        //     index into `backend_addrs` (`< backend_addrs.len() <= size`), so a
1510        //     later `table[slot]` lookup can never index out of `backend_addrs`.
1511        //   * `backend_addrs` is non-empty whenever the table is non-empty (the
1512        //     table maps slots to addresses; an empty address vector would make
1513        //     every lookup resolve to nothing).
1514        #[cfg(debug_assertions)]
1515        {
1516            debug_assert!(
1517                is_prime(self.size),
1518                "Maglev table size {} is not prime",
1519                self.size
1520            );
1521            if self.table.is_empty() {
1522                debug_assert!(
1523                    self.backend_addrs.is_empty(),
1524                    "Maglev: empty table but non-empty backend_addrs"
1525                );
1526            } else {
1527                debug_assert_eq!(
1528                    self.table.len(),
1529                    self.size,
1530                    "Maglev table must have exactly `size` slots"
1531                );
1532                debug_assert!(
1533                    !self.backend_addrs.is_empty(),
1534                    "Maglev: non-empty table but empty backend_addrs"
1535                );
1536                debug_assert!(
1537                    self.table.iter().all(|&idx| idx < self.backend_addrs.len()),
1538                    "Maglev table holds an index out of backend_addrs range"
1539                );
1540            }
1541        }
1542    }
1543}
1544
1545impl LoadBalancingAlgorithm for Maglev {
1546    fn next_available_backend(
1547        &mut self,
1548        key: Option<u64>,
1549        backends: Candidates<'_>,
1550    ) -> Option<Rc<RefCell<Backend>>> {
1551        let Some(key) = key else {
1552            return self.round_robin.next_available_backend(None, backends);
1553        };
1554
1555        if backends.is_empty() {
1556            return None;
1557        }
1558
1559        // Cold start ONLY: the table is empty because no `rebuild` has fired
1560        // yet (the control plane wires `rebuild` into backend-set mutations,
1561        // but a freshly constructed policy or a never-mutated cluster may not
1562        // have one). Build once from the current set. This is NOT a per-packet
1563        // rebuild: once populated, the table is only refreshed by `rebuild` on
1564        // an actual set change — a partial outage (shrunk healthy subset) never
1565        // rebuilds, it is handled by the probe-forward below.
1566        if self.table.is_empty() {
1567            self.rebuild_from(backends.iter());
1568        }
1569
1570        if self.table.is_empty() {
1571            // Still empty (e.g. zero backends captured): nothing to select.
1572            return None;
1573        }
1574
1575        // Probe forward from `key % M` through the table and return the first
1576        // entry whose backend address is present in the passed healthy subset.
1577        // The table was built from the full set; checking membership against
1578        // the subset lets us skip unhealthy backends WITHOUT rebuilding, which
1579        // is what keeps a partial outage off the hot path and keeps healthy
1580        // keys pinned to the same backend.
1581        let start = (key % self.size as u64) as usize;
1582        // `key % size` is always a valid table slot, and the table is exactly
1583        // `size` long here (non-empty checked above, rebuild post-condition).
1584        debug_assert!(start < self.size, "Maglev start slot out of range");
1585        debug_assert_eq!(
1586            self.table.len(),
1587            self.size,
1588            "Maglev lookup on a table whose length != size"
1589        );
1590        for i in 0..self.size {
1591            let slot = (start + i) % self.size;
1592            // Both the running slot and the index it stores are in range: the
1593            // table is `size` long and every entry is a valid `backend_addrs`
1594            // index (rebuild post-condition), so the resolution below is total.
1595            debug_assert!(slot < self.size, "Maglev probe slot out of range");
1596            let idx = self.table[slot];
1597            debug_assert!(
1598                idx < self.backend_addrs.len(),
1599                "Maglev table entry indexes outside the captured address set"
1600            );
1601            // Resolve the table's backend index back to an address captured at
1602            // the last rebuild, then look it up in the live healthy subset.
1603            if let Some(addr) = self.backend_addrs.get(idx)
1604                && let Some(backend) = backends.iter().find(|b| b.borrow().address == *addr)
1605            {
1606                // The chosen backend is, by construction, a member of the
1607                // healthy subset we were handed.
1608                debug_assert!(
1609                    backends
1610                        .iter()
1611                        .any(|b| b.borrow().address == backend.borrow().address),
1612                    "Maglev must return a backend from the healthy subset"
1613                );
1614                return Some(backend.clone());
1615            }
1616        }
1617
1618        // None of the table entries resolved to a healthy backend (every
1619        // backend the table knows about is currently unhealthy/absent). Fall
1620        // back to round-robin over the healthy subset so we still route.
1621        self.round_robin.next_available_backend(None, backends)
1622    }
1623
1624    fn rebuild(&mut self, backends: &[Rc<RefCell<Backend>>]) {
1625        Maglev::rebuild(self, backends);
1626        self.round_robin.rebuild(backends);
1627    }
1628}
1629
1630#[cfg(test)]
1631mod test {
1632    use std::{
1633        net::{IpAddr, Ipv4Addr, SocketAddr},
1634        time::Duration,
1635    };
1636
1637    use super::*;
1638    use crate::{
1639        PeakEWMA,
1640        backends::{BackendStatus, HealthState},
1641        retry::{ExponentialBackoffPolicy, RetryPolicyWrapper},
1642        sozu_command::proto::command::{LoadBalancingParams, LoadMetric},
1643    };
1644
1645    /// Run `policy` over every backend of `backends`, the way a caller whose
1646    /// whole list is available would.
1647    fn pick(
1648        policy: &mut (impl LoadBalancingAlgorithm + ?Sized),
1649        key: Option<u64>,
1650        backends: &[Rc<RefCell<Backend>>],
1651    ) -> Option<Rc<RefCell<Backend>>> {
1652        pick_at(policy, key, backends, Instant::now())
1653    }
1654
1655    /// [`pick`] for a selection at `now`.
1656    fn pick_at(
1657        policy: &mut (impl LoadBalancingAlgorithm + ?Sized),
1658        key: Option<u64>,
1659        backends: &[Rc<RefCell<Backend>>],
1660        now: Instant,
1661    ) -> Option<Rc<RefCell<Backend>>> {
1662        let indices: Vec<usize> = (0..backends.len()).collect();
1663        policy.next_available_backend(key, Candidates::new(backends, &indices, now))
1664    }
1665
1666    fn create_backend(id: String, connections: Option<usize>) -> Backend {
1667        Backend {
1668            sticky_id: None,
1669            backend_id: id,
1670            address: SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 8080),
1671            status: BackendStatus::Normal,
1672            retry_policy: RetryPolicyWrapper::ExponentialBackoff(ExponentialBackoffPolicy::new(1)),
1673            active_connections: connections.unwrap_or(0),
1674            active_requests: 0,
1675            failures: 0,
1676            load_balancing_parameters: None,
1677            backup: false,
1678            connection_time: PeakEWMA::new(Instant::now()),
1679            health: HealthState::default(),
1680            unix_path: None,
1681        }
1682    }
1683
1684    #[test]
1685    fn it_should_find_the_backend_with_least_connections() {
1686        let backend_with_least_connection =
1687            Rc::new(RefCell::new(create_backend("yolo".to_string(), Some(1))));
1688
1689        let backends = vec![
1690            Rc::new(RefCell::new(create_backend("nolo".to_string(), Some(10)))),
1691            Rc::new(RefCell::new(create_backend("philo".to_string(), Some(20)))),
1692            backend_with_least_connection.clone(),
1693        ];
1694
1695        let mut least_connection_algorithm = LeastLoaded {
1696            metric: LoadMetric::Connections,
1697        };
1698
1699        let backend_res = pick(&mut least_connection_algorithm, None, &backends).unwrap();
1700        let backend = backend_res.borrow();
1701
1702        assert!(*backend == *backend_with_least_connection.borrow());
1703    }
1704
1705    #[test]
1706    fn it_shouldnt_find_backend_with_least_connections_when_list_is_empty() {
1707        let backends = vec![];
1708
1709        let mut least_connection_algorithm = LeastLoaded {
1710            metric: LoadMetric::Connections,
1711        };
1712
1713        let backend = pick(&mut least_connection_algorithm, None, &backends);
1714        assert!(backend.is_none());
1715    }
1716
1717    #[test]
1718    fn it_should_find_backend_with_roundrobin_when_some_backends_were_removed() {
1719        let mut backends = vec![
1720            Rc::new(RefCell::new(create_backend("toto".to_string(), None))),
1721            Rc::new(RefCell::new(create_backend("voto".to_string(), None))),
1722            Rc::new(RefCell::new(create_backend("yoto".to_string(), None))),
1723        ];
1724
1725        let mut roundrobin = RoundRobin {
1726            next_backend: 1,
1727            ..RoundRobin::new()
1728        };
1729        let backend = pick(&mut roundrobin, None, &backends);
1730        assert_eq!(backend.as_ref(), backends.get(1));
1731
1732        backends.remove(1);
1733
1734        let backend2 = pick(&mut roundrobin, None, &backends);
1735        assert_eq!(backend2.as_ref(), backends.first());
1736    }
1737
1738    // ----- HRW (Rendezvous) and Maglev affinity tests -----
1739
1740    /// Build a backend with a distinct address so the address-based affinity
1741    /// hashers (HRW / Maglev) see distinct identities, and an optional weight.
1742    fn addr_backend(id: &str, last_octet: u8, port: u16, weight: Option<i32>) -> Backend {
1743        let mut b = create_backend(id.to_string(), None);
1744        b.address = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, last_octet)), port);
1745        b.load_balancing_parameters = weight.map(|weight| LoadBalancingParams { weight });
1746        b
1747    }
1748
1749    fn rc(b: Backend) -> Rc<RefCell<Backend>> {
1750        Rc::new(RefCell::new(b))
1751    }
1752
1753    fn make_backends(n: u8) -> Vec<Rc<RefCell<Backend>>> {
1754        (0..n)
1755            .map(|i| rc(addr_backend(&format!("b{i}"), i + 1, 8000 + i as u16, None)))
1756            .collect()
1757    }
1758
1759    fn chosen_addr(b: &Rc<RefCell<Backend>>) -> SocketAddr {
1760        b.borrow().address
1761    }
1762
1763    #[test]
1764    fn hrw_is_deterministic_for_a_fixed_key() {
1765        let backends = make_backends(5);
1766        let mut hrw = Rendezvous::new();
1767
1768        let first = pick(&mut hrw, Some(42), &backends).map(|b| chosen_addr(&b));
1769        for _ in 0..50 {
1770            let again = pick(&mut hrw, Some(42), &backends).map(|b| chosen_addr(&b));
1771            assert_eq!(first, again, "HRW must be deterministic for a fixed key");
1772        }
1773    }
1774
1775    #[test]
1776    fn hrw_none_key_falls_back_to_round_robin() {
1777        let backends = make_backends(3);
1778        let mut hrw = Rendezvous::new();
1779
1780        // With None it should cycle round-robin: addresses in order then wrap.
1781        let a = chosen_addr(&pick(&mut hrw, None, &backends).unwrap());
1782        let b = chosen_addr(&pick(&mut hrw, None, &backends).unwrap());
1783        let c = chosen_addr(&pick(&mut hrw, None, &backends).unwrap());
1784        let d = chosen_addr(&pick(&mut hrw, None, &backends).unwrap());
1785        assert_eq!(a, chosen_addr(&backends[0]));
1786        assert_eq!(b, chosen_addr(&backends[1]));
1787        assert_eq!(c, chosen_addr(&backends[2]));
1788        assert_eq!(d, a, "round-robin should wrap around");
1789    }
1790
1791    #[test]
1792    fn hrw_minimal_disruption_when_removing_a_non_winner() {
1793        // For every key whose winner is NOT the removed backend, the choice
1794        // must be unchanged after removal (HRW minimal-disruption property).
1795        let mut backends = make_backends(6);
1796        let mut hrw = Rendezvous::new();
1797
1798        // Pick a backend to remove (index 3).
1799        let removed_addr = chosen_addr(&backends[3]);
1800
1801        // Record winners for many keys on the full set.
1802        let mut before = std::collections::HashMap::new();
1803        for key in 0..2000u64 {
1804            let w = chosen_addr(&pick(&mut hrw, Some(key), &backends).unwrap());
1805            before.insert(key, w);
1806        }
1807
1808        // Remove the backend and re-evaluate.
1809        backends.remove(3);
1810        for key in 0..2000u64 {
1811            let after = chosen_addr(&pick(&mut hrw, Some(key), &backends).unwrap());
1812            let prev = before[&key];
1813            if prev != removed_addr {
1814                assert_eq!(
1815                    prev, after,
1816                    "removing a non-winner changed the choice for key {key}"
1817                );
1818            }
1819        }
1820    }
1821
1822    #[test]
1823    fn hrw_distribution_is_roughly_even() {
1824        let n = 5u8;
1825        let backends = make_backends(n);
1826        let mut hrw = Rendezvous::new();
1827
1828        let total = 20_000u64;
1829        let mut counts: std::collections::HashMap<SocketAddr, u64> =
1830            std::collections::HashMap::new();
1831        for key in 0..total {
1832            let w = chosen_addr(&pick(&mut hrw, Some(key), &backends).unwrap());
1833            *counts.entry(w).or_default() += 1;
1834        }
1835
1836        let expected = total / n as u64;
1837        for b in &backends {
1838            let c = counts.get(&chosen_addr(b)).copied().unwrap_or(0);
1839            // Allow generous +/-35% slack: this is a statistical property.
1840            assert!(
1841                c > expected * 65 / 100 && c < expected * 135 / 100,
1842                "HRW distribution skewed: backend got {c}, expected ~{expected}"
1843            );
1844        }
1845    }
1846
1847    #[test]
1848    fn maglev_is_deterministic_for_a_fixed_key() {
1849        let backends = make_backends(7);
1850        let mut mag = Maglev::new();
1851        mag.rebuild(&backends);
1852
1853        let sel = backends.clone();
1854        let first = chosen_addr(&pick(&mut mag, Some(12345), &sel).unwrap());
1855        for _ in 0..50 {
1856            let again = chosen_addr(&pick(&mut mag, Some(12345), &sel).unwrap());
1857            assert_eq!(first, again, "Maglev must be deterministic for a fixed key");
1858        }
1859    }
1860
1861    #[test]
1862    fn maglev_none_key_falls_back_to_round_robin() {
1863        let backends = make_backends(3);
1864        let mut mag = Maglev::new();
1865        mag.rebuild(&backends);
1866
1867        let a = chosen_addr(&pick(&mut mag, None, &backends).unwrap());
1868        let b = chosen_addr(&pick(&mut mag, None, &backends).unwrap());
1869        let c = chosen_addr(&pick(&mut mag, None, &backends).unwrap());
1870        let d = chosen_addr(&pick(&mut mag, None, &backends).unwrap());
1871        assert_eq!(a, chosen_addr(&backends[0]));
1872        assert_eq!(b, chosen_addr(&backends[1]));
1873        assert_eq!(c, chosen_addr(&backends[2]));
1874        assert_eq!(d, a);
1875    }
1876
1877    #[test]
1878    fn maglev_distribution_is_roughly_even() {
1879        let n = 5u8;
1880        let backends = make_backends(n);
1881        // Use a small prime table so the test is fast but still representative.
1882        let mut mag = Maglev::with_seed_and_size(DEFAULT_HASH_SEED, 1009);
1883        mag.rebuild(&backends);
1884
1885        let total = 50_000u64;
1886        let mut counts: std::collections::HashMap<SocketAddr, u64> =
1887            std::collections::HashMap::new();
1888        let sel = backends.clone();
1889        for key in 0..total {
1890            let w = chosen_addr(&pick(&mut mag, Some(key), &sel).unwrap());
1891            *counts.entry(w).or_default() += 1;
1892        }
1893
1894        let expected = total / n as u64;
1895        for b in &backends {
1896            let c = counts.get(&chosen_addr(b)).copied().unwrap_or(0);
1897            // Maglev is more even than HRW; +/-15% is comfortable.
1898            assert!(
1899                c > expected * 85 / 100 && c < expected * 115 / 100,
1900                "Maglev distribution skewed: backend got {c}, expected ~{expected}"
1901            );
1902        }
1903    }
1904
1905    #[test]
1906    fn maglev_table_rebuild_keeps_most_keys_stable() {
1907        // Adding a backend should move only a bounded fraction of keys
1908        // (Maglev disruption bound ~ 1/N of the table, well under 50%).
1909        let backends5 = make_backends(5);
1910        let mut mag = Maglev::with_seed_and_size(DEFAULT_HASH_SEED, 1009);
1911        mag.rebuild(&backends5);
1912
1913        let total = 20_000u64;
1914        let mut before = std::collections::HashMap::new();
1915        let sel = backends5.clone();
1916        for key in 0..total {
1917            before.insert(key, chosen_addr(&pick(&mut mag, Some(key), &sel).unwrap()));
1918        }
1919
1920        // Add a sixth backend and rebuild.
1921        let mut backends6 = backends5.clone();
1922        backends6.push(rc(addr_backend("b5", 6, 8005, None)));
1923        mag.rebuild(&backends6);
1924
1925        let mut moved = 0u64;
1926        let sel6 = backends6.clone();
1927        for key in 0..total {
1928            let after = chosen_addr(&pick(&mut mag, Some(key), &sel6).unwrap());
1929            if after != before[&key] {
1930                moved += 1;
1931            }
1932        }
1933
1934        // Disruption should be well under half: most keys stay on their backend.
1935        assert!(
1936            moved < total / 2,
1937            "Maglev rebuild moved too many keys: {moved}/{total}"
1938        );
1939    }
1940
1941    #[test]
1942    fn maglev_partial_outage_does_not_rebuild_and_stays_stable() {
1943        // FIX 1 guard: a partial outage (one backend missing from the healthy
1944        // subset passed to `next_available_backend`) must NOT rebuild the
1945        // table — the table tracks the full set captured at `rebuild`, and the
1946        // unhealthy backend is simply skipped via probe-forward. Selection for
1947        // keys that did not land on the unhealthy backend stays identical.
1948        let full = make_backends(5);
1949        let mut mag = Maglev::with_seed_and_size(DEFAULT_HASH_SEED, 1009);
1950        mag.rebuild(&full);
1951
1952        // Snapshot the table state built from the FULL set.
1953        let table_before = mag.table.clone();
1954        let addrs_before = mag.backend_addrs.clone();
1955        assert_eq!(addrs_before.len(), 5, "table built from the full set");
1956
1957        // Record winners on the full healthy set.
1958        let total = 4000u64;
1959        let mut before = std::collections::HashMap::new();
1960        let sel_full = full.clone();
1961        for key in 0..total {
1962            before.insert(
1963                key,
1964                chosen_addr(&pick(&mut mag, Some(key), &sel_full).unwrap()),
1965            );
1966        }
1967
1968        // Now mark backend index 2 unhealthy by passing a REDUCED subset
1969        // (everything except index 2). This is exactly what
1970        // `BackendList::next_available_backend_with_key` does when a backend
1971        // fails its health check / enters retry backoff.
1972        let unhealthy_addr = chosen_addr(&full[2]);
1973        let subset: Vec<_> = full
1974            .iter()
1975            .filter(|b| chosen_addr(b) != unhealthy_addr)
1976            .cloned()
1977            .collect();
1978
1979        for key in 0..total {
1980            let after = chosen_addr(&pick(&mut mag, Some(key), &subset).unwrap());
1981            // The unhealthy backend must never be selected.
1982            assert_ne!(
1983                after, unhealthy_addr,
1984                "selection returned the unhealthy backend for key {key}"
1985            );
1986            // Keys that did NOT previously land on the unhealthy backend must
1987            // keep their original backend — probe-forward only reroutes the
1988            // keys that were pinned to the now-missing backend.
1989            if before[&key] != unhealthy_addr {
1990                assert_eq!(
1991                    before[&key], after,
1992                    "a healthy key moved during a partial outage (key {key})"
1993                );
1994            }
1995        }
1996
1997        // The crux of FIX 1: the table was NOT rebuilt by the partial outage.
1998        // Neither the table contents nor the captured address set changed.
1999        assert_eq!(
2000            mag.table, table_before,
2001            "partial outage must not rebuild the Maglev table"
2002        );
2003        assert_eq!(
2004            mag.backend_addrs, addrs_before,
2005            "partial outage must not change the captured backend set"
2006        );
2007    }
2008
2009    #[test]
2010    fn maglev_all_table_backends_unhealthy_falls_back_to_round_robin() {
2011        // When the healthy subset shares NO address with the table (every
2012        // table backend is down), probe-forward finds nothing and we fall back
2013        // to round-robin over the subset rather than returning None.
2014        let table_set = make_backends(3);
2015        let mut mag = Maglev::with_seed_and_size(DEFAULT_HASH_SEED, 1009);
2016        mag.rebuild(&table_set);
2017        let table_before = mag.table.clone();
2018
2019        // A disjoint healthy subset (different addresses than the table).
2020        let fresh = vec![
2021            rc(addr_backend("n0", 50, 9000, None)),
2022            rc(addr_backend("n1", 51, 9001, None)),
2023        ];
2024        let fresh_addrs: Vec<_> = fresh.iter().map(chosen_addr).collect();
2025
2026        let picked = chosen_addr(&pick(&mut mag, Some(7), &fresh).unwrap());
2027        assert!(
2028            fresh_addrs.contains(&picked),
2029            "fallback must route to a backend in the healthy subset"
2030        );
2031        // And still no rebuild happened.
2032        assert_eq!(
2033            mag.table, table_before,
2034            "fallback must not rebuild the table"
2035        );
2036    }
2037
2038    #[test]
2039    fn maglev_cold_start_builds_table_once() {
2040        // A freshly constructed Maglev with no prior `rebuild` must still
2041        // select (one-time cold-start build), then keep the table populated.
2042        let backends = make_backends(4);
2043        let mut mag = Maglev::with_seed_and_size(DEFAULT_HASH_SEED, 1009);
2044        assert!(mag.table.is_empty(), "table starts empty (cold)");
2045
2046        let _ = pick(&mut mag, Some(99), &backends).unwrap();
2047        assert_eq!(mag.table.len(), mag.size, "cold start populated the table");
2048        let table_after_cold = mag.table.clone();
2049
2050        // A subsequent selection (still the full set) must NOT rebuild.
2051        let _ = pick(&mut mag, Some(100), &backends).unwrap();
2052        assert_eq!(
2053            mag.table, table_after_cold,
2054            "selection after cold start must not rebuild"
2055        );
2056    }
2057
2058    #[test]
2059    fn round_robin_empty_set_returns_none_without_panic() {
2060        // FIX 2 guard: `% backends.len()` would panic on an empty Vec.
2061        let empty: Vec<Rc<RefCell<Backend>>> = vec![];
2062        let mut rr = RoundRobin::new();
2063        assert!(pick(&mut rr, None, &empty).is_none());
2064
2065        // Delegators (Rendezvous / Maglev with key == None) route through
2066        // RoundRobin and must be safe on an empty set too.
2067        let mut hrw = Rendezvous::new();
2068        assert!(pick(&mut hrw, None, &empty).is_none());
2069        let mut mag = Maglev::new();
2070        assert!(pick(&mut mag, None, &empty).is_none());
2071    }
2072
2073    #[test]
2074    fn maglev_table_size_one_is_clamped_to_a_prime() {
2075        // FIX 3 guard: requesting size 1 would make `skip = h2 % (m - 1)`
2076        // divide by zero. The constructor must clamp to a prime >= 2.
2077        let mut mag = Maglev::with_seed_and_size(DEFAULT_HASH_SEED, 1);
2078        assert!(
2079            mag.size >= 2,
2080            "size must be clamped to >= 2, got {}",
2081            mag.size
2082        );
2083
2084        let backends = make_backends(3);
2085        // Building and selecting must not panic.
2086        mag.rebuild(&backends);
2087        assert_eq!(mag.table.len(), mag.size);
2088        let _ = pick(&mut mag, Some(1), &backends).unwrap();
2089    }
2090
2091    #[test]
2092    fn next_prime_picks_the_smallest_prime_at_least_n() {
2093        assert_eq!(next_prime(0), 2);
2094        assert_eq!(next_prime(1), 2);
2095        assert_eq!(next_prime(2), 2);
2096        assert_eq!(next_prime(3), 3);
2097        assert_eq!(next_prime(4), 5);
2098        assert_eq!(next_prime(1009), 1009); // already prime — preserved
2099        assert_eq!(next_prime(65537), 65537); // default size is prime
2100    }
2101
2102    #[test]
2103    fn maglev_honors_weight() {
2104        // A backend with 4x weight should win clearly more slots.
2105        let backends = vec![
2106            rc(addr_backend("light", 1, 8001, Some(100))),
2107            rc(addr_backend("heavy", 2, 8002, Some(400))),
2108        ];
2109        let mut mag = Maglev::with_seed_and_size(DEFAULT_HASH_SEED, 1009);
2110        mag.rebuild(&backends);
2111
2112        let heavy_addr = chosen_addr(&backends[1]);
2113        let total = 20_000u64;
2114        let mut heavy = 0u64;
2115        let sel = backends.clone();
2116        for key in 0..total {
2117            if chosen_addr(&pick(&mut mag, Some(key), &sel).unwrap()) == heavy_addr {
2118                heavy += 1;
2119            }
2120        }
2121        // Expected ~80% (400 / 500); allow a margin.
2122        assert!(
2123            heavy > total * 70 / 100,
2124            "weighted Maglev did not favor the heavy backend: {heavy}/{total}"
2125        );
2126    }
2127
2128    /// One address, one share: two ids at one address claim the table slots
2129    /// of ONE backend, weighted by the heaviest of them, not one share each.
2130    ///
2131    /// TO SEE THIS RED: push every backend in `Maglev::rebuild_from` instead
2132    /// of one entry per distinct address; the shared address then holds two
2133    /// thirds of the table.
2134    #[test]
2135    fn maglev_gives_a_shared_address_the_slots_of_one_backend() {
2136        let count_shared = |backends: &[Rc<RefCell<Backend>>]| {
2137            let mut mag = Maglev::with_seed_and_size(DEFAULT_HASH_SEED, 1009);
2138            mag.rebuild(backends);
2139            let shared = chosen_addr(&backends[0]);
2140            mag.table
2141                .iter()
2142                .filter(|&&index| mag.backend_addrs[index] == shared)
2143                .count()
2144        };
2145        let equal = vec![
2146            rc(addr_backend("a", 1, 8001, None)),
2147            rc(addr_backend("b", 1, 8001, None)),
2148            rc(addr_backend("c", 2, 8002, None)),
2149        ];
2150        // 1009 slots between two addresses of equal weight.
2151        assert_eq!(count_shared(&equal), 505);
2152
2153        // The heaviest id weighs for the address: 300 against 100.
2154        let unequal = vec![
2155            rc(addr_backend("a", 1, 8001, Some(100))),
2156            rc(addr_backend("b", 1, 8001, Some(300))),
2157            rc(addr_backend("c", 2, 8002, Some(100))),
2158        ];
2159        assert_eq!(count_shared(&unequal), 1009 * 3 / 4 + 1);
2160    }
2161
2162    // ----- Random / PowerOfTwo: injected, seedable entropy -----
2163
2164    /// Map a test backend's address back to the index `make_backends` gave it
2165    /// (`last_octet - 1`), so a selection sequence can be asserted as a
2166    /// compact `Vec<u8>` instead of a `Vec<SocketAddr>`.
2167    fn addr_index(addr: SocketAddr) -> u8 {
2168        match addr.ip() {
2169            IpAddr::V4(v4) => v4.octets()[3] - 1,
2170            IpAddr::V6(_) => unreachable!("test backends are always IPv4"),
2171        }
2172    }
2173
2174    #[test]
2175    fn random_is_deterministic_and_matches_an_expected_sequence() {
2176        let backends = make_backends(4);
2177        let mut r1 = Random::with_seed(DEFAULT_HASH_SEED);
2178        let mut r2 = Random::with_seed(DEFAULT_HASH_SEED);
2179
2180        let sel1 = backends.clone();
2181        let sel2 = backends.clone();
2182
2183        let seq1: Vec<u8> = (0..20)
2184            .map(|_| addr_index(chosen_addr(&pick(&mut r1, None, &sel1).unwrap())))
2185            .collect();
2186        let seq2: Vec<u8> = (0..20)
2187            .map(|_| addr_index(chosen_addr(&pick(&mut r2, None, &sel2).unwrap())))
2188            .collect();
2189
2190        // Reproducibility: same seed + same inputs + same call sequence must
2191        // give the same picks.
2192        assert_eq!(
2193            seq1, seq2,
2194            "Random::with_seed must reproduce the same sequence for the same seed"
2195        );
2196
2197        // Absolute-value check (paired with the determinism check above so a
2198        // test that is merely "reproducible" but reproducibly wrong still
2199        // fails): the sequence for DEFAULT_HASH_SEED over 4 uniformly
2200        // weighted backends, captured from this implementation.
2201        //
2202        // Brittleness note: this literal is coupled to `rand` 0.10.2's exact
2203        // `StdRng` algorithm (ChaCha12) and `seed_from_u64`'s splitmix64
2204        // expansion. A `rand` version bump that changes either (the type's
2205        // own docs disclaim portability/reproducibility across versions) will
2206        // break this assertion for reasons unrelated to Sōzu's load-balancing
2207        // logic — recapture the sequence, don't "fix" the algorithm.
2208        let expected: Vec<u8> = vec![1, 2, 1, 1, 0, 2, 1, 3, 1, 0, 3, 1, 0, 0, 3, 2, 1, 3, 3, 1];
2209        assert_eq!(
2210            seq1, expected,
2211            "Random sequence for DEFAULT_HASH_SEED regressed"
2212        );
2213    }
2214
2215    #[test]
2216    fn random_distribution_does_not_collapse() {
2217        // A seeded RNG must still spread draws across every backend rather
2218        // than pinning one — determinism is per-(seed, sequence), not
2219        // "always the same backend".
2220        //
2221        // This is a smoke check, not a uniformity test: the +/-35% band is
2222        // wide enough (roughly ±12 sigma at n=5, total=5000) to reliably
2223        // catch full collapse (one backend gets ~0 or ~100%) or a gross bias,
2224        // but a systematic ~20-25% skew would pass it. Tightening the band
2225        // to something a real skew would fail is a legitimate follow-up; the
2226        // name deliberately promises no more than "does not collapse".
2227        let n = 5u8;
2228        let backends = make_backends(n);
2229        let mut r = Random::with_seed(DEFAULT_HASH_SEED);
2230        let sel = backends.clone();
2231
2232        let total = 5_000u64;
2233        let mut counts = [0u64; 5];
2234        for _ in 0..total {
2235            let idx = addr_index(chosen_addr(&pick(&mut r, None, &sel).unwrap()));
2236            counts[idx as usize] += 1;
2237        }
2238
2239        let expected = total / n as u64;
2240        for (idx, &c) in counts.iter().enumerate() {
2241            assert!(
2242                c > expected * 65 / 100 && c < expected * 135 / 100,
2243                "Random distribution skewed for backend {idx}: got {c}, expected ~{expected}"
2244            );
2245        }
2246    }
2247
2248    #[test]
2249    fn power_of_two_tie_break_is_deterministic_and_matches_an_expected_sequence() {
2250        // All backends report the same load (0 active connections), so both
2251        // sampled candidates always tie and every call reaches the coin-flip
2252        // branch.
2253        let backends = make_backends(4);
2254        let mut p1 = PowerOfTwo::with_seed(DEFAULT_HASH_SEED, LoadMetric::Connections);
2255        let mut p2 = PowerOfTwo::with_seed(DEFAULT_HASH_SEED, LoadMetric::Connections);
2256
2257        let sel1 = backends.clone();
2258        let sel2 = backends.clone();
2259
2260        let seq1: Vec<u8> = (0..20)
2261            .map(|_| addr_index(chosen_addr(&pick(&mut p1, None, &sel1).unwrap())))
2262            .collect();
2263        let seq2: Vec<u8> = (0..20)
2264            .map(|_| addr_index(chosen_addr(&pick(&mut p2, None, &sel2).unwrap())))
2265            .collect();
2266
2267        // Reproducibility: `with_seed` is deterministic given the same seed
2268        // and call sequence — the property tests/simulation need. This is
2269        // NOT a claim about production: there the seed comes from the
2270        // backend map's generator, which `BackendMap::new` seeds from OS
2271        // entropy specifically so two independent worker processes do NOT
2272        // draw from the same keystream (see the struct docs' "Seed source"
2273        // section for why sharing a seed here would be a regression, not a
2274        // feature).
2275        assert_eq!(
2276            seq1, seq2,
2277            "PowerOfTwo::with_seed must reproduce the same tie-break sequence for the same seed"
2278        );
2279
2280        // Absolute-value check: the sequence ranges over all four backends,
2281        // because each call samples a fresh uniformly random PAIR and the
2282        // tie then resolves inside that pair. That is the visible signature
2283        // of random sampling — the O(n) fold this replaced could only ever
2284        // emit backends 2 and 3, the last two it happened to retain.
2285        //
2286        // Brittleness note: this literal is coupled to `rand` 0.10.2's exact
2287        // `StdRng` algorithm (ChaCha12), to `seed_from_u64`'s splitmix64
2288        // expansion, AND to the exact order and shape of the draws
2289        // `PowerOfTwo::next_available_backend` makes. TWO kinds of change
2290        // move it legitimately, not one: a `rand` version bump that alters
2291        // either generator (the type's own docs disclaim
2292        // portability/reproducibility across versions), and a deliberate
2293        // change to the selection procedure itself. The second is not
2294        // hypothetical — this literal was recaptured once already, when the
2295        // two-lightest fold became real power-of-two-choices and the
2296        // sequence stopped being confined to backends 2 and 3.
2297        //
2298        // HOW TO TELL A LEGITIMATE RECAPTURE FROM PAPERING OVER A
2299        // REGRESSION. Every other `power_of_two_*` test in this module
2300        // asserts a SEMANTIC property that no RNG stream can shift. There
2301        // are SEVEN of them — `grep -cE '^\s+fn power_of_two_'` in this file
2302        // gives eight, this test included — and the rule below covers all
2303        // seven, not a convenient subset. The pattern is anchored on
2304        // purpose: unanchored it also counts this very comment, reports
2305        // nine, and sends the reader after a test that does not exist. The
2306        // seven are:
2307        // `power_of_two_tie_break_distribution_does_not_collapse`,
2308        // `power_of_two_tie_break_is_decided_by_the_coin_flip_not_by_sample_position`,
2309        // `power_of_two_touches_exactly_two_backends`,
2310        // `power_of_two_always_returns_the_lighter_of_two_backends`,
2311        // `power_of_two_handles_empty_and_singleton_sets_without_panic`,
2312        // `power_of_two_never_returns_the_strictly_heaviest_backend` and
2313        // `power_of_two_sample_size_is_two_not_the_whole_set`.
2314        // The eighth this rule used to name, the seed-correlation guard, is
2315        // `power_of_two_policies_seeded_by_the_os_are_not_correlated` in
2316        // `backends.rs` now that the production seed comes from `BackendMap`
2317        // (#1684); it belongs in the rule too.
2318        // Recapture only when THIS assertion is the only failing one and all
2319        // eight are green with their bodies untouched. If any of them is
2320        // red, or one had to be edited to get green, the algorithm regressed
2321        // and the new sequence is evidence of it — do not capture it. Check
2322        // the grep count first: a `power_of_two_*` test added after this
2323        // comment was written belongs in the rule too.
2324        let expected: Vec<u8> = vec![1, 3, 1, 1, 1, 2, 1, 2, 0, 1, 2, 3, 1, 3, 3, 3, 0, 2, 0, 2];
2325        assert_eq!(
2326            seq1, expected,
2327            "PowerOfTwo selection sequence for DEFAULT_HASH_SEED regressed"
2328        );
2329    }
2330
2331    #[test]
2332    fn power_of_two_tie_break_distribution_does_not_collapse() {
2333        // Every backend reports the same load, so the pair sample is uniform
2334        // and the tie inside it resolves by coin flip: the selection must
2335        // spread over ALL FOUR backends, not pin one — otherwise the seeded
2336        // RNG would have traded "ambient nondeterminism" for "deterministic
2337        // bias", the regression the herding argument in `PowerOfTwo`'s doc
2338        // comment warns about.
2339        //
2340        // This assertion got STRICTER when the O(n) scan became real
2341        // power-of-two-choices: it used to require backends 0 and 1 to win
2342        // exactly ZERO times, because the old fold could only ever retain
2343        // the last two backends it walked. Requiring all four to take a
2344        // roughly equal share is a stronger statement about the same draws.
2345        //
2346        // Smoke check, not a uniformity test: the +/-35% band (~14 sigma at
2347        // p=0.25, total=5000) reliably catches a collapse onto a subset but
2348        // would pass a systematic skew — see the equivalent note on
2349        // `random_distribution_does_not_collapse`.
2350        let backends = make_backends(4);
2351        let mut p = PowerOfTwo::with_seed(DEFAULT_HASH_SEED, LoadMetric::Connections);
2352        let sel = backends.clone();
2353
2354        let total = 5_000u64;
2355        let mut counts = [0u64; 4];
2356        for _ in 0..total {
2357            let idx = addr_index(chosen_addr(&pick(&mut p, None, &sel).unwrap()));
2358            counts[idx as usize] += 1;
2359        }
2360
2361        // Uniform pair sampling over four equally loaded backends makes
2362        // every backend equally likely, so each must take roughly a quarter
2363        // and none may be starved.
2364        let expected = total / 4;
2365        for (idx, &c) in counts.iter().enumerate() {
2366            assert!(
2367                c > expected * 65 / 100 && c < expected * 135 / 100,
2368                "PowerOfTwo distribution skewed for backend {idx}: got {c}, expected ~{expected}"
2369            );
2370        }
2371    }
2372
2373    #[test]
2374    fn power_of_two_tie_break_is_decided_by_the_coin_flip_not_by_sample_position() {
2375        // `power_of_two_tie_break_distribution_does_not_collapse` CANNOT see
2376        // the coin flip, and neither can any other distribution test here.
2377        // The sampler draws an ordered pair uniformly — `first` over
2378        // `0..len`, `second` uniformly over the rest — so the two positions
2379        // are exchangeable, and "always keep `first`" has exactly the same
2380        // marginal distribution as a fair coin. Replacing
2381        // `self.rng.random_bool(0.5)` with `true` leaves every other test in
2382        // this module green except the captured-sequence one, which would
2383        // only fail through RNG-stream drift — the same way a `rand` upgrade
2384        // fails it, and whose own comment invites a recapture. A dropped
2385        // coin flip could ride in behind such a recapture unnoticed.
2386        //
2387        // This test makes the tie-break directly observable instead. Each
2388        // trial is the first selection of a policy seeded with the trial's
2389        // own seed, so the pair is the first two draws of that seed's
2390        // keystream, and a generator seeded alike replays which backend was
2391        // drawn as `first` and which as `second`. The replay stops before the
2392        // coin flip, so it agrees with the policy whether the flip is there
2393        // or not. (It used to read the order off `PeakEWMA` stamps taken on
2394        // the wall clock; a selection now measures both samples at one
2395        // injected instant, #1684.)
2396        //
2397        // The replay mirrors the sampler's exact draws, so it is checked, not
2398        // trusted: with THREE equally loaded backends the pair leaves one out,
2399        // and every pick must be one of the two the replay predicts. A change
2400        // to the sampling desynchronises the replay, and a pick then lands on
2401        // the backend it left out, failing this test loudly instead of
2402        // quietly counting noise.
2403        //
2404        // Every backend carries the same load, so every selection is an exact
2405        // tie and every one of them reaches the tie-break.
2406        const TOTAL: usize = 2_000;
2407        const BACKENDS: u8 = 3;
2408        let backends = make_backends(BACKENDS);
2409
2410        let sel = backends.clone();
2411        let mut second_measured_won = 0usize;
2412        for trial in 0..TOTAL as u64 {
2413            let seed = DEFAULT_HASH_SEED ^ trial;
2414            let mut p = PowerOfTwo::with_seed(seed, LoadMetric::Connections);
2415            let mut replay = StdRng::seed_from_u64(seed);
2416            let len = usize::from(BACKENDS);
2417            let first = replay.random_range(0..len);
2418            let mut second = replay.random_range(0..len - 1);
2419            if second >= first {
2420                second += 1;
2421            }
2422            let picked = usize::from(addr_index(chosen_addr(&pick(&mut p, None, &sel).unwrap())));
2423            assert!(
2424                picked == first || picked == second,
2425                "trial {trial}: picked backend {picked}, outside the pair ({first}, {second}) \
2426                 the replay predicts; the sampler no longer draws the pair this test replays"
2427            );
2428            if picked == second {
2429                second_measured_won += 1;
2430            }
2431        }
2432
2433        // A fair coin gives the second-measured backend half the ties. The
2434        // band is +/-20 percentage points around 50%, so its half-width is
2435        // `0.2 * TOTAL` against a standard deviation of `sqrt(TOTAL * 0.25)`:
2436        // 400 against sigma 22.4, ~17.9 sigma. It cannot flake, and it is
2437        // two-sided: pinning the tie-break to `first` drives this to 0,
2438        // pinning it to `second` drives it to `TOTAL`.
2439        let low = TOTAL * 30 / 100;
2440        let high = TOTAL * 70 / 100;
2441        assert!(
2442            second_measured_won > low && second_measured_won < high,
2443            "power-of-two resolved {second_measured_won} of {TOTAL} exact ties in favour of \
2444             the second-measured sample; a coin flip must land near {}, and a count at either \
2445             end means the tie is being awarded by sample position instead",
2446            TOTAL / 2
2447        );
2448    }
2449
2450    // ----- PowerOfTwo: the sampling itself, not just the outcome -----
2451
2452    #[test]
2453    fn power_of_two_touches_exactly_two_backends() {
2454        // The two-load-reads claim, MEASURED rather than asserted
2455        // structurally. It is a claim about how many backends the POLICY
2456        // reads, not about the cost of a selection: the caller has already
2457        // walked every backend in `BackendList::next_available_backend_with_key`
2458        // to build the candidate set handed in here, so a selection is `O(n)` whatever
2459        // this policy does. `LoadMetric::ConnectionTime` reads a backend's load
2460        // through `Backend::peak_ewma_connection` -> `PeakEWMA::get` ->
2461        // `PeakEWMA::observe`, and `observe` stamps `last_event` with the
2462        // selection's instant. That stamp is an exact per-backend receipt saying
2463        // "this backend's load was measured", so the number of fresh stamps
2464        // after one selection IS the number of backends the algorithm
2465        // consulted: a full scan leaves `n` of them, power-of-two-choices
2466        // leaves exactly 2.
2467        //
2468        // This is the assertion a "did a backend come back?" test cannot
2469        // make: the scan-then-coin-flip shape this replaced also returned a
2470        // backend, it just read all 64 first.
2471        const N: u8 = 64;
2472        let backends = make_backends(N);
2473
2474        // Stamp every backend with one instant and select one nanosecond
2475        // later on the injected clock, so "was this backend measured?" is
2476        // decidable without depending on the wall clock's resolution.
2477        let before = Instant::now();
2478        for backend in &backends {
2479            backend.borrow_mut().connection_time.last_event = before;
2480        }
2481
2482        let mut p = PowerOfTwo::with_seed(DEFAULT_HASH_SEED, LoadMetric::ConnectionTime);
2483        assert!(pick_at(&mut p, None, &backends, before + Duration::from_nanos(1)).is_some());
2484
2485        let touched = backends
2486            .iter()
2487            .filter(|backend| backend.borrow().connection_time.last_event > before)
2488            .count();
2489        assert_eq!(
2490            touched, 2,
2491            "power-of-two must measure exactly 2 of the {N} backends; measuring {touched} means \
2492             the policy reads every backend's load, which the name exists to rule out"
2493        );
2494    }
2495
2496    #[test]
2497    fn power_of_two_always_returns_the_lighter_of_two_backends() {
2498        // With exactly two backends the sample IS the whole set, so P2C is
2499        // fully determined and needs no statistics: the lighter backend must
2500        // win every single call. The scan-then-coin-flip shape this replaced
2501        // computed both measures, asserted their order, and then discarded
2502        // that order for a coin flip — returning the HEAVIER backend about
2503        // half the time.
2504        let backends = make_backends(2);
2505        backends[0].borrow_mut().active_connections = 7;
2506        backends[1].borrow_mut().active_connections = 1;
2507
2508        let mut p = PowerOfTwo::with_seed(DEFAULT_HASH_SEED, LoadMetric::Connections);
2509        let sel = backends.clone();
2510        for call in 0..200 {
2511            let picked = addr_index(chosen_addr(&pick(&mut p, None, &sel).unwrap()));
2512            assert_eq!(
2513                picked, 1,
2514                "power-of-two must keep the lighter of the two sampled backends (call {call} \
2515                 picked backend {picked}, which carries 7 connections against 1)"
2516            );
2517        }
2518    }
2519
2520    #[test]
2521    fn power_of_two_handles_empty_and_singleton_sets_without_panic() {
2522        // Sampling a second DISTINCT index draws from `len - 1` slots, which
2523        // is an empty range for a one-backend set — `random_range` panics on
2524        // an empty range, so the singleton case must short-circuit before the
2525        // draw. A cluster scaled down to one backend is ordinary, not exotic.
2526        let mut p = PowerOfTwo::with_seed(DEFAULT_HASH_SEED, LoadMetric::Connections);
2527
2528        let empty: Vec<Rc<RefCell<Backend>>> = vec![];
2529        assert!(
2530            pick(&mut p, None, &empty).is_none(),
2531            "power-of-two selects nothing from an empty candidate set"
2532        );
2533
2534        let single = make_backends(1);
2535        for _ in 0..10 {
2536            assert_eq!(
2537                addr_index(chosen_addr(&pick(&mut p, None, &single).unwrap())),
2538                0,
2539                "power-of-two returns the only backend of a singleton set"
2540            );
2541        }
2542    }
2543
2544    #[test]
2545    fn power_of_two_never_returns_the_strictly_heaviest_backend() {
2546        // Regression guard for the load profile that broke the O(n) fold
2547        // this replaced. With three strictly ordered loads the heaviest
2548        // backend loses EVERY pairing it can be drawn into, so P2C can never
2549        // return it — a deterministic assertion over any number of calls.
2550        //
2551        // The previous implementation walked every backend keeping a
2552        // `(first, second)` pair, and its "otherwise" branch fired both when
2553        // the new measure was lighter than `first` AND when it was heavier
2554        // than `second`, evicting both candidates in the second case. On
2555        // loads [0, 1, 5] that left `first = 5` and `second = 0`, tripping
2556        // its own `first <= second` invariant (a debug-build panic) and, in
2557        // a release build, coin-flipping the HEAVIEST backend into the
2558        // result half the time.
2559        let backends = make_backends(3);
2560        backends[0].borrow_mut().active_connections = 0;
2561        backends[1].borrow_mut().active_connections = 1;
2562        backends[2].borrow_mut().active_connections = 5;
2563
2564        let mut p = PowerOfTwo::with_seed(DEFAULT_HASH_SEED, LoadMetric::Connections);
2565        let sel = backends.clone();
2566        for call in 0..2_000 {
2567            let picked = addr_index(chosen_addr(&pick(&mut p, None, &sel).unwrap()));
2568            assert_ne!(
2569                picked, 2,
2570                "power-of-two returned the strictly heaviest backend on call {call}: it is \
2571                 heavier than either backend it can be paired against"
2572            );
2573        }
2574    }
2575
2576    #[test]
2577    fn power_of_two_sample_size_is_two_not_the_whole_set() {
2578        // Same two-load-reads property as `power_of_two_touches_exactly_two_backends`,
2579        // measured from the OUTSIDE — through the selection distribution
2580        // alone, with no access to a backend's internals.
2581        //
2582        // Give one backend a uniquely light load among `N`. A policy that
2583        // samples `k` backends uniformly and keeps the lightest returns that
2584        // backend with probability exactly `k / N`, so the observed hit rate
2585        // MEASURES the sample size: `k = hit_rate * N`. Power-of-two-choices
2586        // gives `k = 2`.
2587        //
2588        // The fold this replaced scores worse here than ANY sample size.
2589        // Restored under this exact test it gives `k_est = 0.00`: the unique
2590        // minimum came back 0 times out of 100_000. It was `first` after
2591        // backend 0, but from backend 2 on every remaining measure tied
2592        // `second`, so the fold's "otherwise" branch fired on every step,
2593        // shifting `first` into `second` and evicting the minimum for good.
2594        // What it returned was a coin flip between the last two backends it
2595        // happened to walk.
2596        const N: u8 = 100;
2597        let backends = make_backends(N);
2598        for backend in backends.iter().skip(1) {
2599            backend.borrow_mut().active_connections = 1_000;
2600        }
2601        // `backends[0]` keeps 0 active connections: the unique minimum.
2602
2603        let mut p = PowerOfTwo::with_seed(DEFAULT_HASH_SEED, LoadMetric::Connections);
2604        let sel = backends.clone();
2605        let total = 100_000u32;
2606        let mut hits = 0u32;
2607        for _ in 0..total {
2608            if addr_index(chosen_addr(&pick(&mut p, None, &sel).unwrap())) == 0 {
2609                hits += 1;
2610            }
2611        }
2612
2613        // At k = 2, N = 100 and total = 100_000 the standard deviation of
2614        // `k_est` is `N * sqrt(p * (1 - p) / total)` ~ 0.044, so the
2615        // [1.5, 2.5] band is a ~11-sigma envelope: wide enough never to
2616        // flake, narrow enough to exclude k = 1 (plain random), k = 3, and
2617        // the k_est = 0.00 the old fold produced.
2618        let k_est = f64::from(hits) / f64::from(total) * f64::from(N);
2619        assert!(
2620            (1.5..=2.5).contains(&k_est),
2621            "power-of-two consulted ~{k_est:.2} of {N} backends (unique minimum returned \
2622             {hits}/{total} times); the algorithm must sample exactly 2"
2623        );
2624    }
2625
2626    /// The weighted draw `Random` made before it stopped allocating: collect
2627    /// the weights, build a `WeightedIndex`, and fall back to a uniform
2628    /// `choose` when `WeightedIndex::new` rejects them.
2629    fn weighted_index_reference(rng: &mut StdRng, weights: &[i32]) -> Option<usize> {
2630        use rand::distr::{Distribution, weighted::WeightedIndex};
2631        match WeightedIndex::new(weights.to_vec()) {
2632            Ok(distribution) => Some(distribution.sample(rng)),
2633            Err(_) => {
2634                let positions: Vec<usize> = (0..weights.len()).collect();
2635                positions.choose(rng).copied()
2636            }
2637        }
2638    }
2639
2640    #[test]
2641    fn random_weighted_pick_matches_weighted_index() {
2642        // Weight sets covering every branch of `WeightedIndex::new`: valid,
2643        // zero weights among positive ones, all zero, a negative weight
2644        // (first and later), an `i32` overflow, a singleton and the empty set.
2645        let weight_sets: &[&[i32]] = &[
2646            &[100, 100, 100],
2647            &[1, 2, 3, 4, 5, 6, 7],
2648            &[0, 5, 0, 0, 9, 0],
2649            &[7, 0, 0],
2650            &[0, 0, 0, 3],
2651            &[0, 0, 0],
2652            &[-1, 5, 5],
2653            &[5, 5, -1, 5],
2654            &[i32::MAX, 1],
2655            &[i32::MAX / 2, i32::MAX / 2, 1],
2656            &[i32::MAX / 2, i32::MAX / 2, 2],
2657            &[42],
2658            &[],
2659        ];
2660        for (set, weights) in weight_sets.iter().enumerate() {
2661            let backends: Vec<_> = weights
2662                .iter()
2663                .enumerate()
2664                .map(|(index, &weight)| rc(addr_backend("w", index as u8 + 1, 80, Some(weight))))
2665                .collect();
2666            for seed in 0..64 {
2667                let mut policy = Random::with_seed(seed);
2668                let mut reference = StdRng::seed_from_u64(seed);
2669                for draw in 0..32 {
2670                    let expected = weighted_index_reference(&mut reference, weights);
2671                    let got = pick(&mut policy, None, &backends).map(|backend| {
2672                        backends
2673                            .iter()
2674                            .position(|candidate| Rc::ptr_eq(candidate, &backend))
2675                            .expect("the pick is one of the candidates")
2676                    });
2677                    assert_eq!(
2678                        got, expected,
2679                        "weight set {set} {weights:?}, seed {seed}, draw {draw}"
2680                    );
2681                }
2682            }
2683        }
2684    }
2685
2686    // ----- #524: the client affinity key of HTTP, HTTPS and TCP -----
2687
2688    /// #524: the affinity key is a pure function of the client, pinned to a
2689    /// value.
2690    ///
2691    /// Every worker of every instance must derive the same key for the same
2692    /// client, or `HRW`/`MAGLEV` send one client to different backends
2693    /// depending on which worker accepted it, and an upgrade to a build whose
2694    /// hash moved would remap every client at once. The pinned values make a
2695    /// change to the construction a deliberate, visible edit.
2696    ///
2697    /// TO SEE THIS RED: seed [`affinity_key_from_ip`] with anything other than
2698    /// [`DEFAULT_HASH_SEED`], or drop its domain tag.
2699    #[test]
2700    fn affinity_keys_are_pinned_across_workers_and_releases() {
2701        let ip = IpAddr::V4(Ipv4Addr::new(192, 0, 2, 7));
2702        assert_eq!(affinity_key_from_ip(ip), affinity_key_from_ip(ip));
2703        assert_eq!(affinity_key_from_ip(ip), PINNED_IP_KEY);
2704        assert_eq!(affinity_key_from_value(b"tenant-42"), PINNED_VALUE_KEY);
2705    }
2706    /// [`affinity_key_from_ip`] of `192.0.2.7`, measured when the construction
2707    /// was introduced.
2708    const PINNED_IP_KEY: u64 = 12_659_869_658_761_991_936;
2709    /// [`affinity_key_from_value`] of `tenant-42`, measured likewise.
2710    const PINNED_VALUE_KEY: u64 = 4_927_821_811_038_511_670;
2711
2712    /// #524: a dual-stack listener reports an IPv4 client as the mapped
2713    /// `::ffff:a.b.c.d`; it must key exactly as the same client seen on an
2714    /// IPv4 listener, or one client lands on two backends depending on which
2715    /// listener it reached.
2716    #[test]
2717    fn an_ipv4_mapped_client_keys_as_its_ipv4_address() {
2718        let v4 = Ipv4Addr::new(198, 51, 100, 9);
2719        assert_eq!(
2720            affinity_key_from_ip(IpAddr::V6(v4.to_ipv6_mapped())),
2721            affinity_key_from_ip(IpAddr::V4(v4)),
2722            "a mapped IPv4 client must key as its IPv4 address"
2723        );
2724        assert_ne!(
2725            affinity_key_from_ip(IpAddr::V6(v4.to_ipv6_compatible())),
2726            affinity_key_from_ip(IpAddr::V4(v4)),
2727            "only the mapped form is the same client; a compatible address is a distinct IPv6 one"
2728        );
2729    }
2730
2731    /// #524: a header value spelling an address's bytes is not that address.
2732    /// The two constructions carry distinct domain tags, so a client cannot
2733    /// collide with another client's source IP by choosing its header value.
2734    #[test]
2735    fn a_value_never_keys_as_the_address_it_spells() {
2736        let v4 = Ipv4Addr::new(10, 1, 2, 3);
2737        let mut spelled = vec![4u8];
2738        spelled.extend_from_slice(&v4.octets());
2739        assert_ne!(
2740            affinity_key_from_value(&spelled),
2741            affinity_key_from_ip(IpAddr::V4(v4))
2742        );
2743        assert_ne!(
2744            affinity_key_from_value(b"tenant-a"),
2745            affinity_key_from_value(b"tenant-b")
2746        );
2747    }
2748
2749    /// #524: `k = max(2, ceil(percent × N / 100))`, capped at N, and no
2750    /// shard at all below `min_backends`.
2751    #[test]
2752    fn a_shard_size_follows_the_decided_formula() {
2753        let sharding = |percent, min_backends| ShuffleSharding {
2754            percent,
2755            min_backends,
2756            mode: ShardMode::Fallback,
2757        };
2758        for (percent, min_backends, primaries, expected) in [
2759            (1, 8, 8, Some(2)),     // ceil(0.08) = 1, floored at 2
2760            (25, 8, 8, Some(2)),    // exactly 2
2761            (30, 8, 10, Some(3)),   // exactly 3
2762            (50, 8, 9, Some(5)),    // ceil(4.5)
2763            (34, 8, 100, Some(34)), // exactly 34
2764            (100, 8, 12, Some(12)), // every backend
2765            (25, 8, 7, None),       // below the threshold
2766            (25, 2, 2, Some(2)),    // the smallest cluster that shards
2767            (1, 2, 3, Some(2)),     // floored at 2
2768        ] {
2769            assert_eq!(
2770                sharding(percent, min_backends).shard_size(primaries),
2771                expected,
2772                "percent={percent} min={min_backends} N={primaries}"
2773            );
2774        }
2775        assert_eq!(
2776            ShuffleSharding::from_cluster(Some(10), None, None),
2777            Some(sharding(10, DEFAULT_SHARD_MIN_BACKENDS)),
2778            "min_backends defaults to 8 and the mode to FALLBACK"
2779        );
2780        assert_eq!(ShuffleSharding::from_cluster(None, Some(4), Some(1)), None);
2781    }
2782
2783    // ----- `weight` under ROUND_ROBIN, LEAST_LOADED and POWER_OF_TWO -----
2784
2785    /// Backends at distinct addresses carrying `weights`, `None` for unset.
2786    fn weighted_backends(weights: &[Option<i32>]) -> Vec<Rc<RefCell<Backend>>> {
2787        weights
2788            .iter()
2789            .enumerate()
2790            .map(|(i, &weight)| rc(addr_backend(&format!("w{i}"), i as u8 + 1, 8000, weight)))
2791            .collect()
2792    }
2793
2794    /// The list positions `policy` picks over `count` selections.
2795    fn sequence(
2796        policy: &mut impl LoadBalancingAlgorithm,
2797        backends: &[Rc<RefCell<Backend>>],
2798        count: usize,
2799    ) -> Vec<usize> {
2800        (0..count)
2801            .map(|_| {
2802                let picked = pick(policy, None, backends).expect("a non-empty set selects");
2803                backends
2804                    .iter()
2805                    .position(|backend| Rc::ptr_eq(backend, &picked))
2806                    .expect("the pick is one of the candidates")
2807            })
2808            .collect()
2809    }
2810
2811    /// Smooth weighted round robin interleaves the heavy backend's picks
2812    /// instead of bunching them, and repeats every `Σw` selections.
2813    ///
2814    /// TO SEE THIS RED: drop the `smooth_weighted` branch of
2815    /// `RoundRobin::next_available_backend`; the cursor then yields
2816    /// `0 1 2 0 1 2 0`.
2817    #[test]
2818    fn round_robin_smooth_weighted_interleaves_picks() {
2819        let backends = weighted_backends(&[Some(5), Some(1), Some(1)]);
2820        let mut round_robin = RoundRobin::new();
2821        let cycle = [0, 0, 1, 0, 2, 0, 0];
2822        let expected: Vec<usize> = cycle.iter().chain(cycle.iter()).copied().collect();
2823        assert_eq!(sequence(&mut round_robin, &backends, 14), expected);
2824    }
2825
2826    /// Over one full cycle of `Σw` selections, each backend is picked exactly
2827    /// `w` times. A `0` or negative weight counts as `1`: the backend keeps
2828    /// the smallest share, as under `HRW` and `MAGLEV`.
2829    #[test]
2830    fn round_robin_picks_each_backend_exactly_its_weight_per_cycle() {
2831        let backends = weighted_backends(&[Some(300), None, Some(0), Some(-5), Some(7)]);
2832        // `None` reads as the default weight, 100.
2833        let weights = [300, DEFAULT_WEIGHT as usize, 1, 1, 7];
2834        let total: usize = weights.iter().sum();
2835        let mut round_robin = RoundRobin::new();
2836        for cycle in 0..3 {
2837            let mut counts = [0usize; 5];
2838            for position in sequence(&mut round_robin, &backends, total) {
2839                counts[position] += 1;
2840            }
2841            assert_eq!(counts, weights, "cycle {cycle}");
2842        }
2843    }
2844
2845    /// Equal weights, whatever their value, select exactly as the historical
2846    /// cursor does, including across a removal from the set; so do weights
2847    /// that only clamp to equal (`0` and `1`).
2848    #[test]
2849    fn round_robin_equal_weights_reproduce_the_cursor_sequence() {
2850        let reference = weighted_backends(&[None, None, None, None]);
2851        for weights in [
2852            [Some(50); 4],
2853            [Some(DEFAULT_WEIGHT); 4],
2854            [Some(0), Some(1), Some(-3), Some(1)],
2855        ] {
2856            let mut reference = reference.clone();
2857            let mut backends = weighted_backends(&weights);
2858            let mut expected_policy = RoundRobin::new();
2859            let mut policy = RoundRobin::new();
2860            assert_eq!(
2861                sequence(&mut policy, &backends, 9),
2862                sequence(&mut expected_policy, &reference, 9),
2863                "{weights:?}"
2864            );
2865            reference.remove(1);
2866            backends.remove(1);
2867            assert_eq!(
2868                sequence(&mut policy, &backends, 7),
2869                sequence(&mut expected_policy, &reference, 7),
2870                "{weights:?} after a removal"
2871            );
2872        }
2873    }
2874
2875    /// Equal weights keep the exact sequence the cursor-only `RoundRobin` of
2876    /// `origin/main` (588c4e2f5) produced, across candidate membership
2877    /// transitions. That `RoundRobin` picks view index `next_backend % len` and
2878    /// then stores `(next_backend + 1) % len`; the expected literals below are
2879    /// worked by hand from those two lines, not captured from this
2880    /// implementation. Positions are list positions of a four-backend list:
2881    ///
2882    /// - view `[0, 1, 2, 3]`, six picks: `0 1 2 3 0 1`, cursor `2`;
2883    /// - backend 1 leaves the view (`[0, 2, 3]`): `2 % 3 = 2` is list
2884    ///   position 3, then `0 1 2 0` in view order: list `3 0 2 3 0`, cursor `1`;
2885    /// - backend 1 re-enters (`[0, 1, 2, 3]`): `1 2 3 0 1`, cursor `2`;
2886    /// - the view shrinks to `[0, 3]`: `2 % 2 = 0`, then `1 0 1`: list `0 3 0 3`.
2887    ///
2888    /// TO SEE THIS RED: let equal weights fall through to `smooth_weighted`;
2889    /// its running weights ignore the cursor, so the removal leg moves.
2890    #[test]
2891    fn round_robin_equal_weights_match_the_origin_main_sequence_across_membership_changes() {
2892        for weights in [
2893            [None; 4],
2894            [Some(50); 4],
2895            [Some(DEFAULT_WEIGHT); 4],
2896            [Some(0), Some(1), Some(-3), Some(1)],
2897        ] {
2898            for rebuild_on_change in [false, true] {
2899                let backends = weighted_backends(&weights);
2900                let mut policy = RoundRobin::new();
2901                let now = Instant::now();
2902                let positions = |policy: &mut RoundRobin, view: &[usize], count: usize| {
2903                    (0..count)
2904                        .map(|_| {
2905                            let picked = policy
2906                                .next_available_backend(None, Candidates::new(&backends, view, now))
2907                                .expect("a non-empty view selects");
2908                            backends
2909                                .iter()
2910                                .position(|backend| Rc::ptr_eq(backend, &picked))
2911                                .expect("the pick is one of the candidates")
2912                        })
2913                        .collect::<Vec<_>>()
2914                };
2915                let change = |policy: &mut RoundRobin| {
2916                    if rebuild_on_change {
2917                        policy.rebuild(&backends);
2918                    }
2919                };
2920                let label = format!("{weights:?} rebuild={rebuild_on_change}");
2921
2922                assert_eq!(
2923                    positions(&mut policy, &[0, 1, 2, 3], 6),
2924                    [0, 1, 2, 3, 0, 1],
2925                    "{label}"
2926                );
2927                change(&mut policy);
2928                assert_eq!(
2929                    positions(&mut policy, &[0, 2, 3], 5),
2930                    [3, 0, 2, 3, 0],
2931                    "{label} after a removal"
2932                );
2933                change(&mut policy);
2934                assert_eq!(
2935                    positions(&mut policy, &[0, 1, 2, 3], 5),
2936                    [1, 2, 3, 0, 1],
2937                    "{label} after a re-add"
2938                );
2939                change(&mut policy);
2940                assert_eq!(
2941                    positions(&mut policy, &[0, 3], 4),
2942                    [0, 3, 0, 3],
2943                    "{label} after shrinking to two"
2944                );
2945            }
2946        }
2947    }
2948
2949    /// With no client key, `HRW` and `MAGLEV` fall back to round robin, now
2950    /// weighted like it.
2951    #[test]
2952    fn affinity_fallback_round_robin_honors_weight() {
2953        let backends = weighted_backends(&[Some(3), Some(1)]);
2954        let expected = vec![0, 0, 1, 0, 0, 0, 1, 0];
2955        assert_eq!(
2956            sequence(&mut Rendezvous::new(), &backends, 8),
2957            expected,
2958            "HRW"
2959        );
2960        let mut maglev = Maglev::new();
2961        maglev.rebuild(&backends);
2962        assert_eq!(sequence(&mut maglev, &backends, 8), expected, "MAGLEV");
2963    }
2964
2965    /// `rebuild`, which the backend list runs when a weight changes, starts
2966    /// the running weights afresh: the selection after a reweight is the one
2967    /// a fresh policy makes over the new weights.
2968    #[test]
2969    fn round_robin_rebuild_resets_running_weights() {
2970        use crate::{backends::BackendList, sozu_command::proto::command::LoadBalancingAlgorithms};
2971
2972        let address = |port| SocketAddr::from(([127, 0, 0, 1], port));
2973        let backend = |id: &str, port, weight| {
2974            Backend::new(
2975                id,
2976                address(port),
2977                None,
2978                Some(LoadBalancingParams { weight }),
2979                None,
2980            )
2981        };
2982        let mut list = BackendList::with_seed(1);
2983        list.set_load_balancing_policy(LoadBalancingAlgorithms::RoundRobin, None, 1);
2984        list.add_backend(backend("a", 9200, 3));
2985        list.add_backend(backend("b", 9201, 1));
2986        let now = Instant::now();
2987        let picked = |list: &mut BackendList| {
2988            list.next_available_backend(now)
2989                .expect("a healthy list selects")
2990                .borrow()
2991                .address
2992        };
2993        // Running weights after this pick: a = -1, b = 1.
2994        assert_eq!(picked(&mut list), address(9200));
2995        // Reweighted to (3, 2): kept running weights would make b win
2996        // (a = 2, b = 3); reset ones make a win (a = 3, b = 2).
2997        list.add_backend(backend("b", 9201, 2));
2998        assert_eq!(picked(&mut list), address(9200));
2999    }
3000
3001    /// `LEAST_LOADED` divides each load by its weight: `10` connections on
3002    /// weight `100` is lighter than `4` on weight `10`.
3003    ///
3004    /// TO SEE THIS RED: compare the raw loads in `LeastLoaded`; the second
3005    /// backend then wins.
3006    #[test]
3007    fn least_loaded_compares_load_per_unit_of_weight() {
3008        for metric in [LoadMetric::Connections, LoadMetric::Requests] {
3009            let backends = weighted_backends(&[Some(100), Some(10)]);
3010            for (backend, load) in backends.iter().zip([10, 4]) {
3011                let mut backend = backend.borrow_mut();
3012                backend.active_connections = load;
3013                backend.active_requests = load;
3014            }
3015            let mut policy = LeastLoaded { metric };
3016            assert_eq!(sequence(&mut policy, &backends, 1), [0], "{metric:?}");
3017
3018            // Equal per-weight loads (10/100 against 1/10): the first wins.
3019            backends[1].borrow_mut().active_connections = 1;
3020            backends[1].borrow_mut().active_requests = 1;
3021            assert_eq!(sequence(&mut policy, &backends, 1), [0], "{metric:?} tie");
3022        }
3023    }
3024
3025    /// Equal weights keep `LEAST_LOADED`'s raw-load order and first-minimum
3026    /// tie-break.
3027    #[test]
3028    fn least_loaded_equal_weights_compare_raw_loads() {
3029        let backends = weighted_backends(&[Some(7), Some(7), Some(7)]);
3030        for (backend, load) in backends.iter().zip([5, 3, 3]) {
3031            backend.borrow_mut().active_connections = load;
3032        }
3033        let mut policy = LeastLoaded {
3034            metric: LoadMetric::Connections,
3035        };
3036        assert_eq!(sequence(&mut policy, &backends, 1), [1]);
3037    }
3038
3039    /// `POWER_OF_TWO` divides each sample's load by its weight. With two
3040    /// backends both are sampled on every selection, so the lighter per unit
3041    /// of weight wins every time.
3042    ///
3043    /// TO SEE THIS RED: compare the raw loads in `PowerOfTwo`; the second
3044    /// backend then wins every time.
3045    #[test]
3046    fn power_of_two_compares_load_per_unit_of_weight() {
3047        for metric in [LoadMetric::Connections, LoadMetric::Requests] {
3048            let backends = weighted_backends(&[Some(100), Some(10)]);
3049            for (backend, load) in backends.iter().zip([10, 4]) {
3050                let mut backend = backend.borrow_mut();
3051                backend.active_connections = load;
3052                backend.active_requests = load;
3053            }
3054            let mut policy = PowerOfTwo::with_seed(DEFAULT_HASH_SEED, metric);
3055            assert_eq!(
3056                sequence(&mut policy, &backends, 64),
3057                vec![0; 64],
3058                "{metric:?}"
3059            );
3060        }
3061    }
3062
3063    /// `POWER_OF_TWO` orders count loads as integers: two loads that differ by
3064    /// one above 2^53 (where `f64` cannot tell them apart) still select the
3065    /// lighter, every time, instead of falling to the coin flip.
3066    ///
3067    /// TO SEE THIS RED: convert the counts to `f64` in `PowerOfTwo::compare`;
3068    /// both loads then round to the same value and the tie is a coin flip.
3069    #[test]
3070    fn power_of_two_orders_large_counts_exactly() {
3071        let big = (1usize << 53) + 1;
3072        for metric in [LoadMetric::Connections, LoadMetric::Requests] {
3073            let backends = weighted_backends(&[None, None]);
3074            for (backend, load) in backends.iter().zip([big, big - 1]) {
3075                let mut backend = backend.borrow_mut();
3076                backend.active_connections = load;
3077                backend.active_requests = load;
3078            }
3079            let mut policy = PowerOfTwo::with_seed(DEFAULT_HASH_SEED, metric);
3080            assert_eq!(
3081                sequence(&mut policy, &backends, 64),
3082                vec![1; 64],
3083                "{metric:?}"
3084            );
3085        }
3086    }
3087
3088    /// The per-weight comparison is exact and, with equal weights, is the
3089    /// raw comparison.
3090    #[test]
3091    fn per_weight_comparisons_are_exact() {
3092        assert_eq!(cmp_count_per_weight(10, 100, 4, 10), Ordering::Less);
3093        assert_eq!(cmp_count_per_weight(10, 100, 1, 10), Ordering::Equal);
3094        assert_eq!(
3095            cmp_count_per_weight(usize::MAX, 1, usize::MAX, 2),
3096            Ordering::Greater
3097        );
3098        assert_eq!(cmp_count_per_weight(3, 7, 4, 7), Ordering::Less);
3099        assert_eq!(
3100            cmp_cost_per_weight(2.0, 200, 1.5, 100),
3101            Some(Ordering::Less)
3102        );
3103        assert_eq!(cmp_cost_per_weight(1.0, 5, 1.0, 5), Some(Ordering::Equal));
3104        assert_eq!(cmp_cost_per_weight(f64::NAN, 5, 1.0, 5), None);
3105    }
3106}