Repository: ydb-platform/nbs
Pull request #3266: Wait for in-flight writes before the lagging agent migration
Author: login: komarevtsev-d
Labels: name: large-tests
description: Launch large tests for PR
name: blockstore
description: Add this label to run only cloud/blockstore build and tests on PR
Created: 2025-03-26T19:17:12Z | Merged: 2025-04-04T14:54:10Z

#3
Добавляю новый тип drain: `WaitForInFlightWrites`. Это запрос, который дожидается завершения всех записей в полёте на момент получения сообщения. Нужно это перед началом миграции лагающего диска. 
Мигрировать сразу мы не можем, т.к. миграционный партишн не знает о текущих записях пользователя и начинает вычитывать из этих диапазонов, что может привести к чтению устареших данных.


Diff:
diff --git a/cloud/blockstore/libs/storage/api/partition.h b/cloud/blockstore/libs/storage/api/partition.h
index 3295a0868e0..7450440490a 100644
--- a/cloud/blockstore/libs/storage/api/partition.h
+++ b/cloud/blockstore/libs/storage/api/partition.h
@@ -10,9 +10,13 @@ namespace NCloud::NBlockStore::NStorage::NPartition {
 ////////////////////////////////////////////////////////////////////////////////
 
 #define BLOCKSTORE_PARTITION_REQUESTS(xxx, ...)                                \
-    xxx(WaitReady,              __VA_ARGS__)                                   \
-    xxx(StatPartition,          __VA_ARGS__)                                   \
-    xxx(Drain,                  __VA_ARGS__)                                   \
+    xxx(WaitReady,                                                 __VA_ARGS__)\
+    xxx(StatPartition,                                             __VA_ARGS__)\
+    /* Waits until there are no more in-flight write requests. */              \
+    xxx(Drain,                                                     __VA_ARGS__)\
+    /* Waits for current in-flight writes to finish and does not affect any    \
+     * requests that come after. */                                            \
+    xxx(WaitForInFlightWrites,                                     __VA_ARGS__)\
 // BLOCKSTORE_PARTITION_REQUESTS
 
 // requests forwarded from service to partition
@@ -83,6 +87,18 @@ struct TEvPartition
     {
     };
 
+    //
+    // WaitForInFlightWrites
+    //
+
+    struct TWaitForInFlightWritesRequest
+    {
+    };
+
+    struct TWaitForInFlightWritesResponse
+    {
+    };
+
     //
     // Garbage collector finish report
     //
@@ -119,6 +135,9 @@ struct TEvPartition
         EvAddLaggingAgentRequest = EvBegin + 9,
         EvRemoveLaggingReplicaRequest = EvBegin + 10,
 
+        EvWaitForInFlightWritesRequest = EvBegin + 11,
+        EvWaitForInFlightWritesResponse = EvBegin + 12,
+
         EvEnd
     };
 
diff --git a/cloud/blockstore/libs/storage/model/requests_in_progress.h b/cloud/blockstore/libs/storage/model/requests_in_progress.h
index ed6ea9c3a17..98afb106560 100644
--- a/cloud/blockstore/libs/storage/model/requests_in_progress.h
+++ b/cloud/blockstore/libs/storage/model/requests_in_progress.h
@@ -24,7 +24,12 @@ class TEmptyType
 class IRequestsInProgress
 {
 public:
-    virtual bool WriteRequestInProgress() const = 0;
+    virtual ~IRequestsInProgress() = default;
+
+    [[nodiscard]] virtual bool WriteRequestInProgress() const = 0;
+
+    virtual void WaitForInFlightWrites() = 0;
+    [[nodiscard]] virtual bool IsWaitingForInFlightWrites() const = 0;
 };
 
 ///////////////////////////////////////////////////////////////////////////////
@@ -46,8 +51,10 @@ class TRequestsInProgress: public IRequestsInProgress
     size_t WriteRequestCount = 0;
     TKey RequestIdentityKeyCounter = {};
 
+    THashSet<TKey> WaitingForWriteRequests;
+
 public:
-    TRequestsInProgress(EAllowedRequests allowedRequests)
+    explicit TRequestsInProgress(EAllowedRequests allowedRequests)
         : AllowedRequests(allowedRequests)
     {}
 
@@ -121,6 +128,7 @@ class TRequestsInProgress: public IRequestsInProgress
         }
 
         if (it->second.Write) {
+            WaitingForWriteRequests.erase(key);
             --WriteRequestCount;
         }
         TValue res = std::move(it->second.Value);
@@ -147,6 +155,20 @@ class TRequestsInProgress: public IRequestsInProgress
     {
         return WriteRequestCount != 0;
     }
+
+    void WaitForInFlightWrites() override
+    {
+        for (const auto& [key, value]: RequestsInProgress) {
+            if (value.Write) {
+                WaitingForWriteRequests.insert(key);
+            }
+        }
+    }
+
+    [[nodiscard]] bool IsWaitingForInFlightWrites() const override
+    {
+        return !WaitingForWriteRequests.empty();
+    }
 };
 
 }  // namespace NCloud::NBlockStore::NStorage
diff --git a/cloud/blockstore/libs/storage/model/requests_in_progress_ut.cpp b/cloud/blockstore/libs/storage/model/requests_in_progress_ut.cpp
index 5c7a96f37cc..b0fdae8bb20 100644
--- a/cloud/blockstore/libs/storage/model/requests_in_progress_ut.cpp
+++ b/cloud/blockstore/libs/storage/model/requests_in_progress_ut.cpp
@@ -119,6 +119,43 @@ Y_UNIT_TEST_SUITE(TRequestsInProgressTest)
             UNIT_ASSERT_EQUAL(testItem.Value, item.Value);
         }
     }
+
+    Y_UNIT_TEST(ShouldWaitForInFlightWrites)
+    {
+        using TRequests = TRequestsInProgress<ui32, TString>;
+
+        TMap<ui32, TRequests::TRequest> testData{
+            {0, {"Read Request 1", false}},
+            {1, {"Write Request 1", true}},
+            {10, {"Read Request 2", false}},
+            {20, {"Write Request 2", true}},
+        };
+
+        // When there is no in-flight requests waiting does nothing.
+        TRequests requestsInProgress{EAllowedRequests::ReadWrite};
+        UNIT_ASSERT(!requestsInProgress.IsWaitingForInFlightWrites());
+        requestsInProgress.WaitForInFlightWrites();
+        UNIT_ASSERT(!requestsInProgress.IsWaitingForInFlightWrites());
+
+        for (const auto& item : testData) {
+            if (item.second.Write) {
+                requestsInProgress.AddWriteRequest(
+                    item.first, TString(item.second.Value));
+            } else {
+                requestsInProgress.AddReadRequest(
+                    item.first, TString(item.second.Value));
+            }
+        }
+
+        UNIT_ASSERT(!requestsInProgress.IsWaitingForInFlightWrites());
+        requestsInProgress.WaitForInFlightWrites();
+        UNIT_ASSERT(requestsInProgress.IsWaitingForInFlightWrites());
+
+        requestsInProgress.RemoveRequest(1);
+        UNIT_ASSERT(requestsInProgress.IsWaitingForInFlightWrites());
+        requestsInProgress.RemoveRequest(20);
+        UNIT_ASSERT(!requestsInProgress.IsWaitingForInFlightWrites());
+    }
 }
 
 }  // namespace NCloud::NBlockStore::NStorage::NPartition
diff --git a/cloud/blockstore/libs/storage/partition/part_actor.cpp b/cloud/blockstore/libs/storage/partition/part_actor.cpp
index 20f05821436..149b4848bbd 100644
--- a/cloud/blockstore/libs/storage/partition/part_actor.cpp
+++ b/cloud/blockstore/libs/storage/partition/part_actor.cpp
@@ -728,6 +728,13 @@ void TPartitionActor::HandleDrain(
     DrainActorCompanion.HandleDrain(ev, ctx);
 }
 
