provide a ranking function via an allocation query

Author: sharpeye

cloud/blockstore/libs/storage/disk_registry/model/device_list.cpp

@@ -428,10 +428,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 +442,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 +481,74 @@ 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) {
+            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,
-        [](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);
-        });
+    // Sort racks by occupied space ASC then by free space DESC
+    auto byRackSpace = [](const TRack& lhs, const TRack& rhs)
+    {
+        return std::tie(lhs.OccupiedSpace, rhs.FreeSpace, lhs.Id) <
+               std::tie(rhs.OccupiedSpace, lhs.FreeSpace, rhs.Id);
+    };
 
-    TVector<TNodeInfo> nodes;
+    if (query.PreferredRacks) {
+        // Move preferred racks to the start of the list and sort each part
+        auto it = std::partition(
+            racks.begin(),
+            racks.end(),
+            [&](const TRack& rack) {
+                return query.PreferredRacks.contains(rack.Id);
+            });
+        Sort(racks.begin(), it, byRackSpace);
+        Sort(it, racks.end(), byRackSpace);
+    } else {
+        Sort(racks, byRackSpace);
+    }
+
+    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 +562,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 +625,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) {

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 {
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -38,7 +42,6 @@ class TDeviceList
         TVector<TNodeInfo> Nodes;
         ui64 FreeSpace = 0;
         ui64 OccupiedSpace = 0;
-        bool Preferred = false;
     };
 
     struct TNodeDevices
@@ -71,7 +74,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 +184,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;
 

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());