Moved commit "Use 1-thread pool executor for subsession event handlers" from ydb repo

qyryq · qyryq · commit 4b100516ab83 · 2024-08-26T17:21:10.000Z
diff --git a/src/client/federated_topic/impl/federated_topic_impl.cpp b/src/client/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);
 }
@@ -48,4 +49,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/src/client/federated_topic/impl/federated_topic_impl.h b/src/client/federated_topic/impl/federated_topic_impl.h
@@ -72,12 +72,20 @@ class TFederatedTopicClient::TImpl {
 
     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, std::unique_ptr<NTopic::ICodec>>> ProvidedCodecs =
          std::make_shared<std::unordered_map<NTopic::ECodec, std::unique_ptr<NTopic::ICodec>>>();
+
+    NTopic::IExecutor::TPtr SubsessionHandlersExecutor;
+
     TAdaptiveLock Lock;
 };
 
diff --git a/src/client/federated_topic/impl/federated_write_session.cpp b/src/client/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, std::unique_ptr<NTopic::ICodec>>> codecs
+    std::shared_ptr<std::unordered_map<NTopic::ECodec, std::unique_ptr<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)
@@ -70,15 +72,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() {
@@ -103,15 +106,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.reset();
-        OldSubsession = std::move(Subsession);
-        OldSubsession->Close(TDuration::Zero());
+        std::swap(oldSubsession, Subsession);
     }
 
     auto clientSettings = SubclientSettings;
@@ -121,22 +125,18 @@ void TFederatedWriteSessionImpl::OpenSubsessionImpl(std::shared_ptr<TDbInfo> db)
     auto subclient = std::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;
                     }
 
                     Y_ABORT_UNLESS(!self->PendingToken.has_value());
                     self->PendingToken = std::move(ev.ContinuationToken);
-                    self->PrepareDeferredWriteImpl(deferred);
+                    self->MaybeWriteImpl();
                 }
-
-                deferred.DoWrite();
             }
         })
         .AcksHandler([selfCtx = SelfContext](NTopic::TWriteSessionEvent::TAcksEvent& ev) {
@@ -181,6 +181,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(
@@ -265,13 +267,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());
@@ -290,19 +292,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) {
@@ -312,8 +317,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());
                 }
             }
         }
@@ -376,24 +383,20 @@ 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_.has_value()) {
             wrapped.Message.CreateTimestamp_ = TInstant::Now();
         }
         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.has_value()) {
         return false;
@@ -403,9 +406,7 @@ bool TFederatedWriteSessionImpl::PrepareDeferredWriteImpl(TDeferredWrite& deferr
     }
     OriginalMessagesToGetAck.push_back(std::move(OriginalMessagesToPassDown.front()));
     OriginalMessagesToPassDown.pop_front();
-    deferred.Writer = Subsession;
-    deferred.Token.emplace(std::move(*PendingToken));
-    deferred.Message.emplace(std::move(OriginalMessagesToGetAck.back().Message));
+    Subsession->Write(std::move(*PendingToken), std::move(OriginalMessagesToGetAck.back().Message));
     PendingToken.reset();
     return true;
 }
diff --git a/src/client/federated_topic/impl/federated_write_session.h b/src/client/federated_topic/impl/federated_write_session.h
@@ -26,7 +26,8 @@ class TFederatedWriteSessionImpl : public NTopic::TContinuationTokenIssuer,
                                std::shared_ptr<TGRpcConnectionsImpl> connections,
                                const TFederatedTopicClientSettings& clientSetttings,
                                std::shared_ptr<TFederatedDbObserver> observer,
-                               std::shared_ptr<std::unordered_map<NTopic::ECodec, std::unique_ptr<NTopic::ICodec>>> codecs);
+                               std::shared_ptr<std::unordered_map<NTopic::ECodec, std::unique_ptr<NTopic::ICodec>>> codecs,
+                               NTopic::IExecutor::TPtr subsessionHandlersExecutor);
 
     ~TFederatedWriteSessionImpl() = default;
 
@@ -61,40 +62,22 @@ class TFederatedWriteSessionImpl : public NTopic::TContinuationTokenIssuer,
         }
     };
 
