[nbs2] ddisk sessions: reject ddisk requests without locked session (#46153)

Co-authored-by: Yudov Maksim <[email protected]>

3 files changed, 215 insertions, 76 deletions

diff --git a/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/direct_block_group_impl.cpp b/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/direct_block_group_impl.cpp
index 3c279105506..245bbc5e67a 100644
--- a/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/direct_block_group_impl.cpp
+++ b/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/direct_block_group_impl.cpp
@@ -31,6 +31,9 @@ constexpr ui64 InitialDDiskSessionSeqNo = 0;
 constexpr size_t MinLockedDDiskSessionsToStart =
     QuorumDirectBlockGroupHostCount;
 
+constexpr TStringBuf DDiskSessionIsNotEstablishedMessage =
+    "DDisk session is not established";
+
 ////////////////////////////////////////////////////////////////////////////////
 
 TListPBufferResponse MakeListPBufferResponse(
@@ -76,6 +79,61 @@ THostSnapshot MakeHostSnapshot(const TOracleHostStat& stat)
     };
 }
 
+// help function for TDirectBlockGroup::SyncWithPBuffer
+std::function<void(const TFuture<NProto::TError>&)>
+CreateWaitSessionCbForSyncWithPBuffer(
+    TPromise<TDBGFlushResponse>&& promise,
+    std::weak_ptr<TDirectBlockGroup>&& weakSelf,
+    ui32 vChunkIndex,
+    THostIndex pbufferHostIndex,
+    THostIndex ddiskHostIndex,
+    const TVector<TPBufferSegment>& segments,
+    std::shared_ptr<NWilson::TSpan> childSpan)
+{
+    using TDBGFlushResponseFuture = NThreading::TFuture<TDBGFlushResponse>;
+    auto cb = [weakSelf = std::move(weakSelf),
+               promise = std::move(promise),
+               vChunkIndex,
+               pbufferHostIndex,
+               ddiskHostIndex,
+               segments = segments,
+               childSpan = std::move(childSpan)]   //
+        (const TFuture<NProto::TError>& f) mutable
+    {
+        TDBGFlushResponse flushResponse;
+        if (HasError(f.GetValue())) {
+            for (size_t i = 0; i < segments.size(); ++i) {
+                flushResponse.Errors.push_back(MakeError(
+                    E_REJECTED,
+                    TString(DDiskSessionIsNotEstablishedMessage)));
+            }
+            promise.SetValue(std::move(flushResponse));
+            return;
+        }
+
+        if (auto self = weakSelf.lock()) {
+            childSpan->Event("ConnectionReady");
+
+            self->SyncWithPBuffer(
+                    vChunkIndex,
+                    pbufferHostIndex,
+                    ddiskHostIndex,
+                    segments,
+                    childSpan->GetTraceId())
+                .Subscribe([promise = std::move(promise)]   //
+                           (const TDBGFlushResponseFuture& f) mutable
+                           { promise.SetValue(f.GetValue()); });
+        } else {
+            for (size_t i = 0; i < segments.size(); ++i) {
+                flushResponse.Errors.push_back(MakeError(E_CANCELLED));
+            }
+        }
+        promise.SetValue(std::move(flushResponse));
+    };
+
+    return cb;
+}
+
 }   // namespace
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -259,19 +317,60 @@ TDirectBlockGroup::ReadBlocksFromDDisk(
 
     using TEvReadResultFuture =
         TFuture<NKikimrBlobStorage::NDDisk::TEvReadResult>;
-
-    if (!WaitForSessionLock(hostIndex)) {
-        return MakeFuture<TDBGReadBlocksResponse>(
-            {.Error =
-                 MakeError(E_REJECTED, "DDisk session is not established")});
-    }
-
-    auto childSpan =
-        CreateChildSpan(traceId, "NbsPartition.ReadBlocks.ReadDDisk");
+    using TDBGReadBlocksResponseFuture =
+        NThreading::TFuture<TDBGReadBlocksResponse>;
 
     auto startAt = TMonotonic::Now();
     auto promise = NewPromise<TDBGReadBlocksResponse>();
     auto result = promise.GetFuture();
+    auto childSpan =
+        CreateChildSpan(traceId, "NbsPartition.ReadBlocks.ReadDDisk");
+
+    if (DDiskConnections[hostIndex].SessionState != EDDiskSessionState::Locked)
+    {
+        childSpan->Event("WaitConnectionReady");
+
+        auto waitReadyCb = [weakSelf = weak_from_this(),
+                            promise = std::move(promise),
+                            vChunkIndex,
+                            hostIndex,
+                            range,
+                            guardedSglist = guardedSglist,
+                            childSpan = std::move(childSpan)]   //
+            (const TFuture<NProto::TError>& f) mutable
+        {
+            if (HasError(f.GetValue())) {
+                promise.SetValue(TDBGReadBlocksResponse{
+                    .Error = MakeError(
+                        E_REJECTED,
+                        TString(DDiskSessionIsNotEstablishedMessage))});
+                return;
+            }
+
+            if (auto self = weakSelf.lock()) {
+                childSpan->Event("ConnectionReady");
+
+                self->ReadBlocksFromDDisk(
+                        vChunkIndex,
+                        hostIndex,
+                        range,
+                        guardedSglist,
+                        childSpan->GetTraceId())
+                    .Subscribe([promise = std::move(promise)]   //
+                               (const TDBGReadBlocksResponseFuture& f) mutable
+                               { promise.SetValue(UnsafeExtractValue(f)); });
+
+            } else {
+                promise.SetValue(
+                    TDBGReadBlocksResponse{.Error = MakeError(E_CANCELLED)});
+            }
+        };
+        DDiskConnections[hostIndex].GetFuture().Subscribe(
+            std::move(waitReadyCb));
+
+        return result;
+    }
+
     OnRequest(hostIndex, EOperation::ReadFromDDisk);
     auto future = StorageTransport->ReadFromDDisk(
         DDiskConnections[hostIndex].HostConnection,
@@ -420,20 +519,60 @@ TDirectBlockGroup::WriteBlocksToDDisk(
 
     using TEvWriterResultFuture =
         NThreading::TFuture<NKikimrBlobStorage::NDDisk::TEvWriteResult>;
+    using TDBGWriteBlocksResponseFuture =
+        NThreading::TFuture<TDBGWriteBlocksResponse>;
 
     const auto startAt = TMonotonic::Now();
-
-    if (!WaitForSessionLock(hostIndex)) {
-        return MakeFuture<TDBGWriteBlocksResponse>(
-            {.Error =
-                 MakeError(E_REJECTED, "DDisk session is not established")});
-    }
-
     auto childSpan =
         CreateChildSpan(traceId, "NbsPartition.WriteBlocksToDDisk");
 
     auto promise = NewPromise<TDBGWriteBlocksResponse>();
     auto result = promise.GetFuture();
+
+    if (DDiskConnections[hostIndex].SessionState != EDDiskSessionState::Locked)
+    {
+        childSpan->Event("WaitConnectionReady");
+
+        auto waitReadyCb = [weakSelf = weak_from_this(),
+                            promise = std::move(promise),
+                            vChunkIndex,
+                            hostIndex,
+                            range,
+                            guardedSglist = guardedSglist,
+                            childSpan = std::move(childSpan)]   //
+            (const TFuture<NProto::TError>& f) mutable
+        {
+            if (HasError(f.GetValue())) {
+                promise.SetValue(TDBGWriteBlocksResponse{
+                    .Error = MakeError(
+                        E_REJECTED,
+                        TString(DDiskSessionIsNotEstablishedMessage))});
+                return;
+            }
+            if (auto self = weakSelf.lock()) {
+                childSpan->Event("ConnectionReady");
+
+                self->WriteBlocksToDDisk(
+                        vChunkIndex,
+                        hostIndex,
+                        range,
+                        guardedSglist,
+                        childSpan->GetTraceId())
+                    .Subscribe([promise = std::move(promise)]   //
+                               (const TDBGWriteBlocksResponseFuture& f) mutable
+                               { promise.SetValue(f.GetValue()); });
+
+            } else {
+                promise.SetValue(
+                    TDBGWriteBlocksResponse{.Error = MakeError(E_CANCELLED)});
+            }
+        };
+        DDiskConnections[hostIndex].GetFuture().Subscribe(
+            std::move(waitReadyCb));
+
+        return result;
+    }
+
     OnRequest(hostIndex, EOperation::WriteToDDisk);
     auto future = StorageTransport->WriteToDDisk(
         DDiskConnections[hostIndex].HostConnection,
@@ -737,14 +876,26 @@ NThreading::TFuture<TDBGFlushResponse> TDirectBlockGroup::SyncWithPBuffer(
 {
     Y_ABORT_UNLESS(ExecutorThreadChecker.Check());
 
-    if (!WaitForSessionLock(ddiskHostIndex)) {
-        TDBGFlushResponse rejectResponse;
-        rejectResponse.Errors.reserve(segments.size());
-        for (size_t i = 0; i < segments.size(); ++i) {
-            rejectResponse.Errors.push_back(
-                MakeError(E_REJECTED, "DDisk session is not established"));
-        }
-        return MakeFuture<TDBGFlushResponse>(std::move(rejectResponse));
+    auto childSpan = CreateChildSpan(traceId, "NbsPartition.SyncWithPBuffer");
+    auto promise = NewPromise<TDBGFlushResponse>();
+    auto flushFuture = promise.GetFuture();
+
+    if (DDiskConnections[ddiskHostIndex].SessionState !=
+        EDDiskSessionState::Locked)
+    {
+        childSpan->Event("WaitConnectionReady");
+
+        auto cb = CreateWaitSessionCbForSyncWithPBuffer(
+            std::move(promise),
+            std::move(weak_from_this()),
+            vChunkIndex,
+            pbufferHostIndex,
+            ddiskHostIndex,
+            segments,
+            childSpan);
+        DDiskConnections[ddiskHostIndex].GetFuture().Subscribe(std::move(cb));
+
+        return flushFuture;
     }
 
     const auto startAt = TMonotonic::Now();
@@ -757,8 +908,6 @@ NThreading::TFuture<TDBGFlushResponse> TDirectBlockGroup::SyncWithPBuffer(
             segment.Range.Size() * DefaultBlockSize));
     }
 
-    auto childSpan = CreateChildSpan(traceId, "NbsPartition.SyncWithPBuffer");
-
     if (pbufferHostIndex == ddiskHostIndex) {
         OnRequest(ddiskHostIndex, EOperation::Flush);
     } else {
@@ -773,9 +922,6 @@ NThreading::TFuture<TDBGFlushResponse> TDirectBlockGroup::SyncWithPBuffer(
         TPBufferSegment::MakeLsnVector(segments),
         childSpan.get());
 
-    auto promise = NewPromise<TDBGFlushResponse>();
-    auto flushFuture = promise.GetFuture();
-
     future.Subscribe(
         [weakSelf = weak_from_this(),
          promise = std::move(promise),
@@ -1734,17 +1880,6 @@ void TDirectBlockGroup::ScheduleOracleThinking()
         });
 }
 
-bool TDirectBlockGroup::WaitForSessionLock(THostIndex hostIndex)
-{
-    const auto& conn = DDiskConnections[hostIndex];
-    if (conn.SessionState != EDDiskSessionState::Locked) {
-        const auto sessionReadyFuture = conn.GetFuture();
-        Executor->WaitFor(sessionReadyFuture);
-        return conn.SessionState == EDDiskSessionState::Locked;
-    }
-    return true;
-}
-
 void TDirectBlockGroup::HandleBlockedGeneration(
     THostIndex hostIndex,
     TStringBuf context)
diff --git a/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/direct_block_group_impl.h b/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/direct_block_group_impl.h
index 341b8a53961..183662ea539 100644
--- a/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/direct_block_group_impl.h
+++ b/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/direct_block_group_impl.h
@@ -249,8 +249,6 @@ private:
     void Thinking();
     void ScheduleOracleThinking();
 
-    [[nodiscard]] bool WaitForSessionLock(THostIndex hostIndex);
-
     void HandleBlockedGeneration(THostIndex hostIndex, TStringBuf context);
 
     [[nodiscard]] TDBGDumpResponse DoDebugPrintDirtyMap() const;
diff --git a/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/direct_block_group_impl_ut.cpp b/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/direct_block_group_impl_ut.cpp
index 85f0122bfa1..a1ad2ff6ede 100644
--- a/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/direct_block_group_impl_ut.cpp
+++ b/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/direct_block_group_impl_ut.cpp
@@ -66,6 +66,13 @@ void ResolveConnects(
     }
 }
 
+NWilson::TTraceId CreateTraceId()
+{
+    return NWilson::TTraceId::NewTraceId(
+        NWilson::TTraceId::MAX_VERBOSITY,
+        NWilson::TTraceId::MAX_TIME_TO_LIVE);
+}
+
 }   // namespace
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -146,7 +153,7 @@ Y_UNIT_TEST_SUITE(TDirectBlockGroupTest)
                         0,
                         range,
                         MakeSgList(readyBuffer),
-                        NWilson::TTraceId());
+                        CreateTraceId());
                 });
             UNIT_ASSERT_VALUES_EQUAL(
                 S_OK,
@@ -165,7 +172,7 @@ Y_UNIT_TEST_SUITE(TDirectBlockGroupTest)
                         1,
                         range,
                         MakeSgList(readyWriteBuffer),
-                        NWilson::TTraceId());
+                        CreateTraceId());
                 });
             UNIT_ASSERT_VALUES_EQUAL(
                 S_OK,
@@ -207,16 +214,17 @@ Y_UNIT_TEST_SUITE(TDirectBlockGroupTest)
                     3,
                     range,
                     MakeSgList(pendingBuffer),
-                    NWilson::TTraceId());
+                    CreateTraceId());
             });