+void TPartitionActor::HandleWaitForInFlightWrites(
+    const TEvPartition::TEvWaitForInFlightWritesRequest::TPtr& ev,
+    const TActorContext& ctx)
+{
+    DrainActorCompanion.HandleWaitForInFlightWrites(ev, ctx);
+}
+
 bool TPartitionActor::HandleRequests(STFUNC_SIG)
 {
     switch (ev->GetTypeRewrite()) {
diff --git a/cloud/blockstore/libs/storage/partition/part_actor.h b/cloud/blockstore/libs/storage/partition/part_actor.h
index e3ded32352e..95d36b80b9d 100644
--- a/cloud/blockstore/libs/storage/partition/part_actor.h
+++ b/cloud/blockstore/libs/storage/partition/part_actor.h
@@ -222,11 +222,22 @@ class TPartitionActor final
 
     void SendStatsToService(const NActors::TActorContext& ctx);
 
+    // IRequestsInProgress implementation:
     bool WriteRequestInProgress() const override
     {
         return WriteAndZeroRequestsInProgress != 0;
     }
 
+    void WaitForInFlightWrites() override
+    {
+        Y_ABORT("Unimplemented");
+    }
+
+    bool IsWaitingForInFlightWrites() const override
+    {
+        Y_ABORT("Unimplemented");
+    }
+
     template <typename TMethod>
     void HandleWriteBlocksRequest(
         const typename TMethod::TRequest::TPtr& ev,
diff --git a/cloud/blockstore/libs/storage/partition2/part2_actor.cpp b/cloud/blockstore/libs/storage/partition2/part2_actor.cpp
index c701c4bd8c9..a31ec4b4e41 100644
--- a/cloud/blockstore/libs/storage/partition2/part2_actor.cpp
+++ b/cloud/blockstore/libs/storage/partition2/part2_actor.cpp
@@ -703,6 +703,13 @@ void TPartitionActor::HandleDrain(
     DrainActorCompanion.HandleDrain(ev, ctx);
 }
 
+void TPartitionActor::HandleWaitForInFlightWrites(
+    const TEvPartition::TEvWaitForInFlightWritesRequest::TPtr& ev,
+    const TActorContext& ctx)
+{
+    DrainActorCompanion.HandleWaitForInFlightWrites(ev, ctx);
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 #define BLOCKSTORE_HANDLE_UNIMPLEMENTED_REQUEST(name, ns)                      \
diff --git a/cloud/blockstore/libs/storage/partition2/part2_actor.h b/cloud/blockstore/libs/storage/partition2/part2_actor.h
index fbcfca59b14..f52a5c35fdd 100644
--- a/cloud/blockstore/libs/storage/partition2/part2_actor.h
+++ b/cloud/blockstore/libs/storage/partition2/part2_actor.h
@@ -241,11 +241,22 @@ class TPartitionActor final
 
     void SendStatsToService(const NActors::TActorContext& ctx);
 
+    // IRequestsInProgress implementation:
     bool WriteRequestInProgress() const override
     {
         return WriteAndZeroRequestsInProgress != 0;
     }
 
+    void WaitForInFlightWrites() override
+    {
+        Y_ABORT("Unimplemented");
+    }
+
+    bool IsWaitingForInFlightWrites() const override
+    {
+        Y_ABORT("Unimplemented");
+    }
+
     template <typename TMethod>
     void HandleWriteBlocksRequest(
         const typename TMethod::TRequest::TPtr& ev,
diff --git a/cloud/blockstore/libs/storage/partition_common/drain_actor_companion.cpp b/cloud/blockstore/libs/storage/partition_common/drain_actor_companion.cpp
index 6467d905079..61d1aa4dac8 100644
--- a/cloud/blockstore/libs/storage/partition_common/drain_actor_companion.cpp
+++ b/cloud/blockstore/libs/storage/partition_common/drain_actor_companion.cpp
@@ -12,6 +12,8 @@
 
 namespace NCloud::NBlockStore::NStorage {
 
+using namespace NActors;
+
 LWTRACE_USING(BLOCKSTORE_STORAGE_PROVIDER);
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -31,7 +33,7 @@ TDrainActorCompanion::TDrainActorCompanion(
 
 void TDrainActorCompanion::HandleDrain(
     const NPartition::TEvPartition::TEvDrainRequest::TPtr& ev,
-    const NActors::TActorContext& ctx)
+    const TActorContext& ctx)
 {
     auto* msg = ev->Get();
 
@@ -46,15 +48,15 @@ void TDrainActorCompanion::HandleDrain(
         "Drain",
         requestInfo->CallContext->RequestId);
 
-    LOG_TRACE(
+    LOG_DEBUG(
         ctx,
         TBlockStoreComponents::PARTITION,
         "[%s] Start drain",
         LoggingId.c_str());
 
-    const ui32 maxDrainRequests = 10;
-    if (RequestsInProgress.WriteRequestInProgress()
-            && DrainRequests.size() >= maxDrainRequests)
+    constexpr ui32 MaxDrainRequests = 10;
+    if (RequestsInProgress.WriteRequestInProgress() &&
+        DrainRequests.size() >= MaxDrainRequests)
     {
         NCloud::Reply(
             ctx,
@@ -65,18 +67,62 @@ void TDrainActorCompanion::HandleDrain(
     }
 
     DrainRequests.push_back(std::move(requestInfo));
-    ProcessDrainRequests(ctx);
+    DoProcessDrainRequests(ctx);
+}
+
+void TDrainActorCompanion::HandleWaitForInFlightWrites(
+    const NPartition::TEvPartition::TEvWaitForInFlightWritesRequest::TPtr& ev,
+    const TActorContext& ctx)
+{
+    auto* msg = ev->Get();
+
+    if (WaitForInFlightWritesRequest) {
+        NCloud::Reply(
+            ctx,
+            *ev,
+            std::make_unique<
+                NPartition::TEvPartition::TEvWaitForInFlightWritesResponse>(
+                MakeError(
+                    E_REJECTED,
+                    "Already waiting for in-flight write requests")));
+        return;
+    }
+
+    WaitForInFlightWritesRequest =
+        CreateRequestInfo(ev->Sender, ev->Cookie, msg->CallContext);
+
+    LOG_DEBUG(
+        ctx,
+        TBlockStoreComponents::PARTITION,
+        "[%s] Start wait for in-flight write requests",
+        LoggingId.c_str());
+
+    Y_DEBUG_ABORT_UNLESS(!RequestsInProgress.IsWaitingForInFlightWrites());
+    RequestsInProgress.WaitForInFlightWrites();
+    DoProcessWaitForInFlightWritesRequests(ctx);
 }
 
 ////////////////////////////////////////////////////////////////////////////////
 
 void TDrainActorCompanion::ProcessDrainRequests(
-    const NActors::TActorContext& ctx)
+    const TActorContext& ctx)
+{
+    DoProcessDrainRequests(ctx);
+    DoProcessWaitForInFlightWritesRequests(ctx);
+}
+
+void TDrainActorCompanion::DoProcessDrainRequests(const TActorContext& ctx)
 {
-    if (RequestsInProgress.WriteRequestInProgress()) {
+    if (DrainRequests.empty() || RequestsInProgress.WriteRequestInProgress()) {
         return;
     }
 
+    LOG_DEBUG(
+        ctx,
+        TBlockStoreComponents::PARTITION,
+        "[%s] Complete drain",
+        LoggingId.c_str());
+
     for (auto& requestInfo: DrainRequests) {
         auto response =
             std::make_unique<NPartition::TEvPartition::TEvDrainResponse>();
@@ -87,16 +133,31 @@ void TDrainActorCompanion::ProcessDrainRequests(
             "Drain",
             requestInfo->CallContext->RequestId);
 
-        LOG_TRACE(
-            ctx,
-            TBlockStoreComponents::PARTITION,
-            "[%s] Complete drain",
-            LoggingId.c_str());
-
         NCloud::Reply(ctx, *requestInfo, std::move(response));
     }
 
     DrainRequests.clear();
 }
 
+void TDrainActorCompanion::DoProcessWaitForInFlightWritesRequests(
+    const TActorContext& ctx)
+{
+    if (!WaitForInFlightWritesRequest ||
+        RequestsInProgress.IsWaitingForInFlightWrites())
+    {
+        return;
+    }
+
+    LOG_DEBUG(
+        ctx,
+        TBlockStoreComponents::PARTITION,
+        "[%s] Complete wait for in-flight write requests",
+        LoggingId.c_str());
+
+    auto response = std::make_unique<
+        NPartition::TEvPartition::TEvWaitForInFlightWritesResponse>();
+    NCloud::Reply(ctx, *WaitForInFlightWritesRequest, std::move(response));
+    WaitForInFlightWritesRequest.Reset();
+}
+
 }  // namespace NCloud::NBlockStore::NStorage
diff --git a/cloud/blockstore/libs/storage/partition_common/drain_actor_companion.h b/cloud/blockstore/libs/storage/partition_common/drain_actor_companion.h
index 87e086e1434..4849f7ead18 100644
--- a/cloud/blockstore/libs/storage/partition_common/drain_actor_companion.h
+++ b/cloud/blockstore/libs/storage/partition_common/drain_actor_companion.h
@@ -9,7 +9,9 @@ namespace NCloud::NBlockStore::NStorage {
 
 class TDrainActorCompanion
 {
+private:
     TVector<TRequestInfoPtr> DrainRequests;
+    TRequestInfoPtr WaitForInFlightWritesRequest;
     IRequestsInProgress& RequestsInProgress;
     const TString LoggingId;
 
@@ -26,7 +28,15 @@ class TDrainActorCompanion
         const NPartition::TEvPartition::TEvDrainRequest::TPtr& ev,
         const NActors::TActorContext& ctx);
 
+    void HandleWaitForInFlightWrites(
+        const NPartition::TEvPartition::TEvWaitForInFlightWritesRequest::TPtr& ev,
+        const NActors::TActorContext& ctx);
+
     void ProcessDrainRequests(const NActors::TActorContext& ctx);
+
+private:
+    void DoProcessDrainRequests(const NActors::TActorContext& ctx);
+    void DoProcessWaitForInFlightWritesRequests(const NActors::TActorContext& ctx);
 };
 
 }  // namespace NCloud::NBlockStore::NStorage
diff --git a/cloud/blockstore/libs/storage/partition_common/drain_actor_companion_ut.cpp b/cloud/blockstore/libs/storage/partition_common/drain_actor_companion_ut.cpp
index 62e4fae8bd1..bd24cb81bc9 100644
--- a/cloud/blockstore/libs/storage/partition_common/drain_actor_companion_ut.cpp
+++ b/cloud/blockstore/libs/storage/partition_common/drain_actor_companion_ut.cpp
@@ -51,11 +51,27 @@ class TMyTestEnv final
 public:
     TMyTestEnv()
     {
+        SetupLogging();
         NKikimr::SetupTabletServices(Runtime);
 
         Sender = Runtime.AllocateEdgeActor();
     }
 
+    void SetupLogging()
+    {
+        Runtime.AppendToLogSettings(
+            TBlockStoreComponents::START,
+            TBlockStoreComponents::END,
+            GetComponentName);
+
+        for (ui32 i = TBlockStoreComponents::START;
+             i < TBlockStoreComponents::END;
+             ++i)
+        {
+            Runtime.SetLogPriority(i, NLog::PRI_TRACE);
+        }
+    }
+
     TActorId Register(IActorPtr actor)
     {
         auto actorId = Runtime.Register(actor.release());
@@ -80,6 +96,14 @@ class TMyTestEnv final
             .GrabEdgeEvent<NPartition::TEvPartition::TEvDrainResponse>(
                 TDuration());
     }
+
+    THolder<NPartition::TEvPartition::TEvWaitForInFlightWritesResponse>
+    GrabWaitForInFlightWritesResponse()
+    {
+        return Runtime.GrabEdgeEvent<
+            NPartition::TEvPartition::TEvWaitForInFlightWritesResponse>(
+            TDuration());
+    }
 };
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -89,6 +113,7 @@ class TActorWithDrain
     , private IRequestsInProgress
 {
     ui32 CurrentWriteRequestInProgress = 0;
+    bool WaitingForInFlightWrites = false;
     TDrainActorCompanion drainCompanion{*this, "LoggingId"};
 
 public:
@@ -100,7 +125,12 @@ class TActorWithDrain
     {
         switch (ev->GetTypeRewrite()) {
             HFunc(NPartition::TEvPartition::TEvDrainRequest, HandleDrain);
+            HFunc(
+                NPartition::TEvPartition::TEvWaitForInFlightWritesRequest,
+                HandleWaitForInFlightWrites);
             HFunc(TEvSetWriteInProgressCount, HandleSetWriteCount);
+            default:
+                HandleUnexpectedEvent(ev, TBlockStoreComponents::PARTITION);
         }
     }
 
@@ -110,6 +140,14 @@ class TActorWithDrain
         return CurrentWriteRequestInProgress != 0;
     }
 
+    void WaitForInFlightWrites() override {
+        WaitingForInFlightWrites = true;
+    }
+
+    bool IsWaitingForInFlightWrites() const override {
+        return WaitingForInFlightWrites && WriteRequestInProgress();
+    }
+
     void HandleDrain(
         const NPartition::TEvPartition::TEvDrainRequest::TPtr& ev,
         const TActorContext& ctx)
@@ -117,11 +155,21 @@ class TActorWithDrain
         drainCompanion.HandleDrain(ev, ctx);
     }
 
+    void HandleWaitForInFlightWrites(
+        const NPartition::TEvPartition::TEvWaitForInFlightWritesRequest::TPtr& ev,
+        const TActorContext& ctx)
+    {
+        drainCompanion.HandleWaitForInFlightWrites(ev, ctx);
+    }
+
     void HandleSetWriteCount(
         const TEvSetWriteInProgressCount::TPtr& ev,
         const TActorContext& ctx)
     {
         CurrentWriteRequestInProgress = ev->Get()->WriteCount;
+        if (!WriteRequestInProgress()) {
+            WaitingForInFlightWrites = false;
+        }
         drainCompanion.ProcessDrainRequests(ctx);
     }
 };
@@ -316,6 +364,68 @@ Y_UNIT_TEST_SUITE(TDrainActorCompanionTest)
             UNIT_ASSERT_VALUES_EQUAL(nullptr, drainResponse);
         }
     }
+
+    Y_UNIT_TEST(BasicWaitForInFlightWrites)
+    {
+        TMyTestEnv testEnv;
+
+        auto actorId = testEnv.Register(std::make_unique<TActorWithDrain>());
+
+        // Waiting for in-flight writes.
+        testEnv.Send(
+            actorId,
+            std::make_unique<
+                NPartition::TEvPartition::TEvWaitForInFlightWritesRequest>());
+        testEnv.DispatchEvents();
+
+        {   // Get the response
+            auto waitResponse = testEnv.GrabWaitForInFlightWritesResponse();
+            UNIT_ASSERT_VALUES_EQUAL(S_OK, waitResponse->GetStatus());
+        }
+    }
+
+    Y_UNIT_TEST(WaitForInFlightWritesWithWriteInProgress)
+    {
+        TMyTestEnv testEnv;
+
+        auto actorId = testEnv.Register(std::make_unique<TActorWithDrain>());
+
+        // Set actor have one write inflight
+        testEnv.Send(actorId, std::make_unique<TEvSetWriteInProgressCount>(1));
+
+        // Waiting for in-flight writes.
+        testEnv.Send(
+            actorId,
+            std::make_unique<
+                NPartition::TEvPartition::TEvWaitForInFlightWritesRequest>());
+        testEnv.DispatchEvents();
+
+        {   // Assert waiting is not finished since the write is in progress
+            auto waitResponse = testEnv.GrabWaitForInFlightWritesResponse();
+            UNIT_ASSERT_VALUES_EQUAL(nullptr, waitResponse);
+        }
+
+        // Sending another wait.
+        testEnv.Send(
+            actorId,
+            std::make_unique<
+                NPartition::TEvPartition::TEvWaitForInFlightWritesRequest>());
+        testEnv.DispatchEvents();
+
+        {   // It should be rejected because the actor is already waiting.
+            auto waitResponse = testEnv.GrabWaitForInFlightWritesResponse();
+            UNIT_ASSERT_VALUES_EQUAL(E_REJECTED, waitResponse->GetStatus());
+        }
+
+        // Finish write
+        testEnv.Send(actorId, std::make_unique<TEvSetWriteInProgressCount>(0));
+        testEnv.DispatchEvents();
+
+        {   // Get the response
+            auto waitResponse = testEnv.GrabWaitForInFlightWritesResponse();
+            UNIT_ASSERT_VALUES_EQUAL(S_OK, waitResponse->GetStatus());
+        }
+    }
 }
 
 }  // namespace NCloud::NBlockStore::NStorage