-    struct TDeferredWrite {
-        TDeferredWrite() {}
-        explicit TDeferredWrite(std::shared_ptr<NTopic::IWriteSession> writer)
-            : Writer(std::move(writer)) {
-        }
-
-        void DoWrite() {
-            if (!Token.has_value() && !Message.has_value()) {
-                return;
-            }
-            Y_ABORT_UNLESS(Token.has_value() && Message.has_value());
-            return Writer->Write(std::move(*Token), std::move(*Message));
-        }
-
-        std::shared_ptr<NTopic::IWriteSession> Writer;
-        std::optional<NTopic::TContinuationToken> Token;
-        std::optional<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();
 
@@ -106,6 +89,7 @@ class TFederatedWriteSessionImpl : public NTopic::TContinuationTokenIssuer,
     std::shared_ptr<TGRpcConnectionsImpl> Connections;
     const NTopic::TTopicClientSettings SubclientSettings;
     std::shared_ptr<std::unordered_map<NTopic::ECodec, std::unique_ptr<NTopic::ICodec>>> ProvidedCodecs;
+    NTopic::IExecutor::TPtr SubsessionHandlersExecutor;
 
     NTopic::IRetryPolicy::IRetryState::TPtr RetryState;
     std::shared_ptr<TFederatedDbObserver> Observer;
@@ -124,7 +108,6 @@ class TFederatedWriteSessionImpl : public NTopic::TContinuationTokenIssuer,
 
     size_t SubsessionGeneration = 0;
     std::shared_ptr<NTopic::IWriteSession> Subsession;
-    std::shared_ptr<NTopic::IWriteSession> OldSubsession;
 
     std::shared_ptr<NTopic::TWriteSessionEventsQueue> ClientEventsQueue;
 
@@ -155,8 +138,9 @@ class TFederatedWriteSession : public NTopic::IWriteSession,
                            std::shared_ptr<TGRpcConnectionsImpl> connections,
                            const TFederatedTopicClientSettings& clientSettings,
                            std::shared_ptr<TFederatedDbObserver> observer,
-                           std::shared_ptr<std::unordered_map<NTopic::ECodec, std::unique_ptr<NTopic::ICodec>>> codecs)
-        : TContextOwner(settings, std::move(connections), clientSettings, std::move(observer), codecs) {}
+                           std::shared_ptr<std::unordered_map<NTopic::ECodec, std::unique_ptr<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();

Original file line number	Diff line number	Diff line change
`@@ -35,7 +35,8 @@ TFederatedTopicClient::TImpl::CreateWriteSession(const TFederatedWriteSessionSet`
`35`	`35`	`splitSettings.EventHandlers_.HandlersExecutor(ClientSettings.DefaultHandlersExecutor_);`
`36`	`36`	`}`
`37`	`37`	`}`
`38`		`- auto session = std::make_shared<TFederatedWriteSession>(splitSettings, Connections, ClientSettings, GetObserver(), ProvidedCodecs);`
	`38`	`+ auto session = std::make_shared<TFederatedWriteSession>(`
	`39`	`+ splitSettings, Connections, ClientSettings, GetObserver(), ProvidedCodecs, GetSubsessionHandlersExecutor());`
`39`	`40`	`session->Start();`
`40`	`41`	`return std::move(session);`
`41`	`42`	`}`
`@@ -48,4 +49,13 @@ void TFederatedTopicClient::TImpl::InitObserver() {`
`48`	`49`	`}`
`49`	`50`	`}`
`50`	`51`
	`52`	`+auto TFederatedTopicClient::TImpl::GetSubsessionHandlersExecutor() -> NTopic::IExecutor::TPtr {`
	`53`	`+ with_lock (Lock) {`
	`54`	`+ if (!SubsessionHandlersExecutor) {`
	`55`	`+ SubsessionHandlersExecutor = NTopic::CreateThreadPoolExecutor(1);`
	`56`	`+ }`
	`57`	`+ return SubsessionHandlersExecutor;`
	`58`	`+ }`
	`59`	`+}`
	`60`	`+`
`51`	`61`	`}`