Add NFederatedTopic::TDeferredCommit implementation (#17430)

Author: qyryq

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

@@ -21,6 +21,8 @@ using TAsyncDescribeTopicResult = NTopic::TAsyncDescribeTopicResult;
 struct TFederatedPartitionSession : public TThrRefBase, public TPrintable<TFederatedPartitionSession> {
     using TPtr = TIntrusivePtr<TFederatedPartitionSession>;
 
+    friend class TDeferredCommit;
+
 public:
     TFederatedPartitionSession(const NTopic::TPartitionSession::TPtr& partitionSession,
                                std::shared_ptr<TDbInfo> db,
@@ -223,10 +225,10 @@ class TDeferredCommit {
     void Add(const TReadSessionEvent::TDataReceivedEvent& dataReceivedEvent);
 
     //! Add offsets range to set.
-    void Add(const TFederatedPartitionSession& partitionSession, ui64 startOffset, ui64 endOffset);
+    void Add(const TFederatedPartitionSession::TPtr& partitionSession, ui64 startOffset, ui64 endOffset);
 
     //! Add offset to set.
-    void Add(const TFederatedPartitionSession& partitionSession, ui64 offset);
+    void Add(const TFederatedPartitionSession::TPtr& partitionSession, ui64 offset);
 
     //! Commit all added offsets.
     void Commit();
@@ -399,7 +401,7 @@ struct TFederatedReadSessionSettings: public NTopic::TReadSessionSettings {
     //! See description in TFederatedEventHandlers class.
     FLUENT_SETTING(TFederatedEventHandlers, FederatedEventHandlers);
 
-    
+
 
     //! Read policy settings
 

src/client/federated_topic/impl/federated_deferred_commit.cpp

@@ -0,0 +1,138 @@
+#include <ydb-cpp-sdk/client/federated_topic/federated_topic.h>
+#include <src/client/topic/impl/read_session_impl.ipp>
+
+#include <library/cpp/containers/disjoint_interval_tree/disjoint_interval_tree.h>
+
+namespace NYdb::inline V3::NFederatedTopic {
+
+std::pair<ui64, ui64> GetMessageOffsetRange(const TReadSessionEvent::TDataReceivedEvent& dataReceivedEvent, ui64 index) {
+    if (dataReceivedEvent.HasCompressedMessages()) {
+        const auto& msg = dataReceivedEvent.GetCompressedMessages()[index];
+        return {msg.GetOffset(), msg.GetOffset() + 1};
+    }
+    const auto& msg = dataReceivedEvent.GetMessages()[index];
+    return {msg.GetOffset(), msg.GetOffset() + 1};
+}
+
+
+////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
+// NFederatedTopic::TDeferredCommit
+
+class TDeferredCommit::TImpl {
+public:
+
+    void Add(const TFederatedPartitionSession::TPtr& partitionStream, ui64 startOffset, ui64 endOffset);
+    void Add(const TFederatedPartitionSession::TPtr& partitionStream, ui64 offset);
+
+    void Add(const TReadSessionEvent::TDataReceivedEvent::TMessage& message);
+    void Add(const TReadSessionEvent::TDataReceivedEvent& dataReceivedEvent);
+
+    void Commit();
+
+private:
+    static void Add(const TFederatedPartitionSession::TPtr& partitionStream, TDisjointIntervalTree<ui64>& offsetSet, ui64 startOffset, ui64 endOffset);
+
+private:
+    // Partition stream -> offsets set.
+    std::unordered_map<TFederatedPartitionSession::TPtr, TDisjointIntervalTree<ui64>, THash<TFederatedPartitionSession::TPtr>> Offsets;
+};
+
+TDeferredCommit::TDeferredCommit() {
+}
+
+TDeferredCommit::TDeferredCommit(TDeferredCommit&&) = default;
+
+TDeferredCommit& TDeferredCommit::operator=(TDeferredCommit&&) = default;
+
+TDeferredCommit::~TDeferredCommit() {
+}
+
+#define GET_IMPL()                              \
+    if (!Impl) {                                \
+        Impl = std::make_unique<TImpl>();       \
+    }                                           \
+    Impl
+
+void TDeferredCommit::Add(const TFederatedPartitionSession::TPtr& partitionStream, ui64 startOffset, ui64 endOffset) {
+    GET_IMPL()->Add(partitionStream, startOffset, endOffset);
+}
+
+void TDeferredCommit::Add(const TFederatedPartitionSession::TPtr& partitionStream, ui64 offset) {
+    GET_IMPL()->Add(partitionStream, offset);
+}
+
+void TDeferredCommit::Add(const TReadSessionEvent::TDataReceivedEvent::TMessage& message) {
+    GET_IMPL()->Add(message);
+}
+
+void TDeferredCommit::Add(const TReadSessionEvent::TDataReceivedEvent& dataReceivedEvent) {
+    GET_IMPL()->Add(dataReceivedEvent);
+}
+
+#undef GET_IMPL
+
+void TDeferredCommit::Commit() {
+    if (Impl) {
+        Impl->Commit();
+    }
+}
+
+void TDeferredCommit::TImpl::Add(const TReadSessionEvent::TDataReceivedEvent::TMessage& message) {
+    Y_ASSERT(message.GetFederatedPartitionSession());
+    Add(message.GetFederatedPartitionSession(), message.GetOffset());
+}
+
+void TDeferredCommit::TImpl::Add(const TFederatedPartitionSession::TPtr& partitionStream, TDisjointIntervalTree<ui64>& offsetSet, ui64 startOffset, ui64 endOffset) {
+    if (offsetSet.Intersects(startOffset, endOffset)) {
+        ThrowFatalError(TStringBuilder() << "Commit set already has some offsets from half-interval ["
+                                         << startOffset << "; " << endOffset
+                                         << ") for partition stream with id " << partitionStream->GetPartitionSessionId());
+    } else {
+        offsetSet.InsertInterval(startOffset, endOffset);
+    }
+}
+
+void TDeferredCommit::TImpl::Add(const TFederatedPartitionSession::TPtr& partitionStream, ui64 startOffset, ui64 endOffset) {
+    Y_ASSERT(partitionStream);
+    Add(partitionStream, Offsets[partitionStream], startOffset, endOffset);
+}
+
+void TDeferredCommit::TImpl::Add(const TFederatedPartitionSession::TPtr& partitionStream, ui64 offset) {
+    Y_ASSERT(partitionStream);
+    auto& offsetSet = Offsets[partitionStream];
+    if (offsetSet.Has(offset)) {
+        ThrowFatalError(TStringBuilder() << "Commit set already has offset " << offset
+                                         << " for partition stream with id " << partitionStream->GetPartitionSessionId());
+    } else {
+        offsetSet.Insert(offset);
+    }
+}
+
+void TDeferredCommit::TImpl::Add(const TReadSessionEvent::TDataReceivedEvent& dataReceivedEvent) {
+    const TFederatedPartitionSession::TPtr& partitionStream = dataReceivedEvent.GetFederatedPartitionSession();
+    Y_ASSERT(partitionStream);
+    auto& offsetSet = Offsets[partitionStream];
+    auto [startOffset, endOffset] = GetMessageOffsetRange(dataReceivedEvent, 0);
+    for (size_t i = 1; i < dataReceivedEvent.GetMessagesCount(); ++i) {
+        auto msgOffsetRange = GetMessageOffsetRange(dataReceivedEvent, i);
+        if (msgOffsetRange.first == endOffset) {
+            endOffset= msgOffsetRange.second;
+        } else {
+            Add(partitionStream, offsetSet, startOffset, endOffset);
+            startOffset = msgOffsetRange.first;
+            endOffset = msgOffsetRange.second;
+        }
+    }
+    Add(partitionStream, offsetSet, startOffset, endOffset);
+}
+
+void TDeferredCommit::TImpl::Commit() {
+    for (auto&& [partitionStream, offsetRanges] : Offsets) {
+        for (auto&& [startOffset, endOffset] : offsetRanges) {
+            static_cast<NTopic::TPartitionStreamImpl<false>*>(partitionStream.Get()->PartitionSession.Get())->Commit(startOffset, endOffset);
+        }
+    }
+    Offsets.clear();
+}
+
+}