Use 1-thread pool executor for subsession event handlers (#6457)

5 files changed, 60 insertions, 55 deletions

diff --git a/ydb/public/sdk/cpp/client/ydb_federated_topic/impl/federated_topic_impl.cpp b/ydb/public/sdk/cpp/client/ydb_federated_topic/impl/federated_topic_impl.cpp
index 143f88d0f1a..5071a31a647 100644
--- a/ydb/public/sdk/cpp/client/ydb_federated_topic/impl/federated_topic_impl.cpp
+++ b/ydb/public/sdk/cpp/client/ydb_federated_topic/impl/federated_topic_impl.cpp
@@ -35,7 +35,8 @@ TFederatedTopicClient::TImpl::CreateWriteSession(const TFederatedWriteSessionSet
             splitSettings.EventHandlers_.HandlersExecutor(ClientSettings.DefaultHandlersExecutor_);
         }
     }
-    auto session = std::make_shared<TFederatedWriteSession>(splitSettings, Connections, ClientSettings, GetObserver(), ProvidedCodecs);
+    auto session = std::make_shared<TFederatedWriteSession>(
+        splitSettings, Connections, ClientSettings, GetObserver(), ProvidedCodecs, GetSubsessionHandlersExecutor());
     session->Start();
     return std::move(session);
 }
@@ -49,4 +50,13 @@ void TFederatedTopicClient::TImpl::InitObserver() {
     }
 }
 
+auto TFederatedTopicClient::TImpl::GetSubsessionHandlersExecutor() -> NTopic::IExecutor::TPtr {
+    with_lock (Lock) {
+        if (!SubsessionHandlersExecutor) {
+            SubsessionHandlersExecutor = NTopic::CreateThreadPoolExecutor(1);
+        }
+        return SubsessionHandlersExecutor;
+    }
+}
+
 }
diff --git a/ydb/public/sdk/cpp/client/ydb_federated_topic/impl/federated_topic_impl.h b/ydb/public/sdk/cpp/client/ydb_federated_topic/impl/federated_topic_impl.h
index 1f8bbcaf066..58320dd23f8 100644
--- a/ydb/public/sdk/cpp/client/ydb_federated_topic/impl/federated_topic_impl.h
+++ b/ydb/public/sdk/cpp/client/ydb_federated_topic/impl/federated_topic_impl.h
@@ -75,11 +75,19 @@ public:
     void InitObserver();
 
 private:
+
+    // Use single-threaded executor to prevent deadlocks inside subsession event handlers.
+    NTopic::IExecutor::TPtr GetSubsessionHandlersExecutor();
+
+private:
     std::shared_ptr<TGRpcConnectionsImpl> Connections;
     const TFederatedTopicClientSettings ClientSettings;
     std::shared_ptr<TFederatedDbObserver> Observer;
     std::shared_ptr<std::unordered_map<NTopic::ECodec, THolder<NTopic::ICodec>>> ProvidedCodecs =
         std::make_shared<std::unordered_map<NTopic::ECodec, THolder<NTopic::ICodec>>>();
+
+    NTopic::IExecutor::TPtr SubsessionHandlersExecutor;
+
     TAdaptiveLock Lock;
 };
 
diff --git a/ydb/public/sdk/cpp/client/ydb_federated_topic/impl/federated_write_session.cpp b/ydb/public/sdk/cpp/client/ydb_federated_topic/impl/federated_write_session.cpp
index 4c07ea0819b..bdfb2abeb81 100644
--- a/ydb/public/sdk/cpp/client/ydb_federated_topic/impl/federated_write_session.cpp
+++ b/ydb/public/sdk/cpp/client/ydb_federated_topic/impl/federated_write_session.cpp
@@ -29,12 +29,14 @@ TFederatedWriteSessionImpl::TFederatedWriteSessionImpl(
     std::shared_ptr<TGRpcConnectionsImpl> connections,
     const TFederatedTopicClientSettings& clientSettings,
     std::shared_ptr<TFederatedDbObserver> observer,
-    std::shared_ptr<std::unordered_map<NTopic::ECodec, THolder<NTopic::ICodec>>> codecs
+    std::shared_ptr<std::unordered_map<NTopic::ECodec, THolder<NTopic::ICodec>>> codecs,
+    NTopic::IExecutor::TPtr subsessionHandlersExecutor
 )
     : Settings(settings)
     , Connections(std::move(connections))
     , SubclientSettings(FromFederated(clientSettings))
     , ProvidedCodecs(std::move(codecs))
+    , SubsessionHandlersExecutor(subsessionHandlersExecutor)
     , Observer(std::move(observer))
     , AsyncInit(Observer->WaitForFirstState())
     , FederationState(nullptr)
