NBS-7459: block PBuffer cleanup until vchunk restore completes (#45705)

4 files changed, 76 insertions, 1 deletions

diff --git a/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/fast_path_service.cpp b/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/fast_path_service.cpp
index 51485f46e84..9d6134fba7a 100644
--- a/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/fast_path_service.cpp
+++ b/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/fast_path_service.cpp
@@ -483,7 +483,9 @@ void TFastPathService::FinishPBufferCleanup()
 
     CleanupGather.Active.store(false);
 
-    if (!globalMin) {
+    if (!globalMin || *globalMin == 0) {
+        // 0 is the blocking bound: some vchunk has not finished restoring its
+        // dirty map, so its records are not accounted for yet. Skip the tick.
         return;
     }
 
diff --git a/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/vchunk.cpp b/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/vchunk.cpp
index dfd1dd04e90..edb908bbe46 100644
--- a/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/vchunk.cpp
+++ b/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/vchunk.cpp
@@ -269,6 +269,14 @@ std::optional<ui64> TVChunk::GetSafeBarrierForErase() const
 {
     Y_ABORT_UNLESS(ExecutorThreadChecker.Check());
 
+    if (!DirtyMapReady.HasValue()) {
+        // Not restored yet: this vchunk's records may still exist only in the
+        // PBuffers and are not inflight, so an empty dirty map does not mean
+        // "no constraint". Report the blocking bound so the tablet-wide
+        // cleanup skips its tick until every vchunk finishes restoring.
+        return 0;
+    }
+
     return BlocksDirtyMap.GetSafeBarrierForErase();
 }
 
diff --git a/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/vchunk.h b/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/vchunk.h
index 6760e1d2c74..49045b6e9fd 100644
--- a/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/vchunk.h
+++ b/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/vchunk.h
@@ -67,6 +67,8 @@ public:
 
     // This vchunk's contribution to the tablet-wide cleanup watermark: the
     // smallest lsn still held in PBuffers, or nullopt when nothing is inflight.
+    // Until the dirty map is restored it returns 0 (the blocking bound), so
+    // the cleanup cannot erase records that are not accounted for yet.
     // Must run on the executor thread.
     [[nodiscard]] std::optional<ui64> GetSafeBarrierForErase() const;
 
diff --git a/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/vchunk_ut.cpp b/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/vchunk_ut.cpp
index 2303ae01aca..343d423cf05 100644
--- a/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/vchunk_ut.cpp
+++ b/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/vchunk_ut.cpp
@@ -191,6 +191,69 @@ Y_UNIT_TEST_SUITE(TVChunkTest)
         vchunk->Stop().GetValue(TDuration::Seconds(10));
     }
 
+    // Until the vchunk finishes restoring its dirty map from the PBuffers,
+    // its pre-flush records exist only in the PBuffers and are not inflight.
+    // Reporting "no constraint" (nullopt) in that window is indistinguishable
+    // from an idle vchunk, so FinishPBufferCleanup would skip it and a
+    // tablet-wide barrier erase could wipe the very records the restore is
+    // about to return. An un-restored vchunk must report the blocking bound
+    // (0) instead; cleanup skips its tick on it.
+    Y_UNIT_TEST_F(
+        ShouldConstrainCleanupBarrierUntilRestoreCompletes,
+        TBaseFixture)
+    {
+        Init();
+
+        // Keep the restore pending: the vchunk stays not-ready.
+        auto neverResolvePromise =
+            NThreading::NewPromise<TDBGRestoreResponse>();
+        DirectBlockGroup->RestoreDBGPBuffersHandler =
+            [neverResolvePromise](const auto& vChunkIndex) mutable
+        {
+            Y_UNUSED(vChunkIndex);
+            return neverResolvePromise.GetFuture();
+        };
+
+        auto vchunk = std::make_shared<TVChunk>(
+            Runtime->GetActorSystem(0),
+            PartitionDirectService.get(),
+            VChunkConfig,
+            DirectBlockGroup,
+            3,   // syncRequestsBatchSize
+            DefaultVChunkSize,
+            Counters);
+        vchunk->Start();
+
+        DrainExecutor(DirectBlockGroup->GetExecutor());
+        UNIT_ASSERT_EQUAL(false, IsDirtyMapReady(*vchunk));
+
+        const auto barrierWhileRestoring =
+            GetSafeBarrierOnExecutor(DirectBlockGroup->GetExecutor(), *vchunk);
+
+        // Resolve the restore; with an empty dirty map and restore complete
+        // the vchunk stops constraining the cleanup.
+        RunOnExecutor(
+            DirectBlockGroup->GetExecutor(),
+            [&]() -> bool
+            {
+                InvokeUpdateDirtyMap(
+                    *vchunk,
+                    TDBGRestoreResponse{.Error = MakeError(S_OK)});
+                return true;
+            });
+        const auto barrierAfterRestore =
+            GetSafeBarrierOnExecutor(DirectBlockGroup->GetExecutor(), *vchunk);
+
+        vchunk->Stop().GetValue(TDuration::Seconds(10));
+
+        UNIT_ASSERT_C(
+            barrierWhileRestoring.has_value(),
+            "vchunk with a pending restore reported 'no constraint' to the "
+            "cleanup barrier gather");
+        UNIT_ASSERT_VALUES_EQUAL(0, *barrierWhileRestoring);
+        UNIT_ASSERT(!barrierAfterRestore.has_value());
+    }
+
     Y_UNIT_TEST_F(ShouldSwitchHostToTemporaryOfflineAndBack, TBaseFixture)
     {
         Init();