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}