Calculate per-group balancing weight

3 files changed, 51 insertions, 23 deletions

diff --git a/ydb/core/blob_depot/agent.cpp b/ydb/core/blob_depot/agent.cpp
index 67c9ec49ba..e1de4260a7 100644
--- a/ydb/core/blob_depot/agent.cpp
+++ b/ydb/core/blob_depot/agent.cpp
@@ -1,5 +1,6 @@
 #include "blob_depot_tablet.h"
 #include "data.h"
+#include "space_monitor.h"
 
 namespace NKikimr::NBlobDepot {
 
@@ -114,32 +115,36 @@ namespace NKikimr::NBlobDepot {
 
             ui64 accum = 0;
             for (const auto& [groupId, group] : groups) {
-                //const ui64 allocatedBytes = Groups[groupId].AllocatedBytes;
-                const ui64 groupWeight = 1;
-                accum += groupWeight;
-                options.emplace_back(accum, &group);
+                if (const ui64 w = SpaceMonitor->GetGroupAllocationWeight(groupId)) {
+                    accum += w;
+                    options.emplace_back(accum, &group);
+                }
             }
 
-            THashMap<ui8, NKikimrBlobDepot::TGivenIdRange::TChannelRange*> issuedRanges;
-            for (ui32 i = 0, count = ev->Get()->Record.GetCount(); i < count; ++i) {
-                const ui64 selection = RandomNumber(accum);
-                const auto it = std::upper_bound(options.begin(), options.end(), selection,
-                    [](ui64 x, const auto& y) { return x < std::get<0>(y); });
-                const auto& [_, group] = *it;
-
-                const size_t channelIndex = RandomNumber(group->Channels.size());
-                TChannelInfo* const channel = group->Channels[channelIndex];
-
-                const ui64 value = channel->NextBlobSeqId++;
-
-                // fill in range item
-                auto& range = issuedRanges[channel->Index];
-                if (!range || range->GetEnd() != value) {
-                    range = givenIdRange->AddChannelRanges();
-                    range->SetChannel(channel->Index);
-                    range->SetBegin(value);
+            if (accum) {
+                THashMap<ui8, NKikimrBlobDepot::TGivenIdRange::TChannelRange*> issuedRanges;
+                for (ui32 i = 0, count = ev->Get()->Record.GetCount(); i < count; ++i) {
+                    const ui64 selection = RandomNumber(accum);
+                    const auto it = std::upper_bound(options.begin(), options.end(), selection,
+                        [](ui64 x, const auto& y) { return x < std::get<0>(y); });
+                    const auto& [_, group] = *it;
+
+                    const size_t channelIndex = RandomNumber(group->Channels.size());
+                    TChannelInfo* const channel = group->Channels[channelIndex];
+
+                    const ui64 value = channel->NextBlobSeqId++;
+
+                    // fill in range item
+                    auto& range = issuedRanges[channel->Index];
+                    if (!range || range->GetEnd() != value) {
+                        range = givenIdRange->AddChannelRanges();
+                        range->SetChannel(channel->Index);
+                        range->SetBegin(value);
+                    }
+                    range->SetEnd(value + 1);
                 }
-                range->SetEnd(value + 1);
+            } else {
+                Y_VERIFY_DEBUG(false); // TODO(alexvru): handle this situation somehow -- agent needs to retry this query?
             }
 
             // register issued ranges in agent and global records
diff --git a/ydb/core/blob_depot/space_monitor.cpp b/ydb/core/blob_depot/space_monitor.cpp
index 43326b5c51..2741fbfd20 100644
--- a/ydb/core/blob_depot/space_monitor.cpp
+++ b/ydb/core/blob_depot/space_monitor.cpp
@@ -52,6 +52,27 @@ namespace NKikimr::NBlobDepot {
         }
     }
 
+    ui64 TSpaceMonitor::GetGroupAllocationWeight(ui32 groupId) const {
+        const auto it = Groups.find(groupId);
+        if (it == Groups.end()) {
+            Y_VERIFY_DEBUG(false);
+            return 0;
+        }
+
+        const TGroupRecord& group = it->second;
+        if (group.StatusFlags.Check(NKikimrBlobStorage::StatusDiskSpaceLightYellowMove)) {
+            return 0; // do not write data to this group
+        }
+
+        const float weight = group.ApproximateFreeSpaceShare < 0.25
+            ? group.ApproximateFreeSpaceShare * 3
+            : (group.ApproximateFreeSpaceShare + 2) / 3;
+
+        const bool isCyan = group.StatusFlags.Check(NKikimrBlobStorage::StatusDiskSpaceCyan);
+
+        return weight * 16'777'216.0f * (isCyan ? 0.5f : 1.0f /* cyan penalty */);
+    }
+
     void TBlobDepot::KickSpaceMonitor() {
         SpaceMonitor->Kick();
     }
diff --git a/ydb/core/blob_depot/space_monitor.h b/ydb/core/blob_depot/space_monitor.h
index af83fd4722..174cecdc89 100644
--- a/ydb/core/blob_depot/space_monitor.h
+++ b/ydb/core/blob_depot/space_monitor.h
@@ -25,6 +25,8 @@ namespace NKikimr::NBlobDepot {
         void Handle(TEvBlobStorage::TEvStatusResult::TPtr ev);
         void Kick();
 
+        ui64 GetGroupAllocationWeight(ui32 groupId) const;
+
     private:
         void Init();
     };