EXT-1082 Add replication aware vdisk mapper option to the bs controller (#25063)

SammyVimes · web-flow · commit 91d4b005be84 · 2025-09-18T01:15:11.000+04:00
(cherry picked from commit c6e1031)
diff --git a/ydb/apps/dstool/lib/dstool_cmd_cluster_balance.py b/ydb/apps/dstool/lib/dstool_cmd_cluster_balance.py
@@ -10,6 +10,8 @@
 def add_options(p):
     p.add_argument('--max-replicating-pdisks', type=int, help='Limit number of maximum replicating PDisks in the cluster')
     p.add_argument('--only-from-overpopulated-pdisks', action='store_true', help='Move vdisks out only from pdisks with over expected slot count')
+    p.add_argument('--prefer-less-occupied-rack', action='store_true', help='Take into account racks\' free slots picking pdisk from rack with more free slots first')
+    p.add_argument('--with-attention-to-replication', action='store_true', help='Take into account replicating vdisks picking node and pdisk with less amount of them')
     common.add_basic_format_options(p)
 
 
@@ -122,6 +124,8 @@ def add_reassign_cmd(request, vslot):
                 cmd.FailRealmIdx = vslot.FailRealmIdx
                 cmd.FailDomainIdx = vslot.FailDomainIdx
                 cmd.VDiskIdx = vslot.VDiskIdx
+                cmd.PreferLessOccupiedRack = args.prefer_less_occupied_rack
+                cmd.WithAttentionToReplication = args.with_attention_to_replication
 
             request = common.kikimr_bsconfig.TConfigRequest(Rollback=True)
             index = len(request.Command)
diff --git a/ydb/core/mind/bscontroller/cmds_storage_pool.cpp b/ydb/core/mind/bscontroller/cmds_storage_pool.cpp
@@ -382,6 +382,14 @@ namespace NKikimr::NBsController {
             Fit.OnlyToLessOccupiedPDisk = true;
         }
 
+        if (cmd.GetPreferLessOccupiedRack()) {
+            Fit.PreferLessOccupiedRack = true;
+        }
+
+        if (cmd.GetWithAttentionToReplication()) {
+            Fit.WithAttentionToReplication = true;
+        }
+
         Fit.PoolsAndGroups.emplace(group->StoragePoolId, group->ID);
     }
 