+        auto inFlightRead = WaitFuture(executor, pendingRead, WaitTimeout);
         DrainExecutor(executor);
-        UNIT_ASSERT(!pendingRead.HasValue());
+        UNIT_ASSERT(!inFlightRead.HasValue());
 
         // Establishing the session unblocks the read.
         pendingHost3.SetValue(TStorageTransportMock::MakeConnectResult());
-        UNIT_ASSERT_VALUES_EQUAL(
-            S_OK,
-            GetResponse(pendingRead).Error.GetCode());
+        DrainExecutor(executor);
+        auto response = WaitFuture(executor, inFlightRead, WaitTimeout);
+        UNIT_ASSERT_VALUES_EQUAL(S_OK, response.Error.GetCode());
     }
 
     // The tablet-wide "all DBGs ready" gate (WaitAll over per-DBG
@@ -296,7 +304,7 @@ Y_UNIT_TEST_SUITE(TDirectBlockGroupTest)
                     ddiskHost,   // host
                     TBlockRange64::WithLength(0, 3),
                     guardedSglist,
-                    NWilson::TTraceId());
+                    CreateTraceId());
             });
         auto readyReadResponse =
             pendingRead.GetValue(WaitTimeout).GetValue(WaitTimeout);
@@ -359,7 +367,7 @@ Y_UNIT_TEST_SUITE(TDirectBlockGroupTest)
                     TBlockRange64::WithLength(0, 3),
                     TDuration::Seconds(1),
                     guardedSglist,