diff --git a/cloud/blockstore/libs/storage/partition_nonrepl/lagging_agent_migration_actor.cpp b/cloud/blockstore/libs/storage/partition_nonrepl/lagging_agent_migration_actor.cpp
index 06f127233b8..555152a83df 100644
--- a/cloud/blockstore/libs/storage/partition_nonrepl/lagging_agent_migration_actor.cpp
+++ b/cloud/blockstore/libs/storage/partition_nonrepl/lagging_agent_migration_actor.cpp
@@ -59,7 +59,6 @@ void TLaggingAgentMigrationActor::OnBootstrap(const TActorContext& ctx)
             Config->GetMaxMigrationBandwidth(),
             Config->GetExpectedDiskAgentSize(),
             PartConfig));
-    StartWork(ctx);
 }
 
 bool TLaggingAgentMigrationActor::OnMessage(
@@ -67,8 +66,28 @@ bool TLaggingAgentMigrationActor::OnMessage(
     TAutoPtr<IEventHandle>& ev)
 {
     Y_UNUSED(ctx);
+    switch (ev->GetTypeRewrite()) {
+        HFunc(
+            TEvNonreplPartitionPrivate::TEvStartLaggingAgentMigration,
+            HandleStartLaggingAgentMigration);
+
+        default:
+            // Message processing by the base class is required.
+            return false;
+    }
+
+    // We get here if we have processed an incoming message. And its processing
+    // by the base class is not required.
+    return true;
+}
+
+void TLaggingAgentMigrationActor::HandleStartLaggingAgentMigration(
+    const TEvNonreplPartitionPrivate::TEvStartLaggingAgentMigration::TPtr& ev,
+    const NActors::TActorContext& ctx)
+{
     Y_UNUSED(ev);
-    return false;
+    Y_ABORT_IF(IsMigrationAllowed());
+    StartWork(ctx);
 }
 
 void TLaggingAgentMigrationActor::OnMigrationProgress(
diff --git a/cloud/blockstore/libs/storage/partition_nonrepl/lagging_agent_migration_actor.h b/cloud/blockstore/libs/storage/partition_nonrepl/lagging_agent_migration_actor.h
index b170bc2aa0b..1bd6199ff4a 100644
--- a/cloud/blockstore/libs/storage/partition_nonrepl/lagging_agent_migration_actor.h
+++ b/cloud/blockstore/libs/storage/partition_nonrepl/lagging_agent_migration_actor.h
@@ -53,6 +53,12 @@ class TLaggingAgentMigrationActor final
         ui64 migrationIndex) override;
     void OnMigrationFinished(const NActors::TActorContext& ctx) override;
     void OnMigrationError(const NActors::TActorContext& ctx) override;
+
+private:
+    void HandleStartLaggingAgentMigration(
+        const TEvNonreplPartitionPrivate::TEvStartLaggingAgentMigration::TPtr&
+            ev,
+        const NActors::TActorContext& ctx);
 };
 
 }   // namespace NCloud::NBlockStore::NStorage
