Write changes using WriteTxId in consistent mode (#11374)

Author: CyberROFL

ydb/core/base/row_version.cpp

@@ -6,6 +6,10 @@
 
 namespace NKikimr {
 
+TRowVersion TRowVersion::Parse(const NKikimrProto::TRowVersion& proto) {
+    return TRowVersion(proto.GetStep(), proto.GetTxId());
+}
+
 void TRowVersion::Serialize(NKikimrProto::TRowVersion& proto) const {
     proto.SetStep(Step);
     proto.SetTxId(TxId);

ydb/core/base/row_version.h

@@ -75,6 +75,7 @@ namespace NKikimr {
             return ++copy;
         }
 
+        static TRowVersion Parse(const NKikimrProto::TRowVersion& proto);
         void Serialize(NKikimrProto::TRowVersion& proto) const;
 
         friend constexpr bool operator==(const TRowVersion& a, const TRowVersion& b) {

ydb/core/change_exchange/change_record.h

@@ -37,6 +37,8 @@ class IChangeRecord: public TThrRefBase {
 
     virtual void Accept(IVisitor& visitor) const = 0;
 
+    virtual void RewriteTxId(ui64 value) = 0;
+
     virtual TString ToString() const = 0;
     virtual void Out(IOutputStream& out) const = 0;
 
@@ -56,6 +58,8 @@ class TChangeRecordBase: public IChangeRecord {
     const TString& GetSourceId() const override { return SourceId; }
     bool IsBroadcast() const override { return false; }
 
+    void RewriteTxId(ui64) override { Y_ABORT("not implemented"); }
+
     TString ToString() const override;
     void Out(IOutputStream& out) const override;
 

ydb/core/protos/replication.proto

@@ -1,4 +1,5 @@
 import "ydb/core/scheme/protos/pathid.proto";
+import "ydb/core/protos/base.proto";
 import "ydb/public/api/protos/annotations/sensitive.proto";
 import "ydb/public/api/protos/ydb_issue_message.proto";
 
@@ -257,3 +258,22 @@ message TEvWorkerDataEnd {
     repeated uint64 AdjacentPartitionsIds = 3;
     repeated uint64 ChildPartitionsIds = 4;
 }
+
+message TEvGetTxId {
+    repeated NKikimrProto.TRowVersion Versions = 1;
+}
+
+message TEvTxIdResult {
+    message TVersionTxId {
+        optional NKikimrProto.TRowVersion Version = 1;
+        optional uint64 TxId = 2;
+    }
+
+    optional TControllerIdentity Controller = 1;
+    repeated TVersionTxId VersionTxIds = 2;
+}
+
+message TEvHeartbeat {
+    optional TWorkerIdentity Worker = 1;
+    optional NKikimrProto.TRowVersion Version = 2;
+}

ydb/core/tx/replication/service/base_table_writer.cpp

@@ -1,5 +1,6 @@
 #include "base_table_writer.h"
 #include "logging.h"
+#include "service.h"
 #include "worker.h"
 
 #include <ydb/core/base/tablet_pipecache.h>
@@ -16,6 +17,7 @@
 
 #include <util/generic/map.h>
 #include <util/generic/maybe.h>
+#include <util/generic/set.h>
 #include <util/string/builder.h>
 
 namespace NKikimr::NReplication::NService {
@@ -426,16 +428,71 @@ class TLocalTableWriter
 
         Y_ABORT_UNLESS(PendingRecords.empty());
         TVector<NChangeExchange::TEvChangeExchange::TEvEnqueueRecords::TRecordInfo> records(::Reserve(ev->Get()->Records.size()));
+        TSet<TRowVersion> versionsWithoutTxId;
 
-        for (auto& record : ev->Get()->Records) {
-            records.emplace_back(record.Offset, TablePathId, record.Data.size());
-            auto res = PendingRecords.emplace(
-                record.Offset, Parser->Parse(ev->Get()->Source, record.Offset, std::move(record.Data))
-            );
-            Y_ABORT_UNLESS(res.second);
+        for (auto& [offset, data, _] : ev->Get()->Records) {
+            auto record = Parser->Parse(ev->Get()->Source, offset, std::move(data));
+
+            if (Mode == EWriteMode::Consistent) {
+                const auto version = TRowVersion(record->GetStep(), record->GetTxId());
+
+                if (record->GetKind() == NChangeExchange::IChangeRecord::EKind::CdcHeartbeat) {
+                    TxIds.erase(TxIds.begin(), TxIds.upper_bound(version));
+                    Send(Worker, new TEvService::TEvHeartbeat(version));
+                    continue;
+                } else if (record->GetKind() != NChangeExchange::IChangeRecord::EKind::CdcDataChange) {
+                    Y_ABORT("Unexpected record kind");
+                }
+
+                if (auto it = TxIds.upper_bound(version); it != TxIds.end()) {
+                    record->RewriteTxId(it->second);
+                } else {
+                    versionsWithoutTxId.insert(version);
+                    PendingTxId[version].push_back(std::move(record));
+                    continue;
+                }
+            }
+
+            records.emplace_back(record->GetOrder(), TablePathId, record->GetBody().size());
+            Y_ABORT_UNLESS(PendingRecords.emplace(record->GetOrder(), std::move(record)).second);
+        }
+
+        if (versionsWithoutTxId) {
+            Send(Worker, new TEvService::TEvGetTxId(versionsWithoutTxId));
+        }
+
+        if (records) {
+            EnqueueRecords(std::move(records));
+        }
+    }
+
+    void Handle(TEvService::TEvTxIdResult::TPtr& ev) {
+        LOG_D("Handle " << ev->Get()->ToString());
+
+        TVector<NChangeExchange::TEvChangeExchange::TEvEnqueueRecords::TRecordInfo> records;
+
+        for (const auto& kv : ev->Get()->Record.GetVersionTxIds()) {
+            const auto version = TRowVersion::Parse(kv.GetVersion());
+            TxIds.emplace(version, kv.GetTxId());
+
+            for (auto it = PendingTxId.begin(); it != PendingTxId.end();) {
+                if (it->first >= version) {
+                    break;
+                }
+
+                for (auto& record : it->second) {
+                    record->RewriteTxId(kv.GetTxId());
+                    records.emplace_back(record->GetOrder(), TablePathId, record->GetBody().size());
+                    Y_ABORT_UNLESS(PendingRecords.emplace(record->GetOrder(), std::move(record)).second);
+                }
+
+                PendingTxId.erase(it++);
+            }
         }
 
-        EnqueueRecords(std::move(records));
+        if (records) {
+            EnqueueRecords(std::move(records));
+        }
     }
 
     void Handle(NChangeExchange::TEvChangeExchange::TEvRequestRecords::TPtr& ev) {
@@ -502,12 +559,14 @@ class TLocalTableWriter
     }
 
     explicit TLocalTableWriter(
+            EWriteMode mode,
             const TPathId& tablePathId,
             THolder<IChangeRecordParser>&& parser,
             THolder<IChangeRecordSerializer>&& serializer,
             std::function<NChangeExchange::IPartitionResolverVisitor*(const NKikimr::TKeyDesc&)>&& createResolverFn)
         : TActor(&TThis::StateWork)
         , TChangeSender(this, this, this, this, TActorId())
+        , Mode(mode)
         , TablePathId(tablePathId)
         , Parser(std::move(parser))
         , Serializer(std::move(serializer))
@@ -523,6 +582,7 @@ class TLocalTableWriter
         switch (ev->GetTypeRewrite()) {
             hFunc(TEvWorker::TEvHandshake, Handle);
             hFunc(TEvWorker::TEvData, Handle);
+            hFunc(TEvService::TEvTxIdResult, Handle);
             hFunc(NChangeExchange::TEvChangeExchange::TEvRequestRecords, Handle);
             hFunc(NChangeExchange::TEvChangeExchange::TEvRemoveRecords, Handle);
             hFunc(NChangeExchange::TEvChangeExchangePrivate::TEvReady, Handle);
@@ -536,6 +596,7 @@ class TLocalTableWriter
 
 private:
     mutable TMaybe<TString> LogPrefix;
+    const EWriteMode Mode;
     const TPathId TablePathId;
     THolder<IChangeRecordParser> Parser;
     THolder<IChangeRecordSerializer> Serializer;
@@ -549,16 +610,19 @@ class TLocalTableWriter
     bool Initialized = false;
 
     TMap<ui64, NChangeExchange::IChangeRecord::TPtr> PendingRecords;
+    TMap<TRowVersion, ui64> TxIds; // key is non-inclusive right hand edge
+    TMap<TRowVersion, TVector<NChangeExchange::IChangeRecord::TPtr>> PendingTxId;
 
 }; // TLocalTableWriter
 
 IActor* CreateLocalTableWriter(
         const TPathId& tablePathId,
         THolder<IChangeRecordParser>&& parser,
         THolder<IChangeRecordSerializer>&& serializer,
-        std::function<NChangeExchange::IPartitionResolverVisitor*(const NKikimr::TKeyDesc&)>&& createResolverFn)
+        std::function<NChangeExchange::IPartitionResolverVisitor*(const NKikimr::TKeyDesc&)>&& createResolverFn,
+        EWriteMode mode)
 {
-    return new TLocalTableWriter(tablePathId, std::move(parser), std::move(serializer), std::move(createResolverFn));
+    return new TLocalTableWriter(mode, tablePathId, std::move(parser), std::move(serializer), std::move(createResolverFn));
 }
 
 } // namespace NKikimr::NReplication::NService

ydb/core/tx/replication/service/base_table_writer.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include "lightweight_schema.h"
+#include "table_writer.h"
 
 #include <ydb/core/base/defs.h>
 #include <ydb/core/change_exchange/change_record.h>
@@ -31,6 +32,7 @@ IActor* CreateLocalTableWriter(
     const TPathId& tablePathId,
     THolder<IChangeRecordParser>&& parser,
     THolder<IChangeRecordSerializer>&& serializer,
-    std::function<NChangeExchange::IPartitionResolverVisitor*(const NKikimr::TKeyDesc&)>&& createResolverFn);
+    std::function<NChangeExchange::IPartitionResolverVisitor*(const NKikimr::TKeyDesc&)>&& createResolverFn,
+    EWriteMode mode = EWriteMode::Simple);
 
 }

ydb/core/tx/replication/service/json_change_record.cpp

@@ -91,7 +91,9 @@ static bool ParseValue(TVector<NTable::TTag>& tags, TVector<TCell>& cells,
 void TChangeRecord::Serialize(NKikimrTxDataShard::TEvApplyReplicationChanges_TChange& record, TMemoryPool& pool) const {
     pool.Clear();
     record.SetSourceOffset(GetOrder());
-    // TODO: fill WriteTxId
+    if (WriteTxId) {
+        record.SetWriteTxId(WriteTxId);
+    }
 
     TString error;
 
@@ -155,4 +157,8 @@ void TChangeRecord::Accept(NChangeExchange::IVisitor& visitor) const {
     return visitor.Visit(*this);
 }
 
+void TChangeRecord::RewriteTxId(ui64 value) {
+    WriteTxId = value;
+}
+
 }

ydb/core/tx/replication/service/json_change_record.h

@@ -33,10 +33,12 @@ class TChangeRecord: public NChangeExchange::TChangeRecordBase {
     TConstArrayRef<TCell> GetKey() const;
 
     void Accept(NChangeExchange::IVisitor& visitor) const override;
+    void RewriteTxId(ui64 value) override;
 
 private:
     NJson::TJsonValue JsonBody;
     TLightweightSchema::TCPtr Schema;
+    ui64 WriteTxId = 0;
 
     mutable TMaybe<TOwnedCellVec> Key;