diff --git a/ydb/core/mind/bscontroller/config.h b/ydb/core/mind/bscontroller/config.h
@@ -13,6 +13,8 @@ namespace NKikimr {
             std::set<TBoxId> Boxes;
             std::multiset<std::tuple<TBoxStoragePoolId, std::optional<TGroupId>>> PoolsAndGroups; // nullopt goes first and means 'cover all groups in the pool'
             bool OnlyToLessOccupiedPDisk = false;
+            bool PreferLessOccupiedRack = false;
+            bool WithAttentionToReplication = false;
 
             operator bool() const {
                 return !Boxes.empty() || !PoolsAndGroups.empty();
diff --git a/ydb/core/mind/bscontroller/config_fit_groups.cpp b/ydb/core/mind/bscontroller/config_fit_groups.cpp
@@ -450,24 +450,25 @@ namespace NKikimr {
                 State.CheckConsistency();
             }
 
-        private:            
+        private:
             template<typename T>
             std::invoke_result_t<T, TGroupGeometryInfo&, TGroupMapper&, TGroupId, TGroupMapper::TGroupDefinition&, TGroupMapper::TGroupConstraintsDefinition&,
                     const THashMap<TVDiskIdShort, TPDiskId>&, TGroupMapper::TForbiddenPDisks, i64> AllocateOrSanitizeGroup(
                     TGroupId groupId, TGroupMapper::TGroupDefinition& group, TGroupMapper::TGroupConstraintsDefinition& constraints,
                     const THashMap<TVDiskIdShort, TPDiskId>& replacedDisks, TGroupMapper::TForbiddenPDisks forbid,
                     i64 requiredSpace, bool addExistingDisks, T&& func) {
                 if (!Mapper) {
-                    Mapper.emplace(Geometry, StoragePool.RandomizeGroupMapping);
+                    Mapper.emplace(Geometry, StoragePool.RandomizeGroupMapping, State.Fit.PreferLessOccupiedRack, State.Fit.WithAttentionToReplication);
                     PopulateGroupMapper();
                 }
+                TPDiskSlotTracker& pdiskSlotTracker= Mapper->GetPDiskSlotTracker();
                 TStackVec<TPDiskId, 32> removeQ;
                 if (addExistingDisks) {
                     for (const auto& realm : group) {
                         for (const auto& domain : realm) {
                             for (const TPDiskId id : domain) {
                                 if (id != TPDiskId()) {
-                                    if (auto *info = State.PDisks.Find(id); info && RegisterPDisk(id, *info, false, "X")) {
+                                    if (auto *info = State.PDisks.Find(id); info && RegisterPDisk(id, *info, false, pdiskSlotTracker, "X")) {
                                         removeQ.push_back(id);
                                     }
                                 }
@@ -476,14 +477,14 @@ namespace NKikimr {
                     }
                 }
                 struct TUnregister {
-                    TGroupMapper& Mapper;
+                    TBlobStorageController::TGroupFitter& Self;
                     TStackVec<TPDiskId, 32>& RemoveQ;
                     ~TUnregister() {
                         for (const TPDiskId pdiskId : RemoveQ) {
-                            Mapper.UnregisterPDisk(pdiskId);
+                            Self.UnregisterPDisk(pdiskId);
                         }
                     }
-                } unregister{*Mapper, removeQ};
+                } unregister{*this, removeQ};
                 return std::invoke(func, Geometry, *Mapper, groupId, group, constraints, replacedDisks, std::move(forbid), requiredSpace);
             }
 
@@ -500,6 +501,22 @@ namespace NKikimr {
             void PopulateGroupMapper() {
                 const TBoxId boxId = std::get<0>(StoragePoolId);
 
+                TPDiskSlotTracker pdiskSlotTracker;
+
+                bool populateSlotTracker = State.Fit.PreferLessOccupiedRack || State.Fit.WithAttentionToReplication;
+
+                if (populateSlotTracker) {
+                    State.VSlots.ForEach([&](const TVSlotId& id, const TVSlotInfo& info) {
+                        if (info.IsBeingDeleted()) {
+                            return; // ignore slots being deleted
+                        }
+                        if (info.GetStatus() == NKikimrBlobStorage::EVDiskStatus::REPLICATING) {
+                            TPDiskId pdiskId = id.ComprisingPDiskId();
+                            pdiskSlotTracker.AddReplicatingVSlot(pdiskId);
+                        }
+                    });
+                }
+
                 State.PDisks.ForEach([&](const TPDiskId& id, const TPDiskInfo& info) {
                     if (info.BoxId != boxId) {
                         return; // ignore disks not from desired box
@@ -511,15 +528,17 @@ namespace NKikimr {
 
                     for (const auto& filter : StoragePool.PDiskFilters) {
                         if (filter.MatchPDisk(info)) {
-                            const bool inserted = RegisterPDisk(id, info, true);
+                            const bool inserted = RegisterPDisk(id, info, true, pdiskSlotTracker);
                             Y_ABORT_UNLESS(inserted);
                             break;
                         }
                     }
                 });
+
+                Mapper->SetPDiskSlotTracker(std::move(pdiskSlotTracker));
             }
 
-            bool RegisterPDisk(TPDiskId id, const TPDiskInfo& info, bool usable, TString whyUnusable = {}) {
+            bool RegisterPDisk(TPDiskId id, const TPDiskInfo& info, bool usable, TPDiskSlotTracker& pdiskSlotTracker, TString whyUnusable = {}) {
                 // calculate number of used slots on this PDisk, also counting the static ones
                 ui32 numSlots = info.NumActiveSlots + info.StaticSlotUsage;
 
@@ -574,19 +593,42 @@ namespace NKikimr {
                     whyUnusable.append('D');
                 }
 
+                ui32 maxSlots = info.ExpectedSlotCount;
+                auto location = State.HostRecords->GetLocation(id.NodeId);
+
                 // register PDisk in the mapper
-                return Mapper->RegisterPDisk({
+                bool registered = Mapper->RegisterPDisk({
                     .PDiskId = id,
-                    .Location = State.HostRecords->GetLocation(id.NodeId),
+                    .Location = location,
                     .Usable = usable,
                     .NumSlots = numSlots,
-                    .MaxSlots = info.ExpectedSlotCount,
+                    .MaxSlots = maxSlots,
                     .Groups = std::move(groups),
                     .SpaceAvailable = availableSpace,
                     .Operational = info.Operational,
                     .Decommitted = info.Decommitted(),
                     .WhyUnusable = std::move(whyUnusable),
                 });
+
+                bool populateSlotTracker = State.Fit.PreferLessOccupiedRack || State.Fit.WithAttentionToReplication;
+
+                if (registered && populateSlotTracker) {
+                    i32 freeSlots = i32(maxSlots) - numSlots;
+                    pdiskSlotTracker.AddFreeSlotsForRack(location.GetRackId(), freeSlots);
+                }
+
+                return registered;
+            }
+
+            void UnregisterPDisk(TPDiskId id) {
+                TGroupMapper::TPDiskRecord rec = Mapper->UnregisterPDisk(id);
+
+                bool populatedSlotTracker = State.Fit.PreferLessOccupiedRack || State.Fit.WithAttentionToReplication;
+
+                if (populatedSlotTracker) {
+                    i32 freeSlots = i32(rec.MaxSlots) - rec.NumSlots;
+                    Mapper->GetPDiskSlotTracker().AddFreeSlotsForRack(rec.Location.GetRackId(), -freeSlots);
+                }
             }
 
             std::map<TVDiskIdShort, TVSlotInfo*> CreateVSlotsForGroup(TGroupInfo *groupInfo,
diff --git a/ydb/core/mind/bscontroller/group_mapper.cpp b/ydb/core/mind/bscontroller/group_mapper.cpp
@@ -250,9 +250,50 @@ namespace NKikimr::NBsController {
             }
 
             bool DiskIsBetter(const TPDiskInfo& pretender, const TPDiskInfo& king) const {
+                if (Self.PreferLessOccupiedRack) {
+                    Y_ABORT_UNLESS(Self.PDiskSlotTracker.has_value());
+
+                    auto& pdiskSlotTracker = *Self.PDiskSlotTracker;
+
+                    // Compare by number of free slots in PDisk's rack.
+                    i32 freeSlotsPretender = pdiskSlotTracker.GetFreeSlotsOnRack(pretender.Location.GetRackId());
+                    i32 freeSlotsKing = pdiskSlotTracker.GetFreeSlotsOnRack(king.Location.GetRackId());
+
+                    if (freeSlotsPretender != freeSlotsKing) {
+                        return freeSlotsPretender > freeSlotsKing;
+                    }
+                }
+
+                if (Self.WithAttentionToReplication) {
+                    auto pretenderNode = pretender.PDiskId.NodeId;
+                    auto kingNode = king.PDiskId.NodeId;
+
+                    Y_ABORT_UNLESS(Self.PDiskSlotTracker.has_value());
+
+                    auto& pdiskSlotTracker = *Self.PDiskSlotTracker;
+
+                    // Compare by number of replicating VDisks on the PDisk's node.
+                    auto pretenderNodeRepls = pdiskSlotTracker.GetReplicatingVDisksOnNode(pretenderNode);
+                    auto kingNodeRepls = pdiskSlotTracker.GetReplicatingVDisksOnNode(kingNode);
+
+                    if (pretenderNodeRepls != kingNodeRepls) {
+                        return pretenderNodeRepls < kingNodeRepls;
+                    }
+
+                    // Compare by number of replicating VDisks on the PDisk.
+                    auto pretenderPDiskRepls = pdiskSlotTracker.GetReplicatingVDisksOnPDisk(pretender.PDiskId);
+                    auto kingPDiskRepls = pdiskSlotTracker.GetReplicatingVDisksOnPDisk(king.PDiskId);
+
+                    if (pretenderPDiskRepls != kingPDiskRepls) {
+                        return pretenderPDiskRepls < kingPDiskRepls;
+                    }
+                }
+
                 if (pretender.FreeSlots() != king.FreeSlots()) {
                     return pretender.FreeSlots() > king.FreeSlots();
-                } else if (GivesLocalityBoost(pretender, king) || BetterQuotaMatch(pretender, king)) {
+                }
+
+                if (GivesLocalityBoost(pretender, king) || BetterQuotaMatch(pretender, king)) {
                     return true;
                 } else {
                     if (pretender.NumDomainMatchingDisks != king.NumDomainMatchingDisks) {
@@ -849,13 +890,26 @@ namespace NKikimr::NBsController {
         TPDisks PDisks;
         TPDiskByPosition PDiskByPosition;
         bool Dirty = false;
+        bool PreferLessOccupiedRack;
+        bool WithAttentionToReplication;
+        std::optional<TPDiskSlotTracker> PDiskSlotTracker;
 
     public:
-        TImpl(TGroupGeometryInfo geom, bool randomize)
+        TImpl(TGroupGeometryInfo geom, bool randomize, bool preferLessOccupiedRack, bool withAttentionToReplication)
             : Geom(std::move(geom))
             , Randomize(randomize)
+            , PreferLessOccupiedRack(preferLessOccupiedRack)
+            , WithAttentionToReplication(withAttentionToReplication)
         {}
 
+        void SetPDiskSlotTracker(TPDiskSlotTracker&& tracker) {
+            PDiskSlotTracker = std::move(tracker);
+        }
+
+        TPDiskSlotTracker& GetPDiskSlotTracker() {
+            return PDiskSlotTracker.value();
+        }
+
         bool RegisterPDisk(const TPDiskRecord& pdisk) {
             // calculate disk position
             const TPDiskLayoutPosition p(DomainMapper, pdisk.Location, pdisk.PDiskId, Geom);
@@ -872,13 +926,15 @@ namespace NKikimr::NBsController {
             return inserted;
         }
 
-        void UnregisterPDisk(TPDiskId pdiskId) {
+        TPDiskRecord UnregisterPDisk(TPDiskId pdiskId) {
             const auto it = PDisks.find(pdiskId);
             Y_ABORT_UNLESS(it != PDisks.end());
             auto x = std::remove(PDiskByPosition.begin(), PDiskByPosition.end(), std::make_pair(it->second.Position, &it->second));
             Y_ABORT_UNLESS(x + 1 == PDiskByPosition.end());
             PDiskByPosition.pop_back();
+            TPDiskRecord ret = it->second;
             PDisks.erase(it);
+            return ret;
         }
 
         void AdjustSpaceAvailable(TPDiskId pdiskId, i64 increment) {
@@ -1126,17 +1182,25 @@ namespace NKikimr::NBsController {
         }
     };
 
-    TGroupMapper::TGroupMapper(TGroupGeometryInfo geom, bool randomize)
-        : Impl(new TImpl(std::move(geom), randomize))
+    TGroupMapper::TGroupMapper(TGroupGeometryInfo geom, bool randomize, bool preferLessOccupiedRack, bool withAttentionToReplication)
+        : Impl(new TImpl(std::move(geom), randomize, preferLessOccupiedRack, withAttentionToReplication))
     {}
 
     TGroupMapper::~TGroupMapper() = default;
 
+    void TGroupMapper::SetPDiskSlotTracker(TPDiskSlotTracker&& tracker) {
+        Impl->SetPDiskSlotTracker(std::move(tracker));
+    }
+
+    TPDiskSlotTracker& TGroupMapper::GetPDiskSlotTracker() {
+        return Impl->GetPDiskSlotTracker();
+    }
+
     bool TGroupMapper::RegisterPDisk(const TPDiskRecord& pdisk) {
         return Impl->RegisterPDisk(pdisk);
     }
 
-    void TGroupMapper::UnregisterPDisk(TPDiskId pdiskId) {
+    TGroupMapper::TPDiskRecord TGroupMapper::UnregisterPDisk(TPDiskId pdiskId) {
         return Impl->UnregisterPDisk(pdiskId);
     }
 
diff --git a/ydb/core/mind/bscontroller/group_mapper.h b/ydb/core/mind/bscontroller/group_mapper.h
@@ -8,6 +8,42 @@ namespace NKikimr {
 
         class TGroupGeometryInfo;
 
+        class TPDiskSlotTracker {
+            absl::flat_hash_map<ui32, ui16> ReplicatingVDisksByNode;
+            absl::flat_hash_map<TPDiskId, ui8> ReplicatingVDisksByPDisk;
+            absl::flat_hash_map<TString, i32> FreeSlotsPerRack;
+        public:
+            ui16 GetReplicatingVDisksOnNode(ui32 nodeId) const {
+                if (const auto it = ReplicatingVDisksByNode.find(nodeId); it != ReplicatingVDisksByNode.end()) {
+                    return it->second;
+                }
+                return 0;
+            }
+
+            ui8 GetReplicatingVDisksOnPDisk(TPDiskId pdiskId) const {
+                if (const auto it = ReplicatingVDisksByPDisk.find(pdiskId); it != ReplicatingVDisksByPDisk.end()) {
+                    return it->second;
+                }
+                return 0;
+            }
+
+            i32 GetFreeSlotsOnRack(const TString& rack) const {
+                if (const auto it = FreeSlotsPerRack.find(rack); it != FreeSlotsPerRack.end()) {
+                    return it->second;
+                }
+                return 0;
+            }
+
+            void AddReplicatingVSlot(TPDiskId pdiskId) {
+                ++ReplicatingVDisksByNode[pdiskId.NodeId];
+                ++ReplicatingVDisksByPDisk[pdiskId];
+            }
+            
+            void AddFreeSlotsForRack(const TString& rack, i32 freeSlots) {
+                FreeSlotsPerRack[rack] += freeSlots;
+            }
+        };
+
         // TGroupMapper is a helper class used to create groups from a set of PDisks with their respective locations
         // over physical hardware
         class TGroupMapper {
@@ -65,14 +101,18 @@ namespace NKikimr {
             };
 
         public:
-            TGroupMapper(TGroupGeometryInfo geom, bool randomize = false);
+            TGroupMapper(TGroupGeometryInfo geom, bool randomize = false, bool preferLessOccupiedRack = false, bool withAttentionToReplication = false);
             ~TGroupMapper();
 
+            void SetPDiskSlotTracker(TPDiskSlotTracker&& state);
+
+            TPDiskSlotTracker& GetPDiskSlotTracker();
+
             // Register PDisk inside mapper to use it in subsequent map operations
             bool RegisterPDisk(const TPDiskRecord& pdisk);
 
             // Remove PDisk from the table.
-            void UnregisterPDisk(TPDiskId pdiskId);
+            TPDiskRecord UnregisterPDisk(TPDiskId pdiskId);
 
             // Adjust VDisk space quota.
             void AdjustSpaceAvailable(TPDiskId pdiskId, i64 increment);
diff --git a/ydb/core/mind/bscontroller/group_mapper_ut.cpp b/ydb/core/mind/bscontroller/group_mapper_ut.cpp
diff --git a/ydb/core/protos/blobstorage_config.proto b/ydb/core/protos/blobstorage_config.proto

Original file line number	Diff line number	Diff line change
`@@ -382,6 +382,14 @@ namespace NKikimr::NBsController {`
`382`	`382`	`Fit.OnlyToLessOccupiedPDisk = true;`
`383`	`383`	`}`
`384`	`384`
	`385`	`+ if (cmd.GetPreferLessOccupiedRack()) {`
	`386`	`+ Fit.PreferLessOccupiedRack = true;`
	`387`	`+ }`
	`388`	`+`
	`389`	`+ if (cmd.GetWithAttentionToReplication()) {`
	`390`	`+ Fit.WithAttentionToReplication = true;`
	`391`	`+ }`
	`392`	`+`
`385`	`393`	`Fit.PoolsAndGroups.emplace(group->StoragePoolId, group->ID);`
`386`	`394`	`}`
`387`	`395`