diff --git a/cloud/blockstore/libs/storage/partition_nonrepl/lagging_agents_replica_proxy_actor.cpp b/cloud/blockstore/libs/storage/partition_nonrepl/lagging_agents_replica_proxy_actor.cpp
index 591d699e7b7..d5bc75cf61b 100644
--- a/cloud/blockstore/libs/storage/partition_nonrepl/lagging_agents_replica_proxy_actor.cpp
+++ b/cloud/blockstore/libs/storage/partition_nonrepl/lagging_agents_replica_proxy_actor.cpp
@@ -296,8 +296,7 @@ TLaggingAgentsReplicaProxyActor::TLaggingAgentsReplicaProxyActor(
     , PoisonPillHelper(this)
 {}
 
-TLaggingAgentsReplicaProxyActor::
-    ~TLaggingAgentsReplicaProxyActor() = default;
+TLaggingAgentsReplicaProxyActor::~TLaggingAgentsReplicaProxyActor() = default;
 
 void TLaggingAgentsReplicaProxyActor::Bootstrap(const TActorContext& ctx)
 {
@@ -317,25 +316,29 @@ void TLaggingAgentsReplicaProxyActor::MarkBlocksAsDirty(
     const TString& unavailableAgentId,
     TBlockRange64 range)
 {
+    // We artificially lower the resolution of the "CleanBlocksMap". This is
+    // done because the migration actor migrates by ranges of
+    // "ProcessingRangeSize" anyway. The full resolution here would lead to
+    // migrating not aligned ranges which is not ideal.
+    const ui64 blocksPerRange =
+        ProcessingRangeSize / PartConfig->GetBlockSize();
+    auto alignedStart = AlignDown<ui64>(range.Start, blocksPerRange);
+    auto alignedEnd = AlignUp<ui64>(range.End + 1, blocksPerRange);
+
     LOG_TRACE(
         ctx,
         TBlockStoreComponents::PARTITION_WORKER,
-        "[%s] Marking block range %s as dirty for agent %s",
+        "[%s] Lower block range resolution %s -> %s and mark blocks as dirty "
+        "for lagging agent %s",
         PartConfig->GetName().c_str(),
         range.Print().c_str(),
+        TBlockRange64::MakeClosedInterval(alignedStart, alignedEnd)
+            .Print()
+            .c_str(),
         unavailableAgentId.Quote().c_str());
 
     auto& state = AgentState[unavailableAgentId];
     Y_ABORT_UNLESS(state.CleanBlocksMap);
-
-    // We artificially lower the resolution of the "CleanBlocksMap". This is
-    // done because the migration actor migrates by ranges of
-    // "ProcessingRangeSize" anyway. The full resolution here would lead to
-    // migrating not aligned ranges which is not ideal.
-    const ui64 blocksPerRange =
-        ProcessingRangeSize / PartConfig->GetBlockSize();
-    auto alignedStart = AlignDown<ui64>(range.Start, blocksPerRange);
-    auto alignedEnd = AlignUp<ui64>(range.End + 1, blocksPerRange);
     state.CleanBlocksMap->Unset(alignedStart, alignedEnd);
 }
 
@@ -520,7 +523,8 @@ TVector<TSplitRequest> TLaggingAgentsReplicaProxyActor::DoSplitRequest(
         result.emplace_back(
             std::move(request),
             forkedCallContext,
-            GetRecipientActorId(deviceRequest.Device.GetAgentId()));
+            GetRecipientActorId(deviceRequest.Device.GetAgentId()),
+            deviceRequest.Device.GetDeviceUUID());
     }
 
     return result;
@@ -598,8 +602,9 @@ NActors::TActorId TLaggingAgentsReplicaProxyActor::GetRecipientActorId(
     switch (agentState->State) {
         case EAgentState::Unavailable:
             return {};
+        case EAgentState::WaitingForDrain:
         case EAgentState::Resyncing:
-            Y_DEBUG_ABORT_UNLESS(agentState->MigrationActorId);
+            Y_ABORT_UNLESS(agentState->MigrationActorId);
             return agentState->MigrationActorId;
     }
     Y_ABORT("Unknown enum value: %u", static_cast<ui8>(agentState->State));
@@ -670,15 +675,26 @@ void TLaggingAgentsReplicaProxyActor::HandleAgentIsBackOnline(
 {
     const auto* msg = ev->Get();
 
-    Y_DEBUG_ABORT_UNLESS(AgentState.contains(msg->AgentId));
-    if (!AgentState.contains(msg->AgentId)) {
+    auto* state = AgentState.FindPtr(msg->AgentId);
+    Y_DEBUG_ABORT_UNLESS(state);
+    if (!state) {
         return;
     }
+    switch (state->State) {
+        case EAgentState::Unavailable:
+            break;
+        case EAgentState::WaitingForDrain:
+            return;
+        case EAgentState::Resyncing:
+            Y_DEBUG_ABORT_UNLESS(false);
+            return;
+    }
+    state->State = EAgentState::WaitingForDrain;
 
     LOG_INFO(
         ctx,
         TBlockStoreComponents::PARTITION_WORKER,
-        "[%s] Lagging agent %s is back online. Starting migration actor",
+        "[%s] Lagging agent %s is back online. Draining other partitions",
         PartConfig->GetName().c_str(),
         msg->AgentId.Quote().c_str());
 
@@ -688,30 +704,33 @@ void TLaggingAgentsReplicaProxyActor::HandleAgentIsBackOnline(
         std::make_unique<TEvNonreplPartitionPrivate::TEvAgentIsBackOnline>(
             msg->AgentId));
 
-    auto& state = AgentState[msg->AgentId];
-    if (state.State == EAgentState::Resyncing) {
-        return;
-    }
-    Y_DEBUG_ABORT_UNLESS(!state.MigrationActorId);
-    Y_DEBUG_ABORT_UNLESS(state.AvailabilityMonitoringActorId);
+    DrainRequestCounter++;
+    Y_ABORT_UNLESS(!CurrentDrainingAgents.contains(DrainRequestCounter));
+    CurrentDrainingAgents[DrainRequestCounter] = msg->AgentId;
+    NCloud::Send<NPartition::TEvPartition::TEvWaitForInFlightWritesRequest>(
+        ctx,
+        MirrorPartitionActorId,
+        DrainRequestCounter);
+
+    Y_DEBUG_ABORT_UNLESS(!state->MigrationActorId);
+    Y_DEBUG_ABORT_UNLESS(state->AvailabilityMonitoringActorId);
 
-    state.State = EAgentState::Resyncing;
-    DestroyChildActor(ctx, &state.AvailabilityMonitoringActorId);
+    DestroyChildActor(ctx, &state->AvailabilityMonitoringActorId);
 
     TCompressedBitmap cleanBlocksCopy(PartConfig->GetBlockCount());
-    cleanBlocksCopy.Update(*state.CleanBlocksMap, 0);
+    cleanBlocksCopy.Update(*state->CleanBlocksMap, 0);
 
     LOG_INFO(
         ctx,
         TBlockStoreComponents::PARTITION_WORKER,
-        "[%s] Starting lagging agent %s migration. Block count: %lu, dirty "
+        "[%s] Preparing lagging agent %s migration. Block count: %lu, dirty "
         "block count: %lu",
         PartConfig->GetName().c_str(),
         msg->AgentId.Quote().c_str(),
         PartConfig->GetBlockCount(),
         PartConfig->GetBlockCount() - cleanBlocksCopy.Count());
 
-    state.MigrationActorId = NCloud::Register(
+    state->MigrationActorId = NCloud::Register(
         ctx,
         std::make_unique<TLaggingAgentMigrationActor>(
             Config,
@@ -725,7 +744,7 @@ void TLaggingAgentsReplicaProxyActor::HandleAgentIsBackOnline(
             MirrorPartitionActorId,
             std::move(cleanBlocksCopy),
             msg->AgentId));
-    PoisonPillHelper.TakeOwnership(ctx, state.MigrationActorId);
+    PoisonPillHelper.TakeOwnership(ctx, state->MigrationActorId);
 }
 
 void TLaggingAgentsReplicaProxyActor::HandleLaggingAgentMigrationFinished(
@@ -744,6 +763,9 @@ void TLaggingAgentsReplicaProxyActor::HandleLaggingAgentMigrationFinished(
     switch (state.State) {
         case EAgentState::Unavailable:
             return;
+        case EAgentState::WaitingForDrain:
+            Y_DEBUG_ABORT_UNLESS(false);
+            return;
         case EAgentState::Resyncing:
             break;
     }
@@ -766,6 +788,82 @@ void TLaggingAgentsReplicaProxyActor::HandleLaggingAgentMigrationFinished(
         ev->ReleaseBase().Release()));
 }
 
+void TLaggingAgentsReplicaProxyActor::HandleWaitForInFlightWrites(
+    const NPartition::TEvPartition::TEvWaitForInFlightWritesRequest::TPtr& ev,
+    const NActors::TActorContext& ctx)
+{
+    ForwardMessageToActor(ev, ctx, NonreplPartitionActorId);
+}
+
+void TLaggingAgentsReplicaProxyActor::HandleWaitForInFlightWritesResponse(
+    const NPartition::TEvPartition::TEvWaitForInFlightWritesResponse::TPtr& ev,
+    const NActors::TActorContext& ctx)
+{
+    auto it = CurrentDrainingAgents.find(ev->Cookie);
+    if (it == CurrentDrainingAgents.end()) {
+        return;
+    }
+
+    TString agentId = std::move(it->second);
+    CurrentDrainingAgents.erase(it);
+
+    const auto* msg = ev->Get();
+    if (HasError(msg->GetError())) {
+        LOG_WARN(
+            ctx,
+            TBlockStoreComponents::PARTITION_WORKER,
+            "[%s] Failed to drain partitions for lagging agent %s: %s",
+            PartConfig->GetName().c_str(),
+            agentId.Quote().c_str(),
+            FormatError(msg->GetError()).c_str());
+
+        DrainRequestCounter++;
+        Y_ABORT_UNLESS(!CurrentDrainingAgents.contains(DrainRequestCounter));
+        CurrentDrainingAgents[DrainRequestCounter] = agentId;
+
+        ctx.ExecutorThread.Schedule(
+            TDuration::Seconds(1),
+            new IEventHandle(
+                MirrorPartitionActorId,
+                SelfId(),
+                new NPartition::TEvPartition::
+                    TEvWaitForInFlightWritesRequest()));
+        return;
+    }
+
+    auto* state = AgentState.FindPtr(agentId);
+    Y_DEBUG_ABORT_UNLESS(state);
+    if (!state) {
+        return;
+    }
+
+    switch (state->State) {
+        case EAgentState::Unavailable:
+            return;
+        case EAgentState::WaitingForDrain:
+            break;
+        case EAgentState::Resyncing:
+            Y_DEBUG_ABORT_UNLESS(false);
+            return;
+    }
+
+    state->State = EAgentState::Resyncing;
+    LOG_DEBUG(
+        ctx,
+        TBlockStoreComponents::PARTITION_WORKER,
+        "[%s] Starting lagging agent %s migration. Block count: %lu, dirty "
+        "block count: %lu",
+        PartConfig->GetName().c_str(),
+        agentId.Quote().c_str(),
+        PartConfig->GetBlockCount(),
+        PartConfig->GetBlockCount() - state->CleanBlocksMap->Count());
+
+    Y_ABORT_UNLESS(state->MigrationActorId);
+    NCloud::Send<TEvNonreplPartitionPrivate::TEvStartLaggingAgentMigration>(
+        ctx,
+        state->MigrationActorId);
+}
+
 void TLaggingAgentsReplicaProxyActor::HandleWriteBlocks(
     const TEvService::TEvWriteBlocksRequest::TPtr& ev,
     const TActorContext& ctx)
@@ -831,8 +929,7 @@ void TLaggingAgentsReplicaProxyActor::HandlePoisonPill(
     AgentState.clear();
 }
 
-void TLaggingAgentsReplicaProxyActor::Die(
-    const NActors::TActorContext& ctx)
+void TLaggingAgentsReplicaProxyActor::Die(const NActors::TActorContext& ctx)
 {
     TBase::Die(ctx);
 }
@@ -857,6 +954,12 @@ STFUNC(TLaggingAgentsReplicaProxyActor::StateWork)
         HFunc(
             TEvVolumePrivate::TEvLaggingAgentMigrationFinished,
             HandleLaggingAgentMigrationFinished);
+        HFunc(
+            NPartition::TEvPartition::TEvWaitForInFlightWritesRequest,
+            HandleWaitForInFlightWrites);
+        HFunc(
+            NPartition::TEvPartition::TEvWaitForInFlightWritesResponse,
+            HandleWaitForInFlightWritesResponse);
         HFunc(TEvVolume::TEvRWClientIdChanged, HandleRWClientIdChanged);
 
         HFunc(TEvents::TEvPoisonPill, HandlePoisonPill);
@@ -883,7 +986,11 @@ STFUNC(TLaggingAgentsReplicaProxyActor::StateZombie)
         HFunc(TEvService::TEvWriteBlocksLocalRequest, RejectWriteBlocksLocal);
         HFunc(TEvService::TEvReadBlocksRequest, RejectReadBlocks);
         HFunc(TEvService::TEvReadBlocksLocalRequest, RejectReadBlocksLocal);
+        HFunc(
+            NPartition::TEvPartition::TEvWaitForInFlightWritesRequest,
+            RejectWaitForInFlightWrites);
 
+        IgnoreFunc(NPartition::TEvPartition::TEvWaitForInFlightWritesResponse);
         IgnoreFunc(TEvVolume::TEvRWClientIdChanged);
         IgnoreFunc(TEvents::TEvPoisonPill);
         HFunc(TEvents::TEvPoisonTaken, PoisonPillHelper.HandlePoisonTaken);
diff --git a/cloud/blockstore/libs/storage/partition_nonrepl/lagging_agents_replica_proxy_actor.h b/cloud/blockstore/libs/storage/partition_nonrepl/lagging_agents_replica_proxy_actor.h
index 076e831048e..b9657e1f023 100644
--- a/cloud/blockstore/libs/storage/partition_nonrepl/lagging_agents_replica_proxy_actor.h
+++ b/cloud/blockstore/libs/storage/partition_nonrepl/lagging_agents_replica_proxy_actor.h
@@ -7,6 +7,7 @@
 
 #include <cloud/blockstore/libs/diagnostics/public.h>
 #include <cloud/blockstore/libs/kikimr/helpers.h>
+#include <cloud/blockstore/libs/storage/api/partition.h>
 #include <cloud/blockstore/libs/storage/api/service.h>
 #include <cloud/blockstore/libs/storage/api/volume.h>
 #include <cloud/blockstore/libs/storage/volume/volume_events_private.h>
@@ -47,6 +48,7 @@ class TLaggingAgentsReplicaProxyActor final
     enum class EAgentState : ui8
     {
         Unavailable,
+        WaitingForDrain,
         Resyncing
     };
     struct TAgentState
@@ -59,6 +61,11 @@ class TLaggingAgentsReplicaProxyActor final
     };
     THashMap<TString, TAgentState> AgentState;
 
+    ui64 DrainRequestCounter = 0;
+    // To determine which agent has drained in-flight writes by cookie in the
+    // response.
+    THashMap<ui64, TString> CurrentDrainingAgents;
+
     TPoisonPillHelper PoisonPillHelper;
 
 public:
@@ -135,6 +142,11 @@ class TLaggingAgentsReplicaProxyActor final
         const TEvVolumePrivate::TEvLaggingAgentMigrationFinished::TPtr& ev,
         const NActors::TActorContext& ctx);
 
+    void HandleWaitForInFlightWritesResponse(
+        const NPartition::TEvPartition::TEvWaitForInFlightWritesResponse::TPtr&
+            ev,
+        const NActors::TActorContext& ctx);
+
     void HandleWriteOrZeroCompleted(
         const TEvNonreplPartitionPrivate::TEvWriteOrZeroCompleted::TPtr& ev,
         const NActors::TActorContext& ctx);
@@ -170,6 +182,9 @@ class TLaggingAgentsReplicaProxyActor final
     BLOCKSTORE_IMPLEMENT_REQUEST(ZeroBlocks, TEvService);
     BLOCKSTORE_IMPLEMENT_REQUEST(ReadBlocks, TEvService);
     BLOCKSTORE_IMPLEMENT_REQUEST(ReadBlocksLocal, TEvService);
+    BLOCKSTORE_IMPLEMENT_REQUEST(
+        WaitForInFlightWrites,
+        NPartition::TEvPartition);
 };
 
 }   // namespace NCloud::NBlockStore::NStorage
diff --git a/cloud/blockstore/libs/storage/partition_nonrepl/lagging_agents_replica_proxy_ut.cpp b/cloud/blockstore/libs/storage/partition_nonrepl/lagging_agents_replica_proxy_ut.cpp
index a2d2e76a7ab..1a692d03729 100644
--- a/cloud/blockstore/libs/storage/partition_nonrepl/lagging_agents_replica_proxy_ut.cpp
+++ b/cloud/blockstore/libs/storage/partition_nonrepl/lagging_agents_replica_proxy_ut.cpp
@@ -173,6 +173,8 @@ struct TTestEnv
         storageConfig.SetMaxMigrationIoDepth(1);
         storageConfig.SetLaggingDevicePingInterval(
             TDuration::Seconds(5).MilliSeconds());
+        storageConfig.SetLaggingDeviceTimeoutThreshold(
+            TDuration::Minutes(5).MilliSeconds());
         Config = std::make_shared<TStorageConfig>(
             std::move(storageConfig),
             std::make_shared<NFeatures::TFeaturesConfig>(
@@ -274,7 +276,7 @@ struct TTestEnv
              i < TBlockStoreComponents::END;
              ++i)
         {
-            Runtime.SetLogPriority(i, NLog::PRI_DEBUG);
+            Runtime.SetLogPriority(i, NLog::PRI_TRACE);
         }
     }
 
@@ -302,6 +304,18 @@ struct TTestEnv
         }
     }
 
