[NBS-7153][nbs2] ddisk session handle connect errors (#46817)

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

7 files changed, 242 insertions, 32 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 c124311c3d5..f56d55568fb 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
@@ -112,14 +112,18 @@ CreateWaitSessionCbForSyncWithPBuffer(
         }
 
         if (auto self = weakSelf.lock()) {
-            childSpan->Event("ConnectionReady");
+            NWilson::TTraceId traceId;
+            if (childSpan) {
+                childSpan->Event("ConnectionReady");
+                traceId = childSpan->GetTraceId();
+            }
 
             self->SyncWithPBuffer(
                     vChunkIndex,
                     pbufferHostIndex,
                     ddiskHostIndex,
                     segments,
-                    childSpan->GetTraceId())
+                    traceId)
                 .Subscribe([promise = std::move(promise)]   //
                            (const TDBGFlushResponseFuture& f) mutable
                            { promise.SetValue(f.GetValue()); });
@@ -127,8 +131,8 @@ CreateWaitSessionCbForSyncWithPBuffer(
             for (size_t i = 0; i < segments.size(); ++i) {
                 flushResponse.Errors.push_back(MakeError(E_CANCELLED));
             }
+            promise.SetValue(std::move(flushResponse));
         }
-        promise.SetValue(std::move(flushResponse));
     };
 
     return cb;
@@ -315,14 +319,18 @@ TDirectBlockGroup::ReadBlocksFromDDisk(
             }
 
             if (auto self = weakSelf.lock()) {
-                childSpan->Event("ConnectionReady");
+                NWilson::TTraceId traceId;
+                if (childSpan) {
+                    childSpan->Event("ConnectionReady");
+                    traceId = childSpan->GetTraceId();
+                }
 
                 self->ReadBlocksFromDDisk(
                         vChunkIndex,
                         hostIndex,
                         range,
                         guardedSglist,
-                        childSpan->GetTraceId())
+                        traceId)
                     .Subscribe([promise = std::move(promise)]   //
                                (const TDBGReadBlocksResponseFuture& f) mutable
                                { promise.SetValue(UnsafeExtractValue(f)); });
@@ -519,14 +527,18 @@ TDirectBlockGroup::WriteBlocksToDDisk(
                 return;
             }
             if (auto self = weakSelf.lock()) {
-                childSpan->Event("ConnectionReady");
+                NWilson::TTraceId traceId;
+                if (childSpan) {
+                    childSpan->Event("ConnectionReady");
+                    traceId = childSpan->GetTraceId();
+                }
 
                 self->WriteBlocksToDDisk(
                         vChunkIndex,
                         hostIndex,
                         range,
                         guardedSglist,
-                        childSpan->GetTraceId())
+                        traceId)
                     .Subscribe([promise = std::move(promise)]   //
                                (const TDBGWriteBlocksResponseFuture& f) mutable
                                { promise.SetValue(f.GetValue()); });
@@ -1211,9 +1223,11 @@ NThreading::TFuture<TListPBufferResponse> TDirectBlockGroup::ListPBuffers(
     }
 
     const auto& connection = PBufferConnections[hostIndex];
+    // Hold a local copy of the connect future,
+    // do not put an address of changeable field into a wait.
+    auto connectFuture = connection.GetFuture();
     // Switch co-routine context if needed.
-    const NProto::TError& connectError =
-        Executor->WaitFor(connection.GetFuture());
+    const NProto::TError& connectError = Executor->WaitFor(connectFuture);
     if (HasError(connectError)) {
         return MakeFuture(TListPBufferResponse{.Error = connectError});
     }
@@ -1551,18 +1565,16 @@ void TDirectBlockGroup::OnConnectionEstablished(
         // Unblock waiters on ConnectFuture with the error.
         connection.ConnectPromise.SetValue(error);
         return;
-        // TODO (future phase): handle non-BLOCKED connect errors
-        // (ERROR/unavailability) via reconnect.
-    } else if (IsInitialized()) {
+    } else {
         LOG_ERROR(
             *ActorSystem,
             NKikimrServices::NBS_PARTITION,
-            "%s connection failed for %s (post-init): %s",
+            "%s connection failed %s: %s",
             LogTitle.GetWithTime().c_str(),
             PrintHostIndex(hostIndex).c_str(),
             FormatError(error).c_str());
-    } else {
-        Y_ABORT("Unhandled branch of connect error");
+        ReEstablishConnection(connectionType, hostIndex);
+        return;
     }
 
     // ConnectPromise resolves both "connection ready" and "session ready" in
