Skip to main content

sozu_lib/
load_balancing.rs

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