Skip to main content

reddb_server/cluster/
placement.rs

1//! Weighted placement and the multi-signal rebalancer planner (issue #1003,
2//! PRD #987, ADR 0037).
3//!
4//! Where the [`supervisor`](super::supervisor) reacts to a *failed* owner, this
5//! module is the proactive counterpart: it decides where ranges *should* live so
6//! the cluster's storage and traffic stay balanced as members come, go, and grow
7//! their disks. It is the glossary's **weighted placement** policy
8//! (`clustering.md`) — *"shard/range placement policy that accounts for advertised
9//! node capacity such as usable disk … and operator weights. Expanding a node's
10//! disk changes its placement weight; data moves only through explicit rebalancing
11//! transitions"* — driven by the **multi-signal rebalancer** — *"Cluster
12//! Supervisor policy that plans ownership transitions using bytes-used versus
13//! weighted capacity as the primary safety signal and read/write load as a
14//! secondary hotspot signal."*
15//!
16//! ## Two signals, one a safety floor and one a hint
17//!
18//! * **Primary — bytes-used vs weighted capacity.** Every member advertises its
19//!   usable disk and an operator weight ([`MemberCapacity`]); the product is its
20//!   [`weighted_capacity`](MemberCapacity::weighted_capacity), the member's share
21//!   of the cluster it is *meant* to hold. The planner compares each member's
22//!   bytes-used against its **fair share** (cluster bytes apportioned by weighted
23//!   capacity) and proposes moving ranges off members that are over their share
24//!   onto members that are under it. This is the safety signal: a member running
25//!   out of disk is an availability risk, so capacity balance is what the planner
26//!   acts on.
27//! * **Secondary — read/write load.** A range can be perfectly placed by bytes
28//!   yet still be a **hotspot**: it absorbs a disproportionate share of the
29//!   cluster's read/write traffic. The planner surfaces hotspots
30//!   ([`HotspotRange`]) and, when capacity allows, proposes spreading them off
31//!   their over-loaded owner. This is a hint layered on top of the capacity
32//!   floor, never in place of it — a hotspot move is only taken when it does not
33//!   itself create a capacity problem.
34//!
35//! ## Planning, not moving
36//!
37//! [`WeightedPlacementPlanner::plan_rebalance`] reads the membership catalog, the
38//! ownership catalog, and the live signals, and returns a [`RebalancePlan`] of
39//! [`PlannedMove`]s. It takes the ownership catalog by shared reference and
40//! **never mutates it** — *nothing moves implicitly*. Each [`PlannedMove`] is the
41//! intent for one rebalancing transition; executing it (copy the range to the
42//! target, let it catch up to the range commit watermark, then cut over through
43//! the fenced [`Handoff`](super::ownership_transition::TransitionKind::Handoff)
44//! transition machine) is a separate, explicit step. This is why *expanding a
45//! member's disk changes its placement weight but moves no data*: the new weight
46//! changes what the *next* plan proposes, and data only relocates when that plan
47//! is run.
48//!
49//! ## Purity
50//!
51//! All live state — per-member advertised capacity and per-range bytes/traffic —
52//! is read through the [`PlacementSignals`] trait, injected by the caller.
53//! Production backs it onto the disk-usage reporter and the per-range traffic
54//! counters; tests back it onto a scripted fake. The planner itself is a pure
55//! policy over the two catalogs plus those signals, so the whole weighting,
56//! balancing, and hotspot story is exercised deterministically — no disk, no
57//! clock, no network.
58
59use std::collections::{BTreeMap, BTreeSet};
60
61use super::identity::NodeIdentity;
62use super::membership::MembershipCatalog;
63use super::ownership::{CollectionId, PlacementMetadata, RangeId, ShardOwnershipCatalog};
64
65/// The neutral operator weight: a member with this weight is placed strictly by
66/// its usable disk. The weight is expressed in hundredths, so `100` means a 1.0×
67/// multiplier; `200` doubles a member's placement weight and `50` halves it. An
68/// operator nudges placement without lying about disk by tuning this.
69pub const NEUTRAL_OPERATOR_WEIGHT: u32 = 100;
70
71/// Authority-sharding unit for placement. Related small collections may share a
72/// group; a large collection can be isolated in its own group.
73#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
74pub struct CollectionGroupId(String);
75
76impl CollectionGroupId {
77    pub fn new(value: impl Into<String>) -> Result<Self, PlacementAuthorityError> {
78        let value = value.into();
79        if value.trim().is_empty() {
80            return Err(PlacementAuthorityError::EmptyCollectionGroup);
81        }
82        Ok(Self(value))
83    }
84}
85
86impl std::fmt::Display for CollectionGroupId {
87    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
88        f.write_str(&self.0)
89    }
90}
91
92/// The Placement Authority responsible for one Collection group and the
93/// collections whose ownership-catalog slice belongs to it.
94#[derive(Debug, Clone, PartialEq, Eq)]
95pub struct CollectionGroupPlacementAuthority {
96    collection_group: CollectionGroupId,
97    authority: NodeIdentity,
98    collections: BTreeSet<CollectionId>,
99}
100
101impl CollectionGroupPlacementAuthority {
102    pub fn new(
103        collection_group: CollectionGroupId,
104        authority: NodeIdentity,
105        collections: impl IntoIterator<Item = CollectionId>,
106    ) -> Result<Self, PlacementAuthorityError> {
107        let collections: BTreeSet<_> = collections.into_iter().collect();
108        if collections.is_empty() {
109            return Err(PlacementAuthorityError::EmptyCollectionSet {
110                collection_group,
111                authority,
112            });
113        }
114        Ok(Self {
115            collection_group,
116            authority,
117            collections,
118        })
119    }
120
121    pub fn collection_group(&self) -> &CollectionGroupId {
122        &self.collection_group
123    }
124
125    pub fn authority(&self) -> &NodeIdentity {
126        &self.authority
127    }
128
129    pub fn covers(&self, collection: &CollectionId) -> bool {
130        self.collections.contains(collection)
131    }
132}
133
134/// Pure in-memory authority index used by placement planning.
135#[derive(Debug, Clone, Default, PartialEq, Eq)]
136pub struct PlacementAuthorityCatalog {
137    by_collection: BTreeMap<CollectionId, CollectionGroupPlacementAuthority>,
138}
139
140impl PlacementAuthorityCatalog {
141    pub fn new() -> Self {
142        Self::default()
143    }
144
145    pub fn register(
146        &mut self,
147        authority: CollectionGroupPlacementAuthority,
148    ) -> Result<(), PlacementAuthorityError> {
149        for collection in &authority.collections {
150            if let Some(existing) = self.by_collection.get(collection) {
151                return Err(PlacementAuthorityError::OverlappingCollection {
152                    collection: collection.clone(),
153                    existing_group: existing.collection_group.clone(),
154                    new_group: authority.collection_group.clone(),
155                });
156            }
157        }
158        for collection in &authority.collections {
159            self.by_collection
160                .insert(collection.clone(), authority.clone());
161        }
162        Ok(())
163    }
164
165    pub fn authority_for(
166        &self,
167        collection: &CollectionId,
168    ) -> Option<&CollectionGroupPlacementAuthority> {
169        self.by_collection.get(collection)
170    }
171}
172
173#[derive(Debug, Clone, PartialEq, Eq)]
174pub enum PlacementAuthorityError {
175    EmptyCollectionGroup,
176    EmptyCollectionSet {
177        collection_group: CollectionGroupId,
178        authority: NodeIdentity,
179    },
180    OverlappingCollection {
181        collection: CollectionId,
182        existing_group: CollectionGroupId,
183        new_group: CollectionGroupId,
184    },
185    MissingCollectionAuthority {
186        collection: CollectionId,
187    },
188}
189
190impl std::fmt::Display for PlacementAuthorityError {
191    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
192        match self {
193            Self::EmptyCollectionGroup => write!(f, "collection group id must not be empty"),
194            Self::EmptyCollectionSet {
195                collection_group,
196                authority,
197            } => write!(
198                f,
199                "placement authority {authority} for collection group {collection_group} has no collections"
200            ),
201            Self::OverlappingCollection {
202                collection,
203                existing_group,
204                new_group,
205            } => write!(
206                f,
207                "collection {collection} is already assigned to collection group {existing_group}, cannot also assign it to {new_group}"
208            ),
209            Self::MissingCollectionAuthority { collection } => {
210                write!(f, "no placement authority for collection {collection}")
211            }
212        }
213    }
214}
215
216impl std::error::Error for PlacementAuthorityError {}
217
218/// A member's advertised placement capacity: how much usable disk it offers and
219/// the operator's weight multiplier on top of it.
220///
221/// The two combine into the member's [`weighted_capacity`](Self::weighted_capacity)
222/// — its share of the cluster it is meant to hold. Advertising more usable disk,
223/// or a higher operator weight, raises that share; the planner then apportions
224/// ranges toward it on the *next* plan. The struct is pure advertised state: it
225/// records what a member *offers*, never moves anything by itself.
226#[derive(Debug, Clone, Copy, PartialEq, Eq)]
227pub struct MemberCapacity {
228    /// Usable disk the member advertises for user ranges, in bytes.
229    pub usable_disk_bytes: u64,
230    /// Operator weight in hundredths ([`NEUTRAL_OPERATOR_WEIGHT`] = 1.0×). Lets an
231    /// operator bias placement on or off a member without misreporting disk.
232    pub operator_weight: u32,
233}
234
235impl MemberCapacity {
236    /// Capacity with an explicit usable disk and operator weight.
237    pub fn new(usable_disk_bytes: u64, operator_weight: u32) -> Self {
238        Self {
239            usable_disk_bytes,
240            operator_weight,
241        }
242    }
243
244    /// Capacity from usable disk alone, at the neutral operator weight — the
245    /// common case where the operator has expressed no preference.
246    pub fn with_disk(usable_disk_bytes: u64) -> Self {
247        Self::new(usable_disk_bytes, NEUTRAL_OPERATOR_WEIGHT)
248    }
249
250    /// The member's **placement weight**: usable disk scaled by the operator
251    /// weight. This is the value the rebalancer apportions the cluster's bytes by,
252    /// and it is exactly what *expanding a member's disk changes* — a larger disk
253    /// (or a higher operator weight) yields a larger weighted capacity and so a
254    /// larger fair share on the next plan. Computed in `u128` so a large disk
255    /// times a large weight cannot overflow.
256    pub fn weighted_capacity(&self) -> u128 {
257        self.usable_disk_bytes as u128 * self.operator_weight as u128
258            / NEUTRAL_OPERATOR_WEIGHT as u128
259    }
260
261    /// Whether this member can hold any ranges at all — a member advertising no
262    /// usable disk (or a zero operator weight) has zero weighted capacity and is
263    /// never a placement target.
264    pub fn is_placeable(&self) -> bool {
265        self.weighted_capacity() > 0
266    }
267}
268
269/// The live load on one range: its on-disk size and its recent read/write
270/// traffic.
271///
272/// `bytes_used` feeds the **primary** capacity signal (it is what a member's
273/// bytes-used is summed from); `read_ops`/`write_ops` feed the **secondary**
274/// hotspot signal. Keeping both on one struct lets a single
275/// [`PlacementSignals::range_load`] call answer everything the planner needs about
276/// a range.
277#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
278pub struct RangeLoad {
279    /// The range's on-disk size in bytes — the primary capacity signal.
280    pub bytes_used: u64,
281    /// Read operations served in the recent observation window.
282    pub read_ops: u64,
283    /// Write operations served in the recent observation window.
284    pub write_ops: u64,
285}
286
287impl RangeLoad {
288    /// A range that occupies `bytes_used` but serves no traffic — handy when only
289    /// the capacity signal matters.
290    pub fn idle(bytes_used: u64) -> Self {
291        Self {
292            bytes_used,
293            read_ops: 0,
294            write_ops: 0,
295        }
296    }
297
298    /// Total read + write traffic — the hotspot signal. A range with high traffic
299    /// relative to the cluster mean is a hotspot candidate regardless of its size.
300    pub fn traffic(&self) -> u64 {
301        self.read_ops.saturating_add(self.write_ops)
302    }
303}
304
305/// The live cluster state the planner reads but does not own: each member's
306/// advertised capacity and each range's bytes/traffic.
307///
308/// Production backs this onto the disk-usage reporter and the per-range traffic
309/// counters; tests back it onto a scripted fake. Keeping it behind a trait is what
310/// makes the planner a pure policy.
311pub trait PlacementSignals {
312    /// The capacity `member` currently advertises. A member that advertises
313    /// nothing (or is unknown) should report a zero-disk [`MemberCapacity`], which
314    /// makes it un-placeable rather than a div-by-zero hazard.
315    fn member_capacity(&self, member: &NodeIdentity) -> MemberCapacity;
316
317    /// Placement attribute advertised by `member`, such as a `zone` or `region`.
318    /// Implementations that do not track topology attributes can leave this
319    /// absent; spread rules then degrade through operator warnings instead of
320    /// inventing topology.
321    fn member_attribute(&self, _member: &NodeIdentity, _key: &str) -> Option<&str> {
322        None
323    }
324
325    /// The current load on `(collection, range_id)` — its bytes and its recent
326    /// read/write traffic.
327    fn range_load(&self, collection: &CollectionId, range_id: RangeId) -> RangeLoad;
328}
329
330/// Why the planner proposed moving a range.
331#[derive(Debug, Clone, Copy, PartialEq, Eq)]
332pub enum MoveReason {
333    /// The **primary** signal: the source member is over its weighted-capacity
334    /// fair share and the target is under its own. Moving the range relieves a
335    /// disk-pressure (availability) risk.
336    CapacityBalance,
337    /// The **secondary** signal: the range is a read/write hotspot on an
338    /// over-loaded owner, and a target with both load and capacity headroom can
339    /// absorb it. Taken only when it does not create a capacity problem.
340    HotspotRelief,
341}
342
343/// One proposed rebalancing transition: move authority for a range from its
344/// current owner to a target member.
345///
346/// A [`PlannedMove`] is *intent*, not an executed transition. Carrying it out
347/// means copying the range to `to`, letting it catch up to the range commit
348/// watermark, and then cutting over through the fenced
349/// [`Handoff`](super::ownership_transition::TransitionKind::Handoff) machine — a
350/// separate, explicit step. The planner only ever produces these; it moves no
351/// data.
352#[derive(Debug, Clone, PartialEq, Eq)]
353pub struct PlannedMove {
354    pub collection: CollectionId,
355    pub range_id: RangeId,
356    /// The range's current owner in the catalog — the move's source.
357    pub from: NodeIdentity,
358    /// The proposed new owner — an active data member with capacity headroom.
359    pub to: NodeIdentity,
360    /// The range's size in bytes at planning time (what the move relocates).
361    pub bytes: u64,
362    pub reason: MoveReason,
363}
364
365/// One proposed replica-list transition to restore failure-domain spread without
366/// moving write authority.
367#[derive(Debug, Clone, PartialEq, Eq)]
368pub struct PlannedReplicaReplacement {
369    pub collection: CollectionId,
370    pub range_id: RangeId,
371    /// The current replica whose domain duplicates another copy of the range.
372    pub replace: NodeIdentity,
373    /// The eligible member in a distinct failure domain.
374    pub with: NodeIdentity,
375    pub domain_key: String,
376    pub domain_value: String,
377}
378
379/// A range the **secondary** signal flagged as a read/write hotspot: it serves
380/// traffic well above the cluster mean. Surfaced whether or not a relief move was
381/// possible, so an operator can see a hotspot even when there is no headroom to
382/// relieve it.
383#[derive(Debug, Clone, PartialEq, Eq)]
384pub struct HotspotRange {
385    pub collection: CollectionId,
386    pub range_id: RangeId,
387    /// The range's current owner — the member bearing the hot traffic.
388    pub owner: NodeIdentity,
389    /// The range's read + write traffic in the observation window.
390    pub traffic: u64,
391}
392
393/// The planner's decision for one pass: the moves to schedule and the hotspots it
394/// observed.
395///
396/// A cluster already balanced by capacity, with no hotspot, yields an empty plan
397/// ([`is_empty`](Self::is_empty)) — the no-op a stable cluster must produce.
398#[derive(Debug, Clone, Default, PartialEq, Eq)]
399pub struct RebalancePlan {
400    /// Proposed moves, in a deterministic order (capacity moves first, then
401    /// hotspot-relief moves, each in `(collection, range_id)` order).
402    pub moves: Vec<PlannedMove>,
403    /// Proposed replica-list changes that restore failure-domain spread without
404    /// changing the range owner.
405    pub replica_replacements: Vec<PlannedReplicaReplacement>,
406    /// Ranges observed to be hotspots this pass, hottest first.
407    pub hotspots: Vec<HotspotRange>,
408    /// Loud operator warnings for placement rules the topology cannot satisfy.
409    pub warnings: Vec<PlacementWarning>,
410}
411
412#[derive(Debug, Clone, PartialEq, Eq)]
413pub enum PlacementWarning {
414    UnsatisfiedFailureDomainSpread {
415        collection: CollectionId,
416        range_id: RangeId,
417        domain_key: String,
418        domain_value: String,
419    },
420}
421
422impl std::fmt::Display for PlacementWarning {
423    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
424        match self {
425            Self::UnsatisfiedFailureDomainSpread {
426                collection,
427                range_id,
428                domain_key,
429                domain_value,
430            } => write!(
431                f,
432                "WARNING: range {collection}/{range_id} cannot satisfy failure-domain spread for {domain_key}={domain_value}"
433            ),
434        }
435    }
436}
437
438#[derive(Debug, Clone, PartialEq, Eq)]
439pub struct AuthorityScopedPlannedMove {
440    pub movement: PlannedMove,
441    pub placement_authority: CollectionGroupPlacementAuthority,
442}
443
444#[derive(Debug, Clone, Default, PartialEq, Eq)]
445pub struct AuthorityScopedRebalancePlan {
446    pub moves: Vec<AuthorityScopedPlannedMove>,
447    pub replica_replacements: Vec<PlannedReplicaReplacement>,
448    pub hotspots: Vec<HotspotRange>,
449    pub warnings: Vec<PlacementWarning>,
450}
451
452impl RebalancePlan {
453    /// Nothing to schedule *and* nothing hot — a fully balanced, evenly-loaded
454    /// cluster. Distinct from [`no_moves`](Self::no_moves): a balanced cluster can
455    /// still have an *observed* hotspot it cannot relieve.
456    pub fn is_empty(&self) -> bool {
457        self.moves.is_empty()
458            && self.replica_replacements.is_empty()
459            && self.hotspots.is_empty()
460            && self.warnings.is_empty()
461    }
462
463    /// Whether the plan proposes any actual range movement. False on a cluster
464    /// that is balanced by capacity and has no relievable hotspot, even if a
465    /// hotspot was *observed*.
466    pub fn no_moves(&self) -> bool {
467        self.moves.is_empty() && self.replica_replacements.is_empty()
468    }
469
470    /// The capacity-balance moves only (the primary signal).
471    pub fn capacity_moves(&self) -> impl Iterator<Item = &PlannedMove> {
472        self.moves
473            .iter()
474            .filter(|m| m.reason == MoveReason::CapacityBalance)
475    }
476
477    /// The hotspot-relief moves only (the secondary signal).
478    pub fn hotspot_moves(&self) -> impl Iterator<Item = &PlannedMove> {
479        self.moves
480            .iter()
481            .filter(|m| m.reason == MoveReason::HotspotRelief)
482    }
483}
484
485/// The tunables that gate when imbalance and traffic are worth a move.
486///
487/// The defaults are deliberately slack: a cluster within 10% of its fair share is
488/// "balanced enough" not to churn ownership, and a hotspot must run at 2× the
489/// cluster-mean traffic before it is worth spreading. Tight thresholds would make
490/// the planner thrash on noise.
491#[derive(Debug, Clone, Copy, PartialEq)]
492pub struct PlacementPolicy {
493    /// Fractional tolerance around a member's fair share. A member is "over" only
494    /// when its bytes-used exceeds `fair * (1 + balance_tolerance)`; a move's
495    /// target must have room under `fair * (1 + balance_tolerance)`. Larger values
496    /// tolerate more imbalance for less churn.
497    pub balance_tolerance: f64,
498    /// How many times the cluster-mean range traffic a range must serve to count
499    /// as a hotspot. `2.0` means "twice the average".
500    pub hotspot_load_factor: f64,
501}
502
503impl Default for PlacementPolicy {
504    fn default() -> Self {
505        Self {
506            balance_tolerance: 0.10,
507            hotspot_load_factor: 2.0,
508        }
509    }
510}
511
512/// The fair share of `total_bytes` that a member with `member_capacity` deserves,
513/// out of the cluster's `total_capacity`. Apportions bytes strictly by weighted
514/// capacity; `u128` math keeps a large cluster from overflowing.
515fn fair_share(total_bytes: u64, member_capacity: u128, total_capacity: u128) -> u64 {
516    if total_capacity == 0 {
517        return 0;
518    }
519    let share = total_bytes as u128 * member_capacity / total_capacity;
520    share.min(u64::MAX as u128) as u64
521}
522
523/// The weighted-placement, multi-signal rebalancer planner.
524///
525/// Holds only the [`PlacementPolicy`]; all live state is read through
526/// [`PlacementSignals`] at plan time, so one planner instance serves the whole
527/// cluster lifetime.
528#[derive(Debug, Clone, Default)]
529pub struct WeightedPlacementPlanner {
530    policy: PlacementPolicy,
531}
532
533impl WeightedPlacementPlanner {
534    /// A planner with the given policy.
535    pub fn new(policy: PlacementPolicy) -> Self {
536        Self { policy }
537    }
538
539    pub fn policy(&self) -> &PlacementPolicy {
540        &self.policy
541    }
542
543    /// Plan a rebalance across the whole ownership catalog **without** mutating
544    /// it. Runs the primary capacity-balance pass, then the secondary
545    /// hotspot-relief pass on top, and returns the combined [`RebalancePlan`].
546    /// Ranges owned by members that are not placement-eligible (draining members,
547    /// witnesses) are left to the drain flow and never moved here.
548    pub fn plan_rebalance(
549        &self,
550        membership: &MembershipCatalog,
551        ownership: &ShardOwnershipCatalog,
552        signals: &impl PlacementSignals,
553    ) -> RebalancePlan {
554        let mut state = ClusterState::observe(membership, ownership, signals, &self.policy);
555        let (replica_replacements, warnings) =
556            state.plan_failure_domain_spread(membership, ownership, signals);
557        let mut moves = state.plan_capacity_moves(&self.policy);
558        let (hotspots, hotspot_moves) = state.plan_hotspot_moves(&self.policy);
559        moves.extend(hotspot_moves);
560        RebalancePlan {
561            moves,
562            replica_replacements,
563            hotspots,
564            warnings,
565        }
566    }
567
568    pub fn plan_rebalance_scoped(
569        &self,
570        membership: &MembershipCatalog,
571        ownership: &ShardOwnershipCatalog,
572        signals: &impl PlacementSignals,
573        authorities: &PlacementAuthorityCatalog,
574    ) -> Result<AuthorityScopedRebalancePlan, PlacementAuthorityError> {
575        let plan = self.plan_rebalance(membership, ownership, signals);
576        let mut moves = Vec::with_capacity(plan.moves.len());
577        for movement in plan.moves {
578            let placement_authority = authorities
579                .authority_for(&movement.collection)
580                .ok_or_else(|| PlacementAuthorityError::MissingCollectionAuthority {
581                    collection: movement.collection.clone(),
582                })?
583                .clone();
584            moves.push(AuthorityScopedPlannedMove {
585                movement,
586                placement_authority,
587            });
588        }
589        Ok(AuthorityScopedRebalancePlan {
590            moves,
591            replica_replacements: plan.replica_replacements,
592            hotspots: plan.hotspots,
593            warnings: plan.warnings,
594        })
595    }
596}
597
598/// The mutable simulation the planner balances against. Built once from the live
599/// catalogs and signals, then evolved as moves are chosen so each successive move
600/// sees the effect of the ones before it. Crucially this is a *copy* of the live
601/// state — evolving it changes nothing in the real catalog, which is what makes
602/// planning side-effect free.
603struct ClusterState {
604    /// Placement-eligible members (active data members) with non-zero weighted
605    /// capacity, in stable identity order.
606    eligible: Vec<NodeIdentity>,
607    weighted_capacity: BTreeMap<NodeIdentity, u128>,
608    /// Total weighted capacity across `eligible` — the denominator of fair share.
609    total_capacity: u128,
610    /// Total bytes across all movable ranges — the numerator of fair share.
611    total_bytes: u64,
612    /// Per-range size and traffic, keyed in `(collection, range_id)` order.
613    ranges: BTreeMap<(CollectionId, RangeId), RangeFacts>,
614    /// Simulated current owner of each movable range (evolves as moves are taken).
615    owner_of: BTreeMap<(CollectionId, RangeId), NodeIdentity>,
616    /// The range's true catalog owner — the `from` every move records, even if the
617    /// simulation has since reassigned it (a range moves at most once per plan).
618    origin_owner: BTreeMap<(CollectionId, RangeId), NodeIdentity>,
619    /// Simulated bytes-used per member.
620    used: BTreeMap<NodeIdentity, u64>,
621    /// Simulated read/write load per member.
622    load: BTreeMap<NodeIdentity, u64>,
623    /// Ranges already scheduled to move — never moved twice in one plan.
624    moved: std::collections::BTreeSet<(CollectionId, RangeId)>,
625}
626
627#[derive(Clone, Copy)]
628struct RangeFacts {
629    bytes: u64,
630    traffic: u64,
631}
632
633impl ClusterState {
634    fn observe(
635        membership: &MembershipCatalog,
636        ownership: &ShardOwnershipCatalog,
637        signals: &impl PlacementSignals,
638        _policy: &PlacementPolicy,
639    ) -> Self {
640        let mut weighted_capacity = BTreeMap::new();
641        let mut eligible = Vec::new();
642        let mut total_capacity: u128 = 0;
643        for member in membership.placement_eligible_members() {
644            let id = member.identity().clone();
645            let cap = signals.member_capacity(&id).weighted_capacity();
646            if cap == 0 {
647                // A placement-eligible member advertising no usable disk is not a
648                // valid target; exclude it so it is never apportioned bytes.
649                continue;
650            }
651            total_capacity += cap;
652            weighted_capacity.insert(id.clone(), cap);
653            eligible.push(id);
654        }
655
656        let eligible_set: std::collections::BTreeSet<&NodeIdentity> = eligible.iter().collect();
657
658        let mut ranges = BTreeMap::new();
659        let mut owner_of = BTreeMap::new();
660        let mut origin_owner = BTreeMap::new();
661        let mut used: BTreeMap<NodeIdentity, u64> =
662            eligible.iter().map(|id| (id.clone(), 0)).collect();
663        let mut load: BTreeMap<NodeIdentity, u64> =
664            eligible.iter().map(|id| (id.clone(), 0)).collect();
665        let mut total_bytes: u64 = 0;
666
667        for entry in ownership.entries() {
668            let owner = entry.owner().clone();
669            // Only ranges owned by an eligible member are movable here; a draining
670            // owner's ranges belong to the drain flow.
671            if !eligible_set.contains(&owner) {
672                continue;
673            }
674            let key = (entry.collection().clone(), entry.range_id());
675            let load_facts = signals.range_load(entry.collection(), entry.range_id());
676            ranges.insert(
677                key.clone(),
678                RangeFacts {
679                    bytes: load_facts.bytes_used,
680                    traffic: load_facts.traffic(),
681                },
682            );
683            *used.get_mut(&owner).unwrap() += load_facts.bytes_used;
684            *load.get_mut(&owner).unwrap() += load_facts.traffic();
685            total_bytes = total_bytes.saturating_add(load_facts.bytes_used);
686            owner_of.insert(key.clone(), owner.clone());
687            origin_owner.insert(key, owner);
688        }
689
690        Self {
691            eligible,
692            weighted_capacity,
693            total_capacity,
694            total_bytes,
695            ranges,
696            owner_of,
697            origin_owner,
698            used,
699            load,
700            moved: std::collections::BTreeSet::new(),
701        }
702    }
703
704    fn plan_failure_domain_spread(
705        &self,
706        membership: &MembershipCatalog,
707        ownership: &ShardOwnershipCatalog,
708        signals: &impl PlacementSignals,
709    ) -> (Vec<PlannedReplicaReplacement>, Vec<PlacementWarning>) {
710        let mut replacements = Vec::new();
711        let mut warnings = Vec::new();
712
713        for range in ownership.entries() {
714            let Some(domain_key) = failure_domain_key(range.placement()) else {
715                continue;
716            };
717            let mut used_domains: BTreeMap<String, NodeIdentity> = BTreeMap::new();
718            let owner_domain = signals.member_attribute(range.owner(), domain_key);
719            if let Some(domain) = owner_domain {
720                used_domains.insert(domain.to_string(), range.owner().clone());
721            }
722            let mut selected_replacements = BTreeSet::new();
723
724            for replica in range.replicas() {
725                let Some(domain) = signals.member_attribute(replica, domain_key) else {
726                    continue;
727                };
728                if used_domains.contains_key(domain) {
729                    let mut retained_domains =
730                        copy_domains_except(range, replica, domain_key, signals);
731                    retained_domains.extend(used_domains.keys().cloned());
732                    if let Some(candidate) = select_spread_replacement(
733                        range,
734                        &retained_domains,
735                        &selected_replacements,
736                        domain_key,
737                        membership,
738                        signals,
739                    ) {
740                        if let Some(candidate_domain) =
741                            signals.member_attribute(&candidate, domain_key)
742                        {
743                            used_domains.insert(candidate_domain.to_string(), candidate.clone());
744                        }
745                        selected_replacements.insert(candidate.clone());
746                        replacements.push(PlannedReplicaReplacement {
747                            collection: range.collection().clone(),
748                            range_id: range.range_id(),
749                            replace: replica.clone(),
750                            with: candidate,
751                            domain_key: domain_key.to_string(),
752                            domain_value: domain.to_string(),
753                        });
754                    } else {
755                        warnings.push(PlacementWarning::UnsatisfiedFailureDomainSpread {
756                            collection: range.collection().clone(),
757                            range_id: range.range_id(),
758                            domain_key: domain_key.to_string(),
759                            domain_value: domain.to_string(),
760                        });
761                    }
762                    continue;
763                }
764                used_domains.insert(domain.to_string(), replica.clone());
765            }
766        }
767
768        (replacements, warnings)
769    }
770
771    fn fair(&self, member: &NodeIdentity) -> u64 {
772        let cap = self.weighted_capacity.get(member).copied().unwrap_or(0);
773        fair_share(self.total_bytes, cap, self.total_capacity)
774    }
775
776    /// Ranges currently owned by `member` in the simulation, in `(collection,
777    /// range_id)` order, that have not already been moved this plan.
778    fn ranges_owned_by(&self, member: &NodeIdentity) -> Vec<(CollectionId, RangeId)> {
779        self.owner_of
780            .iter()
781            .filter(|(key, owner)| *owner == member && !self.moved.contains(*key))
782            .map(|(key, _)| key.clone())
783            .collect()
784    }
785
786    fn apply_move(&mut self, key: &(CollectionId, RangeId), to: &NodeIdentity) {
787        let facts = self.ranges[key];
788        let from = self.owner_of[key].clone();
789        *self.used.get_mut(&from).unwrap() -= facts.bytes;
790        *self.load.get_mut(&from).unwrap() -= facts.traffic;
791        *self.used.get_mut(to).unwrap() += facts.bytes;
792        *self.load.get_mut(to).unwrap() += facts.traffic;
793        self.owner_of.insert(key.clone(), to.clone());
794        self.moved.insert(key.clone());
795    }
796
797    /// The **primary** pass: greedily move ranges off members over their
798    /// weighted-capacity fair share onto members under theirs, until no member is
799    /// over tolerance or no move strictly improves the worst imbalance.
800    fn plan_capacity_moves(&mut self, policy: &PlacementPolicy) -> Vec<PlannedMove> {
801        let mut planned = Vec::new();
802        if self.total_capacity == 0 || self.eligible.len() < 2 {
803            return planned;
804        }
805
806        // Each range moves at most once, so the loop is bounded by the range count.
807        // Pick the member most over its fair share (beyond tolerance) each round,
808        // then the member most under its own — the pair whose rebalance helps most.
809        while let Some(source) = self.most_over(policy) {
810            let Some(target) = self.most_under(&source) else {
811                break;
812            };
813
814            let dev_src = self.deviation(&source);
815            let dev_tgt = self.deviation(&target);
816            let worst_before = dev_src.abs().max(dev_tgt.abs());
817
818            // Among the source's still-movable ranges, choose the one that most
819            // reduces the worse of the two deviations after the move.
820            let mut best: Option<((CollectionId, RangeId), f64)> = None;
821            for key in self.ranges_owned_by(&source) {
822                let s = self.ranges[&key].bytes as f64;
823                let after = (dev_src - s).abs().max((dev_tgt + s).abs());
824                let better = match &best {
825                    None => true,
826                    Some((_, best_after)) => after < *best_after,
827                };
828                if better {
829                    best = Some((key, after));
830                }
831            }
832
833            let Some((key, worst_after)) = best else {
834                break;
835            };
836            // Only take the move if it strictly improves the worst imbalance —
837            // otherwise we would churn ownership for nothing.
838            if worst_after >= worst_before {
839                break;
840            }
841
842            let bytes = self.ranges[&key].bytes;
843            let from = self.origin_owner[&key].clone();
844            self.apply_move(&key, &target);
845            planned.push(PlannedMove {
846                collection: key.0,
847                range_id: key.1,
848                from,
849                to: target,
850                bytes,
851                reason: MoveReason::CapacityBalance,
852            });
853        }
854
855        planned
856    }
857
858    /// A member's deviation from its fair share in bytes: positive is over-full,
859    /// negative is under-full.
860    fn deviation(&self, member: &NodeIdentity) -> f64 {
861        self.used.get(member).copied().unwrap_or(0) as f64 - self.fair(member) as f64
862    }
863
864    /// The eligible member furthest over its fair share, beyond the tolerance
865    /// band, or `None` if everyone is within tolerance.
866    fn most_over(&self, policy: &PlacementPolicy) -> Option<NodeIdentity> {
867        self.eligible
868            .iter()
869            .filter(|id| {
870                let used = self.used.get(*id).copied().unwrap_or(0) as f64;
871                let fair = self.fair(id) as f64;
872                used > fair * (1.0 + policy.balance_tolerance) && used > fair
873            })
874            .max_by(|a, b| {
875                self.deviation(a)
876                    .partial_cmp(&self.deviation(b))
877                    .unwrap()
878                    // Tie-break by identity so the plan is deterministic.
879                    .then_with(|| b.cmp(a))
880            })
881            .cloned()
882    }
883
884    /// The eligible member furthest *under* its fair share (the best target),
885    /// excluding `source`.
886    fn most_under(&self, source: &NodeIdentity) -> Option<NodeIdentity> {
887        self.eligible
888            .iter()
889            .filter(|id| *id != source && self.deviation(id) < 0.0)
890            .min_by(|a, b| {
891                self.deviation(a)
892                    .partial_cmp(&self.deviation(b))
893                    .unwrap()
894                    // Tie-break by identity so the plan is deterministic.
895                    .then_with(|| a.cmp(b))
896            })
897            .cloned()
898    }
899
900    /// The **secondary** pass: identify hotspot ranges (traffic well above the
901    /// cluster mean) and, for each, propose spreading it to a member with both
902    /// load and capacity headroom — but only when that strictly lowers the owner's
903    /// load concentration and respects the capacity tolerance. Returns the
904    /// observed hotspots (hottest first) and any relief moves.
905    fn plan_hotspot_moves(
906        &mut self,
907        policy: &PlacementPolicy,
908    ) -> (Vec<HotspotRange>, Vec<PlannedMove>) {
909        let mut hotspots = Vec::new();
910        let mut moves = Vec::new();
911
912        let range_count = self.ranges.len();
913        if range_count == 0 {
914            return (hotspots, moves);
915        }
916        let total_traffic: u64 = self.ranges.values().map(|f| f.traffic).sum();
917        let mean = total_traffic as f64 / range_count as f64;
918        let threshold = mean * policy.hotspot_load_factor;
919        if mean <= 0.0 {
920            return (hotspots, moves);
921        }
922
923        // Collect hotspots, hottest first; tie-break by key for determinism.
924        let mut hot: Vec<((CollectionId, RangeId), u64)> = self
925            .ranges
926            .iter()
927            .filter(|(_, f)| f.traffic as f64 > threshold)
928            .map(|(key, f)| (key.clone(), f.traffic))
929            .collect();
930        hot.sort_by(|a, b| b.1.cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
931
932        for (key, traffic) in hot {
933            let owner = self.owner_of[&key].clone();
934            hotspots.push(HotspotRange {
935                collection: key.0.clone(),
936                range_id: key.1,
937                // Report the range's true catalog owner — the member actually
938                // bearing the hot traffic — independent of any simulated move.
939                owner: self.origin_owner[&key].clone(),
940                traffic,
941            });
942
943            // A hotspot already scheduled to move (by the capacity pass) needs no
944            // second move, and an owner holding only this one range cannot be
945            // relieved by moving it — that just relocates the hotspot.
946            if self.moved.contains(&key) || self.ranges_owned_by(&owner).len() < 2 {
947                continue;
948            }
949
950            let facts = self.ranges[&key];
951            let owner_load = self.load.get(&owner).copied().unwrap_or(0);
952
953            // Pick the eligible member with the least load that can take the range
954            // without breaching its capacity tolerance and ends up less loaded than
955            // the owner is now — otherwise the move does not spread load.
956            let target = self
957                .eligible
958                .iter()
959                .filter(|id| **id != owner)
960                .filter(|id| {
961                    let used = self.used.get(*id).copied().unwrap_or(0);
962                    let fair = self.fair(id) as f64;
963                    (used + facts.bytes) as f64 <= fair * (1.0 + policy.balance_tolerance)
964                })
965                .filter(|id| {
966                    let tgt_load = self.load.get(*id).copied().unwrap_or(0);
967                    tgt_load + facts.traffic < owner_load
968                })
969                .min_by(|a, b| {
970                    let la = self.load.get(*a).copied().unwrap_or(0);
971                    let lb = self.load.get(*b).copied().unwrap_or(0);
972                    la.cmp(&lb).then_with(|| a.cmp(b))
973                })
974                .cloned();
975
976            if let Some(target) = target {
977                let from = self.origin_owner[&key].clone();
978                self.apply_move(&key, &target);
979                moves.push(PlannedMove {
980                    collection: key.0,
981                    range_id: key.1,
982                    from,
983                    to: target,
984                    bytes: facts.bytes,
985                    reason: MoveReason::HotspotRelief,
986                });
987            }
988        }
989
990        (hotspots, moves)
991    }
992}
993
994fn failure_domain_key(placement: &PlacementMetadata) -> Option<&str> {
995    placement
996        .attribute("failure_domain")
997        .or_else(|| placement.attribute("failure-domain"))
998}
999
1000fn copy_domains_except(
1001    range: &super::ownership::RangeOwnership,
1002    excluded: &NodeIdentity,
1003    domain_key: &str,
1004    signals: &impl PlacementSignals,
1005) -> BTreeSet<String> {
1006    let mut domains = BTreeSet::new();
1007    if range.owner() != excluded {
1008        if let Some(domain) = signals.member_attribute(range.owner(), domain_key) {
1009            domains.insert(domain.to_string());
1010        }
1011    }
1012    for replica in range.replicas() {
1013        if replica == excluded {
1014            continue;
1015        }
1016        if let Some(domain) = signals.member_attribute(replica, domain_key) {
1017            domains.insert(domain.to_string());
1018        }
1019    }
1020    domains
1021}
1022
1023fn select_spread_replacement(
1024    range: &super::ownership::RangeOwnership,
1025    retained_domains: &BTreeSet<String>,
1026    selected_replacements: &BTreeSet<NodeIdentity>,
1027    domain_key: &str,
1028    membership: &MembershipCatalog,
1029    signals: &impl PlacementSignals,
1030) -> Option<NodeIdentity> {
1031    membership
1032        .placement_eligible_members()
1033        .map(super::membership::ClusterMember::identity)
1034        .filter(|id| {
1035            range.owner() != *id
1036                && !range.replicas().contains(id)
1037                && !selected_replacements.contains(id)
1038        })
1039        .find(|id| {
1040            signals
1041                .member_attribute(id, domain_key)
1042                .is_some_and(|domain| !retained_domains.contains(domain))
1043        })
1044        .cloned()
1045}
1046
1047#[cfg(test)]
1048mod tests {
1049    use super::*;
1050    use crate::cluster::membership::{ClusterId, ClusterMember, MemberKind};
1051    use crate::cluster::ownership::{PlacementMetadata, RangeBounds, RangeOwnership, ShardKeyMode};
1052    use std::collections::HashMap;
1053
1054    fn ident(cn: &str) -> NodeIdentity {
1055        NodeIdentity::from_certificate_subject(cn).unwrap()
1056    }
1057
1058    fn collection(name: &str) -> CollectionId {
1059        CollectionId::new(name).unwrap()
1060    }
1061
1062    fn data_member(cn: &str) -> ClusterMember {
1063        ClusterMember::joined_empty(ident(cn), MemberKind::Data)
1064    }
1065
1066    fn membership(members: &[&str]) -> MembershipCatalog {
1067        MembershipCatalog::new(
1068            ClusterId::new("cluster-x").unwrap(),
1069            members.iter().map(|m| data_member(m)),
1070        )
1071    }
1072
1073    /// A catalog of `n` single-owner ranges in `orders`, assigning range `i` to
1074    /// `owners[i]`. Each range is a distinct hash partition so they never overlap.
1075    fn catalog(owners: &[&str]) -> (ShardOwnershipCatalog, CollectionId) {
1076        let orders = collection("orders");
1077        let mut catalog = ShardOwnershipCatalog::new();
1078        for (i, owner) in owners.iter().enumerate() {
1079            let lower = vec![i as u8];
1080            let upper = vec![i as u8 + 1];
1081            let bounds = RangeBounds::new(
1082                crate::cluster::ownership::RangeBound::key(lower),
1083                crate::cluster::ownership::RangeBound::key(upper),
1084            )
1085            .unwrap();
1086            catalog
1087                .apply_update(RangeOwnership::establish(
1088                    orders.clone(),
1089                    RangeId::new(i as u64 + 1),
1090                    ShardKeyMode::Hash,
1091                    bounds,
1092                    ident(owner),
1093                    Vec::<NodeIdentity>::new(),
1094                    PlacementMetadata::with_replication_factor(1),
1095                ))
1096                .unwrap();
1097        }
1098        (catalog, orders)
1099    }
1100
1101    fn catalog_with_replicas(
1102        owner: &str,
1103        replicas: &[&str],
1104        placement: PlacementMetadata,
1105    ) -> (ShardOwnershipCatalog, CollectionId) {
1106        let orders = collection("orders");
1107        let mut catalog = ShardOwnershipCatalog::new();
1108        catalog
1109            .apply_update(RangeOwnership::establish(
1110                orders.clone(),
1111                RangeId::new(1),
1112                ShardKeyMode::Hash,
1113                RangeBounds::full(),
1114                ident(owner),
1115                replicas.iter().map(|replica| ident(replica)),
1116                placement,
1117            ))
1118            .unwrap();
1119        (catalog, orders)
1120    }
1121
1122    /// A scripted [`PlacementSignals`]: per-member capacity (defaulting to a
1123    /// uniform disk) and per-range load keyed by range id.
1124    struct FakeSignals {
1125        default_capacity: MemberCapacity,
1126        capacity: HashMap<NodeIdentity, MemberCapacity>,
1127        load: HashMap<u64, RangeLoad>,
1128        member_attributes: HashMap<(NodeIdentity, String), String>,
1129        default_bytes: u64,
1130    }
1131
1132    impl FakeSignals {
1133        fn uniform(disk: u64, default_bytes: u64) -> Self {
1134            Self {
1135                default_capacity: MemberCapacity::with_disk(disk),
1136                capacity: HashMap::new(),
1137                load: HashMap::new(),
1138                member_attributes: HashMap::new(),
1139                default_bytes,
1140            }
1141        }
1142
1143        fn with_capacity(mut self, cn: &str, cap: MemberCapacity) -> Self {
1144            self.capacity.insert(ident(cn), cap);
1145            self
1146        }
1147
1148        fn with_load(mut self, range_id: u64, load: RangeLoad) -> Self {
1149            self.load.insert(range_id, load);
1150            self
1151        }
1152
1153        fn with_member_attribute(mut self, cn: &str, key: &str, value: &str) -> Self {
1154            self.member_attributes
1155                .insert((ident(cn), key.to_string()), value.to_string());
1156            self
1157        }
1158    }
1159
1160    impl PlacementSignals for FakeSignals {
1161        fn member_capacity(&self, member: &NodeIdentity) -> MemberCapacity {
1162            self.capacity
1163                .get(member)
1164                .copied()
1165                .unwrap_or(self.default_capacity)
1166        }
1167
1168        fn range_load(&self, _collection: &CollectionId, range_id: RangeId) -> RangeLoad {
1169            self.load
1170                .get(&range_id.value())
1171                .copied()
1172                .unwrap_or_else(|| RangeLoad::idle(self.default_bytes))
1173        }
1174
1175        fn member_attribute(&self, member: &NodeIdentity, key: &str) -> Option<&str> {
1176            self.member_attributes
1177                .get(&(member.clone(), key.to_string()))
1178                .map(String::as_str)
1179        }
1180    }
1181
1182    // --- weighted capacity model -----------------------------------------
1183
1184    #[test]
1185    fn weighted_capacity_scales_disk_by_operator_weight() {
1186        // Neutral weight places strictly by disk.
1187        assert_eq!(MemberCapacity::with_disk(1_000).weighted_capacity(), 1_000);
1188        // A 2.0x operator weight doubles the placement weight; 0.5x halves it.
1189        assert_eq!(MemberCapacity::new(1_000, 200).weighted_capacity(), 2_000);
1190        assert_eq!(MemberCapacity::new(1_000, 50).weighted_capacity(), 500);
1191        // No disk -> not placeable.
1192        assert!(!MemberCapacity::with_disk(0).is_placeable());
1193        assert!(MemberCapacity::with_disk(1).is_placeable());
1194    }
1195
1196    // --- acceptance scenario: homogeneous placement ----------------------
1197
1198    #[test]
1199    fn homogeneous_cluster_is_balanced_and_plans_nothing() {
1200        // Three members, equal disk, three equal-sized ranges one each: already
1201        // perfectly balanced, so the planner proposes no move.
1202        let planner = WeightedPlacementPlanner::default();
1203        let members = membership(&["CN=node-a", "CN=node-b", "CN=node-c"]);
1204        let (catalog, _orders) = catalog(&["CN=node-a", "CN=node-b", "CN=node-c"]);
1205        let signals = FakeSignals::uniform(1_000_000, 100);
1206
1207        let plan = planner.plan_rebalance(&members, &catalog, &signals);
1208        assert!(plan.is_empty(), "balanced homogeneous cluster is a no-op");
1209    }
1210
1211    #[test]
1212    fn homogeneous_cluster_with_skew_spreads_ranges() {
1213        // All three ranges sit on node-a while node-b and node-c are empty. With
1214        // equal capacity the fair share is one range each, so the planner moves two
1215        // ranges off node-a.
1216        let planner = WeightedPlacementPlanner::default();
1217        let members = membership(&["CN=node-a", "CN=node-b", "CN=node-c"]);
1218        let (catalog, _orders) = catalog(&["CN=node-a", "CN=node-a", "CN=node-a"]);
1219        let signals = FakeSignals::uniform(1_000_000, 100);
1220
1221        let plan = planner.plan_rebalance(&members, &catalog, &signals);
1222        assert_eq!(
1223            plan.capacity_moves().count(),
1224            2,
1225            "two ranges move off node-a"
1226        );
1227        for mv in plan.capacity_moves() {
1228            assert_eq!(mv.from, ident("CN=node-a"));
1229            assert_ne!(mv.to, ident("CN=node-a"));
1230            assert_eq!(mv.reason, MoveReason::CapacityBalance);
1231        }
1232        // node-b and node-c each receive exactly one range.
1233        let targets: std::collections::BTreeSet<_> =
1234            plan.capacity_moves().map(|m| m.to.clone()).collect();
1235        assert_eq!(targets.len(), 2);
1236    }
1237
1238    #[test]
1239    fn scoped_plan_identifies_the_collection_group_placement_authority() {
1240        let planner = WeightedPlacementPlanner::default();
1241        let members = membership(&["CN=node-a", "CN=node-b"]);
1242        let (catalog, _orders) = catalog(&["CN=node-a", "CN=node-a", "CN=node-a"]);
1243        let signals = FakeSignals::uniform(1_000, 100);
1244        let mut authorities = PlacementAuthorityCatalog::new();
1245        let group = CollectionGroupId::new("commerce").unwrap();
1246        let authority = CollectionGroupPlacementAuthority::new(
1247            group.clone(),
1248            ident("CN=pa-commerce"),
1249            [collection("orders"), collection("payments")],
1250        )
1251        .unwrap();
1252        authorities.register(authority).unwrap();
1253
1254        let plan = planner
1255            .plan_rebalance_scoped(&members, &catalog, &signals, &authorities)
1256            .unwrap();
1257
1258        assert!(
1259            !plan.moves.is_empty(),
1260            "skewed ownership should plan movement"
1261        );
1262        for planned in &plan.moves {
1263            assert_eq!(planned.placement_authority.collection_group(), &group);
1264            assert_eq!(
1265                planned.placement_authority.authority(),
1266                &ident("CN=pa-commerce")
1267            );
1268            assert_eq!(planned.movement.collection, collection("orders"));
1269        }
1270    }
1271
1272    // --- acceptance scenario: heterogeneous disk weights -----------------
1273
1274    #[test]
1275    fn heterogeneous_disk_weights_apportion_by_capacity() {
1276        // node-big advertises 4x the disk of node-small. Six equal ranges all start
1277        // on node-small; fair shares are big≈4.8, small≈1.2 ranges, so the planner
1278        // moves the bulk onto node-big.
1279        let planner = WeightedPlacementPlanner::default();
1280        let members = membership(&["CN=node-big", "CN=node-small"]);
1281        let (catalog, _orders) = catalog(&[
1282            "CN=node-small",
1283            "CN=node-small",
1284            "CN=node-small",
1285            "CN=node-small",
1286            "CN=node-small",
1287            "CN=node-small",
1288        ]);
1289        let signals = FakeSignals::uniform(1_000, 100)
1290            .with_capacity("CN=node-big", MemberCapacity::with_disk(4_000))
1291            .with_capacity("CN=node-small", MemberCapacity::with_disk(1_000));
1292
1293        let plan = planner.plan_rebalance(&members, &catalog, &signals);
1294        assert!(!plan.no_moves(), "imbalanced cluster must plan moves");
1295        // Every move goes from small to big, and big ends with ~4-5 of the 6 ranges.
1296        let to_big = plan
1297            .capacity_moves()
1298            .filter(|m| m.to == ident("CN=node-big"))
1299            .count();
1300        assert!(
1301            (4..=5).contains(&to_big),
1302            "node-big should receive ~4/5 of 6 ranges, got {to_big}"
1303        );
1304        for mv in plan.capacity_moves() {
1305            assert_eq!(mv.from, ident("CN=node-small"));
1306            assert_eq!(mv.to, ident("CN=node-big"));
1307        }
1308    }
1309
1310    #[test]
1311    fn operator_weight_biases_placement_without_more_disk() {
1312        // Same disk on both, but node-pref carries a 3x operator weight, so it
1313        // deserves the larger share of four ranges that all start on node-plain.
1314        let planner = WeightedPlacementPlanner::default();
1315        let members = membership(&["CN=node-pref", "CN=node-plain"]);
1316        let (catalog, _orders) = catalog(&[
1317            "CN=node-plain",
1318            "CN=node-plain",
1319            "CN=node-plain",
1320            "CN=node-plain",
1321        ]);
1322        let signals = FakeSignals::uniform(1_000, 100)
1323            .with_capacity("CN=node-pref", MemberCapacity::new(1_000, 300));
1324
1325        let plan = planner.plan_rebalance(&members, &catalog, &signals);
1326        let to_pref = plan
1327            .capacity_moves()
1328            .filter(|m| m.to == ident("CN=node-pref"))
1329            .count();
1330        assert!(
1331            to_pref >= 2,
1332            "higher operator weight pulls more ranges, got {to_pref}"
1333        );
1334    }
1335
1336    #[test]
1337    fn failure_domain_spread_wins_over_marginal_weight_difference() {
1338        let planner = WeightedPlacementPlanner::default();
1339        let members = membership(&["CN=node-a", "CN=node-b", "CN=node-c"]);
1340        let placement =
1341            PlacementMetadata::with_replication_factor(2).with_attribute("failure_domain", "zone");
1342        let (catalog, _orders) = catalog_with_replicas("CN=node-a", &["CN=node-b"], placement);
1343        let signals = FakeSignals::uniform(1_000, 100)
1344            .with_capacity("CN=node-b", MemberCapacity::new(1_000, 110))
1345            .with_capacity("CN=node-c", MemberCapacity::new(1_000, 100))
1346            .with_member_attribute("CN=node-a", "zone", "zone-1")
1347            .with_member_attribute("CN=node-b", "zone", "zone-1")
1348            .with_member_attribute("CN=node-c", "zone", "zone-2");
1349
1350        let plan = planner.plan_rebalance(&members, &catalog, &signals);
1351
1352        assert_eq!(plan.replica_replacements.len(), 1);
1353        let replacement = &plan.replica_replacements[0];
1354        assert_eq!(replacement.collection, collection("orders"));
1355        assert_eq!(replacement.range_id, RangeId::new(1));
1356        assert_eq!(replacement.replace, ident("CN=node-b"));
1357        assert_eq!(replacement.with, ident("CN=node-c"));
1358        assert!(
1359            plan.warnings.is_empty(),
1360            "satisfiable spread is not degraded"
1361        );
1362    }
1363
1364    #[test]
1365    fn unsatisfiable_failure_domain_spread_warns_and_keeps_coplacement() {
1366        let planner = WeightedPlacementPlanner::default();
1367        let members = membership(&["CN=node-a", "CN=node-b"]);
1368        let placement =
1369            PlacementMetadata::with_replication_factor(2).with_attribute("failure_domain", "zone");
1370        let (catalog, _orders) = catalog_with_replicas("CN=node-a", &["CN=node-b"], placement);
1371        let signals = FakeSignals::uniform(1_000, 100)
1372            .with_member_attribute("CN=node-a", "zone", "zone-1")
1373            .with_member_attribute("CN=node-b", "zone", "zone-1");
1374
1375        let plan = planner.plan_rebalance(&members, &catalog, &signals);
1376
1377        assert!(
1378            plan.replica_replacements.is_empty(),
1379            "no distinct-domain member exists"
1380        );
1381        assert_eq!(plan.warnings.len(), 1);
1382        assert_eq!(
1383            plan.warnings[0],
1384            PlacementWarning::UnsatisfiedFailureDomainSpread {
1385                collection: collection("orders"),
1386                range_id: RangeId::new(1),
1387                domain_key: "zone".to_string(),
1388                domain_value: "zone-1".to_string(),
1389            }
1390        );
1391        let warning = plan.warnings[0].to_string();
1392        assert!(
1393            warning.contains("orders/1"),
1394            "warning names range: {warning}"
1395        );
1396        assert!(
1397            warning.contains("zone=zone-1"),
1398            "warning names violated domain: {warning}"
1399        );
1400    }
1401
1402    #[test]
1403    fn failure_domain_spread_is_restored_after_distinct_domain_join() {
1404        let planner = WeightedPlacementPlanner::default();
1405        let placement =
1406            PlacementMetadata::with_replication_factor(2).with_attribute("failure_domain", "zone");
1407        let (catalog, _orders) = catalog_with_replicas("CN=node-a", &["CN=node-b"], placement);
1408        let signals = FakeSignals::uniform(1_000, 100)
1409            .with_member_attribute("CN=node-a", "zone", "zone-1")
1410            .with_member_attribute("CN=node-b", "zone", "zone-1")
1411            .with_member_attribute("CN=node-c", "zone", "zone-2");
1412
1413        let before_join =
1414            planner.plan_rebalance(&membership(&["CN=node-a", "CN=node-b"]), &catalog, &signals);
1415        assert!(before_join.replica_replacements.is_empty());
1416        assert_eq!(before_join.warnings.len(), 1);
1417
1418        let after_join = planner.plan_rebalance(
1419            &membership(&["CN=node-a", "CN=node-b", "CN=node-c"]),
1420            &catalog,
1421            &signals,
1422        );
1423
1424        assert!(after_join.warnings.is_empty());
1425        assert_eq!(after_join.replica_replacements.len(), 1);
1426        assert_eq!(
1427            after_join.replica_replacements[0].replace,
1428            ident("CN=node-b")
1429        );
1430        assert_eq!(after_join.replica_replacements[0].with, ident("CN=node-c"));
1431    }
1432
1433    // --- acceptance scenario: capacity expansion -------------------------
1434
1435    #[test]
1436    fn expanding_disk_changes_weight_and_next_plan_without_moving_data() {
1437        // Start heterogeneous: node-a small, node-b large, all six ranges on node-a.
1438        let planner = WeightedPlacementPlanner::default();
1439        let members = membership(&["CN=node-a", "CN=node-b"]);
1440        let (catalog, orders) = catalog(&[
1441            "CN=node-a",
1442            "CN=node-a",
1443            "CN=node-a",
1444            "CN=node-a",
1445            "CN=node-a",
1446            "CN=node-a",
1447        ]);
1448
1449        // Before expansion: node-b has only modest disk, so it receives a modest
1450        // share.
1451        let before_signals = FakeSignals::uniform(1_000, 100)
1452            .with_capacity("CN=node-a", MemberCapacity::with_disk(3_000))
1453            .with_capacity("CN=node-b", MemberCapacity::with_disk(1_000));
1454        let before = planner.plan_rebalance(&members, &catalog, &before_signals);
1455        let before_to_b = before
1456            .capacity_moves()
1457            .filter(|m| m.to == ident("CN=node-b"))
1458            .count();
1459
1460        // Operator expands node-b's disk 8x. Its placement weight jumps...
1461        let small = MemberCapacity::with_disk(1_000);
1462        let expanded = MemberCapacity::with_disk(8_000);
1463        assert!(
1464            expanded.weighted_capacity() > small.weighted_capacity(),
1465            "expanding disk raises placement weight",
1466        );
1467        let after_signals = FakeSignals::uniform(1_000, 100)
1468            .with_capacity("CN=node-a", MemberCapacity::with_disk(3_000))
1469            .with_capacity("CN=node-b", expanded);
1470        let after = planner.plan_rebalance(&members, &catalog, &after_signals);
1471        let after_to_b = after
1472            .capacity_moves()
1473            .filter(|m| m.to == ident("CN=node-b"))
1474            .count();
1475
1476        // ...so the *next* plan apportions more ranges to node-b than before.
1477        assert!(
1478            after_to_b > before_to_b,
1479            "expanded disk pulls more ranges on the next plan ({before_to_b} -> {after_to_b})",
1480        );
1481
1482        // But planning never moved data: the catalog still shows all six ranges on
1483        // node-a. Data only relocates when a transition plan is executed.
1484        for i in 1..=6 {
1485            let range = catalog.range(&orders, RangeId::new(i)).unwrap();
1486            assert_eq!(
1487                range.owner(),
1488                &ident("CN=node-a"),
1489                "range {i} stayed on node-a; planning moved nothing",
1490            );
1491        }
1492    }
1493
1494    // --- acceptance scenario: hotspot signal influence -------------------
1495
1496    #[test]
1497    fn hotspot_traffic_identifies_secondary_candidate() {
1498        // Capacity is *balanced* — node-a owns two ranges but has twice the disk,
1499        // so every member sits exactly on its fair share and the primary pass
1500        // proposes nothing. Yet range 1 on node-a serves a huge read/write load (a
1501        // small, read-hammered range) while the others are quiet. The secondary
1502        // signal flags it as a hotspot and, because node-a also carries other
1503        // traffic and a quiet member has both load and capacity headroom, proposes
1504        // spreading it.
1505        let planner = WeightedPlacementPlanner::default();
1506        let members = membership(&["CN=node-a", "CN=node-b", "CN=node-c"]);
1507        // node-a owns ranges 1 and 2; node-b owns 3; node-c owns 4.
1508        let (catalog, _orders) = catalog(&["CN=node-a", "CN=node-a", "CN=node-b", "CN=node-c"]);
1509        // node-a has 2x disk so its fair share covers both its ranges (40 bytes);
1510        // node-b and node-c each match their single 20-byte range.
1511        let signals = FakeSignals::uniform(0, 0)
1512            .with_capacity("CN=node-a", MemberCapacity::with_disk(2_000))
1513            .with_capacity("CN=node-b", MemberCapacity::with_disk(1_000))
1514            .with_capacity("CN=node-c", MemberCapacity::with_disk(1_000))
1515            // The hot range is tiny on disk but hammered; node-a keeps real
1516            // residual traffic on range 2.
1517            .with_load(
1518                1,
1519                RangeLoad {
1520                    bytes_used: 2,
1521                    read_ops: 1_000,
1522                    write_ops: 1_000,
1523                },
1524            )
1525            .with_load(
1526                2,
1527                RangeLoad {
1528                    bytes_used: 38,
1529                    read_ops: 300,
1530                    write_ops: 0,
1531                },
1532            )
1533            .with_load(
1534                3,
1535                RangeLoad {
1536                    bytes_used: 20,
1537                    read_ops: 100,
1538                    write_ops: 0,
1539                },
1540            )
1541            .with_load(
1542                4,
1543                RangeLoad {
1544                    bytes_used: 20,
1545                    read_ops: 100,
1546                    write_ops: 0,
1547                },
1548            );
1549
1550        let plan = planner.plan_rebalance(&members, &catalog, &signals);
1551        // Capacity is balanced, so no capacity-balance move is proposed.
1552        assert_eq!(plan.capacity_moves().count(), 0, "capacity is balanced");
1553        // Range 1 is identified as a hotspot, attributed to its real owner.
1554        assert_eq!(plan.hotspots.len(), 1, "the hot range is surfaced");
1555        assert_eq!(plan.hotspots[0].range_id, RangeId::new(1));
1556        assert_eq!(plan.hotspots[0].owner, ident("CN=node-a"));
1557        assert_eq!(plan.hotspots[0].traffic, 2_000);
1558        // And a hotspot-relief move spreads it off node-a onto the quietest member.
1559        let relief: Vec<_> = plan.hotspot_moves().collect();
1560        assert_eq!(relief.len(), 1, "a relief move is planned");
1561        assert_eq!(relief[0].range_id, RangeId::new(1));
1562        assert_eq!(relief[0].from, ident("CN=node-a"));
1563        assert_eq!(
1564            relief[0].to,
1565            ident("CN=node-b"),
1566            "quietest target, tie -> lowest id"
1567        );
1568        assert_eq!(relief[0].reason, MoveReason::HotspotRelief);
1569    }
1570
1571    #[test]
1572    fn no_hotspot_when_traffic_is_even() {
1573        // Balanced capacity (one range each, equal disk) and equal traffic: nothing
1574        // is a hotspot and the secondary signal proposes nothing.
1575        let planner = WeightedPlacementPlanner::default();
1576        let members = membership(&["CN=node-a", "CN=node-b", "CN=node-c"]);
1577        let (catalog, _orders) = catalog(&["CN=node-a", "CN=node-b", "CN=node-c"]);
1578        let signals = FakeSignals::uniform(1_000_000, 100)
1579            .with_load(
1580                1,
1581                RangeLoad {
1582                    bytes_used: 10,
1583                    read_ops: 100,
1584                    write_ops: 100,
1585                },
1586            )
1587            .with_load(
1588                2,
1589                RangeLoad {
1590                    bytes_used: 10,
1591                    read_ops: 100,
1592                    write_ops: 100,
1593                },
1594            )
1595            .with_load(
1596                3,
1597                RangeLoad {
1598                    bytes_used: 10,
1599                    read_ops: 100,
1600                    write_ops: 100,
1601                },
1602            );
1603
1604        let plan = planner.plan_rebalance(&members, &catalog, &signals);
1605        assert!(plan.is_empty(), "balanced, even-traffic cluster is a no-op");
1606    }
1607
1608    // --- acceptance scenario: no implicit data movement ------------------
1609
1610    #[test]
1611    fn planning_never_mutates_the_catalog() {
1612        // A deliberately skewed cluster yields a non-empty plan, yet the ownership
1613        // catalog is byte-for-byte identical before and after planning: the planner
1614        // only *describes* moves, it never performs them.
1615        let planner = WeightedPlacementPlanner::default();
1616        let members = membership(&["CN=node-a", "CN=node-b"]);
1617        let (catalog, orders) = catalog(&["CN=node-a", "CN=node-a", "CN=node-a", "CN=node-a"]);
1618        let signals = FakeSignals::uniform(1_000, 100);
1619
1620        // Snapshot every range's owner/epoch/version before planning.
1621        let before: Vec<_> = (1..=4)
1622            .map(|i| {
1623                let r = catalog.range(&orders, RangeId::new(i)).unwrap();
1624                (r.owner().clone(), r.epoch(), r.version())
1625            })
1626            .collect();
1627
1628        let plan = planner.plan_rebalance(&members, &catalog, &signals);
1629        assert!(!plan.no_moves(), "skewed cluster does plan moves");
1630
1631        // The catalog is unchanged: same owners, same epochs, same versions.
1632        for (i, snap) in before.iter().enumerate() {
1633            let r = catalog.range(&orders, RangeId::new(i as u64 + 1)).unwrap();
1634            assert_eq!(&(r.owner().clone(), r.epoch(), r.version()), snap);
1635        }
1636    }
1637
1638    #[test]
1639    fn draining_owner_ranges_are_left_to_the_drain_flow() {
1640        // node-a is draining (not placement-eligible). Its ranges are not moved by
1641        // the rebalancer — drain owns evacuating them — so no plan targets or
1642        // sources it for placement balancing.
1643        let planner = WeightedPlacementPlanner::default();
1644        let mut members = membership(&["CN=node-a", "CN=node-b"]);
1645        members.begin_drain(&ident("CN=node-a"));
1646        let (catalog, _orders) = catalog(&["CN=node-a", "CN=node-a", "CN=node-a"]);
1647        let signals = FakeSignals::uniform(1_000, 100);
1648
1649        let plan = planner.plan_rebalance(&members, &catalog, &signals);
1650        // node-a's ranges are not movable here, and node-b is the only eligible
1651        // member, so there is nothing to balance.
1652        assert!(
1653            plan.no_moves(),
1654            "draining owner's ranges are not rebalanced"
1655        );
1656    }
1657
1658    #[test]
1659    fn single_member_cluster_plans_nothing() {
1660        let planner = WeightedPlacementPlanner::default();
1661        let members = membership(&["CN=node-a"]);
1662        let (catalog, _orders) = catalog(&["CN=node-a", "CN=node-a"]);
1663        let signals = FakeSignals::uniform(1_000, 100);
1664
1665        let plan = planner.plan_rebalance(&members, &catalog, &signals);
1666        assert!(
1667            plan.no_moves(),
1668            "nowhere to move ranges in a one-member cluster"
1669        );
1670    }
1671}