@@ -1579,12 +1591,16 @@ void TDirectBlockGroup::OnConnectionEstablished(
     }
 }
 
-void TDirectBlockGroup::ReEstablishDDiskConnection(
-    THostIndex hostIndex,
-    TDuration reconnectDelay)
+void TDirectBlockGroup::ReEstablishConnection(
+    EConnectionType connectionType,
+    THostIndex hostIndex)
 {
     Y_ABORT_UNLESS(ExecutorThreadChecker.Check());
-    Y_ABORT_UNLESS(hostIndex < DDiskConnections.size());
+    auto& connections = connectionType == EConnectionType::DDisk
+                            ? DDiskConnections
+                            : PBufferConnections;
+    Y_ABORT_UNLESS(hostIndex < connections.size());
+    TDDiskConnection& connection = connections[hostIndex];
 
     if (BlockedGenerationDetected) {
         LOG_WARN(
@@ -1595,13 +1611,14 @@ void TDirectBlockGroup::ReEstablishDDiskConnection(
         return;
     }
 
-    DDiskConnections[hostIndex].ResetSession();
+    connection.ResetSession();
+    TDuration reconnectDelay = Oracle.GetHostReconnectDelay(hostIndex);
     Schedule(
         reconnectDelay,
-        [hostIndex, weakSelf = weak_from_this()]()
+        [hostIndex, weakSelf = weak_from_this(), connectionType]()
         {
             if (auto self = weakSelf.lock()) {
-                self->DoEstablishConnection(hostIndex, EConnectionType::DDisk);
+                self->DoEstablishConnection(hostIndex, connectionType);
             }
         });
 }
@@ -1619,9 +1636,9 @@ void TDirectBlockGroup::OnNodeDisconnected(THostIndex hostIndex, ui32 nodeId)
         nodeId);
 
     Oracle.OnDDiskDisconnected(hostIndex, TInstant::Now());
+
     // OnNodeDisconnected may be called only for DDisk
-    TDuration reconnectDelay = Oracle.GetDDiskReconnectDelay(hostIndex);
-    ReEstablishDDiskConnection(hostIndex, reconnectDelay);
+    ReEstablishConnection(EConnectionType::DDisk, hostIndex);
 }
 
 bool TDirectBlockGroup::HasPBufferQuorum() const
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 602c15a9f72..366615e0718 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
@@ -194,9 +194,9 @@ private:
         THostIndex hostIndex,
         ui64 seqNo,
         const NKikimrBlobStorage::NDDisk::TEvConnectResult& result);
-    void ReEstablishDDiskConnection(
-        THostIndex hostIndex,
-        TDuration reconnectDelay);
+    void ReEstablishConnection(
+        EConnectionType connectionType,
+        THostIndex hostIndex);
     void OnNodeDisconnected(THostIndex hostIndex, ui32 nodeId);
 
     [[nodiscard]] bool HasPBufferQuorum() 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 79af0b61f78..2ce680d6282 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
@@ -979,6 +979,199 @@ Y_UNIT_TEST_SUITE(TDirectBlockGroupTest)
         UNIT_ASSERT_VALUES_EQUAL(connectsBefore, connectsAfter);
     }
 
