Moved "TFederatedWriteSession: Wait for messages to be sent in Close method" commit from ydb repo

Author: qyryq

include/ydb-cpp-sdk/client/topic/topic.h

@@ -1557,8 +1557,8 @@ class IWriteSession {
                               std::optional<ui64> seqNo = std::nullopt, std::optional<TInstant> createTimestamp = std::nullopt) = 0;
 
 
-    //! Wait for all writes to complete (no more that closeTimeout()), than close. Empty maybe - means infinite timeout.
-    //! return - true if all writes were completed and acked. false if timeout was reached and some writes were aborted.
+    //! Wait for all writes to complete (no more that closeTimeout()), then close.
+    //! Return true if all writes were completed and acked, false if timeout was reached and some writes were aborted.
     virtual bool Close(TDuration closeTimeout = TDuration::Max()) = 0;
 
     //! Writer counters with different stats (see TWriterConuters).

src/client/federated_topic/impl/federated_write_session.cpp

@@ -27,13 +27,13 @@ NTopic::TTopicClientSettings FromFederated(const TFederatedTopicClientSettings&
 TFederatedWriteSessionImpl::TFederatedWriteSessionImpl(
     const TFederatedWriteSessionSettings& settings,
     std::shared_ptr<TGRpcConnectionsImpl> connections,
-    const TFederatedTopicClientSettings& clientSetttings,
+    const TFederatedTopicClientSettings& clientSettings,
     std::shared_ptr<TFederatedDbObserver> observer,
     std::shared_ptr<std::unordered_map<NTopic::ECodec, THolder<NTopic::ICodec>>> codecs
 )
     : Settings(settings)
     , Connections(std::move(connections))
-    , SubClientSetttings(FromFederated(clientSetttings))
+    , SubclientSettings(FromFederated(clientSettings))
     , ProvidedCodecs(std::move(codecs))
     , Observer(std::move(observer))
     , AsyncInit(Observer->WaitForFirstState())
@@ -44,37 +44,72 @@ TFederatedWriteSessionImpl::TFederatedWriteSessionImpl(
 {
 }
 
-TStringBuilder TFederatedWriteSessionImpl::GetLogPrefix() const {
-    return TStringBuilder() << GetDatabaseLogPrefix(SubClientSetttings.Database_.value_or("")) << "[" << SessionId << "] ";
+TStringBuilder TFederatedWriteSessionImpl::GetLogPrefixImpl() const {
+     return TStringBuilder() << GetDatabaseLogPrefix(SubclientSettings.Database_.value_or("")) << "[" << SessionId << "] ";
+}
+
+
+bool TFederatedWriteSessionImpl::MessageQueuesAreEmptyImpl() const {
+    Y_ABORT_UNLESS(Lock.IsLocked());
+    return OriginalMessagesToGetAck.empty() && OriginalMessagesToPassDown.empty();
+}
+
+void TFederatedWriteSessionImpl::IssueTokenIfAllowed() {
+    // The session should not issue tokens after it has transitioned to CLOSING or CLOSE state.
+    // A user may have one spare token, so at most one additional message
+    // could be written to the internal queue after the transition.
+    bool issue = false;
+    with_lock(Lock) {
+        if (BufferFreeSpace > 0 && !ClientHasToken && SessionState < State::CLOSING) {
+            ClientHasToken = true;
+            issue = true;
+        }
+    }
+    if (issue) {
+        ClientEventsQueue->PushEvent(NTopic::TWriteSessionEvent::TReadyToAcceptEvent{IssueContinuationToken()});
+    }
+}
+
+void 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();
+    }
 }
 
 void TFederatedWriteSessionImpl::Start() {
     // TODO validate settings?
-    Settings.EventHandlers_.HandlersExecutor_->Start();
     with_lock(Lock){
-        ClientEventsQueue->PushEvent(NTopic::TWriteSessionEvent::TReadyToAcceptEvent{IssueContinuationToken()});
-        ClientHasToken = true;
+        if (SessionState != State::CREATED) {
+            return;
+        }
+        SessionState = State::WORKING;
+        Settings.EventHandlers_.HandlersExecutor_->Start();
     }
 
-    AsyncInit.Subscribe([selfCtx = SelfContext](const auto& f){
+    IssueTokenIfAllowed();
+
+    AsyncInit.Subscribe([selfCtx = SelfContext](const auto& f) {
         Y_UNUSED(f);
         if (auto self = selfCtx->LockShared()) {
             with_lock(self->Lock) {
-                if (!self->Closing) {
-                    self->FederationState = self->Observer->GetState();
-                    self->OnFederatedStateUpdateImpl();
-                }
+                self->UpdateFederationStateImpl();
             }
         }
     });
 }
 
-void TFederatedWriteSessionImpl::OpenSubSessionImpl(std::shared_ptr<TDbInfo> db) {
+void TFederatedWriteSessionImpl::OpenSubsessionImpl(std::shared_ptr<TDbInfo> db) {
+    Y_ABORT_UNLESS(Lock.IsLocked());
     if (Subsession) {
         PendingToken.reset();
         Subsession->Close(TDuration::Zero());
     }
-    NTopic::TTopicClientSettings clientSettings = SubClientSetttings;
+    auto clientSettings = SubclientSettings;
     clientSettings
         .Database(db->path())
         .DiscoveryEndpoint(db->endpoint());
@@ -89,7 +124,7 @@ void TFederatedWriteSessionImpl::OpenSubSessionImpl(std::shared_ptr<TDbInfo> db)
                 with_lock(self->Lock) {
                     Y_ABORT_UNLESS(!self->PendingToken.has_value());
                     self->PendingToken = std::move(ev.ContinuationToken);
-                    self->PrepareDeferredWrite(deferred);
+                    self->PrepareDeferredWriteImpl(deferred);
                 }
                 deferred.DoWrite();
             }
@@ -98,24 +133,25 @@ void TFederatedWriteSessionImpl::OpenSubSessionImpl(std::shared_ptr<TDbInfo> db)
             if (auto self = selfCtx->LockShared()) {
                 with_lock(self->Lock) {
                     Y_ABORT_UNLESS(ev.Acks.size() <= self->OriginalMessagesToGetAck.size());
+
                     for (size_t i = 0; i < ev.Acks.size(); ++i) {
                         self->BufferFreeSpace += self->OriginalMessagesToGetAck.front().Data.size();
                         self->OriginalMessagesToGetAck.pop_front();
                     }
-                    self->ClientEventsQueue->PushEvent(std::move(ev));
-                    if (self->BufferFreeSpace > 0 && !self->ClientHasToken) {
-                        self->ClientEventsQueue->PushEvent(NTopic::TWriteSessionEvent::TReadyToAcceptEvent{IssueContinuationToken()});
-                        self->ClientHasToken = true;
+
+                    if (self->MessageQueuesAreEmptyImpl() && self->MessageQueuesHaveBeenEmptied.Initialized() && !self->MessageQueuesHaveBeenEmptied.HasValue()) {
+                        self->MessageQueuesHaveBeenEmptied.SetValue();
                     }
                 }
+
+                self->ClientEventsQueue->PushEvent(std::move(ev));
+                self->IssueTokenIfAllowed();
             }
         })
         .SessionClosedHandler([selfCtx = SelfContext](const NTopic::TSessionClosedEvent & ev) {
             if (auto self = selfCtx->LockShared()) {
                 with_lock(self->Lock) {
-                    if (!self->Closing) {
-                        self->CloseImpl(ev);
-                    }
+                    self->CloseImpl(ev);
                 }
             }
         });
@@ -129,7 +165,7 @@ void TFederatedWriteSessionImpl::OpenSubSessionImpl(std::shared_ptr<TDbInfo> db)
     CurrentDatabase = db;
 }
 
-std::pair<std::shared_ptr<TDbInfo>, EStatus> SelectDatabaseByHash(
+std::pair<std::shared_ptr<TDbInfo>, EStatus> SelectDatabaseByHashImpl(
     NTopic::TFederatedWriteSessionSettings const& settings,
     std::vector<std::shared_ptr<TDbInfo>> const& dbInfos
 ) {
@@ -163,7 +199,7 @@ std::pair<std::shared_ptr<TDbInfo>, EStatus> SelectDatabaseByHash(
     Y_UNREACHABLE();
 }
 
-std::pair<std::shared_ptr<TDbInfo>, EStatus> SelectDatabase(
+std::pair<std::shared_ptr<TDbInfo>, EStatus> SelectDatabaseImpl(
     NTopic::TFederatedWriteSessionSettings const& settings,
     std::vector<std::shared_ptr<TDbInfo>> const& dbInfos, std::string const& selfLocation
 ) {
@@ -197,7 +233,7 @@ std::pair<std::shared_ptr<TDbInfo>, EStatus> SelectDatabase(
             if (!settings.AllowFallback_) {
                 return {nullptr, EStatus::NOT_FOUND};
             }
-            return SelectDatabaseByHash(settings, dbInfos);
+            return SelectDatabaseByHashImpl(settings, dbInfos);
         }
     }
 
@@ -208,69 +244,64 @@ std::pair<std::shared_ptr<TDbInfo>, EStatus> SelectDatabase(
     if (!settings.AllowFallback_) {
         return {nullptr, EStatus::UNAVAILABLE};
     }
-    return SelectDatabaseByHash(settings, dbInfos);
-}
-
-std::pair<std::shared_ptr<TDbInfo>, EStatus> TFederatedWriteSessionImpl::SelectDatabaseImpl() {
-    return SelectDatabase(Settings, FederationState->DbInfos, FederationState->SelfLocation);
+    return SelectDatabaseByHashImpl(settings, dbInfos);
 }
 
-void TFederatedWriteSessionImpl::OnFederatedStateUpdateImpl() {
+void 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;
     }
 
     Y_ABORT_UNLESS(!FederationState->DbInfos.empty());
 
-    auto [preferrableDb, status] = SelectDatabaseImpl();
+    auto [preferrableDb, status] = SelectDatabaseImpl(Settings, FederationState->DbInfos, FederationState->SelfLocation);
 
     if (!preferrableDb) {
         if (!RetryState) {
             RetryState = Settings.RetryPolicy_->CreateRetryState();
         }
         if (auto delay = RetryState->GetNextRetryDelay(status)) {
-            LOG_LAZY(Log, TLOG_NOTICE, GetLogPrefix() << "Retry to update federation state in " << delay);
-            ScheduleFederatedStateUpdateImpl(*delay);
+            LOG_LAZY(Log, TLOG_NOTICE, GetLogPrefixImpl() << "Retry to update federation state in " << delay);
+            ScheduleFederationStateUpdateImpl(*delay);
         } else {
             std::string message = "Failed to select database: no available database";
-            LOG_LAZY(Log, TLOG_ERR, GetLogPrefix() << message);
+            LOG_LAZY(Log, TLOG_ERR, GetLogPrefixImpl() << message << ". Status: " << status);
             CloseImpl(status, NYql::TIssues{NYql::TIssue(message)});
         }
         return;
     }
     RetryState.reset();
 
     if (!DatabasesAreSame(preferrableDb, CurrentDatabase)) {
-        LOG_LAZY(Log, TLOG_INFO, GetLogPrefix()
+        LOG_LAZY(Log, TLOG_INFO, GetLogPrefixImpl()
             << "Start federated write session to database '" << preferrableDb->name()
             << "' (previous was " << (CurrentDatabase ? CurrentDatabase->name() : "<empty>") << ")"
             << " FederationState: " << *FederationState);
-        OpenSubSessionImpl(preferrableDb);
+        OpenSubsessionImpl(preferrableDb);
     }
 
-    ScheduleFederatedStateUpdateImpl(UPDATE_FEDERATION_STATE_DELAY);
+    ScheduleFederationStateUpdateImpl(UPDATE_FEDERATION_STATE_DELAY);
 }
 
-void TFederatedWriteSessionImpl::ScheduleFederatedStateUpdateImpl(TDuration delay) {
+void TFederatedWriteSessionImpl::ScheduleFederationStateUpdateImpl(TDuration delay) {
     Y_ABORT_UNLESS(Lock.IsLocked());
     auto cb = [selfCtx = SelfContext](bool ok) {
         if (ok) {
             if (auto self = selfCtx->LockShared()) {
                 with_lock(self->Lock) {
-                    if (self->Closing) {
-                        return;
-                    }
-                    self->FederationState = self->Observer->GetState();
-                    self->OnFederatedStateUpdateImpl();
+                    self->UpdateFederationStateImpl();
                 }
             }
         }
     };
 
     UpdateStateDelayContext = Connections->CreateContext();
     if (!UpdateStateDelayContext) {
-        Closing = true;
+        CloseImpl(EStatus::TRANSPORT_UNAVAILABLE, NYql::TIssues{NYql::TIssue("Could not update federation state")});
         // TODO log DRIVER_IS_STOPPING_DESCRIPTION
         return;
     }
@@ -325,28 +356,25 @@ void TFederatedWriteSessionImpl::WriteEncoded(NTopic::TContinuationToken&& token
 }
 
 void TFederatedWriteSessionImpl::WriteInternal(NTopic::TContinuationToken&&, TWrappedWriteMessage&& wrapped) {
-    ClientHasToken = false;
-    if (!wrapped.Message.CreateTimestamp_.has_value()) {
-        wrapped.Message.CreateTimestamp_ = TInstant::Now();
-    }
+    TDeferredWrite deferred(Subsession);
 
-    {
-        TDeferredWrite deferred(Subsession);
-        with_lock(Lock) {
-            BufferFreeSpace -= wrapped.Message.Data.size();
-            OriginalMessagesToPassDown.emplace_back(std::move(wrapped));
-
-            PrepareDeferredWrite(deferred);
+    with_lock(Lock) {
+        ClientHasToken = false;
+        if (!wrapped.Message.CreateTimestamp_.has_value()) {
+            wrapped.Message.CreateTimestamp_ = TInstant::Now();
         }
-        deferred.DoWrite();
-    }
-    if (BufferFreeSpace > 0) {
-        ClientEventsQueue->PushEvent(NTopic::TWriteSessionEvent::TReadyToAcceptEvent{IssueContinuationToken()});
-        ClientHasToken = true;
+        BufferFreeSpace -= wrapped.Message.Data.size();
+        OriginalMessagesToPassDown.emplace_back(std::move(wrapped));
+        PrepareDeferredWriteImpl(deferred);
     }
+
+    deferred.DoWrite();
+
+    IssueTokenIfAllowed();
 }
 
-bool TFederatedWriteSessionImpl::PrepareDeferredWrite(TDeferredWrite& deferred) {
+bool TFederatedWriteSessionImpl::PrepareDeferredWriteImpl(TDeferredWrite& deferred) {
+    Y_ABORT_UNLESS(Lock.IsLocked());
     if (!PendingToken.has_value()) {
         return false;
     }
@@ -361,27 +389,47 @@ bool TFederatedWriteSessionImpl::PrepareDeferredWrite(TDeferredWrite& deferred)
     return true;
 }
 
-void TFederatedWriteSessionImpl::CloseImpl(EStatus statusCode, NYql::TIssues&& issues, TDuration timeout) {
-    CloseImpl(TPlainStatus(statusCode, std::move(issues)), timeout);
+void TFederatedWriteSessionImpl::CloseImpl(EStatus statusCode, NYql::TIssues&& issues) {
+    CloseImpl(TPlainStatus(statusCode, std::move(issues)));
 }
 
-void TFederatedWriteSessionImpl::CloseImpl(NTopic::TSessionClosedEvent const& ev, TDuration timeout) {
-    if (Closing) {
+void TFederatedWriteSessionImpl::CloseImpl(NTopic::TSessionClosedEvent const& ev) {
+    Y_ABORT_UNLESS(Lock.IsLocked());
+    if (SessionState == State::CLOSED) {
         return;
     }
-    Closing = true;
-    if (Subsession) {
-        Subsession->Close(timeout);
-    }
-    ClientEventsQueue->Close(ev);
+    SessionState = State::CLOSED;
     NTopic::Cancel(UpdateStateDelayContext);
+    if (!HasBeenClosed.HasValue()) {
+        HasBeenClosed.SetValue();
+    }
+    {
+        auto unguard = Unguard(Lock);
+        ClientEventsQueue->Close(ev);
+    }
 }
 
 bool TFederatedWriteSessionImpl::Close(TDuration timeout) {
     with_lock(Lock) {
-        CloseImpl(EStatus::SUCCESS, {}, timeout);
+        if (SessionState == State::CLOSED) {
+            return MessageQueuesAreEmptyImpl();
+        }
+        SessionState = State::CLOSING;
+        if (!MessageQueuesHaveBeenEmptied.Initialized()) {
+            MessageQueuesHaveBeenEmptied = NThreading::NewPromise();
+            if (MessageQueuesAreEmptyImpl()) {
+                MessageQueuesHaveBeenEmptied.SetValue();
+            }
+        }
+    }
+
+    std::vector<NThreading::TFuture<void>> futures{MessageQueuesHaveBeenEmptied.GetFuture(), HasBeenClosed.GetFuture()};
+    NThreading::WaitAny(futures).Wait(timeout);
+
+    with_lock(Lock) {
+        CloseImpl(EStatus::SUCCESS, NYql::TIssues{});
+        return MessageQueuesAreEmptyImpl();
     }
-    return true;
 }
 
 }  // namespace NYdb::NFederatedTopic