1use std::collections::{BTreeMap, BTreeSet};
60
61use super::identity::NodeIdentity;
62use super::membership::MembershipCatalog;
63use super::ownership::{CollectionId, PlacementMetadata, 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 member_attribute(&self, _member: &NodeIdentity, _key: &str) -> Option<&str> {
322 None
323 }
324
325 fn range_load(&self, collection: &CollectionId, range_id: RangeId) -> RangeLoad;
328}
329
330#[derive(Debug, Clone, Copy, PartialEq, Eq)]
332pub enum MoveReason {
333 CapacityBalance,
337 HotspotRelief,
341}
342
343#[derive(Debug, Clone, PartialEq, Eq)]
353pub struct PlannedMove {
354 pub collection: CollectionId,
355 pub range_id: RangeId,
356 pub from: NodeIdentity,
358 pub to: NodeIdentity,
360 pub bytes: u64,
362 pub reason: MoveReason,
363}
364
365#[derive(Debug, Clone, PartialEq, Eq)]
368pub struct PlannedReplicaReplacement {
369 pub collection: CollectionId,
370 pub range_id: RangeId,
371 pub replace: NodeIdentity,
373 pub with: NodeIdentity,
375 pub domain_key: String,
376 pub domain_value: String,
377}
378
379#[derive(Debug, Clone, PartialEq, Eq)]
384pub struct HotspotRange {
385 pub collection: CollectionId,
386 pub range_id: RangeId,
387 pub owner: NodeIdentity,
389 pub traffic: u64,
391}
392
393#[derive(Debug, Clone, Default, PartialEq, Eq)]
399pub struct RebalancePlan {
400 pub moves: Vec<PlannedMove>,
403 pub replica_replacements: Vec<PlannedReplicaReplacement>,
406 pub hotspots: Vec<HotspotRange>,
408 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 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 pub fn no_moves(&self) -> bool {
467 self.moves.is_empty() && self.replica_replacements.is_empty()
468 }
469
470 pub fn capacity_moves(&self) -> impl Iterator<Item = &PlannedMove> {
472 self.moves
473 .iter()
474 .filter(|m| m.reason == MoveReason::CapacityBalance)
475 }
476
477 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#[derive(Debug, Clone, Copy, PartialEq)]
492pub struct PlacementPolicy {
493 pub balance_tolerance: f64,
498 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
512fn 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#[derive(Debug, Clone, Default)]
529pub struct WeightedPlacementPlanner {
530 policy: PlacementPolicy,
531}
532
533impl WeightedPlacementPlanner {
534 pub fn new(policy: PlacementPolicy) -> Self {
536 Self { policy }
537 }
538
539 pub fn policy(&self) -> &PlacementPolicy {
540 &self.policy
541 }
542
543 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
598struct ClusterState {
604 eligible: Vec<NodeIdentity>,
607 weighted_capacity: BTreeMap<NodeIdentity, u128>,
608 total_capacity: u128,
610 total_bytes: u64,
612 ranges: BTreeMap<(CollectionId, RangeId), RangeFacts>,
614 owner_of: BTreeMap<(CollectionId, RangeId), NodeIdentity>,
616 origin_owner: BTreeMap<(CollectionId, RangeId), NodeIdentity>,
619 used: BTreeMap<NodeIdentity, u64>,
621 load: BTreeMap<NodeIdentity, u64>,
623 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 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 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 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 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 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 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 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 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 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 .then_with(|| b.cmp(a))
880 })
881 .cloned()
882 }
883
884 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 .then_with(|| a.cmp(b))
896 })
897 .cloned()
898 }
899
900 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 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 owner: self.origin_owner[&key].clone(),
940 traffic,
941 });
942
943 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 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 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 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 #[test]
1185 fn weighted_capacity_scales_disk_by_operator_weight() {
1186 assert_eq!(MemberCapacity::with_disk(1_000).weighted_capacity(), 1_000);
1188 assert_eq!(MemberCapacity::new(1_000, 200).weighted_capacity(), 2_000);
1190 assert_eq!(MemberCapacity::new(1_000, 50).weighted_capacity(), 500);
1191 assert!(!MemberCapacity::with_disk(0).is_placeable());
1193 assert!(MemberCapacity::with_disk(1).is_placeable());
1194 }
1195
1196 #[test]
1199 fn homogeneous_cluster_is_balanced_and_plans_nothing() {
1200 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 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 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 #[test]
1275 fn heterogeneous_disk_weights_apportion_by_capacity() {
1276 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 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 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 #[test]
1436 fn expanding_disk_changes_weight_and_next_plan_without_moving_data() {
1437 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 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 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 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 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 #[test]
1497 fn hotspot_traffic_identifies_secondary_candidate() {
1498 let planner = WeightedPlacementPlanner::default();
1506 let members = membership(&["CN=node-a", "CN=node-b", "CN=node-c"]);
1507 let (catalog, _orders) = catalog(&["CN=node-a", "CN=node-a", "CN=node-b", "CN=node-c"]);
1509 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 .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 assert_eq!(plan.capacity_moves().count(), 0, "capacity is balanced");
1553 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 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 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 #[test]
1611 fn planning_never_mutates_the_catalog() {
1612 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 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 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 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 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}