+    // A non-BLOCKED connect failure on DDisk (e.g. ERROR/unavailable) must NOT
+    // be terminal: the session is reset and a fresh Connect is retried with the
+    // next seq_no.
+    Y_UNIT_TEST_F(ShouldReconnectDDiskOnNonBlockedConnectError, TDBGFixture)
+    {
+        auto executor = MakeExecutor();
+        auto transport = std::make_unique<TStorageTransportMock>();
+        auto* transportPtr = transport.get();
+
+        const auto& ddisks = transportPtr->GetDDiskIds();
+        // host[0] connect is deferred; hosts 1..4 connect immediately -> the
+        // quorum (3 of 5) is reached without host[0].
+        auto pendingConnect0 =
+            transportPtr->SetPendingConnect(EConnectionType::DDisk, ddisks[0]);
+
+        auto dbg = MakeDirectBlockGroup(executor, std::move(transport));
+        auto initialReady = RunAndGetInitialReady(dbg, false);
+        WaitReady(executor, initialReady);
+
+        // Only the first (initial) connect attempt was issued for host[0].
+        const size_t connectsBefore =
+            transportPtr
+                ->GetConnectCredentials(EConnectionType::DDisk, ddisks[0])
+                .size();
+        UNIT_ASSERT_VALUES_EQUAL(1, connectsBefore);
+
+        // Arm the next (reconnect) Connect for host[0] before delivering the
+        // failure.
+        auto reconnectPending =
+            transportPtr->SetPendingConnect(EConnectionType::DDisk, ddisks[0]);
+
+        // Resolve the in-flight connect with a non-BLOCKED error.
+        pendingConnect0.SetValue(TStorageTransportMock::MakeConnectResult(
+            1,
+            NKikimrBlobStorage::NDDisk::TReplyStatus::ERROR));
+
+        // Drain OnConnectionEstablished (queues the reconnect) then the
+        // reconnect itself (issues a fresh Connect).
+        DrainExecutor(executor);
+        DrainExecutor(executor);
+
+        // A reconnect Connect was issued and carries the next seq_no (=2).
+        const auto credentials = transportPtr->GetConnectCredentials(
+            EConnectionType::DDisk,
+            ddisks[0]);
+        UNIT_ASSERT_VALUES_EQUAL(connectsBefore + 1, credentials.size());
+        UNIT_ASSERT_VALUES_EQUAL(2, credentials.back().DDiskSessionSeqNo);
+
+        // No suicide on a non-BLOCKED connect error.
+        UNIT_ASSERT_VALUES_EQUAL(0, Service->BlockedGenerationCount);
+        {
+            const auto state =
+                GetBlockedDetected(executor, dbg, 0, WaitTimeout);
+            UNIT_ASSERT(!state.DDiskSessionBroken);
+            UNIT_ASSERT(!state.BlockedGenerationDetected);
+        }
+
+        // The retry succeeds -> the session is confirmed at seq_no=2.
+        reconnectPending.SetValue(TStorageTransportMock::MakeConnectResult());
+        DrainExecutor(executor);
+
+        const auto seqNos = ReadAllDDiskSeqNos(executor, dbg, WaitTimeout);
+        UNIT_ASSERT_VALUES_EQUAL(2, seqNos[0]);
+        UNIT_ASSERT_VALUES_EQUAL(0, Service->BlockedGenerationCount);
+    }
+
+    // A non-BLOCKED connect failure on PBuffer is retried the same way. PBuffer
+    // has no session/lock, so only the reconnect (a fresh Connect) is asserted,
+    // and no suicide happens.
+    Y_UNIT_TEST_F(ShouldReconnectPBufferOnNonBlockedConnectError, TDBGFixture)
+    {
+        auto executor = MakeExecutor();
+        auto transport = std::make_unique<TStorageTransportMock>();
+        auto* transportPtr = transport.get();
+
+        const auto& pbuffers = transportPtr->GetPBufferIds();
+        // PBuffer[0] connect is deferred; the rest connect immediately -> the
+        // PBuffer quorum is reached without PBuffer[0]. All DDisks connect
+        // immediately -> locked quorum is reached too.
+        auto pendingConnect0 = transportPtr->SetPendingConnect(
+            EConnectionType::PBuffer,
+            pbuffers[0]);
+
+        auto dbg = MakeDirectBlockGroup(executor, std::move(transport));
+        auto initialReady = RunAndGetInitialReady(dbg, false);
+        WaitReady(executor, initialReady);
+
+        // Only the first (initial) connect attempt was issued for PBuffer[0].
+        const size_t connectsBefore =
+            transportPtr
+                ->GetConnectCredentials(EConnectionType::PBuffer, pbuffers[0])
+                .size();
+        UNIT_ASSERT_VALUES_EQUAL(1, connectsBefore);
+
+        // Arm the next (reconnect) Connect for PBuffer[0] before delivering the
+        // failure.
+        auto reconnectPending = transportPtr->SetPendingConnect(
+            EConnectionType::PBuffer,
+            pbuffers[0]);
+
+        // Resolve the in-flight PBuffer connect with a non-BLOCKED error.
+        pendingConnect0.SetValue(TStorageTransportMock::MakeConnectResult(
+            1,
+            NKikimrBlobStorage::NDDisk::TReplyStatus::ERROR));
+
+        // Drain OnConnectionEstablished (queues the reconnect) then the
+        // reconnect itself (issues a fresh Connect).
+        DrainExecutor(executor);
+        DrainExecutor(executor);
+
+        // A reconnect Connect was issued for PBuffer[0].
+        const size_t connectsAfter =
+            transportPtr
+                ->GetConnectCredentials(EConnectionType::PBuffer, pbuffers[0])
+                .size();
+        UNIT_ASSERT_VALUES_EQUAL(connectsBefore + 1, connectsAfter);
+
+        // A PBuffer connect error never triggers a suicide.
+        UNIT_ASSERT_VALUES_EQUAL(0, Service->BlockedGenerationCount);
+
+        // The retry succeeds without any crash/suicide.
+        reconnectPending.SetValue(TStorageTransportMock::MakeConnectResult());
+        DrainExecutor(executor);
+        UNIT_ASSERT_VALUES_EQUAL(0, Service->BlockedGenerationCount);
+    }
+
+    // Startup must not fail permanently on a non-BLOCKED connect
+    // error: initial-ready stays unresolved while the DDisk connects keep
+    // failing, and resolves once the retried connects reach the locked quorum.
+    Y_UNIT_TEST_F(
+        ShouldReachInitialReadyAfterTransientConnectError,
+        TDBGFixture)
+    {
+        auto executor = MakeExecutor();
+        auto transport = std::make_unique<TStorageTransportMock>();
+        auto* transportPtr = transport.get();
+
+        const auto& ddisks = transportPtr->GetDDiskIds();
+        // Defer every DDisk connect so we control when the quorum is reached.
+        // PBuffers connect immediately -> PBuffer quorum is satisfied.
+        TVector<TStorageTransportMock::TConnectPromise> firstConnects;
+        for (const auto& ddiskId: ddisks) {
+            firstConnects.push_back(transportPtr->SetPendingConnect(
+                EConnectionType::DDisk,
+                ddiskId));
+        }
+
+        auto dbg = MakeDirectBlockGroup(executor, std::move(transport));
+        auto initialReady = RunAndGetInitialReady(dbg, false);
+
+        DrainExecutor(executor);
+        UNIT_ASSERT(!initialReady.HasValue());
+
+        // Arm the reconnect Connect for the first three DDisks (the quorum).
+        constexpr size_t quorum = QuorumDirectBlockGroupHostCount;
+        TVector<TStorageTransportMock::TConnectPromise> reconnects;
+        for (size_t i = 0; i < quorum; ++i) {
+            reconnects.push_back(transportPtr->SetPendingConnect(
+                EConnectionType::DDisk,
+                ddisks[i]));
+        }
+
+        // First attempt fails on all three with a non-BLOCKED error.
+        for (size_t i = 0; i < quorum; ++i) {
+            firstConnects[i].SetValue(TStorageTransportMock::MakeConnectResult(
+                1,
+                NKikimrBlobStorage::NDDisk::TReplyStatus::ERROR));
+        }
+        // Drain OnConnectionEstablished + queued reconnects, then the
+        // reconnects.
+        DrainExecutor(executor);
+        DrainExecutor(executor);
+
+        // Still not ready: no DDisk session is locked yet.
+        UNIT_ASSERT(!initialReady.HasValue());
+        UNIT_ASSERT_VALUES_EQUAL(0, Service->BlockedGenerationCount);
+
+        // The retried connects succeed -> the locked quorum is reached and
+        // startup completes.
+        for (auto& reconnect: reconnects) {
+            reconnect.SetValue(TStorageTransportMock::MakeConnectResult());
+        }
+        WaitReady(executor, initialReady);
+
+        // The recovered sessions are confirmed at seq_no=2 (attempt + retry),
+        // and no suicide happened.
+        const auto seqNos = ReadAllDDiskSeqNos(executor, dbg, WaitTimeout);
+        for (size_t i = 0; i < quorum; ++i) {
+            UNIT_ASSERT_VALUES_EQUAL(2, seqNos[i]);
+        }
+        UNIT_ASSERT_VALUES_EQUAL(0, Service->BlockedGenerationCount);
+    }
+
     // A PBuffer error must NOT trigger the suicide: the BLOCKED check is
     // only wired into DDisk-addressed responses.
     Y_UNIT_TEST_F(ShouldNotSuicideOnPBufferError, TDBGFixture)
