Add UT

serbel324 · serbel324 · commit d44df6c54eb1 · 2025-04-03T13:47:52.000Z
diff --git a/ydb/core/blobstorage/pdisk/mock/pdisk_mock.cpp b/ydb/core/blobstorage/pdisk/mock/pdisk_mock.cpp
@@ -3,6 +3,7 @@
 #include <ydb/core/util/stlog.h>
 #include <ydb/core/util/interval_set.h>
 
+#include <ydb/core/blobstorage/pdisk/blobstorage_pdisk_quota_record.h>
 #include <ydb/core/blobstorage/pdisk/blobstorage_pdisk_util_space_color.h>
 
 namespace NKikimr {
@@ -50,7 +51,12 @@ struct TPDiskMockState::TImpl {
     NPDisk::EDeviceType DeviceType;
     std::optional<TRcBuf> Metadata;
 
-    TImpl(ui32 nodeId, ui32 pdiskId, ui64 pdiskGuid, ui64 size, ui32 chunkSize, bool isDiskReadOnly, NPDisk::EDeviceType deviceType)
+    ESpaceColorPolicy SpaceColorPolicy;
+    std::shared_ptr<NPDisk::TQuotaRecord> ChunkSharedQuota;
+    double Occupancy = 0;
+
+    TImpl(ui32 nodeId, ui32 pdiskId, ui64 pdiskGuid, ui64 size, ui32 chunkSize, bool isDiskReadOnly, NPDisk::EDeviceType deviceType,
+            ESpaceColorPolicy spaceColorPolicy)
         : NodeId(nodeId)
         , PDiskId(pdiskId)
         , PDiskGuid(pdiskGuid)
@@ -62,7 +68,20 @@ struct TPDiskMockState::TImpl {
         , NextFreeChunk(1)
         , StatusFlags(NPDisk::TStatusFlags{})
         , DeviceType(deviceType)
-    {}
+        , SpaceColorPolicy(spaceColorPolicy)
+    {
+        switch (SpaceColorPolicy) {
+            case ESpaceColorPolicy::SharedQuota: {
+                ChunkSharedQuota = std::make_shared<NPDisk::TQuotaRecord>();
+                // 13% for CYAN is default value in prod
+                ChunkSharedQuota->ForceHardLimit(TotalChunks, NPDisk::TColorLimits::MakeChunkLimits(130));
+                break;
+            }
+            case ESpaceColorPolicy::None: 
+            default:
+                break;
+        }
+    }
 
     TImpl(const TImpl&) = default;
 
@@ -76,6 +95,25 @@ struct TPDiskMockState::TImpl {
         }
     }
 
+    void UpdateStatusFlags(i64 chunksDelta) {
+        switch (SpaceColorPolicy) {
+            case ESpaceColorPolicy::SharedQuota: {
+                NKikimrBlobStorage::TPDiskSpaceColor::E newColor =
+                        ChunkSharedQuota->EstimateSpaceColor(chunksDelta, &Occupancy);
+                if (chunksDelta < 0) {
+                    ChunkSharedQuota->Release(-chunksDelta);
+                } else if (chunksDelta > 0) {
+                    ChunkSharedQuota->ForceAllocate(chunksDelta);
+                }
+                SetStatusFlags(SpaceColorToStatusFlag(newColor));
+                break;
+            }
+            case ESpaceColorPolicy::None: 
+            default:
+                break;
+        }
+    }
+
     ui32 AllocateChunk(TOwner& to) {
         ui32 chunkIdx = TotalChunks;
 
@@ -89,6 +127,8 @@ struct TPDiskMockState::TImpl {
         }
 
         Y_ABORT_UNLESS(chunkIdx != TotalChunks);
+
+        UpdateStatusFlags(1);
         return chunkIdx;
     }
 
@@ -172,6 +212,8 @@ struct TPDiskMockState::TImpl {
         for (const TChunkIdx chunkIdx : owner.ReservedChunks) {
             owner.ChunkData.erase(chunkIdx);
         }
+
+        UpdateStatusFlags(-(i64)owner.ReservedChunks.size());
         FreeChunks.merge(owner.ReservedChunks);
         AdjustFreeChunks();
     }
@@ -189,6 +231,8 @@ struct TPDiskMockState::TImpl {
         owner.ChunkData.erase(chunkIdx);
         const bool inserted = FreeChunks.insert(chunkIdx).second;
         Y_ABORT_UNLESS(inserted);
+
+        UpdateStatusFlags(-1);
         AdjustFreeChunks();
     }
 
@@ -286,8 +330,9 @@ struct TPDiskMockState::TImpl {
 };
 
 TPDiskMockState::TPDiskMockState(ui32 nodeId, ui32 pdiskId, ui64 pdiskGuid, ui64 size, ui32 chunkSize, bool isDiskReadOnly,
-        NPDisk::EDeviceType deviceType)
-    : TPDiskMockState(std::make_unique<TImpl>(nodeId, pdiskId, pdiskGuid, size, chunkSize, isDiskReadOnly, deviceType))
+        NPDisk::EDeviceType deviceType, ESpaceColorPolicy spaceColorPolicy)
+    : TPDiskMockState(std::make_unique<TImpl>(nodeId, pdiskId, pdiskGuid, size, chunkSize, isDiskReadOnly, deviceType,
+            spaceColorPolicy))
 {}
 
 TPDiskMockState::TPDiskMockState(std::unique_ptr<TImpl>&& impl)
@@ -475,6 +520,8 @@ class TPDiskMockActor : public TActorBootstrapped<TPDiskMockActor> {
                 break;
             } else {
                 owner.Slain = true;
+                Impl.UpdateStatusFlags(-(i64)owner.ReservedChunks.size());
+                Impl.UpdateStatusFlags(-(i64)owner.CommittedChunks.size());
                 Impl.FreeChunks.merge(owner.ReservedChunks);
                 Impl.FreeChunks.merge(owner.CommittedChunks);
                 Impl.AdjustFreeChunks();
@@ -676,6 +723,7 @@ class TPDiskMockActor : public TActorBootstrapped<TPDiskMockActor> {
                 for (ui32 i = 0; i < msg->SizeChunks; ++i) {
                     res->ChunkIds.push_back(Impl.AllocateChunk(*owner));
                 }
+                res->StatusFlags = GetStatusFlags();
                 PDISK_MOCK_LOG(DEBUG, PDM10, "sending TEvChunkReserveResult", (Msg, res->ToString()));
             }
         }
@@ -848,7 +896,7 @@ class TPDiskMockActor : public TActorBootstrapped<TPDiskMockActor> {
         auto res = std::make_unique<NPDisk::TEvCheckSpaceResult>(NKikimrProto::OK, GetStatusFlags(),
             Impl.GetNumFreeChunks(), Impl.TotalChunks, Impl.TotalChunks - Impl.GetNumFreeChunks(),
             Impl.Owners.size(), TString());
-        res->Occupancy = (double)res->UsedChunks / res->TotalChunks;
+        res->Occupancy = GetOccupancy();
         Impl.FindOwner(msg, res); // to ensure correct owner/round
         Send(ev->Sender, res.release());
     }
@@ -886,6 +934,12 @@ class TPDiskMockActor : public TActorBootstrapped<TPDiskMockActor> {
         return Impl.StatusFlags;
     }
 
+    double GetOccupancy() {
+        return (Impl.Occupancy == 0)
+            ? ((double)(Impl.TotalChunks - Impl.GetNumFreeChunks()) / Impl.TotalChunks)
+            : Impl.Occupancy;
+    }
+
     void ErrorHandle(NPDisk::TEvYardInit::TPtr &ev) {
         Send(ev->Sender, new NPDisk::TEvYardInitResult(NKikimrProto::CORRUPTED, State->GetStateErrorReason()));
     }
diff --git a/ydb/core/blobstorage/pdisk/mock/pdisk_mock.h b/ydb/core/blobstorage/pdisk/mock/pdisk_mock.h
@@ -17,12 +17,19 @@ namespace NKikimr {
         std::unique_ptr<TImpl> Impl;
         friend class TPDiskMockActor;
 
+    public:
+        enum class ESpaceColorPolicy {
+            None = 0,
+            SharedQuota,
+        };
+
     public:
         using TPtr = TIntrusivePtr<TPDiskMockState>;
 
     public:
         TPDiskMockState(ui32 nodeId, ui32 pdiskId, ui64 pdiskGuid, ui64 size, ui32 chunkSize = 128 << 20,
-                bool isDiskReadOnly = false, NPDisk::EDeviceType deviceType = NPDisk::EDeviceType::DEVICE_TYPE_NVME);
+                bool isDiskReadOnly = false, NPDisk::EDeviceType deviceType = NPDisk::EDeviceType::DEVICE_TYPE_NVME,
+                ESpaceColorPolicy spaceColorPolicy = ESpaceColorPolicy::None);
         TPDiskMockState(std::unique_ptr<TImpl>&& impl);
         ~TPDiskMockState();
 
diff --git a/ydb/core/blobstorage/ut_blobstorage/lib/env.h b/ydb/core/blobstorage/ut_blobstorage/lib/env.h
@@ -56,6 +56,8 @@ struct TEnvironmentSetup {
         const ui32 MaxNumOfSlowDisks = 2;
         const ui32 ReplMaxQuantumBytes = 0;
         const ui32 ReplMaxDonorNotReadyCount = 0;
+        const ui64 PDiskSize = 10_TB;
+        const bool TrackSharedQuotaInPDiskMock = false;
     };
 
     const TSettings Settings;
@@ -73,8 +75,11 @@ struct TEnvironmentSetup {
             const auto key = std::make_pair(nodeId, pdiskId);
             TIntrusivePtr<TPDiskMockState>& state = Env.PDiskMockStates[key];
             if (!state) {
-                state.Reset(new TPDiskMockState(nodeId, pdiskId, cfg->PDiskGuid, ui64(10) << 40, cfg->ChunkSize,
-                        cfg->ReadOnly, Env.Settings.DiskType));
+                TPDiskMockState::ESpaceColorPolicy spaceColorPolicy = Env.Settings.TrackSharedQuotaInPDiskMock
+                        ? TPDiskMockState::ESpaceColorPolicy::SharedQuota
+                        : TPDiskMockState::ESpaceColorPolicy::None;
+                state.Reset(new TPDiskMockState(nodeId, pdiskId, cfg->PDiskGuid, Env.Settings.PDiskSize, cfg->ChunkSize,
+                        cfg->ReadOnly, Env.Settings.DiskType, spaceColorPolicy));
             }
             const TActorId& actorId = ctx.Register(CreatePDiskMockActor(state), TMailboxType::HTSwap, poolId);
             const TActorId& serviceId = MakeBlobStoragePDiskID(nodeId, pdiskId);
diff --git a/ydb/core/blobstorage/ut_blobstorage/replication.cpp b/ydb/core/blobstorage/ut_blobstorage/replication.cpp
@@ -1,10 +1,16 @@
 #include <ydb/core/blobstorage/ut_blobstorage/lib/env.h>
 #include <ydb/core/blobstorage/ut_blobstorage/lib/common.h>
 #include <ydb/core/blobstorage/vdisk/hulldb/base/hullbase_barrier.h>
+
 #include <util/system/info.h>
+#include <util/stream/null.h>
+
+#include "ut_helpers.h"
 
 #define SINGLE_THREAD 1
 
+#define Ctest Cnull
+
 enum class EState {
     OK,
     FORMAT,
@@ -350,3 +356,118 @@ Y_UNIT_TEST_SUITE(Replication) {
         DoTestCase(TBlobStorageGroupType::ErasureMirror3dc, {E::OK, E::FORMAT, E::OK, E::OK, E::OFFLINE, E::OK, E::OK, E::OFFLINE, E::OK}, true);
     }
 }
+
+struct TTestCtx : public TTestCtxBase {
+public:
+    TTestCtx(TBlobStorageGroupType erasure, ui64 pdiskSize, ui32 groupsCount = 3)
+        : TTestCtxBase(TEnvironmentSetup::TSettings{
+            .NodeCount = erasure.BlobSubgroupSize(),
+            .Erasure = erasure,
+            .PDiskSize = pdiskSize,
+            .TrackSharedQuotaInPDiskMock = true,
+        })
+        , PDiskSize(pdiskSize)
+        , GroupsCount(groupsCount)
+    {}
+
+    void Initialize() override {
+        Env->CreateBoxAndPool(GroupsCount, GroupsCount);
+        Env->Sim(TDuration::Minutes(1));
+
+        BaseConfig = Env->FetchBaseConfig();
+        UNIT_ASSERT_VALUES_EQUAL(BaseConfig.GroupSize(), GroupsCount);
+        for (const auto& group : BaseConfig.GetGroup()) {
+            Groups.push_back(group.GetGroupId());
+        }
+
+        AllocateEdgeActor();
+        for (const ui32 groupId : Groups) {
+            GetGroupStatus(groupId);
+        }
+    }
+
+public:
+    ui64 PDiskSize;
+    ui32 GroupsCount;
+    std::vector<ui32> Groups;
+};
+
+Y_UNIT_TEST_SUITE(ReplicationSpace) {
+    void TestSpace() {
+        TBlobStorageGroupType erasure = TBlobStorageGroupType::ErasureMirror3dc;
+        ui64 diskSize = 3_GB;
+        TTestCtx ctx(erasure, diskSize);
+        ctx.Initialize();
+
+        ui32 chosenNodeId = 0;
+        ui32 chosenPDiskId = 0;
+
+        for (const auto& vslot : ctx.BaseConfig.GetVSlot()) {
+            if (vslot.GetGroupId() == ctx.Groups[0]) {
+                chosenNodeId = vslot.GetVSlotId().GetNodeId();
+                chosenPDiskId = vslot.GetVSlotId().GetPDiskId();
+                break;
+            }
+        }
+
+        UNIT_ASSERT(chosenNodeId != 0);
+
+        // disable self-heal
+        ctx.Env->UpdateSettings(false, true, false);
+
+        ui64 perDiskDataSize = diskSize * 0.6;
+        ui64 dataSize = perDiskDataSize * ctx.NodeCount / 3;
+
+        for (ui32 groupId : ctx.Groups) {
+            ctx.WriteCompressedData(groupId, dataSize, 8_MB);
+            Ctest << "DATA WRITTEN FOR GROUP " << groupId << Endl;
+        }
+
+        Ctest << "REASSIGN DISK" << Endl;
+
+        // find vdisk from another node and move it to the chosen
+        {
+            NKikimrBlobStorage::TConfigRequest request;
+            ui32 groupToMove = 1;
+            for (const auto& vslot : ctx.BaseConfig.GetVSlot()) {
+                if (vslot.GetGroupId() != ctx.Groups[groupToMove]) {
+                    continue;
+                }
+                const ui32 nodeId = vslot.GetVSlotId().GetNodeId();
+                if (nodeId != chosenNodeId) {
+                    NKikimrBlobStorage::TReassignGroupDisk* cmd = request.AddCommand()->MutableReassignGroupDisk();
+                    cmd->SetGroupId(vslot.GetGroupId());
+                    cmd->SetGroupGeneration(vslot.GetGroupGeneration());
+                    cmd->SetFailRealmIdx(vslot.GetFailRealmIdx());
+                    cmd->SetFailDomainIdx(vslot.GetFailDomainIdx());
+                    cmd->SetVDiskIdx(vslot.GetVDiskIdx());
+                    auto* target = cmd->MutableTargetPDiskId();
+                    target->SetNodeId(chosenNodeId);
+                    target->SetPDiskId(chosenPDiskId);
+                    if (++groupToMove == ctx.GroupsCount) {
+                        break;
+                    }
+                }
+            }
+            auto res = ctx.Env->Invoke(request);
+            UNIT_ASSERT_C(res.GetSuccess(), res.GetErrorDescription());
+            UNIT_ASSERT_C(res.GetStatus(0).GetSuccess(), res.GetStatus(0).GetErrorDescription());
+        }
+        Ctest << "DISK REASSIGNED" << Endl;
+
+        ctx.Env->Sim(TDuration::Minutes(3600));
+
+        ctx.Env->Runtime->WrapInActorContext(ctx.Edge, [&] {
+            SendToBSProxy(ctx.Edge, ctx.GroupId, new TEvBlobStorage::TEvStatus(TInstant::Max()));
+        });
+        auto res = ctx.Env->WaitForEdgeActorEvent<TEvBlobStorage::TEvStatusResult>(ctx.Edge,
+                false, TInstant::Max());
+        UNIT_ASSERT(res->Get()->Status == NKikimrProto::OK);
+        Ctest << "FLAGS " << res->Get()->ToString() << Endl;
+        UNIT_ASSERT(!res->Get()->StatusFlags.Check(NKikimrBlobStorage::StatusDiskSpaceRed));
+    }
+
+    Y_UNIT_TEST(Mirror3dc) {
+        TestSpace();
+    }
+}
diff --git a/ydb/core/blobstorage/ut_blobstorage/ut_helpers.h b/ydb/core/blobstorage/ut_blobstorage/ut_helpers.h