+    void WriteBlocksToReplicaAsync(
+        ui32 replicaIndex,
+        TBlockRange64 range,
+        char content)
+    {
+        if (LocalRequests) {
+            WriteBlocksLocalAsync(ReplicaActors[replicaIndex], range, content);
+        } else {
+            WriteBlocksAsync(ReplicaActors[replicaIndex], range, content);
+        }
+    }
+
     void WriteBlocksToController(
         ui32 replicaIndex,
         TBlockRange64 range,
@@ -437,7 +451,6 @@ struct TTestEnv
     {
         TPartitionClient client(Runtime, MirrorPartActorId);
         auto request = client.CreateWriteBlocksRequest(range, content);
-        // request->MutableHeaders()->SetClientId("")
         client.SendRequest(actorId, std::move(request));
         auto response =
             client.RecvResponse<TEvService::TEvWriteBlocksResponse>();
@@ -462,6 +475,22 @@ struct TTestEnv
         return response;
     }
 
+    void WriteBlocksAsync(TActorId actorId, TBlockRange64 range, char content)
+    {
+        TPartitionClient client(Runtime, MirrorPartActorId);
+        auto request = client.CreateWriteBlocksRequest(range, content);
+        client.SendRequest(actorId, std::move(request));
+    }
+
+    void
+    WriteBlocksLocalAsync(TActorId actorId, TBlockRange64 range, char content)
+    {
+        TPartitionClient client(Runtime, MirrorPartActorId);
+        const TString data(DefaultBlockSize, content);
+        auto request = client.CreateWriteBlocksLocalRequest(range, data);
+        client.SendRequest(actorId, std::move(request));
+    }
+
     void
     ReadAndCheckContents(TActorId actorId, TBlockRange64 range, char content)
     {
@@ -513,9 +542,9 @@ Y_UNIT_TEST_SUITE(TLaggingAgentsReplicaProxyActorTest)
         const auto firstAndSecondDevices =
             TBlockRange64::WithLength(DeviceBlockCount - 1, 2);
         const auto firstDevice =
-            TBlockRange64::WithLength(DeviceBlockCount - 1, 1);
+            TBlockRange64::MakeOneBlock(DeviceBlockCount - 1);
         const auto secondDeviceOneBlock =
-            TBlockRange64::WithLength(DeviceBlockCount, 1);
+            TBlockRange64::MakeOneBlock(DeviceBlockCount);
         env.WriteBlocksToController(0, firstAndSecondDevices, 'B');
         env.ReadFromReplicaAndCheckContents(0, firstDevice, 'B');
         env.ReadFromReplicaAndCheckContents(0, secondDeviceOneBlock, 'A');
@@ -602,11 +631,11 @@ Y_UNIT_TEST_SUITE(TLaggingAgentsReplicaProxyActorTest)
             });
 
         runtime.AdvanceCurrentTime(env.Config->GetLaggingDevicePingInterval());
-        runtime.DispatchEvents({}, TDuration::MilliSeconds(10));
-        UNIT_ASSERT(seenHealthCheck);
-        UNIT_ASSERT(seenMigrationReads);
-        UNIT_ASSERT(seenMigrationWrites);
-        UNIT_ASSERT(seenMigrationFinish);
+        runtime.DispatchEvents(
+            {.CustomFinalCondition = [&]()
+             {
+                 return seenMigrationFinish;
+             }});
         env.ReadFromReplicaAndCheckContents(0, firstAndSecondDevices, 'B');
     }
 
@@ -869,6 +898,114 @@ Y_UNIT_TEST_SUITE(TLaggingAgentsReplicaProxyActorTest)
         MultipleLaggingAgentsOnOneReplica(false);
         MultipleLaggingAgentsOnOneReplica(true);
     }
+
+    void ShouldDrainBeforeMigration(bool localRequests)
+    {
+        TTestBasicRuntime runtime(AgentCount);
+        TTestEnv env(runtime, localRequests);
+        TPartitionClient client(runtime, env.MirrorPartActorId);
+
+        const auto fullDiskRange = TBlockRange64::WithLength(
+            0,
+            DeviceBlockCount * DeviceCountPerReplica);
+        env.WriteBlocksToPartition(fullDiskRange, 'A');
+
+        // Second row in the first column is lagging.
+        env.AddLaggingAgent(runtime.GetNodeId(1), 0);
+        runtime.DispatchEvents({}, TDuration::MilliSeconds(10));
+
+        // Writes across the first and second devices should write only to the
+        // first one.
+        const auto firstAndSecondDevices =
+            TBlockRange64::WithLength(DeviceBlockCount - 1, 2);
+        env.WriteBlocksToController(0, firstAndSecondDevices, 'B');
+
+        bool seenWaitForInFlightWritesRequest = false;
+        bool seenMigrationReads = false;
+        // bool seenMigrationWrites = false;
+        TVector<TRequestInfoPtr> writeDeviceBlocksRequestInfos;
+
+        bool interceptWrites = true;
+        runtime.SetEventFilter(
+            [&](TTestActorRuntimeBase&, TAutoPtr<IEventHandle>& event)
+            {
+                switch (event->GetTypeRewrite()) {
+                    case NPartition::TEvPartition::
+                        EvWaitForInFlightWritesRequest: {
+                        if (event->Sender == env.GetControllerActorId(0)) {
+                            seenWaitForInFlightWritesRequest = true;
+                        }
+                        break;
+                    }
+                    case TEvDiskAgent::EvWriteDeviceBlocksRequest: {
+                        if (!interceptWrites) {
+                            break;
+                        }
+
+                        auto* msg = event->Get<
+                            TEvDiskAgent::TEvWriteDeviceBlocksRequest>();
+                        writeDeviceBlocksRequestInfos.push_back(
+                            CreateRequestInfo(
+                                event->Sender,
+                                event->Cookie,
+                                msg->CallContext));
+                        return true;
+                    }
+                    case TEvService::EvReadBlocksRequest: {
+                        if (event->Recipient != env.ReplicaActors[0]) {
+                            break;
+                        }
+                        auto* msg =
+                            event->Get<TEvService::TEvReadBlocksRequest>();
+                        const auto& clientId =
+                            msg->Record.GetHeaders().GetClientId();
+                        if (clientId == BackgroundOpsClientId) {
+                            seenMigrationReads = true;
+                        }
+                        break;
+                    }
+                }
+                return false;
+            });
+
+        env.WriteBlocksToReplicaAsync(1, TBlockRange64::MakeOneBlock(0), 'C');
+        env.WriteBlocksToReplicaAsync(2, TBlockRange64::MakeOneBlock(0), 'D');
+        runtime.DispatchEvents(
+            {.CustomFinalCondition = [&]()
+             {
+                 return writeDeviceBlocksRequestInfos.size() == 2;
+             }});
+
+        interceptWrites = false;
+        runtime.DispatchEvents(
+            {.CustomFinalCondition = [&]()
+             {
+                 return seenWaitForInFlightWritesRequest;
+             }});
+        runtime.DispatchEvents({}, TDuration::MilliSeconds(10));
+        UNIT_ASSERT(!seenMigrationReads);
+
+        for (const auto& requestInfo: writeDeviceBlocksRequestInfos) {
+            auto response =
+                std::make_unique<TEvDiskAgent::TEvWriteDeviceBlocksResponse>();
+            runtime.Send(
+                new NActors::IEventHandle(
+                    requestInfo->Sender,
+                    TActorId(),
+                    response.release(),
+                    0,   // flags
+                    requestInfo->Cookie),
+                0);
+        }
+        writeDeviceBlocksRequestInfos.clear();
+        env.WaitForMigrationFinishEvent();
+    }
+
+    Y_UNIT_TEST(ShouldDrainBeforeMigration)
+    {
+        ShouldDrainBeforeMigration(false);
+        ShouldDrainBeforeMigration(true);
+    }
 }
 
 }   // namespace NCloud::NBlockStore::NStorage
