1use std::collections::{BTreeMap, BTreeSet};
60
61use super::identity::NodeIdentity;
62use super::membership::MembershipCatalog;
63use super::ownership::{CollectionId, RangeId, ShardOwnershipCatalog};
64
65pub const NEUTRAL_OPERATOR_WEIGHT: u32 = 100;
70
71#[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#[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#[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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
227pub struct MemberCapacity {
228 pub usable_disk_bytes: u64,
230 pub operator_weight: u32,
233}
234
235impl MemberCapacity {
236 pub fn new(usable_disk_bytes: u64, operator_weight: u32) -> Self {
238 Self {
239 usable_disk_bytes,
240 operator_weight,
241 }
242 }
243
244 pub fn with_disk(usable_disk_bytes: u64) -> Self {
247 Self::new(usable_disk_bytes, NEUTRAL_OPERATOR_WEIGHT)
248 }
249
250 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 pub fn is_placeable(&self) -> bool {
265 self.weighted_capacity() > 0
266 }
267}
268
269#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
278pub struct RangeLoad {
279 pub bytes_used: u64,
281 pub read_ops: u64,
283 pub write_ops: u64,
285}
286
287impl RangeLoad {
288 pub fn idle(bytes_used: u64) -> Self {
291 Self {
292 bytes_used,
293 read_ops: 0,
294 write_ops: 0,
295 }
296 }
297
298 pub fn traffic(&self) -> u64 {
301 self.read_ops.saturating_add(self.write_ops)
302 }
303}
304
305pub trait PlacementSignals {
312 fn member_capacity(&self, member: &NodeIdentity) -> MemberCapacity;
316
317 fn range_load(&self, collection: &CollectionId, range_id: RangeId) -> RangeLoad;
320}
321
322#[derive(Debug, Clone, Copy, PartialEq, Eq)]
324pub enum MoveReason {
325 CapacityBalance,
329 HotspotRelief,
333}
334
335#[derive(Debug, Clone, PartialEq, Eq)]
345pub struct PlannedMove {
346 pub collection: CollectionId,
347 pub range_id: RangeId,
348 pub from: NodeIdentity,
350 pub to: NodeIdentity,
352 pub bytes: u64,
354 pub reason: MoveReason,
355}
356
357#[derive(Debug, Clone, PartialEq, Eq)]
362pub struct HotspotRange {
363 pub collection: CollectionId,
364 pub range_id: RangeId,
365 pub owner: NodeIdentity,
367 pub traffic: u64,
369}
370
371#[derive(Debug, Clone, Default, PartialEq, Eq)]
377pub struct RebalancePlan {
378 pub moves: Vec<PlannedMove>,
381 pub hotspots: Vec<HotspotRange>,
383}
384
385#[derive(Debug, Clone, PartialEq, Eq)]
386pub struct AuthorityScopedPlannedMove {
387 pub movement: PlannedMove,
388 pub placement_authority: CollectionGroupPlacementAuthority,
389}
390
391#[derive(Debug, Clone, Default, PartialEq, Eq)]
392pub struct AuthorityScopedRebalancePlan {
393 pub moves: Vec<AuthorityScopedPlannedMove>,
394 pub hotspots: Vec<HotspotRange>,
395}
396
397impl RebalancePlan {
398 pub fn is_empty(&self) -> bool {
402 self.moves.is_empty() && self.hotspots.is_empty()
403 }
404
405 pub fn no_moves(&self) -> bool {
409 self.moves.is_empty()
410 }
411
412 pub fn capacity_moves(&self) -> impl Iterator<Item = &PlannedMove> {
414 self.moves
415 .iter()
416 .filter(|m| m.reason == MoveReason::CapacityBalance)
417 }
418
419 pub fn hotspot_moves(&self) -> impl Iterator<Item = &PlannedMove> {
421 self.moves
422 .iter()
423 .filter(|m| m.reason == MoveReason::HotspotRelief)
424 }
425}
426
427#[derive(Debug, Clone, Copy, PartialEq)]
434pub struct PlacementPolicy {
435 pub balance_tolerance: f64,
440 pub hotspot_load_factor: f64,
443}
444
445impl Default for PlacementPolicy {
446 fn default() -> Self {
447 Self {
448 balance_tolerance: 0.10,
449 hotspot_load_factor: 2.0,
450 }
451 }
452}
453
454fn fair_share(total_bytes: u64, member_capacity: u128, total_capacity: u128) -> u64 {
458 if total_capacity == 0 {
459 return 0;
460 }
461 let share = total_bytes as u128 * member_capacity / total_capacity;
462 share.min(u64::MAX as u128) as u64
463}
464
465#[derive(Debug, Clone, Default)]
471pub struct WeightedPlacementPlanner {
472 policy: PlacementPolicy,
473}
474
475impl WeightedPlacementPlanner {
476 pub fn new(policy: PlacementPolicy) -> Self {
478 Self { policy }
479 }
480
481 pub fn policy(&self) -> &PlacementPolicy {
482 &self.policy
483 }
484
485 pub fn plan_rebalance(
491 &self,
492 membership: &MembershipCatalog,
493 ownership: &ShardOwnershipCatalog,
494 signals: &impl PlacementSignals,
495 ) -> RebalancePlan {
496 let mut state = ClusterState::observe(membership, ownership, signals, &self.policy);
497 let mut moves = state.plan_capacity_moves(&self.policy);
498 let (hotspots, hotspot_moves) = state.plan_hotspot_moves(&self.policy);
499 moves.extend(hotspot_moves);
500 RebalancePlan { moves, hotspots }
501 }
502
503 pub fn plan_rebalance_scoped(
504 &self,
505 membership: &MembershipCatalog,
506 ownership: &ShardOwnershipCatalog,
507 signals: &impl PlacementSignals,
508 authorities: &PlacementAuthorityCatalog,
509 ) -> Result<AuthorityScopedRebalancePlan, PlacementAuthorityError> {
510 let plan = self.plan_rebalance(membership, ownership, signals);
511 let mut moves = Vec::with_capacity(plan.moves.len());
512 for movement in plan.moves {
513 let placement_authority = authorities
514 .authority_for(&movement.collection)
515 .ok_or_else(|| PlacementAuthorityError::MissingCollectionAuthority {
516 collection: movement.collection.clone(),
517 })?
518 .clone();
519 moves.push(AuthorityScopedPlannedMove {
520 movement,
521 placement_authority,
522 });
523 }
524 Ok(AuthorityScopedRebalancePlan {
525 moves,
526 hotspots: plan.hotspots,
527 })
528 }
529}
530
531struct ClusterState {
537 eligible: Vec<NodeIdentity>,
540 weighted_capacity: BTreeMap<NodeIdentity, u128>,
541 total_capacity: u128,
543 total_bytes: u64,
545 ranges: BTreeMap<(CollectionId, RangeId), RangeFacts>,
547 owner_of: BTreeMap<(CollectionId, RangeId), NodeIdentity>,
549 origin_owner: BTreeMap<(CollectionId, RangeId), NodeIdentity>,
552 used: BTreeMap<NodeIdentity, u64>,
554 load: BTreeMap<NodeIdentity, u64>,
556 moved: std::collections::BTreeSet<(CollectionId, RangeId)>,
558}
559
560#[derive(Clone, Copy)]
561struct RangeFacts {
562 bytes: u64,
563 traffic: u64,
564}
565
566impl ClusterState {
567 fn observe(
568 membership: &MembershipCatalog,
569 ownership: &ShardOwnershipCatalog,
570 signals: &impl PlacementSignals,
571 _policy: &PlacementPolicy,
572 ) -> Self {
573 let mut weighted_capacity = BTreeMap::new();
574 let mut eligible = Vec::new();
575 let mut total_capacity: u128 = 0;
576 for member in membership.placement_eligible_members() {
577 let id = member.identity().clone();
578 let cap = signals.member_capacity(&id).weighted_capacity();
579 if cap == 0 {
580 continue;
583 }
584 total_capacity += cap;
585 weighted_capacity.insert(id.clone(), cap);
586 eligible.push(id);
587 }
588
589 let eligible_set: std::collections::BTreeSet<&NodeIdentity> = eligible.iter().collect();
590
591 let mut ranges = BTreeMap::new();
592 let mut owner_of = BTreeMap::new();
593 let mut origin_owner = BTreeMap::new();
594 let mut used: BTreeMap<NodeIdentity, u64> =
595 eligible.iter().map(|id| (id.clone(), 0)).collect();
596 let mut load: BTreeMap<NodeIdentity, u64> =
597 eligible.iter().map(|id| (id.clone(), 0)).collect();
598 let mut total_bytes: u64 = 0;
599
600 for entry in ownership.entries() {
601 let owner = entry.owner().clone();
602 if !eligible_set.contains(&owner) {
605 continue;
606 }
607 let key = (entry.collection().clone(), entry.range_id());
608 let load_facts = signals.range_load(entry.collection(), entry.range_id());
609 ranges.insert(
610 key.clone(),
611 RangeFacts {
612 bytes: load_facts.bytes_used,
613 traffic: load_facts.traffic(),
614 },
615 );
616 *used.get_mut(&owner).unwrap() += load_facts.bytes_used;
617 *load.get_mut(&owner).unwrap() += load_facts.traffic();
618 total_bytes = total_bytes.saturating_add(load_facts.bytes_used);
619 owner_of.insert(key.clone(), owner.clone());
620 origin_owner.insert(key, owner);
621 }
622
623 Self {
624 eligible,
625 weighted_capacity,
626 total_capacity,
627 total_bytes,
628 ranges,
629 owner_of,
630 origin_owner,
631 used,
632 load,
633 moved: std::collections::BTreeSet::new(),
634 }
635 }
636
637 fn fair(&self, member: &NodeIdentity) -> u64 {
638 let cap = self.weighted_capacity.get(member).copied().unwrap_or(0);
639 fair_share(self.total_bytes, cap, self.total_capacity)
640 }
641
642 fn ranges_owned_by(&self, member: &NodeIdentity) -> Vec<(CollectionId, RangeId)> {
645 self.owner_of
646 .iter()
647 .filter(|(key, owner)| *owner == member && !self.moved.contains(*key))
648 .map(|(key, _)| key.clone())
649 .collect()
650 }
651
652 fn apply_move(&mut self, key: &(CollectionId, RangeId), to: &NodeIdentity) {
653 let facts = self.ranges[key];
654 let from = self.owner_of[key].clone();
655 *self.used.get_mut(&from).unwrap() -= facts.bytes;
656 *self.load.get_mut(&from).unwrap() -= facts.traffic;
657 *self.used.get_mut(to).unwrap() += facts.bytes;
658 *self.load.get_mut(to).unwrap() += facts.traffic;
659 self.owner_of.insert(key.clone(), to.clone());
660 self.moved.insert(key.clone());
661 }
662
663 fn plan_capacity_moves(&mut self, policy: &PlacementPolicy) -> Vec<PlannedMove> {
667 let mut planned = Vec::new();
668 if self.total_capacity == 0 || self.eligible.len() < 2 {
669 return planned;
670 }
671
672 while let Some(source) = self.most_over(policy) {
676 let Some(target) = self.most_under(&source) else {
677 break;
678 };
679
680 let dev_src = self.deviation(&source);
681 let dev_tgt = self.deviation(&target);
682 let worst_before = dev_src.abs().max(dev_tgt.abs());
683
684 let mut best: Option<((CollectionId, RangeId), f64)> = None;
687 for key in self.ranges_owned_by(&source) {
688 let s = self.ranges[&key].bytes as f64;
689 let after = (dev_src - s).abs().max((dev_tgt + s).abs());
690 let better = match &best {
691 None => true,
692 Some((_, best_after)) => after < *best_after,
693 };
694 if better {
695 best = Some((key, after));
696 }
697 }
698
699 let Some((key, worst_after)) = best else {
700 break;
701 };
702 if worst_after >= worst_before {
705 break;
706 }
707
708 let bytes = self.ranges[&key].bytes;
709 let from = self.origin_owner[&key].clone();
710 self.apply_move(&key, &target);
711 planned.push(PlannedMove {
712 collection: key.0,
713 range_id: key.1,
714 from,
715 to: target,
716 bytes,
717 reason: MoveReason::CapacityBalance,
718 });
719 }
720
721 planned
722 }
723
724 fn deviation(&self, member: &NodeIdentity) -> f64 {
727 self.used.get(member).copied().unwrap_or(0) as f64 - self.fair(member) as f64
728 }
729
730 fn most_over(&self, policy: &PlacementPolicy) -> Option<NodeIdentity> {
733 self.eligible
734 .iter()
735 .filter(|id| {
736 let used = self.used.get(*id).copied().unwrap_or(0) as f64;
737 let fair = self.fair(id) as f64;
738 used > fair * (1.0 + policy.balance_tolerance) && used > fair
739 })
740 .max_by(|a, b| {
741 self.deviation(a)
742 .partial_cmp(&self.deviation(b))
743 .unwrap()
744 .then_with(|| b.cmp(a))
746 })
747 .cloned()
748 }
749
750 fn most_under(&self, source: &NodeIdentity) -> Option<NodeIdentity> {
753 self.eligible
754 .iter()
755 .filter(|id| *id != source && self.deviation(id) < 0.0)
756 .min_by(|a, b| {
757 self.deviation(a)
758 .partial_cmp(&self.deviation(b))
759 .unwrap()
760 .then_with(|| a.cmp(b))
762 })
763 .cloned()
764 }
765
766 fn plan_hotspot_moves(
772 &mut self,
773 policy: &PlacementPolicy,
774 ) -> (Vec<HotspotRange>, Vec<PlannedMove>) {
775 let mut hotspots = Vec::new();
776 let mut moves = Vec::new();
777
778 let range_count = self.ranges.len();
779 if range_count == 0 {
780 return (hotspots, moves);
781 }
782 let total_traffic: u64 = self.ranges.values().map(|f| f.traffic).sum();
783 let mean = total_traffic as f64 / range_count as f64;
784 let threshold = mean * policy.hotspot_load_factor;
785 if mean <= 0.0 {
786 return (hotspots, moves);
787 }
788
789 let mut hot: Vec<((CollectionId, RangeId), u64)> = self
791 .ranges
792 .iter()
793 .filter(|(_, f)| f.traffic as f64 > threshold)
794 .map(|(key, f)| (key.clone(), f.traffic))
795 .collect();
796 hot.sort_by(|a, b| b.1.cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
797
798 for (key, traffic) in hot {
799 let owner = self.owner_of[&key].clone();
800 hotspots.push(HotspotRange {
801 collection: key.0.clone(),
802 range_id: key.1,
803 owner: self.origin_owner[&key].clone(),
806 traffic,
807 });
808
809 if self.moved.contains(&key) || self.ranges_owned_by(&owner).len() < 2 {
813 continue;
814 }
815
816 let facts = self.ranges[&key];
817 let owner_load = self.load.get(&owner).copied().unwrap_or(0);
818
819 let target = self
823 .eligible
824 .iter()
825 .filter(|id| **id != owner)
826 .filter(|id| {
827 let used = self.used.get(*id).copied().unwrap_or(0);
828 let fair = self.fair(id) as f64;
829 (used + facts.bytes) as f64 <= fair * (1.0 + policy.balance_tolerance)
830 })
831 .filter(|id| {
832 let tgt_load = self.load.get(*id).copied().unwrap_or(0);
833 tgt_load + facts.traffic < owner_load
834 })
835 .min_by(|a, b| {
836 let la = self.load.get(*a).copied().unwrap_or(0);
837 let lb = self.load.get(*b).copied().unwrap_or(0);
838 la.cmp(&lb).then_with(|| a.cmp(b))
839 })
840 .cloned();
841
842 if let Some(target) = target {
843 let from = self.origin_owner[&key].clone();
844 self.apply_move(&key, &target);
845 moves.push(PlannedMove {
846 collection: key.0,
847 range_id: key.1,
848 from,
849 to: target,
850 bytes: facts.bytes,
851 reason: MoveReason::HotspotRelief,
852 });
853 }
854 }
855
856 (hotspots, moves)
857 }
858}
859
860#[cfg(test)]
861mod tests {
862 use super::*;
863 use crate::cluster::membership::{ClusterId, ClusterMember, MemberKind};
864 use crate::cluster::ownership::{PlacementMetadata, RangeBounds, RangeOwnership, ShardKeyMode};
865 use std::collections::HashMap;
866
867 fn ident(cn: &str) -> NodeIdentity {
868 NodeIdentity::from_certificate_subject(cn).unwrap()
869 }
870
871 fn collection(name: &str) -> CollectionId {
872 CollectionId::new(name).unwrap()
873 }
874
875 fn data_member(cn: &str) -> ClusterMember {
876 ClusterMember::joined_empty(ident(cn), MemberKind::Data)
877 }
878
879 fn membership(members: &[&str]) -> MembershipCatalog {
880 MembershipCatalog::new(
881 ClusterId::new("cluster-x").unwrap(),
882 members.iter().map(|m| data_member(m)),
883 )
884 }
885
886 fn catalog(owners: &[&str]) -> (ShardOwnershipCatalog, CollectionId) {
889 let orders = collection("orders");
890 let mut catalog = ShardOwnershipCatalog::new();
891 for (i, owner) in owners.iter().enumerate() {
892 let lower = vec![i as u8];
893 let upper = vec![i as u8 + 1];
894 let bounds = RangeBounds::new(
895 crate::cluster::ownership::RangeBound::key(lower),
896 crate::cluster::ownership::RangeBound::key(upper),
897 )
898 .unwrap();
899 catalog
900 .apply_update(RangeOwnership::establish(
901 orders.clone(),
902 RangeId::new(i as u64 + 1),
903 ShardKeyMode::Hash,
904 bounds,
905 ident(owner),
906 Vec::<NodeIdentity>::new(),
907 PlacementMetadata::with_replication_factor(1),
908 ))
909 .unwrap();
910 }
911 (catalog, orders)
912 }
913
914 struct FakeSignals {
917 default_capacity: MemberCapacity,
918 capacity: HashMap<NodeIdentity, MemberCapacity>,
919 load: HashMap<u64, RangeLoad>,
920 default_bytes: u64,
921 }
922
923 impl FakeSignals {
924 fn uniform(disk: u64, default_bytes: u64) -> Self {
925 Self {
926 default_capacity: MemberCapacity::with_disk(disk),
927 capacity: HashMap::new(),
928 load: HashMap::new(),
929 default_bytes,
930 }
931 }
932
933 fn with_capacity(mut self, cn: &str, cap: MemberCapacity) -> Self {
934 self.capacity.insert(ident(cn), cap);
935 self
936 }
937
938 fn with_load(mut self, range_id: u64, load: RangeLoad) -> Self {
939 self.load.insert(range_id, load);
940 self
941 }
942 }
943
944 impl PlacementSignals for FakeSignals {
945 fn member_capacity(&self, member: &NodeIdentity) -> MemberCapacity {
946 self.capacity
947 .get(member)
948 .copied()
949 .unwrap_or(self.default_capacity)
950 }
951
952 fn range_load(&self, _collection: &CollectionId, range_id: RangeId) -> RangeLoad {
953 self.load
954 .get(&range_id.value())
955 .copied()
956 .unwrap_or_else(|| RangeLoad::idle(self.default_bytes))
957 }
958 }
959
960 #[test]
963 fn weighted_capacity_scales_disk_by_operator_weight() {
964 assert_eq!(MemberCapacity::with_disk(1_000).weighted_capacity(), 1_000);
966 assert_eq!(MemberCapacity::new(1_000, 200).weighted_capacity(), 2_000);
968 assert_eq!(MemberCapacity::new(1_000, 50).weighted_capacity(), 500);
969 assert!(!MemberCapacity::with_disk(0).is_placeable());
971 assert!(MemberCapacity::with_disk(1).is_placeable());
972 }
973
974 #[test]
977 fn homogeneous_cluster_is_balanced_and_plans_nothing() {
978 let planner = WeightedPlacementPlanner::default();
981 let members = membership(&["CN=node-a", "CN=node-b", "CN=node-c"]);
982 let (catalog, _orders) = catalog(&["CN=node-a", "CN=node-b", "CN=node-c"]);
983 let signals = FakeSignals::uniform(1_000_000, 100);
984
985 let plan = planner.plan_rebalance(&members, &catalog, &signals);
986 assert!(plan.is_empty(), "balanced homogeneous cluster is a no-op");
987 }
988
989 #[test]
990 fn homogeneous_cluster_with_skew_spreads_ranges() {
991 let planner = WeightedPlacementPlanner::default();
995 let members = membership(&["CN=node-a", "CN=node-b", "CN=node-c"]);
996 let (catalog, _orders) = catalog(&["CN=node-a", "CN=node-a", "CN=node-a"]);
997 let signals = FakeSignals::uniform(1_000_000, 100);
998
999 let plan = planner.plan_rebalance(&members, &catalog, &signals);
1000 assert_eq!(
1001 plan.capacity_moves().count(),
1002 2,
1003 "two ranges move off node-a"
1004 );
1005 for mv in plan.capacity_moves() {
1006 assert_eq!(mv.from, ident("CN=node-a"));
1007 assert_ne!(mv.to, ident("CN=node-a"));
1008 assert_eq!(mv.reason, MoveReason::CapacityBalance);
1009 }
1010 let targets: std::collections::BTreeSet<_> =
1012 plan.capacity_moves().map(|m| m.to.clone()).collect();
1013 assert_eq!(targets.len(), 2);
1014 }
1015
1016 #[test]
1017 fn scoped_plan_identifies_the_collection_group_placement_authority() {
1018 let planner = WeightedPlacementPlanner::default();
1019 let members = membership(&["CN=node-a", "CN=node-b"]);
1020 let (catalog, _orders) = catalog(&["CN=node-a", "CN=node-a", "CN=node-a"]);
1021 let signals = FakeSignals::uniform(1_000, 100);
1022 let mut authorities = PlacementAuthorityCatalog::new();
1023 let group = CollectionGroupId::new("commerce").unwrap();
1024 let authority = CollectionGroupPlacementAuthority::new(
1025 group.clone(),
1026 ident("CN=pa-commerce"),
1027 [collection("orders"), collection("payments")],
1028 )
1029 .unwrap();
1030 authorities.register(authority).unwrap();
1031
1032 let plan = planner
1033 .plan_rebalance_scoped(&members, &catalog, &signals, &authorities)
1034 .unwrap();
1035
1036 assert!(
1037 !plan.moves.is_empty(),
1038 "skewed ownership should plan movement"
1039 );
1040 for planned in &plan.moves {
1041 assert_eq!(planned.placement_authority.collection_group(), &group);
1042 assert_eq!(
1043 planned.placement_authority.authority(),
1044 &ident("CN=pa-commerce")
1045 );
1046 assert_eq!(planned.movement.collection, collection("orders"));
1047 }
1048 }
1049
1050 #[test]
1053 fn heterogeneous_disk_weights_apportion_by_capacity() {
1054 let planner = WeightedPlacementPlanner::default();
1058 let members = membership(&["CN=node-big", "CN=node-small"]);
1059 let (catalog, _orders) = catalog(&[
1060 "CN=node-small",
1061 "CN=node-small",
1062 "CN=node-small",
1063 "CN=node-small",
1064 "CN=node-small",
1065 "CN=node-small",
1066 ]);
1067 let signals = FakeSignals::uniform(1_000, 100)
1068 .with_capacity("CN=node-big", MemberCapacity::with_disk(4_000))
1069 .with_capacity("CN=node-small", MemberCapacity::with_disk(1_000));
1070
1071 let plan = planner.plan_rebalance(&members, &catalog, &signals);
1072 assert!(!plan.no_moves(), "imbalanced cluster must plan moves");
1073 let to_big = plan
1075 .capacity_moves()
1076 .filter(|m| m.to == ident("CN=node-big"))
1077 .count();
1078 assert!(
1079 (4..=5).contains(&to_big),
1080 "node-big should receive ~4/5 of 6 ranges, got {to_big}"
1081 );
1082 for mv in plan.capacity_moves() {
1083 assert_eq!(mv.from, ident("CN=node-small"));
1084 assert_eq!(mv.to, ident("CN=node-big"));
1085 }
1086 }
1087
1088 #[test]
1089 fn operator_weight_biases_placement_without_more_disk() {
1090 let planner = WeightedPlacementPlanner::default();
1093 let members = membership(&["CN=node-pref", "CN=node-plain"]);
1094 let (catalog, _orders) = catalog(&[
1095 "CN=node-plain",
1096 "CN=node-plain",
1097 "CN=node-plain",
1098 "CN=node-plain",
1099 ]);
1100 let signals = FakeSignals::uniform(1_000, 100)
1101 .with_capacity("CN=node-pref", MemberCapacity::new(1_000, 300));
1102
1103 let plan = planner.plan_rebalance(&members, &catalog, &signals);
1104 let to_pref = plan
1105 .capacity_moves()
1106 .filter(|m| m.to == ident("CN=node-pref"))
1107 .count();
1108 assert!(
1109 to_pref >= 2,
1110 "higher operator weight pulls more ranges, got {to_pref}"
1111 );
1112 }
1113
1114 #[test]
1117 fn expanding_disk_changes_weight_and_next_plan_without_moving_data() {
1118 let planner = WeightedPlacementPlanner::default();
1120 let members = membership(&["CN=node-a", "CN=node-b"]);
1121 let (catalog, orders) = catalog(&[
1122 "CN=node-a",
1123 "CN=node-a",
1124 "CN=node-a",
1125 "CN=node-a",
1126 "CN=node-a",
1127 "CN=node-a",
1128 ]);
1129
1130 let before_signals = FakeSignals::uniform(1_000, 100)
1133 .with_capacity("CN=node-a", MemberCapacity::with_disk(3_000))
1134 .with_capacity("CN=node-b", MemberCapacity::with_disk(1_000));
1135 let before = planner.plan_rebalance(&members, &catalog, &before_signals);
1136 let before_to_b = before
1137 .capacity_moves()
1138 .filter(|m| m.to == ident("CN=node-b"))
1139 .count();
1140
1141 let small = MemberCapacity::with_disk(1_000);
1143 let expanded = MemberCapacity::with_disk(8_000);
1144 assert!(
1145 expanded.weighted_capacity() > small.weighted_capacity(),
1146 "expanding disk raises placement weight",
1147 );
1148 let after_signals = FakeSignals::uniform(1_000, 100)
1149 .with_capacity("CN=node-a", MemberCapacity::with_disk(3_000))
1150 .with_capacity("CN=node-b", expanded);
1151 let after = planner.plan_rebalance(&members, &catalog, &after_signals);
1152 let after_to_b = after
1153 .capacity_moves()
1154 .filter(|m| m.to == ident("CN=node-b"))
1155 .count();
1156
1157 assert!(
1159 after_to_b > before_to_b,
1160 "expanded disk pulls more ranges on the next plan ({before_to_b} -> {after_to_b})",
1161 );
1162
1163 for i in 1..=6 {
1166 let range = catalog.range(&orders, RangeId::new(i)).unwrap();
1167 assert_eq!(
1168 range.owner(),
1169 &ident("CN=node-a"),
1170 "range {i} stayed on node-a; planning moved nothing",
1171 );
1172 }
1173 }
1174
1175 #[test]
1178 fn hotspot_traffic_identifies_secondary_candidate() {
1179 let planner = WeightedPlacementPlanner::default();
1187 let members = membership(&["CN=node-a", "CN=node-b", "CN=node-c"]);
1188 let (catalog, _orders) = catalog(&["CN=node-a", "CN=node-a", "CN=node-b", "CN=node-c"]);
1190 let signals = FakeSignals::uniform(0, 0)
1193 .with_capacity("CN=node-a", MemberCapacity::with_disk(2_000))
1194 .with_capacity("CN=node-b", MemberCapacity::with_disk(1_000))
1195 .with_capacity("CN=node-c", MemberCapacity::with_disk(1_000))
1196 .with_load(
1199 1,
1200 RangeLoad {
1201 bytes_used: 2,
1202 read_ops: 1_000,
1203 write_ops: 1_000,
1204 },
1205 )
1206 .with_load(
1207 2,
1208 RangeLoad {
1209 bytes_used: 38,
1210 read_ops: 300,
1211 write_ops: 0,
1212 },
1213 )
1214 .with_load(
1215 3,
1216 RangeLoad {
1217 bytes_used: 20,
1218 read_ops: 100,
1219 write_ops: 0,
1220 },
1221 )
1222 .with_load(
1223 4,
1224 RangeLoad {
1225 bytes_used: 20,
1226 read_ops: 100,
1227 write_ops: 0,
1228 },
1229 );
1230
1231 let plan = planner.plan_rebalance(&members, &catalog, &signals);
1232 assert_eq!(plan.capacity_moves().count(), 0, "capacity is balanced");
1234 assert_eq!(plan.hotspots.len(), 1, "the hot range is surfaced");
1236 assert_eq!(plan.hotspots[0].range_id, RangeId::new(1));
1237 assert_eq!(plan.hotspots[0].owner, ident("CN=node-a"));
1238 assert_eq!(plan.hotspots[0].traffic, 2_000);
1239 let relief: Vec<_> = plan.hotspot_moves().collect();
1241 assert_eq!(relief.len(), 1, "a relief move is planned");
1242 assert_eq!(relief[0].range_id, RangeId::new(1));
1243 assert_eq!(relief[0].from, ident("CN=node-a"));
1244 assert_eq!(
1245 relief[0].to,
1246 ident("CN=node-b"),
1247 "quietest target, tie -> lowest id"
1248 );
1249 assert_eq!(relief[0].reason, MoveReason::HotspotRelief);
1250 }
1251
1252 #[test]
1253 fn no_hotspot_when_traffic_is_even() {
1254 let planner = WeightedPlacementPlanner::default();
1257 let members = membership(&["CN=node-a", "CN=node-b", "CN=node-c"]);
1258 let (catalog, _orders) = catalog(&["CN=node-a", "CN=node-b", "CN=node-c"]);
1259 let signals = FakeSignals::uniform(1_000_000, 100)
1260 .with_load(
1261 1,
1262 RangeLoad {
1263 bytes_used: 10,
1264 read_ops: 100,
1265 write_ops: 100,
1266 },
1267 )
1268 .with_load(
1269 2,
1270 RangeLoad {
1271 bytes_used: 10,
1272 read_ops: 100,
1273 write_ops: 100,
1274 },
1275 )
1276 .with_load(
1277 3,
1278 RangeLoad {
1279 bytes_used: 10,
1280 read_ops: 100,
1281 write_ops: 100,
1282 },
1283 );
1284
1285 let plan = planner.plan_rebalance(&members, &catalog, &signals);
1286 assert!(plan.is_empty(), "balanced, even-traffic cluster is a no-op");
1287 }
1288
1289 #[test]
1292 fn planning_never_mutates_the_catalog() {
1293 let planner = WeightedPlacementPlanner::default();
1297 let members = membership(&["CN=node-a", "CN=node-b"]);
1298 let (catalog, orders) = catalog(&["CN=node-a", "CN=node-a", "CN=node-a", "CN=node-a"]);
1299 let signals = FakeSignals::uniform(1_000, 100);
1300
1301 let before: Vec<_> = (1..=4)
1303 .map(|i| {
1304 let r = catalog.range(&orders, RangeId::new(i)).unwrap();
1305 (r.owner().clone(), r.epoch(), r.version())
1306 })
1307 .collect();
1308
1309 let plan = planner.plan_rebalance(&members, &catalog, &signals);
1310 assert!(!plan.no_moves(), "skewed cluster does plan moves");
1311
1312 for (i, snap) in before.iter().enumerate() {
1314 let r = catalog.range(&orders, RangeId::new(i as u64 + 1)).unwrap();
1315 assert_eq!(&(r.owner().clone(), r.epoch(), r.version()), snap);
1316 }
1317 }
1318
1319 #[test]
1320 fn draining_owner_ranges_are_left_to_the_drain_flow() {
1321 let planner = WeightedPlacementPlanner::default();
1325 let mut members = membership(&["CN=node-a", "CN=node-b"]);
1326 members.begin_drain(&ident("CN=node-a"));
1327 let (catalog, _orders) = catalog(&["CN=node-a", "CN=node-a", "CN=node-a"]);
1328 let signals = FakeSignals::uniform(1_000, 100);
1329
1330 let plan = planner.plan_rebalance(&members, &catalog, &signals);
1331 assert!(
1334 plan.no_moves(),
1335 "draining owner's ranges are not rebalanced"
1336 );
1337 }
1338
1339 #[test]
1340 fn single_member_cluster_plans_nothing() {
1341 let planner = WeightedPlacementPlanner::default();
1342 let members = membership(&["CN=node-a"]);
1343 let (catalog, _orders) = catalog(&["CN=node-a", "CN=node-a"]);
1344 let signals = FakeSignals::uniform(1_000, 100);
1345
1346 let plan = planner.plan_rebalance(&members, &catalog, &signals);
1347 assert!(
1348 plan.no_moves(),
1349 "nowhere to move ranges in a one-member cluster"
1350 );
1351 }
1352}