Moved commits "Optimization of the UpdateOffsetsInTransaction calls" and "Commit offsets in a transaction and out of a transaction" from ydb repo

Author: Alek5andr-Kotov

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

@@ -200,6 +200,10 @@ struct TReadSessionEvent {
             return CompressedMessages;
         }
 
+        void SetReadInTransaction() {
+            ReadInTransaction = true;
+        }
+
         //! Commits all messages in batch.
         void Commit();
 
@@ -218,6 +222,7 @@ struct TReadSessionEvent {
         std::vector<TMessage> Messages;
         std::vector<TCompressedMessage> CompressedMessages;
         std::vector<std::pair<uint64_t, uint64_t>> OffsetRanges;
+        bool ReadInTransaction = false;
     };
 
     //! Acknowledgement for commit request.

src/client/table/impl/transaction.cpp

@@ -14,12 +14,19 @@ TAsyncStatus TTransaction::TImpl::Precommit() const
     auto result = NThreading::MakeFuture(TStatus(EStatus::SUCCESS, {}));
 
     for (auto& callback : PrecommitCallbacks) {
-        auto action = [curr = callback()](const TAsyncStatus& prev) {
+        if (!callback) {
+            continue;
+        }
+
+        // If you send multiple requests in parallel, the `KQP` service can respond with `SESSION_BUSY`.
+        // Therefore, precommit operations are performed sequentially. Here we capture the closure to
+        // trigger it later.
+        auto action = [callback = std::move(callback)](const TAsyncStatus& prev) {
             if (const TStatus& status = prev.GetValue(); !status.IsSuccess()) {
                 return prev;
             }
 
-            return curr;
+            return callback();
         };
 
         result = result.Apply(action);
@@ -39,6 +46,8 @@ TAsyncCommitTransactionResult TTransaction::TImpl::Commit(const TCommitTxSetting
             return NThreading::MakeFuture(TCommitTransactionResult(TStatus(status), std::nullopt));
         }
 
+        PrecommitCallbacks.clear();
+
         return Session_.Client_->CommitTransaction(Session_,
                                                    TxId_,
                                                    settings);

src/client/table/impl/transaction.h

@@ -31,7 +31,7 @@ class TTransaction::TImpl {
     std::string TxId_;
 
     bool ChangesAreAccepted = true; // haven't called Commit or Rollback yet
-    std::vector<TPrecommitTransactionCallback> PrecommitCallbacks;
+    mutable std::vector<TPrecommitTransactionCallback> PrecommitCallbacks;
 };
 
 }

src/client/topic/impl/CMakeLists.txt

@@ -27,12 +27,14 @@ target_sources(client-ydb_topic-impl
     common.cpp
     deferred_commit.cpp
     event_handlers.cpp
+    offsets_collector.cpp
     read_session_event.cpp
     read_session.cpp
-    write_session.cpp
-    write_session_impl.cpp
     topic_impl.cpp
     topic.cpp
+    transaction.cpp
+    write_session_impl.cpp
+    write_session.cpp
 )
 
 _ydb_sdk_install_targets(TARGETS client-ydb_topic-impl)

src/client/topic/impl/offsets_collector.cpp

@@ -0,0 +1,72 @@
+#include "offsets_collector.h"
+
+namespace NYdb::NTopic {
+
+std::vector<TTopicOffsets> TOffsetsCollector::GetOffsets() const
+{
+    std::vector<TTopicOffsets> topics;
+
+    for (auto& [path, partitions] : Ranges) {
+        TTopicOffsets topic;
+        topic.Path = path;
+
+        topics.push_back(std::move(topic));
+
+        for (auto& [id, ranges] : partitions) {
+            TPartitionOffsets partition;
+            partition.PartitionId = id;
+
+            TTopicOffsets& t = topics.back();
+            t.Partitions.push_back(std::move(partition));
+
+            for (auto& range : ranges) {
+                TPartitionOffsets& p = t.Partitions.back();
+
+                TOffsetsRange r;
+                r.Start = range.first;
+                r.End = range.second;
+
+                p.Offsets.push_back(r);
+            }
+        }
+    }
+
+    return topics;
+}
+
+void TOffsetsCollector::CollectOffsets(const std::vector<TReadSessionEvent::TEvent>& events)
+{
+    for (auto& event : events) {
+        if (auto* e = std::get_if<TReadSessionEvent::TDataReceivedEvent>(&event)) {
+            CollectOffsets(*e);
+        }
+    }
+}
+
+void TOffsetsCollector::CollectOffsets(const TReadSessionEvent::TEvent& event)
+{
+    if (auto* e = std::get_if<TReadSessionEvent::TDataReceivedEvent>(&event)) {
+        CollectOffsets(*e);
+    }
+}
+
+void TOffsetsCollector::CollectOffsets(const TReadSessionEvent::TDataReceivedEvent& event)
+{
+    const auto& session = *event.GetPartitionSession();
+    const std::string& topicPath = session.GetTopicPath();
+    uint32_t partitionId = session.GetPartitionId();
+
+    if (event.HasCompressedMessages()) {
+        for (auto& message : event.GetCompressedMessages()) {
+            uint64_t offset = message.GetOffset();
+            Ranges[topicPath][partitionId].InsertInterval(offset, offset + 1);
+        }
+    } else {
+        for (auto& message : event.GetMessages()) {
+            uint64_t offset = message.GetOffset();
+            Ranges[topicPath][partitionId].InsertInterval(offset, offset + 1);
+        }
+    }
+}
+
+}

src/client/topic/impl/offsets_collector.h

@@ -0,0 +1,35 @@
+#pragma once
+
+#include "topic_impl.h"
+#include "transaction.h"
+
+#include <ydb-cpp-sdk/client/table/table.h>
+#include <ydb-cpp-sdk/client/topic/read_events.h>
+
+#include <library/cpp/containers/disjoint_interval_tree/disjoint_interval_tree.h>
+
+#include <util/generic/hash.h>
+#include <util/generic/string.h>
+#include <util/generic/vector.h>
+
+#include <memory>
+
+namespace NYdb::NTopic {
+
+class TOffsetsCollector {
+public:
+    std::vector<TTopicOffsets> GetOffsets() const;
+
+    void CollectOffsets(const std::vector<TReadSessionEvent::TEvent>& events);
+    void CollectOffsets(const TReadSessionEvent::TEvent& event);
+
+private:
+    // topic -> partition -> (begin, end)
+    using TOffsetRanges = std::unordered_map<std::string, std::unordered_map<uint64_t, TDisjointIntervalTree<uint64_t>>>;
+
+    void CollectOffsets(const TReadSessionEvent::TDataReceivedEvent& event);
+
+    TOffsetRanges Ranges;
+};
+
+}

src/client/topic/impl/read_session.cpp

@@ -11,6 +11,13 @@ namespace NYdb::NTopic {
 
 static const std::string DRIVER_IS_STOPPING_DESCRIPTION = "Driver is stopping";
 
+void SetReadInTransaction(TReadSessionEvent::TEvent& event)
+{
+    if (auto* e = std::get_if<TReadSessionEvent::TDataReceivedEvent>(&event)) {
+        e->SetReadInTransaction();
+    }
+}
+
 TReadSession::TReadSession(const TReadSessionSettings& settings,
              std::shared_ptr<TTopicClient::TImpl> client,
              std::shared_ptr<TGRpcConnectionsImpl> connections,
@@ -136,12 +143,10 @@ std::vector<TReadSessionEvent::TEvent> TReadSession::GetEvents(const TReadSessio
     auto events = GetEvents(settings.Block_, settings.MaxEventsCount_, settings.MaxByteSize_);
     if (!events.empty() && settings.Tx_) {
         auto& tx = settings.Tx_->get();
+        CbContext->TryGet()->CollectOffsets(tx, events, Client);
         for (auto& event : events) {
-            if (auto* dataEvent = std::get_if<TReadSessionEvent::TDataReceivedEvent>(&event)) {
-                CollectOffsets(tx, *dataEvent);
-            }
+            SetReadInTransaction(event);
         }
-        UpdateOffsets(tx);
     }
     return events;
 }
@@ -157,92 +162,14 @@ std::optional<TReadSessionEvent::TEvent> TReadSession::GetEvent(bool block, size
 std::optional<TReadSessionEvent::TEvent> TReadSession::GetEvent(const TReadSessionGetEventSettings& settings)
 {
     auto event = GetEvent(settings.Block_, settings.MaxByteSize_);
-    if (event) {
+    if (event && settings.Tx_) {
         auto& tx = settings.Tx_->get();
-        if (auto* dataEvent = std::get_if<TReadSessionEvent::TDataReceivedEvent>(&*event)) {
-            CollectOffsets(tx, *dataEvent);
-        }
-        UpdateOffsets(tx);
+        CbContext->TryGet()->CollectOffsets(tx, *event, Client);
+        SetReadInTransaction(*event);
     }
     return event;
 }
 
-void TReadSession::CollectOffsets(NTable::TTransaction& tx,
-                                  const TReadSessionEvent::TDataReceivedEvent& event)
-{
-    const auto& session = *event.GetPartitionSession();
-
-    if (event.HasCompressedMessages()) {
-        for (auto& message : event.GetCompressedMessages()) {
-            CollectOffsets(tx, session.GetTopicPath(), session.GetPartitionId(), message.GetOffset());
-        }
-    } else {
-        for (auto& message : event.GetMessages()) {
-            CollectOffsets(tx, session.GetTopicPath(), session.GetPartitionId(), message.GetOffset());
-        }
-    }
-}
-
-void TReadSession::CollectOffsets(NTable::TTransaction& tx,
-                                  const std::string& topicPath, ui32 partitionId, ui64 offset)
-{
-    const std::string& sessionId = tx.GetSession().GetId();
-    const std::string& txId = tx.GetId();
-    TOffsetRanges& ranges = OffsetRanges[std::make_pair(sessionId, txId)];
-    ranges[topicPath][partitionId].InsertInterval(offset, offset + 1);
-}
-
-void TReadSession::UpdateOffsets(const NTable::TTransaction& tx)
-{
-    const std::string& sessionId = tx.GetSession().GetId();
-    const std::string& txId = tx.GetId();
-
-    auto p = OffsetRanges.find(std::make_pair(sessionId, txId));
-    if (p == OffsetRanges.end()) {
-        return;
-    }
-
-    std::vector<TTopicOffsets> topics;
-    for (auto& [path, partitions] : p->second) {
-        TTopicOffsets topic;
-        topic.Path = path;
-
-        topics.push_back(std::move(topic));
-
-        for (auto& [id, ranges] : partitions) {
-            TPartitionOffsets partition;
-            partition.PartitionId = id;
-
-            TTopicOffsets& t = topics.back();
-            t.Partitions.push_back(std::move(partition));
-
-            for (auto& range : ranges) {
-                TPartitionOffsets& p = t.Partitions.back();
-
-                TOffsetsRange r;
-                r.Start = range.first;
-                r.End = range.second;
-
-                p.Offsets.push_back(r);
-            }
-        }
-    }
-
-    Y_ABORT_UNLESS(!topics.empty());
-
-    auto result = Client->UpdateOffsetsInTransaction(tx,
-                                                     topics,
-                                                     Settings.ConsumerName_,
-                                                     {}).GetValueSync();
-    Y_ABORT_UNLESS(!result.IsTransportError());
-
-    if (!result.IsSuccess()) {
-        ythrow yexception() << "error on update offsets: " << result;
-    }
-
-    OffsetRanges.erase(std::make_pair(sessionId, txId));
-}
-
 bool TReadSession::Close(TDuration timeout) {
     LOG_LAZY(Log, TLOG_INFO, GetLogPrefix() << "Closing read session. Close timeout: " << timeout);
     // Log final counters.

src/client/topic/impl/read_session.h

@@ -63,18 +63,6 @@ class TReadSession : public IReadSession {
     void AbortImpl(EStatus statusCode, const std::string& message, TDeferredActions<false>& deferred);
 
 private:
-    using TOffsetRanges = std::unordered_map<std::string, std::unordered_map<ui64, TDisjointIntervalTree<ui64>>>;
-
-    void CollectOffsets(NTable::TTransaction& tx,
-                        const TReadSessionEvent::TDataReceivedEvent& event);
-    void CollectOffsets(NTable::TTransaction& tx,
-                        const std::string& topicPath, ui32 partitionId, ui64 offset);
-    void UpdateOffsets(const NTable::TTransaction& tx);
-
-    //
-    // (session, tx) -> topic -> partition -> (begin, end)
-    //
-    std::unordered_map<std::pair<std::string, std::string>, TOffsetRanges, THash<std::pair<std::string, std::string>>> OffsetRanges;
 
     TReadSessionSettings Settings;
     const std::string SessionId;

src/client/topic/impl/read_session_event.cpp

@@ -248,6 +248,10 @@ TDataReceivedEvent::TDataReceivedEvent(std::vector<TMessage> messages, std::vect
 }
 
 void TDataReceivedEvent::Commit() {
+    if (ReadInTransaction) {
+        ythrow yexception() << "Offsets cannot be commited explicitly when reading in a transaction";
+    }
+
     for (auto [from, to] : OffsetRanges) {
         static_cast<TPartitionStreamImpl<false>*>(PartitionSession.Get())->Commit(from, to);
     }