diff --git a/cloud/blockstore/libs/storage/partition_nonrepl/part_mirror_actor.cpp b/cloud/blockstore/libs/storage/partition_nonrepl/part_mirror_actor.cpp
index 9cd3cc4d18f..fe1ced03ddc 100644
--- a/cloud/blockstore/libs/storage/partition_nonrepl/part_mirror_actor.cpp
+++ b/cloud/blockstore/libs/storage/partition_nonrepl/part_mirror_actor.cpp
@@ -691,6 +691,9 @@ STFUNC(TMirrorPartitionActor::StateWork)
         HFunc(TEvNonreplPartitionPrivate::TEvAddLaggingAgentRequest, HandleAddLaggingAgent);
         HFunc(TEvNonreplPartitionPrivate::TEvRemoveLaggingAgentRequest, HandleRemoveLaggingAgent);
         HFunc(NPartition::TEvPartition::TEvDrainRequest, DrainActorCompanion.HandleDrain);
+        HFunc(
+            NPartition::TEvPartition::TEvWaitForInFlightWritesRequest,
+            DrainActorCompanion.HandleWaitForInFlightWrites);
         HFunc(TEvService::TEvGetChangedBlocksRequest, DeclineGetChangedBlocks);
         HFunc(
             TEvNonreplPartitionPrivate::TEvGetDeviceForRangeRequest,
@@ -750,6 +753,9 @@ STFUNC(TMirrorPartitionActor::StateZombie)
         HFunc(TEvService::TEvWriteBlocksLocalRequest, RejectWriteBlocksLocal);
 
         HFunc(NPartition::TEvPartition::TEvDrainRequest, RejectDrain);
+        HFunc(
+            NPartition::TEvPartition::TEvWaitForInFlightWritesRequest,
+            RejectWaitForInFlightWrites);
         HFunc(TEvService::TEvGetChangedBlocksRequest, DeclineGetChangedBlocks);
         HFunc(
             TEvNonreplPartitionPrivate::TEvGetDeviceForRangeRequest,
diff --git a/cloud/blockstore/libs/storage/partition_nonrepl/part_mirror_actor.h b/cloud/blockstore/libs/storage/partition_nonrepl/part_mirror_actor.h
index 8fdca9196b2..404e4c8313d 100644
--- a/cloud/blockstore/libs/storage/partition_nonrepl/part_mirror_actor.h
+++ b/cloud/blockstore/libs/storage/partition_nonrepl/part_mirror_actor.h
@@ -175,12 +175,12 @@ class TMirrorPartitionActor final
         const NActors::TActorContext& ctx);
 
     void HandleAddLaggingAgent(
-        const TEvNonreplPartitionPrivate::TEvAddLaggingAgentRequest::
-            TPtr& ev,
+        const TEvNonreplPartitionPrivate::TEvAddLaggingAgentRequest::TPtr& ev,
         const NActors::TActorContext& ctx);
 
     void HandleRemoveLaggingAgent(
-        const TEvNonreplPartitionPrivate::TEvRemoveLaggingAgentRequest::TPtr& ev,
+        const TEvNonreplPartitionPrivate::TEvRemoveLaggingAgentRequest::TPtr&
+            ev,
         const NActors::TActorContext& ctx);
 
     void HandlePoisonPill(
@@ -219,6 +219,9 @@ class TMirrorPartitionActor final
     BLOCKSTORE_IMPLEMENT_REQUEST(ZeroBlocks, TEvService);
     BLOCKSTORE_IMPLEMENT_REQUEST(CheckRange, TEvVolume);
     BLOCKSTORE_IMPLEMENT_REQUEST(Drain, NPartition::TEvPartition);
+    BLOCKSTORE_IMPLEMENT_REQUEST(
+        WaitForInFlightWrites,
+        NPartition::TEvPartition);
 
     BLOCKSTORE_IMPLEMENT_REQUEST(DescribeBlocks, TEvVolume);
     BLOCKSTORE_IMPLEMENT_REQUEST(CompactRange, TEvVolume);
diff --git a/cloud/blockstore/libs/storage/partition_nonrepl/part_mirror_lagging_devices_ut.cpp b/cloud/blockstore/libs/storage/partition_nonrepl/part_mirror_lagging_devices_ut.cpp
index 62b6081f752..1b0c9a1681e 100644
--- a/cloud/blockstore/libs/storage/partition_nonrepl/part_mirror_lagging_devices_ut.cpp
+++ b/cloud/blockstore/libs/storage/partition_nonrepl/part_mirror_lagging_devices_ut.cpp
@@ -310,7 +310,7 @@ struct TTestEnv
              i < TBlockStoreComponents::END;
              ++i)
         {
-            Runtime.SetLogPriority(i, NLog::PRI_DEBUG);
+            Runtime.SetLogPriority(i, NLog::PRI_TRACE);
         }
     }
 
@@ -374,6 +374,27 @@ struct TTestEnv
                 TEvNonreplPartitionPrivate::TEvRemoveLaggingAgentRequest>(
                 *laggingAgent));
     }
+
+    void WaitForMigrationFinishEvent()
+    {
+        Runtime.AdvanceCurrentTime(Config->GetLaggingDevicePingInterval());
+        bool migrationFinished = false;
+        for (int i = 0; i < 100; i++) {
+            migrationFinished = Runtime.DispatchEvents(
+                {.FinalEvents =
+                     {{TEvVolumePrivate::EvLaggingAgentMigrationFinished}}},
+                TDuration::MilliSeconds(10));
+            if (migrationFinished) {
+                break;
+            }
+            Runtime.AdvanceCurrentTime(TDuration::Seconds(4));
+        }
+        if (!migrationFinished) {
+            Runtime.DispatchEvents(
+                {.FinalEvents = {
+                     {TEvVolumePrivate::EvLaggingAgentMigrationFinished}}});
+        }
+    }
 };
 
 }   // namespace