@@ -71,15 +73,16 @@ void TFederatedWriteSessionImpl::IssueTokenIfAllowed() {
     }
 }
 
-void TFederatedWriteSessionImpl::UpdateFederationStateImpl() {
+std::shared_ptr<NTopic::IWriteSession> TFederatedWriteSessionImpl::UpdateFederationStateImpl() {
     Y_ABORT_UNLESS(Lock.IsLocked());
     // Even after the user has called the Close method, transitioning the session to the CLOSING state,
     // we keep updating the federation state, as the session may still have some messages to send in its queues,
     // and for that we need to know the current state of the federation.
     if (SessionState < State::CLOSED) {
         FederationState = Observer->GetState();
-        OnFederationStateUpdateImpl();
+        return OnFederationStateUpdateImpl();
     }
+    return {};
 }
 
 void TFederatedWriteSessionImpl::Start() {
@@ -104,15 +107,16 @@ void TFederatedWriteSessionImpl::Start() {
     });
 }
 
-void TFederatedWriteSessionImpl::OpenSubsessionImpl(std::shared_ptr<TDbInfo> db) {
+std::shared_ptr<NTopic::IWriteSession> TFederatedWriteSessionImpl::OpenSubsessionImpl(std::shared_ptr<TDbInfo> db) {
     Y_ABORT_UNLESS(Lock.IsLocked());
 
     ++SubsessionGeneration;
 
+    std::shared_ptr<NTopic::IWriteSession> oldSubsession;
+
     if (Subsession) {
         PendingToken.Clear();
-        OldSubsession = std::move(Subsession);
-        OldSubsession->Close(TDuration::Zero());
+        std::swap(oldSubsession, Subsession);
     }
 
     auto clientSettings = SubclientSettings;
@@ -122,11 +126,9 @@ void TFederatedWriteSessionImpl::OpenSubsessionImpl(std::shared_ptr<TDbInfo> db)
     auto subclient = make_shared<NTopic::TTopicClient::TImpl>(Connections, clientSettings);
 
     auto handlers = NTopic::TWriteSessionSettings::TEventHandlers()
-        .HandlersExecutor(Settings.EventHandlers_.HandlersExecutor_)
+        .HandlersExecutor(SubsessionHandlersExecutor)
         .ReadyToAcceptHandler([selfCtx = SelfContext, generation = SubsessionGeneration](NTopic::TWriteSessionEvent::TReadyToAcceptEvent& ev) {
             if (auto self = selfCtx->LockShared()) {
-                TDeferredWrite deferred;
-
                 with_lock(self->Lock) {
                     if (generation != self->SubsessionGeneration) {
                         return;
@@ -134,10 +136,8 @@ void TFederatedWriteSessionImpl::OpenSubsessionImpl(std::shared_ptr<TDbInfo> db)
 
                     Y_ABORT_UNLESS(self->PendingToken.Empty());
                     self->PendingToken = std::move(ev.ContinuationToken);
-                    self->PrepareDeferredWriteImpl(deferred);
+                    self->MaybeWriteImpl();
                 }
-
-                deferred.DoWrite();
             }
         })
         .AcksHandler([selfCtx = SelfContext](NTopic::TWriteSessionEvent::TAcksEvent& ev) {
@@ -182,6 +182,8 @@ void TFederatedWriteSessionImpl::OpenSubsessionImpl(std::shared_ptr<TDbInfo> db)
 
     Subsession = subclient->CreateWriteSession(wsSettings);
     CurrentDatabase = db;
+
+    return oldSubsession;
 }
 
 std::pair<std::shared_ptr<TDbInfo>, EStatus> SelectDatabaseByHashImpl(
@@ -267,13 +269,13 @@ std::pair<std::shared_ptr<TDbInfo>, EStatus> SelectDatabaseImpl(
     return SelectDatabaseByHashImpl(settings, dbInfos);
 }
 
-void TFederatedWriteSessionImpl::OnFederationStateUpdateImpl() {
+std::shared_ptr<NTopic::IWriteSession> TFederatedWriteSessionImpl::OnFederationStateUpdateImpl() {
     Y_ABORT_UNLESS(Lock.IsLocked());
     if (!FederationState->Status.IsSuccess()) {
         // The observer became stale, it won't try to get federation state anymore due to retry policy,
         // so there's no reason to keep the write session alive.
         CloseImpl(FederationState->Status.GetStatus(), NYql::TIssues(FederationState->Status.GetIssues()));
-        return;
+        return {};
     }
 
     Y_ABORT_UNLESS(!FederationState->DbInfos.empty());
@@ -292,19 +294,22 @@ void TFederatedWriteSessionImpl::OnFederationStateUpdateImpl() {
             LOG_LAZY(Log, TLOG_ERR, GetLogPrefixImpl() << message << ". Status: " << status);
             CloseImpl(status, NYql::TIssues{NYql::TIssue(message)});
         }
-        return;
+        return {};
     }
     RetryState.reset();
 
+    std::shared_ptr<NTopic::IWriteSession> oldSubsession;
     if (!DatabasesAreSame(preferrableDb, CurrentDatabase)) {
         LOG_LAZY(Log, TLOG_INFO, GetLogPrefixImpl()
             << "Start federated write session to database '" << preferrableDb->name()
             << "' (previous was " << (CurrentDatabase ? CurrentDatabase->name() : "<empty>") << ")"
             << " FederationState: " << *FederationState);
-        OpenSubsessionImpl(preferrableDb);
+        oldSubsession = OpenSubsessionImpl(preferrableDb);
     }
 
     ScheduleFederationStateUpdateImpl(UPDATE_FEDERATION_STATE_DELAY);
+
+    return oldSubsession;
 }
 
 void TFederatedWriteSessionImpl::ScheduleFederationStateUpdateImpl(TDuration delay) {
@@ -314,8 +319,10 @@ void TFederatedWriteSessionImpl::ScheduleFederationStateUpdateImpl(TDuration del
             if (auto self = selfCtx->LockShared()) {
                 std::shared_ptr<NTopic::IWriteSession> old;
                 with_lock(self->Lock) {
-                    self->UpdateFederationStateImpl();
-                    old = std::move(self->OldSubsession);
+                    old = self->UpdateFederationStateImpl();
+                }
+                if (old) {
+                    old->Close(TDuration::Zero());
                 }
             }
         }
@@ -378,8 +385,6 @@ void TFederatedWriteSessionImpl::WriteEncoded(NTopic::TContinuationToken&& token
 }
 
 void TFederatedWriteSessionImpl::WriteInternal(NTopic::TContinuationToken&&, TWrappedWriteMessage&& wrapped) {
-    TDeferredWrite deferred(Subsession);
-
     with_lock(Lock) {
         ClientHasToken = false;
         if (!wrapped.Message.CreateTimestamp_.Defined()) {
@@ -387,15 +392,13 @@ void TFederatedWriteSessionImpl::WriteInternal(NTopic::TContinuationToken&&, TWr
         }
         BufferFreeSpace -= wrapped.Message.Data.size();
         OriginalMessagesToPassDown.emplace_back(std::move(wrapped));
-        PrepareDeferredWriteImpl(deferred);
+        MaybeWriteImpl();
     }
 
-    deferred.DoWrite();
-
     IssueTokenIfAllowed();
 }
 
-bool TFederatedWriteSessionImpl::PrepareDeferredWriteImpl(TDeferredWrite& deferred) {
+bool TFederatedWriteSessionImpl::MaybeWriteImpl() {
     Y_ABORT_UNLESS(Lock.IsLocked());
     if (PendingToken.Empty()) {
         return false;
@@ -405,9 +408,7 @@ bool TFederatedWriteSessionImpl::PrepareDeferredWriteImpl(TDeferredWrite& deferr
     }
     OriginalMessagesToGetAck.push_back(std::move(OriginalMessagesToPassDown.front()));
     OriginalMessagesToPassDown.pop_front();
-    deferred.Writer = Subsession;
-    deferred.Token.ConstructInPlace(std::move(*PendingToken));
-    deferred.Message.ConstructInPlace(std::move(OriginalMessagesToGetAck.back().Message));
+    Subsession->Write(std::move(*PendingToken), std::move(OriginalMessagesToGetAck.back().Message));
     PendingToken.Clear();
     return true;
 }
diff --git a/ydb/public/sdk/cpp/client/ydb_federated_topic/impl/federated_write_session.h b/ydb/public/sdk/cpp/client/ydb_federated_topic/impl/federated_write_session.h
index 5cbc3b01bb6..5b348395887 100644
--- a/ydb/public/sdk/cpp/client/ydb_federated_topic/impl/federated_write_session.h
+++ b/ydb/public/sdk/cpp/client/ydb_federated_topic/impl/federated_write_session.h
@@ -27,7 +27,8 @@ public:
                                std::shared_ptr<TGRpcConnectionsImpl> connections,
                                const TFederatedTopicClientSettings& clientSetttings,
                                std::shared_ptr<TFederatedDbObserver> observer,
-                               std::shared_ptr<std::unordered_map<NTopic::ECodec, THolder<NTopic::ICodec>>> codecs);
+                               std::shared_ptr<std::unordered_map<NTopic::ECodec, THolder<NTopic::ICodec>>> codecs,
+                               NTopic::IExecutor::TPtr subsessionHandlersExecutor);
 
     ~TFederatedWriteSessionImpl() = default;
 
@@ -62,40 +63,22 @@ private:
         }
     };
 
-    struct TDeferredWrite {
-        TDeferredWrite() {}
-        explicit TDeferredWrite(std::shared_ptr<NTopic::IWriteSession> writer)
-            : Writer(std::move(writer)) {
-        }
-
-        void DoWrite() {
-            if (Token.Empty() && Message.Empty()) {
-                return;
-            }
-            Y_ABORT_UNLESS(Token.Defined() && Message.Defined());
-            return Writer->Write(std::move(*Token), std::move(*Message));
-        }
-
-        std::shared_ptr<NTopic::IWriteSession> Writer;
-        TMaybe<NTopic::TContinuationToken> Token;
-        TMaybe<NTopic::TWriteMessage> Message;
-    };
-
 private:
     void Start();
-    void OpenSubsessionImpl(std::shared_ptr<TDbInfo> db);
 
-    void OnFederationStateUpdateImpl();
+    std::shared_ptr<NTopic::IWriteSession> OpenSubsessionImpl(std::shared_ptr<TDbInfo> db);
+    std::shared_ptr<NTopic::IWriteSession> UpdateFederationStateImpl();
+    std::shared_ptr<NTopic::IWriteSession> OnFederationStateUpdateImpl();
+
     void ScheduleFederationStateUpdateImpl(TDuration delay);
 
     void WriteInternal(NTopic::TContinuationToken&&, TWrappedWriteMessage&& message);
-    bool PrepareDeferredWriteImpl(TDeferredWrite& deferred);
+    bool MaybeWriteImpl();
 
     void CloseImpl(EStatus statusCode, NYql::TIssues&& issues);
     void CloseImpl(NTopic::TSessionClosedEvent const& ev);
 
     bool MessageQueuesAreEmptyImpl() const;
-    void UpdateFederationStateImpl();
 
     void IssueTokenIfAllowed();
 
@@ -107,6 +90,7 @@ private:
     std::shared_ptr<TGRpcConnectionsImpl> Connections;
     const NTopic::TTopicClientSettings SubclientSettings;
     std::shared_ptr<std::unordered_map<NTopic::ECodec, THolder<NTopic::ICodec>>> ProvidedCodecs;
+    NTopic::IExecutor::TPtr SubsessionHandlersExecutor;
 
     NTopic::IRetryPolicy::IRetryState::TPtr RetryState;
     std::shared_ptr<TFederatedDbObserver> Observer;
@@ -125,7 +109,6 @@ private:
 
     size_t SubsessionGeneration = 0;
     std::shared_ptr<NTopic::IWriteSession> Subsession;
-    std::shared_ptr<NTopic::IWriteSession> OldSubsession;
 
     std::shared_ptr<NTopic::TWriteSessionEventsQueue> ClientEventsQueue;
 
@@ -156,8 +139,9 @@ public:
                            std::shared_ptr<TGRpcConnectionsImpl> connections,
                            const TFederatedTopicClientSettings& clientSettings,
                            std::shared_ptr<TFederatedDbObserver> observer,
-                           std::shared_ptr<std::unordered_map<NTopic::ECodec, THolder<NTopic::ICodec>>> codecs)
-        : TContextOwner(settings, std::move(connections), clientSettings, std::move(observer), codecs) {}
+                           std::shared_ptr<std::unordered_map<NTopic::ECodec, THolder<NTopic::ICodec>>> codecs,
+                           NTopic::IExecutor::TPtr subsessionHandlersExecutor)
+        : TContextOwner(settings, std::move(connections), clientSettings, std::move(observer), codecs, subsessionHandlersExecutor) {}
 
     NThreading::TFuture<void> WaitEvent() override {
         return TryGetImpl()->WaitEvent();
diff --git a/ydb/public/sdk/cpp/client/ydb_federated_topic/ut/basic_usage_ut.cpp b/ydb/public/sdk/cpp/client/ydb_federated_topic/ut/basic_usage_ut.cpp
index 42eebda5d52..1a98937b287 100644
--- a/ydb/public/sdk/cpp/client/ydb_federated_topic/ut/basic_usage_ut.cpp
+++ b/ydb/public/sdk/cpp/client/ydb_federated_topic/ut/basic_usage_ut.cpp
@@ -1096,6 +1096,8 @@ Y_UNIT_TEST_SUITE(BasicUsage) {
                 ++(ev.IsSuccess() ? successfulSessionClosedEvents : otherSessionClosedEvents);
             });
 
+        writeSettings.EventHandlers_.HandlersExecutor(NTopic::CreateSyncExecutor());
+
         auto WriteSession = topicClient.CreateWriteSession(writeSettings);
 
         TMaybe<NTopic::TContinuationToken> token;