-                    NWilson::TTraceId(),
+                    CreateTraceId(),
                     cb);
                 return future;
             });
@@ -422,7 +430,7 @@ Y_UNIT_TEST_SUITE(TDirectBlockGroupTest)
                     pbufferHost,
                     ddiskHost,
                     segments,
-                    NWilson::TTraceId());
+                    CreateTraceId());
             });
         auto flushResponse =
             pendingFlush.GetValue(WaitTimeout).GetValue(WaitTimeout);
@@ -485,7 +493,7 @@ Y_UNIT_TEST_SUITE(TDirectBlockGroupTest)
                     pbufferHost,
                     ddiskHost,
                     segments,
-                    NWilson::TTraceId());
+                    CreateTraceId());
             });
         auto flushResponse =
             pendingFlush.GetValue(WaitTimeout).GetValue(WaitTimeout);
@@ -564,7 +572,7 @@ Y_UNIT_TEST_SUITE(TDirectBlockGroupTest)
                     TBlockRange64::WithLength(0, 3),
                     TDuration::Seconds(1),
                     guardedSglist,
-                    NWilson::TTraceId(),
+                    CreateTraceId(),
                     cb);
                 return future;
             });
@@ -644,7 +652,7 @@ Y_UNIT_TEST_SUITE(TDirectBlockGroupTest)
                     TBlockRange64::WithLength(0, 3),
                     TDuration::Seconds(1),
                     guardedSglist,
