provide a ranking function via an allocation query

sharpeye · sharpeye · commit cbb8b9d659d3 · 2025-10-27T16:34:00.000+01:00
diff --git a/cloud/blockstore/libs/storage/disk_registry/model/device_list.cpp b/cloud/blockstore/libs/storage/disk_registry/model/device_list.cpp
@@ -275,14 +275,11 @@ NProto::TDeviceConfig TDeviceList::AllocateDevice(
     const TDiskId& diskId,
     const TAllocationQuery& query)
 {
-    for (auto& kv: NodeDevices) {
-        if (!query.NodeIds.empty() && !query.NodeIds.contains(kv.first)) {
+    for (auto& [nodeId, nodeDevices]: NodeDevices) {
+        if (!query.NodeIds.empty() && !query.NodeIds.contains(nodeId)) {
             continue;
         }
 
-        const ui32 nodeId = kv.first;
-        auto& nodeDevices = kv.second;
-
         const auto& currentRack = nodeDevices.Rack;
         auto& devices = nodeDevices.FreeDevices;
 
@@ -428,10 +425,6 @@ void TDeviceList::MarkDeviceAllocated(const TDiskId& diskId, const TDeviceId& id
     AllocatedDevices.emplace(id, diskId);
 }
 
-// returns a list of racks sorted by preference and then by occupied space ASC
-// then by free space DESC
-// the nodes in each rack are sorted by occupied space ASC then by free space
-// DESC
 auto TDeviceList::SelectRacks(
     const TAllocationQuery& query,
     const TString& poolName) const -> TVector<TRack>
@@ -446,7 +439,6 @@ auto TDeviceList::SelectRacks(
         auto& rack = racks[currentRack];
         rack.Id = currentRack;
         rack.Nodes.push_back({nodeId, 0, 0});
-        rack.Preferred = query.PreferredRacks.contains(currentRack);
     };
 
     if (!query.NodeIds.empty()) {
@@ -486,59 +478,64 @@ auto TDeviceList::SelectRacks(
 
     TVector<TRack> result;
     result.reserve(racks.size());
-    for (auto& r: racks) {
-        result.push_back(std::move(r.second));
+    for (auto& [_, r]: racks) {
+        if (r.FreeSpace) {
+            r.Preferred = query.PreferredRacks.contains(r.Id);
+            result.push_back(std::move(r));
+        }
     }
 
     return result;
 }
 
+// Returns a ranked list of nodes. Sorting is performed in stages:
+// - Sort racks by preference, occupied space, free space
+// - Sort nodes within each rack by occupied space, free space
+// - Optionally apply custom NodeRankingFunc to all nodes (if set in query)
 auto TDeviceList::RankNodes(
     const TAllocationQuery& query,
-    TVector<TRack> racks) const -> TVector<TNodeInfo>
+    TVector<TRack> racks) const -> TVector<TNodeId>
 {
+    // Sort racks by preference then by occupied space ASC then by free space
+    // DESC
     Sort(
         racks,
-        [](const TRack& l, const TRack& r)
+        [&](const TRack& l, const TRack& r)
         {
             return std::tie(r.Preferred, l.OccupiedSpace, r.FreeSpace, l.Id) <
                    std::tie(l.Preferred, r.OccupiedSpace, l.FreeSpace, r.Id);
         });
 
-    TVector<TNodeInfo> nodes;
+    TVector<TNodeId> nodeIds;
     {
         size_t size = 0;
         for (const auto& rack: racks) {
             size += rack.Nodes.size();
         }
-        nodes.reserve(size);
+        nodeIds.reserve(size);
     }
 
-    for (auto& rack: racks) {
+    for (TRack& rack: racks) {
+        // Sort rack's nodes by occupied space ASC then by free space DESC
         Sort(
             rack.Nodes,
-            [](const TNodeInfo& l, const TNodeInfo& r)
+            [](const TNodeInfo& lhs, const TNodeInfo& rhs)
             {
-                return std::tie(l.OccupiedSpace, r.FreeSpace) <
-                       std::tie(r.OccupiedSpace, l.FreeSpace);
+                return std::tie(lhs.OccupiedSpace, rhs.FreeSpace) <
+                       std::tie(rhs.OccupiedSpace, lhs.FreeSpace);
             });
 
-        nodes.insert(
-            nodes.end(),
-            std::make_move_iterator(rack.Nodes.begin()),
-            std::make_move_iterator(rack.Nodes.end()));
+        for (const TNodeInfo& node: rack.Nodes) {
+            nodeIds.push_back(node.NodeId);
+        }
     }
 
-    if (query.DownrankedNodeIds) {
-        // move downranked nodes to the end of the list
-        std::stable_partition(
-            nodes.begin(),
-            nodes.end(),
-            [&](const TNodeInfo& node)
-            { return !query.DownrankedNodeIds.contains(node.NodeId); });
+    // Get the final order of nodes
+    if (query.NodeRankingFunc) {
+        query.NodeRankingFunc(nodeIds);
     }
 
-    return nodes;
+    return nodeIds;
 }
 
 TVector<TDeviceList::TDeviceRange> TDeviceList::CollectDevices(
@@ -552,8 +549,8 @@ TVector<TDeviceList::TDeviceRange> TDeviceList::CollectDevices(
     TVector<TDeviceRange> ranges;
     ui64 totalSize = query.GetTotalByteCount();
 
-    for (const auto& node: RankNodes(query, SelectRacks(query, poolName))) {
-        const auto* nodeDevices = NodeDevices.FindPtr(node.NodeId);
+    for (ui32 nodeId: RankNodes(query, SelectRacks(query, poolName))) {
+        const auto* nodeDevices = NodeDevices.FindPtr(nodeId);
         Y_ABORT_UNLESS(nodeDevices);
 
         // finding free devices belonging to this node that match our
@@ -615,7 +612,7 @@ TVector<TDeviceList::TDeviceRange> TDeviceList::CollectDevices(
             }
 
             if (deviceInfo.Range.first != it) {
-                ranges.emplace_back(node.NodeId, deviceInfo.Range.first, it);
+                ranges.emplace_back(nodeId, deviceInfo.Range.first, it);
             }
 
             if (totalSize == 0) {
diff --git a/cloud/blockstore/libs/storage/disk_registry/model/device_list.h b/cloud/blockstore/libs/storage/disk_registry/model/device_list.h
@@ -3,14 +3,18 @@
 #include "public.h"
 
 #include <cloud/blockstore/libs/storage/protos/disk.pb.h>
+
 #include <cloud/storage/core/libs/common/error.h>
 
 #include <util/datetime/base.h>
-#include <util/generic/hash_set.h>
 #include <util/generic/hash.h>
+#include <util/generic/hash_set.h>
 #include <util/generic/string.h>
 #include <util/generic/vector.h>
 
+#include <functional>
+#include <span>
+
 namespace NCloud::NBlockStore::NStorage {
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -71,7 +75,7 @@ class TDeviceList
     {
         THashSet<TString> ForbiddenRacks;
         THashSet<TString> PreferredRacks;
-        THashSet<ui32> DownrankedNodeIds;
+        std::function<void (std::span<ui32> nodeIds)> NodeRankingFunc;
 
         ui32 LogicalBlockSize = 0;
         ui64 BlockCount = 0;
@@ -181,7 +185,7 @@ class TDeviceList
         const TAllocationQuery& query,
         const TString& poolName) const;
 
-    [[nodiscard]] TVector<TNodeInfo> RankNodes(
+    [[nodiscard]] TVector<TNodeId> RankNodes(
         const TAllocationQuery& query,
         TVector<TRack> racks) const;
 
diff --git a/cloud/blockstore/libs/storage/disk_registry/model/device_list_ut.cpp b/cloud/blockstore/libs/storage/disk_registry/model/device_list_ut.cpp
@@ -508,7 +508,17 @@ Y_UNIT_TEST_SUITE(TDeviceListTest)
             return deviceList.AllocateDevices(
                 "disk",
                 {
-                    .DownrankedNodeIds = std::move(downrankedNodeIds),
+                    .NodeRankingFunc =
+                        [downrankedNodeIds](std::span<ui32> nodeIds)
+                    {
+                        auto it = std::stable_partition(
+                            nodeIds.begin(),
+                            nodeIds.end(),
+                            [&](ui32 nodeId)
+                            { return !downrankedNodeIds.contains(nodeId); });
+
+                        std::sort(it, nodeIds.end());
+                    },
                     .LogicalBlockSize = DefaultBlockSize,
                     .BlockCount = n * DefaultBlockCount,
                     .NodeIds = std::move(nodeIds),
@@ -636,21 +646,23 @@ Y_UNIT_TEST_SUITE(TDeviceListTest)
         }
 
         {
-            // rack3 *[agent5: 6]  6
-            // rack1  [agent1: 2]  2
+            // downranked nodes are sorted by nodeId, so agent1 is selected
+            // first, even though rack1 has less space than rack3.
+
+            // rack1 *[agent1: 2]  2
+            // rack3  [agent5: 6]  6
 
             auto devices = allocate(2, {}, downrankedNodes);
             UNIT_ASSERT_VALUES_EQUAL(2, devices.size());
             for (const auto& d: devices) {
-                UNIT_ASSERT_VALUES_EQUAL_C("rack3", d.GetRack(), d);
-                UNIT_ASSERT_VALUES_EQUAL(agent5.GetAgentId(), d.GetAgentId());
-                UNIT_ASSERT_VALUES_EQUAL(agent5.GetNodeId(), d.GetNodeId());
+                UNIT_ASSERT_VALUES_EQUAL_C("rack1", d.GetRack(), d);
+                UNIT_ASSERT_VALUES_EQUAL(agent1.GetAgentId(), d.GetAgentId());
+                UNIT_ASSERT_VALUES_EQUAL(agent1.GetNodeId(), d.GetNodeId());
             }
         }
 
         {
-            // rack3 *[agent5: 4]  4
-            // rack1 *[agent1: 2]  2
+            // rack3  [agent5: 6]  6
 
             auto devices = allocate(6, {}, downrankedNodes);
             UNIT_ASSERT_VALUES_EQUAL(6, devices.size());