@@ -382,6 +403,38 @@ struct TTestEnv
 
 Y_UNIT_TEST_SUITE(TMirrorPartitionLaggingDevicesTest)
 {
+    Y_UNIT_TEST(ShouldMigrateLaggingDevices)
+    {
+        constexpr ui32 AgentCount = 3;
+        TTestBasicRuntime runtime(AgentCount);
+
+        TTestEnv env(runtime);
+
+        TPartitionClient client(runtime, env.MirrorPartActorId);
+
+        const auto fullDiskRange =
+            TBlockRange64::WithLength(0, DeviceBlockCount * 3);
+        client.WriteBlocks(fullDiskRange, 'A');
+
+        // The second device in the second replica is lagging.
+        env.AddLaggingAgent(env.Replicas[0][1].GetAgentId());
+
+        // Write blocks.
+        const auto firstAndSecondDevices =
+            TBlockRange64::WithLength(DeviceBlockCount - 1, 2);
+        client.WriteBlocks(firstAndSecondDevices, 'B');
+
+        // The second replica is lagging and is untouched.
+        env.ReadAndCheckContents(0, firstAndSecondDevices, 'B');
+        env.ReadAndCheckContents(1, firstAndSecondDevices, 'A');
+        env.ReadAndCheckContents(2, firstAndSecondDevices, 'B');
+
+        env.WaitForMigrationFinishEvent();
+
+        // After the migration all replicas are in sync.
+        env.ReadAndCheckContents(firstAndSecondDevices, 'B');
+    }
+
     Y_UNIT_TEST(ShouldDisableIOForOutdatedDevices)
     {
         constexpr ui32 AgentCount = 3;
diff --git a/cloud/blockstore/libs/storage/partition_nonrepl/part_mirror_resync_actor.cpp b/cloud/blockstore/libs/storage/partition_nonrepl/part_mirror_resync_actor.cpp
index d1b478b0ea6..447d3e8d12e 100644
--- a/cloud/blockstore/libs/storage/partition_nonrepl/part_mirror_resync_actor.cpp
+++ b/cloud/blockstore/libs/storage/partition_nonrepl/part_mirror_resync_actor.cpp
@@ -261,6 +261,9 @@ STFUNC(TMirrorPartitionResyncActor::StateWork)
         HFunc(TEvService::TEvWriteBlocksLocalRequest, HandleWriteBlocksLocal);
 
         HFunc(NPartition::TEvPartition::TEvDrainRequest, DrainActorCompanion.HandleDrain);
+        HFunc(
+            NPartition::TEvPartition::TEvWaitForInFlightWritesRequest,
+            DrainActorCompanion.HandleWaitForInFlightWrites);
         HFunc(TEvService::TEvGetChangedBlocksRequest, DeclineGetChangedBlocks);
         HFunc(
             TEvNonreplPartitionPrivate::TEvGetDeviceForRangeRequest,
@@ -315,6 +318,9 @@ STFUNC(TMirrorPartitionResyncActor::StateZombie)
         HFunc(TEvService::TEvWriteBlocksLocalRequest, RejectWriteBlocksLocal);
 
         HFunc(NPartition::TEvPartition::TEvDrainRequest, RejectDrain);
+        HFunc(
+            NPartition::TEvPartition::TEvWaitForInFlightWritesRequest,
+            RejectWaitForInFlightWrites);
         HFunc(TEvService::TEvGetChangedBlocksRequest, DeclineGetChangedBlocks);
         HFunc(
             TEvNonreplPartitionPrivate::TEvGetDeviceForRangeRequest,
diff --git a/cloud/blockstore/libs/storage/partition_nonrepl/part_mirror_resync_actor.h b/cloud/blockstore/libs/storage/partition_nonrepl/part_mirror_resync_actor.h
index 7725f693711..b811046a4e6 100644
--- a/cloud/blockstore/libs/storage/partition_nonrepl/part_mirror_resync_actor.h
+++ b/cloud/blockstore/libs/storage/partition_nonrepl/part_mirror_resync_actor.h
@@ -223,6 +223,9 @@ class TMirrorPartitionResyncActor final
     BLOCKSTORE_IMPLEMENT_REQUEST(WriteBlocksLocal, TEvService);
     BLOCKSTORE_IMPLEMENT_REQUEST(ZeroBlocks, TEvService);
     BLOCKSTORE_IMPLEMENT_REQUEST(Drain, NPartition::TEvPartition);
+    BLOCKSTORE_IMPLEMENT_REQUEST(
+        WaitForInFlightWrites,
+        NPartition::TEvPartition);
 
     BLOCKSTORE_IMPLEMENT_REQUEST(DescribeBlocks, TEvVolume);
     BLOCKSTORE_IMPLEMENT_REQUEST(CompactRange, TEvVolume);
diff --git a/cloud/blockstore/libs/storage/partition_nonrepl/part_nonrepl_actor.cpp b/cloud/blockstore/libs/storage/partition_nonrepl/part_nonrepl_actor.cpp
index 07f8aeda1f0..6d75852ed4b 100644
--- a/cloud/blockstore/libs/storage/partition_nonrepl/part_nonrepl_actor.cpp
+++ b/cloud/blockstore/libs/storage/partition_nonrepl/part_nonrepl_actor.cpp
@@ -655,6 +655,7 @@ STFUNC(TNonreplicatedPartitionActor::StateWork)
         HFunc(TEvService::TEvWriteBlocksLocalRequest, HandleWriteBlocksLocal);
 
         HFunc(NPartition::TEvPartition::TEvDrainRequest, DrainActorCompanion.HandleDrain);
+        HFunc(NPartition::TEvPartition::TEvWaitForInFlightWritesRequest, DrainActorCompanion.HandleWaitForInFlightWrites);
         HFunc(TEvService::TEvGetChangedBlocksRequest, DeclineGetChangedBlocks);
         HFunc(
             TEvNonreplPartitionPrivate::TEvGetDeviceForRangeRequest,
@@ -705,6 +706,7 @@ STFUNC(TNonreplicatedPartitionActor::StateZombie)
         HFunc(TEvService::TEvWriteBlocksLocalRequest, RejectWriteBlocksLocal);
 
         HFunc(NPartition::TEvPartition::TEvDrainRequest, RejectDrain);
+        HFunc(NPartition::TEvPartition::TEvWaitForInFlightWritesRequest, RejectWaitForInFlightWrites);
         HFunc(TEvService::TEvGetChangedBlocksRequest, DeclineGetChangedBlocks);
         HFunc(
             TEvNonreplPartitionPrivate::TEvGetDeviceForRangeRequest,
diff --git a/cloud/blockstore/libs/storage/partition_nonrepl/part_nonrepl_actor.h b/cloud/blockstore/libs/storage/partition_nonrepl/part_nonrepl_actor.h
index d1b2e30306e..a11e266f2bc 100644
--- a/cloud/blockstore/libs/storage/partition_nonrepl/part_nonrepl_actor.h
+++ b/cloud/blockstore/libs/storage/partition_nonrepl/part_nonrepl_actor.h
@@ -210,6 +210,9 @@ class TNonreplicatedPartitionActor final
     BLOCKSTORE_IMPLEMENT_REQUEST(DescribeBlocks, TEvVolume);
     BLOCKSTORE_IMPLEMENT_REQUEST(ChecksumBlocks, TEvNonreplPartitionPrivate);
     BLOCKSTORE_IMPLEMENT_REQUEST(Drain, NPartition::TEvPartition);
+    BLOCKSTORE_IMPLEMENT_REQUEST(
+        WaitForInFlightWrites,
+        NPartition::TEvPartition);
 
     BLOCKSTORE_IMPLEMENT_REQUEST(CompactRange, TEvVolume);
     BLOCKSTORE_IMPLEMENT_REQUEST(GetCompactionStatus, TEvVolume);
diff --git a/cloud/blockstore/libs/storage/partition_nonrepl/part_nonrepl_events_private.h b/cloud/blockstore/libs/storage/partition_nonrepl/part_nonrepl_events_private.h
index 3bb178ab41d..92693d2edbc 100644
--- a/cloud/blockstore/libs/storage/partition_nonrepl/part_nonrepl_events_private.h
+++ b/cloud/blockstore/libs/storage/partition_nonrepl/part_nonrepl_events_private.h
@@ -304,6 +304,10 @@ struct TEvNonreplPartitionPrivate
         {}
     };
 
+    struct TStartLaggingAgentMigration
+    {
+    };
+
     //
     // Events declaration
     //
@@ -332,6 +336,7 @@ struct TEvNonreplPartitionPrivate
         EvRemoveLaggingAgentRequest,
         EvAgentIsUnavailable,
         EvAgentIsBackOnline,
+        EvStartLaggingAgentMigration,
 
         BLOCKSTORE_PARTITION_NONREPL_REQUESTS_PRIVATE(BLOCKSTORE_DECLARE_EVENT_IDS)
 
@@ -414,6 +419,11 @@ struct TEvNonreplPartitionPrivate
         EvAgentIsBackOnline
     >;
 
+    using TEvStartLaggingAgentMigration = TRequestEvent<
+        TStartLaggingAgentMigration,
+        EvStartLaggingAgentMigration
+    >;
+
     BLOCKSTORE_PARTITION_NONREPL_REQUESTS_PRIVATE(BLOCKSTORE_DECLARE_PROTO_EVENTS)
 
 };
diff --git a/cloud/blockstore/libs/storage/partition_nonrepl/part_nonrepl_migration_common_actor.cpp b/cloud/blockstore/libs/storage/partition_nonrepl/part_nonrepl_migration_common_actor.cpp
index 9ac219d0841..71516d60204 100644
--- a/cloud/blockstore/libs/storage/partition_nonrepl/part_nonrepl_migration_common_actor.cpp
+++ b/cloud/blockstore/libs/storage/partition_nonrepl/part_nonrepl_migration_common_actor.cpp
@@ -242,6 +242,9 @@ STFUNC(TNonreplicatedPartitionMigrationCommonActor::StateWork)
         HFunc(
             NPartition::TEvPartition::TEvDrainRequest,
             DrainActorCompanion.HandleDrain);
+        HFunc(
+            NPartition::TEvPartition::TEvWaitForInFlightWritesRequest,
+            DrainActorCompanion.HandleWaitForInFlightWrites);
         HFunc(TEvService::TEvGetChangedBlocksRequest, DeclineGetChangedBlocks);
         HFunc(
             TEvNonreplPartitionPrivate::TEvGetDeviceForRangeRequest,
@@ -306,6 +309,9 @@ STFUNC(TNonreplicatedPartitionMigrationCommonActor::StateZombie)
         HFunc(TEvService::TEvWriteBlocksLocalRequest, RejectWriteBlocksLocal);
 
         HFunc(NPartition::TEvPartition::TEvDrainRequest, RejectDrain);
+        HFunc(
+            NPartition::TEvPartition::TEvWaitForInFlightWritesRequest,
+            RejectWaitForInFlightWrites);
         HFunc(TEvService::TEvGetChangedBlocksRequest, DeclineGetChangedBlocks);
         HFunc(
             TEvNonreplPartitionPrivate::TEvGetDeviceForRangeRequest,
diff --git a/cloud/blockstore/libs/storage/partition_nonrepl/part_nonrepl_migration_common_actor.h b/cloud/blockstore/libs/storage/partition_nonrepl/part_nonrepl_migration_common_actor.h
index 0de50284266..8873535cbf1 100644
--- a/cloud/blockstore/libs/storage/partition_nonrepl/part_nonrepl_migration_common_actor.h
+++ b/cloud/blockstore/libs/storage/partition_nonrepl/part_nonrepl_migration_common_actor.h
@@ -208,6 +208,9 @@ class TNonreplicatedPartitionMigrationCommonActor
     // processed.
     ui64 GetProcessedBlockCount() const;
 
+    // Called from the inheritor to check if migration is allowed.
+    [[nodiscard]] bool IsMigrationAllowed() const;
+
     // IPoisonPillHelperOwner implementation
     void Die(const NActors::TActorContext& ctx) override
     {
@@ -220,7 +223,6 @@ class TNonreplicatedPartitionMigrationCommonActor
     const TDiagnosticsConfigPtr& GetDiagnosticsConfig() const;
 
 private:
-    bool IsMigrationAllowed() const;
     bool IsMigrationFinished() const;
     bool IsIoDepthLimitReached() const;
     bool OverlapsWithInflightWriteAndZero(TBlockRange64 range) const;
@@ -303,6 +305,9 @@ class TNonreplicatedPartitionMigrationCommonActor
     BLOCKSTORE_IMPLEMENT_REQUEST(ZeroBlocks, TEvService);
     BLOCKSTORE_IMPLEMENT_REQUEST(ChecksumBlocks, TEvNonreplPartitionPrivate);
     BLOCKSTORE_IMPLEMENT_REQUEST(Drain, NPartition::TEvPartition);
+    BLOCKSTORE_IMPLEMENT_REQUEST(
+        WaitForInFlightWrites,
+        NPartition::TEvPartition);
 
     BLOCKSTORE_IMPLEMENT_REQUEST(DescribeBlocks, TEvVolume);
     BLOCKSTORE_IMPLEMENT_REQUEST(CompactRange, TEvVolume);


Reviews:
user: login: sharpeye
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-01T19:40:01Z
user: login: sharpeye
state: DISMISSED
author_association: COLLABORATOR
submitted_at: 2025-04-04T11:55:02Z
user: login: drbasic
state: DISMISSED
author_association: COLLABORATOR
submitted_at: 2025-04-04T12:19:42Z
user: login: drbasic
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-04T12:23:58Z
user: login: sharpeye
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-04T12:24:29Z

Review comments:
diff_hunk: @@ -222,11 +222,20 @@ class TPartitionActor final
 
     void SendStatsToService(const NActors::TActorContext& ctx);
 
+    // IRequestsInProgress implementation:
     bool WriteRequestInProgress() const override
     {
         return WriteAndZeroRequestsInProgress != 0;
     }
 
+    void WaitForInFlightWrites() override {
+        Y_ABORT("Unimplemented");
+    }
+
+    bool IsWaitingForInFlightWrites() const override {
+        Y_ABORT("Unimplemented");
+    }
path: cloud/blockstore/libs/storage/partition/part_actor.h
user: login: sharpeye
body: clang-format
created_at: 2025-04-01T19:13:15Z
author_association: COLLABORATOR
diff_hunk: @@ -241,11 +241,20 @@ class TPartitionActor final
 
     void SendStatsToService(const NActors::TActorContext& ctx);
 
+    // IRequestsInProgress implementation:
     bool WriteRequestInProgress() const override
     {
         return WriteAndZeroRequestsInProgress != 0;
     }
 
+    void WaitForInFlightWrites() override {
+        Y_ABORT("Unimplemented");
+    }
+
+    bool IsWaitingForInFlightWrites() const override {
+        Y_ABORT("Unimplemented");
+    }
path: cloud/blockstore/libs/storage/partition2/part2_actor.h
user: login: sharpeye
body: clang-format
created_at: 2025-04-01T19:13:28Z
author_association: COLLABORATOR
diff_hunk: @@ -57,16 +57,35 @@ void TLaggingAgentMigrationActor::OnBootstrap(const TActorContext& ctx)
             Config->GetMaxMigrationBandwidth(),
             Config->GetExpectedDiskAgentSize(),
             PartConfig));
-    StartWork(ctx);
 }
 
 bool TLaggingAgentMigrationActor::OnMessage(
     const TActorContext& ctx,
     TAutoPtr<IEventHandle>& ev)
 {
     Y_UNUSED(ctx);
+    switch (ev->GetTypeRewrite()) {
+        HFunc(
+            TEvNonreplPartitionPrivate::TEvStartLaggingAgentMigration,
+            HandleStartLaggingAgentMigration);
+
+        default:
+            // Message processing by the base class is required.
+            return false;
+    }
+
+    // We get here if we have processed an incoming message. And its processing
+    // by the base class is not required.
+    return true;
+}
+
+void TLaggingAgentMigrationActor::HandleStartLaggingAgentMigration(
+    const TEvNonreplPartitionPrivate::TEvStartLaggingAgentMigration::TPtr& ev,
+    const NActors::TActorContext& ctx)
+{
     Y_UNUSED(ev);
-    return false;
+    Y_ABORT_UNLESS(!IsMigrationAllowed());
path: cloud/blockstore/libs/storage/partition_nonrepl/lagging_agent_migration_actor.cpp
user: login: sharpeye
body: Есть еще Y_ABORT_IF:
```cpp
Y_ABORT_IF(IsMigrationAllowed());
```
created_at: 2025-04-01T19:16:22Z
author_association: COLLABORATOR
diff_hunk: @@ -675,10 +680,22 @@ void TLaggingAgentsReplicaProxyActor::HandleAgentIsBackOnline(
         return;
     }
 
+    auto& state = AgentState[msg->AgentId];
path: cloud/blockstore/libs/storage/partition_nonrepl/lagging_agents_replica_proxy_actor.cpp
user: login: sharpeye
body: ```cpp
auto* state = AgentState.FindPtr(msg->AgentId);
Y_DEBUG_ABORT_UNLESS(state);
if (!state) {
    return;
}
```
created_at: 2025-04-01T19:18:14Z
author_association: COLLABORATOR
diff_hunk: @@ -727,6 +747,82 @@ void TLaggingAgentsReplicaProxyActor::HandleAgentIsBackOnline(
     PoisonPillHelper.TakeOwnership(ctx, state.MigrationActorId);
 }
 
+void TLaggingAgentsReplicaProxyActor::HandleWaitForInFlightWrites(
+    const NPartition::TEvPartition::TEvWaitForInFlightWritesRequest::TPtr& ev,
+    const NActors::TActorContext& ctx)
+{
+    ForwardMessageToActor(ev, ctx, NonreplPartitionActorId);
+}
+
+void TLaggingAgentsReplicaProxyActor::HandleWaitForInFlightWritesResponse(
+    const NPartition::TEvPartition::TEvWaitForInFlightWritesResponse::TPtr& ev,
+    const NActors::TActorContext& ctx)
+{
+    if (!CurrentDrainingAgents.contains(ev->Cookie)) {
+        return;
+    }
+
+    auto it = CurrentDrainingAgents.find(ev->Cookie);
path: cloud/blockstore/libs/storage/partition_nonrepl/lagging_agents_replica_proxy_actor.cpp
user: login: sharpeye
body: Зачем по два раза искать?
created_at: 2025-04-01T19:19:50Z
author_association: COLLABORATOR
diff_hunk: @@ -727,6 +747,82 @@ void TLaggingAgentsReplicaProxyActor::HandleAgentIsBackOnline(
     PoisonPillHelper.TakeOwnership(ctx, state.MigrationActorId);
 }
 
+void TLaggingAgentsReplicaProxyActor::HandleWaitForInFlightWrites(
+    const NPartition::TEvPartition::TEvWaitForInFlightWritesRequest::TPtr& ev,
+    const NActors::TActorContext& ctx)
+{
+    ForwardMessageToActor(ev, ctx, NonreplPartitionActorId);
+}
+
+void TLaggingAgentsReplicaProxyActor::HandleWaitForInFlightWritesResponse(
+    const NPartition::TEvPartition::TEvWaitForInFlightWritesResponse::TPtr& ev,
+    const NActors::TActorContext& ctx)
+{
+    if (!CurrentDrainingAgents.contains(ev->Cookie)) {
+        return;
+    }
+
+    auto it = CurrentDrainingAgents.find(ev->Cookie);
+    TString agentId = std::move(it->second);
+    CurrentDrainingAgents.erase(it);
+
+    const auto* msg = ev->Get();
+    if (HasError(msg->GetError())) {
+        LOG_WARN(
+            ctx,
+            TBlockStoreComponents::PARTITION_WORKER,
+            "[%s] Failed to drain partitions for lagging agent %s: %s",
+            PartConfig->GetName().c_str(),
+            agentId.Quote().c_str(),
+            FormatError(msg->GetError()).c_str());
+
+        DrainRequestCounter++;
+        Y_ABORT_UNLESS(!CurrentDrainingAgents.contains(DrainRequestCounter));
+        CurrentDrainingAgents[DrainRequestCounter] = agentId;
+
+        ctx.ExecutorThread.Schedule(
+            TDuration::Seconds(1),
+            new IEventHandle(
+                MirrorPartitionActorId,
+                SelfId(),
+                new NPartition::TEvPartition::
+                    TEvWaitForInFlightWritesRequest()));
+        return;
+    }
+
+    Y_DEBUG_ABORT_UNLESS(AgentState.contains(agentId));
+    if (!AgentState.contains(agentId)) {
+        return;
+    }
+
+    auto& state = AgentState[agentId];
path: cloud/blockstore/libs/storage/partition_nonrepl/lagging_agents_replica_proxy_actor.cpp
user: login: sharpeye
body: FindPtr
created_at: 2025-04-01T19:23:28Z
author_association: COLLABORATOR
diff_hunk: @@ -57,6 +59,9 @@ class TLaggingAgentsReplicaProxyActor final
     };
     THashMap<TString, TAgentState> AgentState;
 
+    ui64 DrainRequestCounter = 0;
+    THashMap<ui64, TString> CurrentDrainingAgents;
path: cloud/blockstore/libs/storage/partition_nonrepl/lagging_agents_replica_proxy_actor.h
user: login: sharpeye
body: Камента не хватает, как используется CurrentDrainingAgents
created_at: 2025-04-01T19:27:59Z
author_association: COLLABORATOR
diff_hunk: @@ -87,16 +132,31 @@ void TDrainActorCompanion::ProcessDrainRequests(
             "Drain",
             requestInfo->CallContext->RequestId);
 
-        LOG_TRACE(
-            ctx,
-            TBlockStoreComponents::PARTITION,
-            "[%s] Complete drain",
-            LoggingId.c_str());
-
         NCloud::Reply(ctx, *requestInfo, std::move(response));
     }
 
     DrainRequests.clear();
 }
 
+void TDrainActorCompanion::DoProcessWaitForInFlightWritesRequests(
+    const NActors::TActorContext& ctx)
path: cloud/blockstore/libs/storage/partition_common/drain_actor_companion.cpp
user: login: sharpeye
body: Не хватает `using namespace NActors;`
created_at: 2025-04-01T19:30:34Z
author_association: COLLABORATOR

PR comments:
user: login: github-actions[bot]
created_at: 2025-03-27T08:18:53Z
body: <!-- status pr=3266, run=10335, build_preset=linux-x86_64-relwithdebinfo, dry_run=False -->
> [!NOTE]
> This is an automated comment that will be appended during run.

:green_circle: **linux-x86_64-relwithdebinfo**: all tests PASSED for commit 87b48272a66a5dddaaa57ad5245f2a30be0e80cf.
| TESTS | PASSED | ERRORS | FAILED | SKIPPED | MUTED<sup>[?](https://github.com/ydb-platform/nbs/tree/main/.github/config "All mute rules are defined here")</sup> |
| ---: | ---: | ---: | ---: | ---: | ---: |
| [3923](https://github-actions-s3.storage.eu-north1.nebius.cloud/ydb-platform/nbs/PR-check/14100523079/1/nebius-x86-64/summary/ya-test.html) | [3923](https://github-actions-s3.storage.eu-north1.nebius.cloud/ydb-platform/nbs/PR-check/14100523079/1/nebius-x86-64/summary/ya-test.html#PASS) | 0 | 0 | 0 | 0 |
author_association: CONTRIBUTOR
user: login: github-actions[bot]
created_at: 2025-03-28T11:07:04Z
body: <!-- status pr=3266, run=10387, build_preset=linux-x86_64-relwithdebinfo, dry_run=False -->
> [!NOTE]
> This is an automated comment that will be appended during run.

:green_circle: **linux-x86_64-relwithdebinfo**: all tests PASSED for commit 2c70f1e9693236b2df4eed8cb906ef9ad2c0ad48.
| TESTS | PASSED | ERRORS | FAILED | SKIPPED | MUTED<sup>[?](https://github.com/ydb-platform/nbs/tree/main/.github/config "All mute rules are defined here")</sup> |
| ---: | ---: | ---: | ---: | ---: | ---: |
| [3933](https://github-actions-s3.storage.eu-north1.nebius.cloud/ydb-platform/nbs/PR-check/14126070764/1/nebius-x86-64/summary/ya-test.html) | [3933](https://github-actions-s3.storage.eu-north1.nebius.cloud/ydb-platform/nbs/PR-check/14126070764/1/nebius-x86-64/summary/ya-test.html#PASS) | 0 | 0 | 0 | 0 |
author_association: CONTRIBUTOR
user: login: github-actions[bot]
created_at: 2025-04-02T10:22:30Z
body: <!-- status pr=3266, run=10526, build_preset=linux-x86_64-relwithdebinfo, dry_run=False -->
> [!NOTE]
> This is an automated comment that will be appended during run.

:green_circle: **linux-x86_64-relwithdebinfo**: all tests PASSED for commit 730d3354ca5f2459a20c0f2fd59597f7ccecac32.
| TESTS | PASSED | ERRORS | FAILED | SKIPPED | MUTED<sup>[?](https://github.com/ydb-platform/nbs/tree/main/.github/config "All mute rules are defined here")</sup> |
| ---: | ---: | ---: | ---: | ---: | ---: |
| [3949](https://github-actions-s3.storage.eu-north1.nebius.cloud/ydb-platform/nbs/PR-check/14214133800/1/nebius-x86-64/summary/ya-test.html) | [3949](https://github-actions-s3.storage.eu-north1.nebius.cloud/ydb-platform/nbs/PR-check/14214133800/1/nebius-x86-64/summary/ya-test.html#PASS) | 0 | 0 | 0 | 0 |
author_association: CONTRIBUTOR
user: login: github-actions[bot]
created_at: 2025-04-04T14:52:05Z
body: <!-- status pr=3266, run=10630, build_preset=linux-x86_64-relwithdebinfo, dry_run=False -->
> [!NOTE]
> This is an automated comment that will be appended during run.

:green_circle: **linux-x86_64-relwithdebinfo**: all tests PASSED for commit 00e008d3d6fc0a63652d6fbe71b1a30864a5ea6a.
| TESTS | PASSED | ERRORS | FAILED | SKIPPED | MUTED<sup>[?](https://github.com/ydb-platform/nbs/tree/main/.github/config "All mute rules are defined here")</sup> |
| ---: | ---: | ---: | ---: | ---: | ---: |
| [3952](https://github-actions-s3.storage.eu-north1.nebius.cloud/ydb-platform/nbs/PR-check/14265241300/1/nebius-x86-64/summary/ya-test.html) | [3952](https://github-actions-s3.storage.eu-north1.nebius.cloud/ydb-platform/nbs/PR-check/14265241300/1/nebius-x86-64/summary/ya-test.html#PASS) | 0 | 0 | 0 | 0 |
author_association: CONTRIBUTOR