-                    NWilson::TTraceId(),
+                    CreateTraceId(),
                     cb);
                 return future;
             });
@@ -722,10 +730,12 @@ Y_UNIT_TEST_SUITE(TDirectBlockGroupTest)
                     0,
                     range,
                     MakeSgList(pendingBuffer),
-                    NWilson::TTraceId());
+                    CreateTraceId());
             });
-        DrainExecutor(executor);
-        UNIT_ASSERT(!pendingRead.HasValue());
+
+        auto inFlightRead = WaitFuture(executor, pendingRead, WaitTimeout);
+        DoAllExecutorAndRuntimeWork(executor);
+        UNIT_ASSERT(!inFlightRead.HasValue());
 
         // Resolve the connect with BLOCKED -> stale generation, suicide.
         pendingConnect0.SetValue(TStorageTransportMock::MakeConnectResult(
@@ -733,9 +743,8 @@ Y_UNIT_TEST_SUITE(TDirectBlockGroupTest)
             NKikimrBlobStorage::NDDisk::TReplyStatus::BLOCKED));
 
         // The suspended read is rejected once the session becomes Broken.
-        UNIT_ASSERT_VALUES_EQUAL(
-            E_REJECTED,
-            GetResponse(pendingRead).Error.GetCode());
+        auto response = WaitFuture(executor, inFlightRead, WaitTimeout);
+        UNIT_ASSERT_VALUES_EQUAL(E_REJECTED, response.Error.GetCode());
 
         const auto state = GetBlockedDetected(executor, dbg, 0, WaitTimeout);
         UNIT_ASSERT(state.DDiskSessionBroken);
