Specify source of changes (#4680)

CyberROFL · web-flow · commit 6735aec7128a · 2024-05-20T18:55:16.000+03:00
diff --git a/ydb/core/tx/replication/service/json_change_record.cpp b/ydb/core/tx/replication/service/json_change_record.cpp
@@ -33,6 +33,10 @@ NChangeExchange::IChangeRecord::EKind TChangeRecord::GetKind() const {
         : EKind::CdcDataChange;
 }
 
+TString TChangeRecord::GetSourceId() const {
+    return SourceId;
+}
+
 static bool ParseKey(TVector<TCell>& cells,
         const NJson::TJsonValue::TArray& key, TLightweightSchema::TCPtr schema, TMemoryPool& pool, TString& error)
 {
diff --git a/ydb/core/tx/replication/service/json_change_record.h b/ydb/core/tx/replication/service/json_change_record.h
@@ -40,6 +40,7 @@ class TChangeRecord: public NChangeExchange::TChangeRecordBase {
     ui64 GetStep() const override;
     ui64 GetTxId() const override;
     EKind GetKind() const override;
+    TString GetSourceId() const;
 
     void Serialize(NKikimrTxDataShard::TEvApplyReplicationChanges::TChange& record, TMemoryPool& pool) const;
     void Serialize(NKikimrTxDataShard::TEvApplyReplicationChanges::TChange& record) const;
@@ -48,6 +49,7 @@ class TChangeRecord: public NChangeExchange::TChangeRecordBase {
     TConstArrayRef<TCell> GetKey() const;
 
 private:
+    TString SourceId;
     NJson::TJsonValue JsonBody;
     TLightweightSchema::TCPtr Schema;
 
@@ -59,6 +61,11 @@ class TChangeRecordBuilder: public NChangeExchange::TChangeRecordBuilder<TChange
 public:
     using TBase::TBase;
 
+    TSelf& WithSourceId(const TString& sourceId) {
+        GetRecord()->SourceId = sourceId;
+        return static_cast<TSelf&>(*this);
+    }
+
     template <typename T>
     TSelf& WithBody(T&& body) {
         auto res = NJson::ReadJsonTree(body, &GetRecord()->JsonBody);
diff --git a/ydb/core/tx/replication/service/table_writer.cpp b/ydb/core/tx/replication/service/table_writer.cpp
@@ -76,11 +76,21 @@ class TTablePartitionWriter: public TActorBootstrapped<TTablePartitionWriter> {
         tableId.SetTableId(TableId.PathId.LocalPathId);
         tableId.SetSchemaVersion(TableId.SchemaVersion);
 
+        TString source;
         for (auto recordPtr : ev->Get()->Records) {
             MemoryPool.Clear();
             const auto& record = *recordPtr->Get<TChangeRecord>();
             record.Serialize(*event->Record.AddChanges(), MemoryPool);
-            // TODO: set WriteTxId, Source
+
+            if (!source) {
+                source = record.GetSourceId();
+            } else {
+                Y_ABORT_UNLESS(source == record.GetSourceId());
+            }
+        }
+
+        if (source) {
+            event->Record.SetSource(source);
         }
 
         Send(LeaderPipeCache, new TEvPipeCache::TEvForward(event.Release(), TabletId, false));
@@ -427,6 +437,7 @@ class TLocalTableWriter
         for (auto& record : ev->Get()->Records) {
             records.emplace_back(record.Offset, PathId, record.Data.size());
             auto res = PendingRecords.emplace(record.Offset, TChangeRecordBuilder()
+                .WithSourceId(ev->Get()->Source)
                 .WithOrder(record.Offset)
                 .WithBody(std::move(record.Data))
                 .WithSchema(Schema)
diff --git a/ydb/core/tx/replication/service/table_writer_ut.cpp b/ydb/core/tx/replication/service/table_writer_ut.cpp
@@ -31,7 +31,7 @@ Y_UNIT_TEST_SUITE(LocalTableWriter) {
         env.Send<TEvWorker::TEvHandshake>(writer, new TEvWorker::TEvHandshake());
 
         using TRecord = TEvWorker::TEvData::TRecord;
-        env.Send<TEvWorker::TEvPoll>(writer, new TEvWorker::TEvData({
+        env.Send<TEvWorker::TEvPoll>(writer, new TEvWorker::TEvData("TestSource", {
             TRecord(1, R"({"key":[1], "update":{"value":"10"}})"),
             TRecord(2, R"({"key":[2], "update":{"value":"20"}})"),
             TRecord(3, R"({"key":[3], "update":{"value":"30"}})"),
@@ -73,7 +73,7 @@ Y_UNIT_TEST_SUITE(LocalTableWriter) {
         env.Send<TEvWorker::TEvHandshake>(writer, new TEvWorker::TEvHandshake());
 
         using TRecord = TEvWorker::TEvData::TRecord;
-        env.Send<TEvWorker::TEvPoll>(writer, new TEvWorker::TEvData({
+        env.Send<TEvWorker::TEvPoll>(writer, new TEvWorker::TEvData("TestSource", {
             TRecord(1, R"({"key":[1], "update":{"int32_value":-100500}})"),
             TRecord(2, R"({"key":[2], "update":{"uint32_value":100500}})"),
             TRecord(3, R"({"key":[3], "update":{"int64_value":-200500}})"),
diff --git a/ydb/core/tx/replication/service/topic_reader.cpp b/ydb/core/tx/replication/service/topic_reader.cpp
@@ -61,7 +61,7 @@ class TRemoteTopicReader: public TActor<TRemoteTopicReader> {
             records.emplace_back(msg.GetOffset(), std::move(msg.GetData()));
         }
 
-        Send(Worker, new TEvWorker::TEvData(std::move(records)));
+        Send(Worker, new TEvWorker::TEvData(ToString(result.PartitionId), std::move(records)));
     }
 
     void Handle(TEvYdbProxy::TEvTopicReaderGone::TPtr& ev) {
diff --git a/ydb/core/tx/replication/service/worker.cpp b/ydb/core/tx/replication/service/worker.cpp
@@ -23,13 +23,15 @@ TEvWorker::TEvData::TRecord::TRecord(ui64 offset, TString&& data)
 {
 }
 
-TEvWorker::TEvData::TEvData(const TVector<TRecord>& records)
-    : Records(records)
+TEvWorker::TEvData::TEvData(const TString& source, const TVector<TRecord>& records)
+    : Source(source)
+    , Records(records)
 {
 }
 
-TEvWorker::TEvData::TEvData(TVector<TRecord>&& records)
-    : Records(std::move(records))
+TEvWorker::TEvData::TEvData(const TString& source, TVector<TRecord>&& records)
+    : Source(source)
+    , Records(std::move(records))
 {
 }
 
@@ -42,6 +44,7 @@ void TEvWorker::TEvData::TRecord::Out(IOutputStream& out) const {
 
 TString TEvWorker::TEvData::ToString() const {
     return TStringBuilder() << ToStringHeader() << " {"
+        << " Source: " << Source
         << " Records [" << JoinSeq(",", Records) << "]"
     << " }";
 }
@@ -115,16 +118,16 @@ class TWorker: public TActorBootstrapped<TWorker> {
                 << ": sender# " << ev->Sender);
 
             Reader.Registered();
-            if (!InFlightRecords) {
+            if (!InFlightData) {
                 Send(Reader, new TEvWorker::TEvPoll());
             }
         } else if (ev->Sender == Writer) {
             LOG_I("Handshake with writer"
                 << ": sender# " << ev->Sender);
 
             Writer.Registered();
-            if (InFlightRecords) {
-                Send(Writer, new TEvWorker::TEvData(InFlightRecords));
+            if (InFlightData) {
+                Send(Writer, new TEvWorker::TEvData(InFlightData->Source, InFlightData->Records));
             }
         } else {
             LOG_W("Handshake from unknown actor"
@@ -142,7 +145,7 @@ class TWorker: public TActorBootstrapped<TWorker> {
             return;
         }
 
-        InFlightRecords.clear();
+        InFlightData.Reset();
         if (Reader) {
             Send(ev->Forward(Reader));
         }
@@ -157,8 +160,8 @@ class TWorker: public TActorBootstrapped<TWorker> {
             return;
         }
 
-        Y_ABORT_UNLESS(InFlightRecords.empty());
-        InFlightRecords = ev->Get()->Records;
+        Y_ABORT_UNLESS(!InFlightData);
+        InFlightData = MakeHolder<TEvWorker::TEvData>(ev->Get()->Source, ev->Get()->Records);
 
         if (Writer) {
             Send(ev->Forward(Writer));
@@ -239,7 +242,7 @@ class TWorker: public TActorBootstrapped<TWorker> {
     mutable TMaybe<TString> LogPrefix;
     TActorInfo Reader;
     TActorInfo Writer;
-    TVector<TEvWorker::TEvData::TRecord> InFlightRecords;
+    THolder<TEvWorker::TEvData> InFlightData;
 };
 
 IActor* CreateWorker(std::function<IActor*(void)>&& createReaderFn, std::function<IActor*(void)>&& createWriterFn) {
diff --git a/ydb/core/tx/replication/service/worker.h b/ydb/core/tx/replication/service/worker.h
@@ -36,10 +36,11 @@ struct TEvWorker {
             void Out(IOutputStream& out) const;
         };
 
+        TString Source;
         TVector<TRecord> Records;
 
-        explicit TEvData(const TVector<TRecord>& records);
-        explicit TEvData(TVector<TRecord>&& records);
+        explicit TEvData(const TString& source, const TVector<TRecord>& records);
+        explicit TEvData(const TString& source, TVector<TRecord>&& records);
         TString ToString() const override;
     };
 
diff --git a/ydb/core/tx/replication/ydb_proxy/ydb_proxy.cpp b/ydb/core/tx/replication/ydb_proxy/ydb_proxy.cpp
@@ -33,6 +33,7 @@ void TEvYdbProxy::TReadTopicResult::TMessage::Out(IOutputStream& out) const {
 
 void TEvYdbProxy::TReadTopicResult::Out(IOutputStream& out) const {
     out << "{"
+        << " PartitionId: " << PartitionId
         << " Messages [" << JoinSeq(",", Messages) << "]"
     << " }";
 }
diff --git a/ydb/core/tx/replication/ydb_proxy/ydb_proxy.h b/ydb/core/tx/replication/ydb_proxy/ydb_proxy.h
@@ -186,6 +186,7 @@ struct TEvYdbProxy {
         };
 
         explicit TReadTopicResult(const NYdb::NTopic::TReadSessionEvent::TDataReceivedEvent& event) {
+            PartitionId = event.GetPartitionSession()->GetPartitionId();
             Messages.reserve(event.GetMessagesCount());
             if (event.HasCompressedMessages()) {
                 for (const auto& msg : event.GetCompressedMessages()) {
@@ -200,6 +201,7 @@ struct TEvYdbProxy {
 
         void Out(IOutputStream& out) const;
 
+        ui64 PartitionId;
         TVector<TMessage> Messages;
     };
 

Original file line number	Diff line number	Diff line change
`@@ -33,6 +33,10 @@ NChangeExchange::IChangeRecord::EKind TChangeRecord::GetKind() const {`
`33`	`33`	`: EKind::CdcDataChange;`
`34`	`34`	`}`
`35`	`35`
	`36`	`+TString TChangeRecord::GetSourceId() const {`
	`37`	`+ return SourceId;`
	`38`	`+}`
	`39`	`+`
`36`	`40`	`static bool ParseKey(TVector<TCell>& cells,`
`37`	`41`	`const NJson::TJsonValue::TArray& key, TLightweightSchema::TCPtr schema, TMemoryPool& pool, TString& error)`
`38`	`42`	`{`
Original file line number	Diff line number	Diff line change
`@@ -61,7 +61,7 @@ class TRemoteTopicReader: public TActor<TRemoteTopicReader> {`
`61`	`61`	`records.emplace_back(msg.GetOffset(), std::move(msg.GetData()));`
`62`	`62`	`}`
`63`	`63`
`64`		`- Send(Worker, new TEvWorker::TEvData(std::move(records)));`
	`64`	`+ Send(Worker, new TEvWorker::TEvData(ToString(result.PartitionId), std::move(records)));`
`65`	`65`	`}`
`66`	`66`
`67`	`67`	`void Handle(TEvYdbProxy::TEvTopicReaderGone::TPtr& ev) {`
Original file line number	Diff line number	Diff line change
`@@ -23,13 +23,15 @@ TEvWorker::TEvData::TRecord::TRecord(ui64 offset, TString&& data)`
`23`	`23`	`{`
`24`	`24`	`}`
`25`	`25`
`26`		`-TEvWorker::TEvData::TEvData(const TVector<TRecord>& records)`
`27`		`- : Records(records)`
	`26`	`+TEvWorker::TEvData::TEvData(const TString& source, const TVector<TRecord>& records)`
	`27`	`+ : Source(source)`
	`28`	`+ , Records(records)`
`28`	`29`	`{`
`29`	`30`	`}`
`30`	`31`
`31`		`-TEvWorker::TEvData::TEvData(TVector<TRecord>&& records)`
`32`		`- : Records(std::move(records))`
	`32`	`+TEvWorker::TEvData::TEvData(const TString& source, TVector<TRecord>&& records)`
	`33`	`+ : Source(source)`
	`34`	`+ , Records(std::move(records))`
`33`	`35`	`{`
`34`	`36`	`}`
`35`	`37`
`@@ -42,6 +44,7 @@ void TEvWorker::TEvData::TRecord::Out(IOutputStream& out) const {`
`42`	`44`
`43`	`45`	`TString TEvWorker::TEvData::ToString() const {`
`44`	`46`	`return TStringBuilder() << ToStringHeader() << " {"`
	`47`	`+ << " Source: " << Source`
`45`	`48`	`<< " Records [" << JoinSeq(",", Records) << "]"`
`46`	`49`	`<< " }";`
`47`	`50`	`}`
`@@ -115,16 +118,16 @@ class TWorker: public TActorBootstrapped<TWorker> {`
`115`	`118`	`<< ": sender# " << ev->Sender);`
`116`	`119`
`117`	`120`	`Reader.Registered();`
`118`		`- if (!InFlightRecords) {`
	`121`	`+ if (!InFlightData) {`
`119`	`122`	`Send(Reader, new TEvWorker::TEvPoll());`
`120`	`123`	`}`
`121`	`124`	`} else if (ev->Sender == Writer) {`
`122`	`125`	`LOG_I("Handshake with writer"`
`123`	`126`	`<< ": sender# " << ev->Sender);`
`124`	`127`
`125`	`128`	`Writer.Registered();`
`126`		`- if (InFlightRecords) {`
`127`		`- Send(Writer, new TEvWorker::TEvData(InFlightRecords));`
	`129`	`+ if (InFlightData) {`
	`130`	`+ Send(Writer, new TEvWorker::TEvData(InFlightData->Source, InFlightData->Records));`
`128`	`131`	`}`
`129`	`132`	`} else {`
`130`	`133`	`LOG_W("Handshake from unknown actor"`
`@@ -142,7 +145,7 @@ class TWorker: public TActorBootstrapped<TWorker> {`
`142`	`145`	`return;`
`143`	`146`	`}`
`144`	`147`
`145`		`- InFlightRecords.clear();`
	`148`	`+ InFlightData.Reset();`
`146`	`149`	`if (Reader) {`
`147`	`150`	`Send(ev->Forward(Reader));`
`148`	`151`	`}`
`@@ -157,8 +160,8 @@ class TWorker: public TActorBootstrapped<TWorker> {`
`157`	`160`	`return;`
`158`	`161`	`}`
`159`	`162`
`160`		`- Y_ABORT_UNLESS(InFlightRecords.empty());`
`161`		`- InFlightRecords = ev->Get()->Records;`
	`163`	`+ Y_ABORT_UNLESS(!InFlightData);`
	`164`	`+ InFlightData = MakeHolder<TEvWorker::TEvData>(ev->Get()->Source, ev->Get()->Records);`
`162`	`165`
`163`	`166`	`if (Writer) {`
`164`	`167`	`Send(ev->Forward(Writer));`
`@@ -239,7 +242,7 @@ class TWorker: public TActorBootstrapped<TWorker> {`
`239`	`242`	`mutable TMaybe<TString> LogPrefix;`
`240`	`243`	`TActorInfo Reader;`
`241`	`244`	`TActorInfo Writer;`
`242`		`- TVector<TEvWorker::TEvData::TRecord> InFlightRecords;`
	`245`	`+ THolder<TEvWorker::TEvData> InFlightData;`
`243`	`246`	`};`
`244`	`247`
`245`	`248`	`IActor* CreateWorker(std::function<IActor(void)>&& createReaderFn, std::function<IActor(void)>&& createWriterFn) {`
Original file line number	Diff line number	Diff line change
`@@ -33,6 +33,7 @@ void TEvYdbProxy::TReadTopicResult::TMessage::Out(IOutputStream& out) const {`
`33`	`33`
`34`	`34`	`void TEvYdbProxy::TReadTopicResult::Out(IOutputStream& out) const {`
`35`	`35`	`out << "{"`
	`36`	`+ << " PartitionId: " << PartitionId`
`36`	`37`	`<< " Messages [" << JoinSeq(",", Messages) << "]"`
`37`	`38`	`<< " }";`
`38`	`39`	`}`