Report space flags and free space share to BlobDepot clients

10 files changed, 81 insertions, 2 deletions

diff --git a/ydb/core/blob_depot/agent.cpp b/ydb/core/blob_depot/agent.cpp
index e1de4260a7a..b5426c3ff45 100644
--- a/ydb/core/blob_depot/agent.cpp
+++ b/ydb/core/blob_depot/agent.cpp
@@ -60,6 +60,8 @@ namespace NKikimr::NBlobDepot {
         agent.AgentId = ev->Sender;
         agent.ConnectedNodeId = nodeId;
         agent.ExpirationTimestamp = TInstant::Max();
+        agent.LastPushedSpaceColor = SpaceMonitor->GetSpaceColor();
+        agent.LastPushedApproximateFreeSpaceShare = SpaceMonitor->GetApproximateFreeSpaceShare();
 
         if (agent.AgentInstanceId && *agent.AgentInstanceId != req.GetAgentInstanceId()) {
             ResetAgent(agent);
@@ -69,6 +71,8 @@ namespace NKikimr::NBlobDepot {
         OnAgentConnect(agent);
 
         auto [response, record] = TEvBlobDepot::MakeResponseFor(*ev, SelfId(), Executor()->Generation());
+        record->SetSpaceColor(agent.LastPushedSpaceColor);
+        record->SetApproximateFreeSpaceShare(agent.LastPushedApproximateFreeSpaceShare);
 
         for (const auto& [k, v] : ChannelKinds) {
             auto *proto = record->AddChannelKinds();
@@ -272,4 +276,21 @@ namespace NKikimr::NBlobDepot {
         }
     }
 
+    void TBlobDepot::OnSpaceColorChange(NKikimrBlobStorage::TPDiskSpaceColor::E spaceColor, float approximateFreeSpaceShare) {
+        for (auto& [nodeId, agent] : Agents) {
+            if (agent.AgentId && (agent.LastPushedSpaceColor != spaceColor ||
+                    agent.LastPushedApproximateFreeSpaceShare != approximateFreeSpaceShare)) {
+                Y_VERIFY(agent.ConnectedNodeId == nodeId);
+                const ui64 id = ++agent.LastRequestId;
+                agent.PushCallbacks.emplace(id, [](TEvBlobDepot::TEvPushNotifyResult::TPtr) {});
+                auto ev = std::make_unique<TEvBlobDepot::TEvPushNotify>();
+                ev->Record.SetSpaceColor(spaceColor);
+                ev->Record.SetApproximateFreeSpaceShare(approximateFreeSpaceShare);
+                Send(*agent.AgentId, ev.release(), 0, id);
+                agent.LastPushedSpaceColor = spaceColor;
+                agent.LastPushedApproximateFreeSpaceShare = approximateFreeSpaceShare;
+            }
+        }
+    }
+
 } // NKikimr::NBlobDepot
diff --git a/ydb/core/blob_depot/agent/agent_impl.h b/ydb/core/blob_depot/agent/agent_impl.h
index 5bf717cf544..fb5bbe6ae27 100644
--- a/ydb/core/blob_depot/agent/agent_impl.h
+++ b/ydb/core/blob_depot/agent/agent_impl.h
@@ -193,6 +193,8 @@ namespace NKikimr::NBlobDepot {
 
         ui32 BlobDepotGeneration = 0;
         std::optional<ui32> DecommitGroupId;
+        NKikimrBlobStorage::TPDiskSpaceColor::E SpaceColor = {};
+        float ApproximateFreeSpaceShare = 0.0f;
 
         void Handle(TEvTabletPipe::TEvClientConnected::TPtr ev);
         void Handle(TEvTabletPipe::TEvClientDestroyed::TPtr ev);
diff --git a/ydb/core/blob_depot/agent/comm.cpp b/ydb/core/blob_depot/agent/comm.cpp
index 4d2e6dafeee..622b993a8a8 100644
--- a/ydb/core/blob_depot/agent/comm.cpp
+++ b/ydb/core/blob_depot/agent/comm.cpp
@@ -75,6 +75,9 @@ namespace NKikimr::NBlobDepot {
         for (auto& [_, kind] : ChannelKinds) {
             IssueAllocateIdsIfNeeded(kind);
         }
+
+        SpaceColor = msg.GetSpaceColor();
+        ApproximateFreeSpaceShare = msg.GetApproximateFreeSpaceShare();
     }
 
     void TBlobDepotAgent::IssueAllocateIdsIfNeeded(TChannelKind& kind) {
@@ -187,6 +190,13 @@ namespace NKikimr::NBlobDepot {
                 (NumAvailableItemsAfter, kind.GetNumAvailableItems()));
         }
 
+        if (msg.HasSpaceColor()) {
+            SpaceColor = msg.GetSpaceColor();
+        }
+        if (msg.HasApproximateFreeSpaceShare()) {
+            ApproximateFreeSpaceShare = msg.GetApproximateFreeSpaceShare();
+        }
+
         // it is essential to send response through the pipe -- otherwise we can break order with, for example, commits:
         // this message can outrun previously sent commit and lead to data loss
         NTabletPipe::SendData(SelfId(), PipeId, response.release(), ev->Cookie);
diff --git a/ydb/core/blob_depot/agent/status.cpp b/ydb/core/blob_depot/agent/status.cpp
index dba6b1662c6..5fc6fc6b238 100644
--- a/ydb/core/blob_depot/agent/status.cpp
+++ b/ydb/core/blob_depot/agent/status.cpp
@@ -3,11 +3,11 @@
 namespace NKikimr::NBlobDepot {
 
     TStorageStatusFlags TBlobDepotAgent::GetStorageStatusFlags() const {
-        return NKikimrBlobStorage::StatusIsValid; // FIXME: implement
+        return SpaceColorToStatusFlag(SpaceColor);
     }
 
     float TBlobDepotAgent::GetApproximateFreeSpaceShare() const {
-        return 1.0; // FIXME: implement
+        return ApproximateFreeSpaceShare;
     }
 
 } // NKikimr::NBlobDepot
diff --git a/ydb/core/blob_depot/blob_depot_tablet.h b/ydb/core/blob_depot/blob_depot_tablet.h
index a453b21dfb6..368e00f7cf9 100644
--- a/ydb/core/blob_depot/blob_depot_tablet.h
+++ b/ydb/core/blob_depot/blob_depot_tablet.h
@@ -65,6 +65,9 @@ namespace NKikimr::NBlobDepot {
             THashMap<ui64, THashMap<ui8, ui32>> InvalidateStepRequests;
             THashMap<ui64, std::function<void(TEvBlobDepot::TEvPushNotifyResult::TPtr)>> PushCallbacks;
             ui64 LastRequestId = 0;
+
+            NKikimrBlobStorage::TPDiskSpaceColor::E LastPushedSpaceColor = {};
+            float LastPushedApproximateFreeSpaceShare = 0.0f;
         };
 
         THashMap<TActorId, std::optional<ui32>> PipeServerToNode;
@@ -96,6 +99,7 @@ namespace NKikimr::NBlobDepot {
         TAgent& GetAgent(ui32 nodeId);
         void ResetAgent(TAgent& agent);
         void Handle(TEvBlobDepot::TEvPushNotifyResult::TPtr ev);
+        void OnSpaceColorChange(NKikimrBlobStorage::TPDiskSpaceColor::E spaceColor, float approximateFreeSpaceShare);
 
         void ProcessRegisterAgentQ();
 
diff --git a/ydb/core/blob_depot/blocks.cpp b/ydb/core/blob_depot/blocks.cpp
index 369ae4a82ec..e09310651b9 100644
--- a/ydb/core/blob_depot/blocks.cpp
+++ b/ydb/core/blob_depot/blocks.cpp
@@ -64,6 +64,7 @@ namespace NKikimr::NBlobDepot {
         ui32 BlocksPending = 0;
         ui32 RetryCount = 0;
         THashSet<ui32> NodesWaitingForPushResult;
+        std::weak_ptr<TToken> Token;
 
     public:
         static constexpr NKikimrServices::TActivity::EType ActorActivityType() {
@@ -78,9 +79,13 @@ namespace NKikimr::NBlobDepot {
             , NodeId(nodeId)
             , IssuerGuid(issuerGuid)
             , Response(std::move(response))
+            , Token(Self->Token)
         {}
 
         bool CheckIfObsolete() {
+            if (Token.expired()) {
+                return true; // tablet is dead
+            }
             auto& block = Self->BlocksManager->Blocks[TabletId];
             if (block.BlockedGeneration == BlockedGeneration && block.IssuerGuid == IssuerGuid) {
                 return false;
diff --git a/ydb/core/blob_depot/group_metrics_exchange.cpp b/ydb/core/blob_depot/group_metrics_exchange.cpp
index 6dfae1f33ab..2b86d5edcf8 100644
--- a/ydb/core/blob_depot/group_metrics_exchange.cpp
+++ b/ydb/core/blob_depot/group_metrics_exchange.cpp
@@ -1,5 +1,6 @@
 #include "blob_depot_tablet.h"
 #include "data.h"
+#include "space_monitor.h"
 
 namespace NKikimr::NBlobDepot {
 
@@ -91,6 +92,12 @@ namespace NKikimr::NBlobDepot {
 
             params->SetAllocatedSize(Data->GetTotalStoredDataSize());
             Send(MakeBlobStorageNodeWardenID(SelfId().NodeId()), response.release());
+
+            // TODO(alexvru): use a better approach
+            const double approximateFreeSpaceShare = (double)params->GetAvailableSize() / (params->GetAvailableSize() +
+                params->GetAllocatedSize());
+
+            SpaceMonitor->SetSpaceColor(dataColor, approximateFreeSpaceShare); // the best data channel space color works for the whole depot
         }
     }
 
diff --git a/ydb/core/blob_depot/space_monitor.cpp b/ydb/core/blob_depot/space_monitor.cpp
index 2741fbfd20f..27d77a4a7b0 100644
--- a/ydb/core/blob_depot/space_monitor.cpp
+++ b/ydb/core/blob_depot/space_monitor.cpp
@@ -64,6 +64,10 @@ namespace NKikimr::NBlobDepot {
             return 0; // do not write data to this group
         }
 
+        if (!group.ApproximateFreeSpaceShare) { // not collected yet?
+            return 1;
+        }
+
         const float weight = group.ApproximateFreeSpaceShare < 0.25
             ? group.ApproximateFreeSpaceShare * 3
             : (group.ApproximateFreeSpaceShare + 2) / 3;
@@ -73,6 +77,22 @@ namespace NKikimr::NBlobDepot {
         return weight * 16'777'216.0f * (isCyan ? 0.5f : 1.0f /* cyan penalty */);
     }
 
+    void TSpaceMonitor::SetSpaceColor(NKikimrBlobStorage::TPDiskSpaceColor::E spaceColor, float approximateFreeSpaceShare) {
+        if (SpaceColor != spaceColor || ApproximateFreeSpaceShare != approximateFreeSpaceShare) {
+            SpaceColor = spaceColor;
+            ApproximateFreeSpaceShare = approximateFreeSpaceShare;
+            Self->OnSpaceColorChange(SpaceColor, ApproximateFreeSpaceShare);
+        }
+    }
+
+    NKikimrBlobStorage::TPDiskSpaceColor::E TSpaceMonitor::GetSpaceColor() const {
+        return SpaceColor;
+    }
+
+    float TSpaceMonitor::GetApproximateFreeSpaceShare() const {
+        return ApproximateFreeSpaceShare;
+    }
+
     void TBlobDepot::KickSpaceMonitor() {
         SpaceMonitor->Kick();
     }
diff --git a/ydb/core/blob_depot/space_monitor.h b/ydb/core/blob_depot/space_monitor.h
index 174cecdc896..8f4330063ed 100644
--- a/ydb/core/blob_depot/space_monitor.h
+++ b/ydb/core/blob_depot/space_monitor.h
@@ -15,6 +15,8 @@ namespace NKikimr::NBlobDepot {
         };
 
         std::unordered_map<ui32, TGroupRecord> Groups;
+        NKikimrBlobStorage::TPDiskSpaceColor::E SpaceColor = {};
+        float ApproximateFreeSpaceShare = 0.0f;
 
         friend class TBlobDepot;
 
@@ -26,6 +28,9 @@ namespace NKikimr::NBlobDepot {
         void Kick();
 
         ui64 GetGroupAllocationWeight(ui32 groupId) const;
+        void SetSpaceColor(NKikimrBlobStorage::TPDiskSpaceColor::E spaceColor, float approximateFreeSpaceShare);
+        NKikimrBlobStorage::TPDiskSpaceColor::E GetSpaceColor() const;
+        float GetApproximateFreeSpaceShare() const;
 
     private:
         void Init();
diff --git a/ydb/core/protos/blob_depot.proto b/ydb/core/protos/blob_depot.proto
index 9b6a3f2c457..270869d5094 100644
--- a/ydb/core/protos/blob_depot.proto
+++ b/ydb/core/protos/blob_depot.proto
@@ -1,5 +1,6 @@
 import "ydb/core/protos/blob_depot_config.proto";
 import "ydb/core/protos/base.proto";
+import "ydb/core/protos/blobstorage_disk_color.proto";
 
 package NKikimrBlobDepot;
 
@@ -83,6 +84,8 @@ message TEvRegisterAgentResult {
     optional uint32 Generation = 1;
     repeated TChannelKind ChannelKinds = 2;
     optional uint32 DecommitGroupId = 3;
+    optional NKikimrBlobStorage.TPDiskSpaceColor.E SpaceColor = 4;
+    optional float ApproximateFreeSpaceShare = 5;
 }
 
 message TEvAllocateIds {
@@ -120,6 +123,8 @@ message TEvPushNotify {
     }
     repeated TBlockedTablet BlockedTablets = 1;
     repeated TInvalidatedStep InvalidatedSteps = 2;
+    optional NKikimrBlobStorage.TPDiskSpaceColor.E SpaceColor = 3;
+    optional float ApproximateFreeSpaceShare = 4;
 }
 
 message TEvPushNotifyResult {