diff --git a/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/direct_block_group_mock.cpp b/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/direct_block_group_mock.cpp
index 4e0ce88de08..bf0941a86ce 100644
--- a/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/direct_block_group_mock.cpp
+++ b/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/direct_block_group_mock.cpp
@@ -38,7 +38,7 @@ void TOracleMock::OnDDiskDisconnected(THostIndex hostIndex, TInstant now)
     Y_UNUSED(hostIndex, now);
 }
 
-TDuration TOracleMock::GetDDiskReconnectDelay(THostIndex hostIndex)
+TDuration TOracleMock::GetHostReconnectDelay(THostIndex hostIndex)
 {
     Y_UNUSED(hostIndex);
     return TDuration::MilliSeconds(1);
diff --git a/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/direct_block_group_mock.h b/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/direct_block_group_mock.h
index c9d535839ed..214edafc69f 100644
--- a/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/direct_block_group_mock.h
+++ b/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/direct_block_group_mock.h
@@ -41,7 +41,7 @@ struct TOracleMock: public IOracle
 
     void OnDDiskDisconnected(THostIndex hostIndex, TInstant now) override;
     void OnDDiskConnected(THostIndex hostIndex, TInstant now) override;
-    TDuration GetDDiskReconnectDelay(THostIndex hostIndex) override;
+    TDuration GetHostReconnectDelay(THostIndex hostIndex) override;
 
     [[nodiscard]] THostIndex SelectBestPBufferHost(
         THostMask hosts,
diff --git a/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/model/oracle.cpp b/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/model/oracle.cpp
index d79c836bf0f..c29e14c7ce7 100644
--- a/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/model/oracle.cpp
+++ b/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/model/oracle.cpp
@@ -285,7 +285,7 @@ void TOracle::OnDDiskConnected(THostIndex hostIndex, TInstant now)
     HostsReconnectDelays[hostIndex].Reset();
 }
 
-TDuration TOracle::GetDDiskReconnectDelay(THostIndex hostIndex)
+TDuration TOracle::GetHostReconnectDelay(THostIndex hostIndex)
 {
     return HostsReconnectDelays[hostIndex].GetDelayAndIncrease();
 }
diff --git a/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/model/oracle.h b/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/model/oracle.h
index 4e68e4eb54d..b78a53f216c 100644
--- a/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/model/oracle.h
+++ b/ydb/core/nbs/cloud/blockstore/libs/storage/partition_direct/model/oracle.h
@@ -71,7 +71,7 @@ public:
 
     virtual void OnDDiskDisconnected(THostIndex hostIndex, TInstant now) = 0;
     virtual void OnDDiskConnected(THostIndex hostIndex, TInstant now) = 0;
-    virtual TDuration GetDDiskReconnectDelay(THostIndex hostIndex) = 0;
+    virtual TDuration GetHostReconnectDelay(THostIndex hostIndex) = 0;
 
     // Picks the best host (by lowest inflight count) out of the provided set
     // of hosts. Ties are broken uniformly at random.
@@ -135,7 +135,7 @@ public:
 
     void OnDDiskDisconnected(THostIndex hostIndex, TInstant now) override;
     void OnDDiskConnected(THostIndex hostIndex, TInstant now) override;
-    [[nodiscard]] TDuration GetDDiskReconnectDelay(
+    [[nodiscard]] TDuration GetHostReconnectDelay(
         THostIndex hostIndex) override;
 
     [[nodiscard]] THostIndex SelectBestPBufferHost(