@@ -779,7 +788,7 @@ Y_UNIT_TEST_SUITE(TDirectBlockGroupTest)
                     0,
                     range,
                     MakeSgList(writeBuffer),
-                    NWilson::TTraceId());
+                    CreateTraceId());
             });
         UNIT_ASSERT_VALUES_EQUAL(
             E_REJECTED,
@@ -818,7 +827,7 @@ Y_UNIT_TEST_SUITE(TDirectBlockGroupTest)
                     0,
                     range,
                     MakeSgList(readBuffer),
-                    NWilson::TTraceId());
+                    CreateTraceId());
             });
         UNIT_ASSERT_VALUES_EQUAL(
             E_REJECTED,
@@ -863,7 +872,7 @@ Y_UNIT_TEST_SUITE(TDirectBlockGroupTest)
                     pbufferHost,
                     ddiskHost,
                     segments,
-                    NWilson::TTraceId());
+                    CreateTraceId());
             });
         auto flushResponse =
             pendingFlush.GetValue(WaitTimeout).GetValue(WaitTimeout);
@@ -907,7 +916,7 @@ Y_UNIT_TEST_SUITE(TDirectBlockGroupTest)
                         host,
                         range,
                         MakeSgList(writeBuffer),
-                        NWilson::TTraceId());
+                        CreateTraceId());
                 });
             UNIT_ASSERT_VALUES_EQUAL(
                 E_REJECTED,
@@ -1012,7 +1021,7 @@ Y_UNIT_TEST_SUITE(TDirectBlockGroupTest)
                     TBlockRange64::WithLength(0, 3),
                     TDuration::Seconds(1),
                     guardedSglist,
-                    NWilson::TTraceId(),
+                    CreateTraceId(),
                     cb);
                 return future;
             });
@@ -1298,7 +1307,7 @@ Y_UNIT_TEST_SUITE(TSessionsWithRealTransport)
         const auto range = TBlockRange64::WithLength(0, 1);
         TString buffer(DefaultBlockSize, 'r');
 
-        // The read on host 0 suspends inside WaitForSessionLock.
+        // The read on host 0 suspends inside ReadBlocksFromDDisk.
         auto pendingRead = RunOnExecutor(
             executor,
             [&]
@@ -1308,10 +1317,11 @@ Y_UNIT_TEST_SUITE(TSessionsWithRealTransport)
                     0,
                     range,
                     MakeSgList(buffer),
-                    NWilson::TTraceId());
+                    CreateTraceId());
             });
+        auto inFlightRead = WaitFuture(executor, pendingRead, WaitTimeout);
         DoAllExecutorAndRuntimeWork(executor);
-        UNIT_ASSERT(!pendingRead.HasValue());
+        UNIT_ASSERT(!inFlightRead.HasValue());
 
         // The disconnect resets the session: ResetSession wakes the waiter with
         // an error.
@@ -1320,9 +1330,7 @@ Y_UNIT_TEST_SUITE(TSessionsWithRealTransport)
             ddisks[0],
             transportPtr->GetNodeId());
 
-        auto innerRead = WaitFuture(executor, pendingRead, WaitTimeout);
-        UNIT_ASSERT(pendingRead.HasValue());
-        auto response = WaitFuture(executor, innerRead, WaitTimeout);
+        auto response = WaitFuture(executor, inFlightRead, WaitTimeout);
         UNIT_ASSERT_VALUES_UNEQUAL(S_OK, response.Error.GetCode());
     }
 
@@ -1355,12 +1363,10 @@ Y_UNIT_TEST_SUITE(TSessionsWithRealTransport)
                     0,
                     range,
                     MakeSgList(buffer),
-                    NWilson::TTraceId());
+                    CreateTraceId());
             });
 
-        // The session is already established, so the read does not suspend in
-        // WaitForSessionLock: the outer future resolves immediately, carrying
-        // the still-pending in-flight read future.
+        // The session is already established, so the read does not suspend.
         auto inFlightRead = WaitFuture(executor, pendingRead, WaitTimeout);
         DoAllExecutorAndRuntimeWork(executor);
         UNIT_ASSERT(!inFlightRead.HasValue());