Process heartbeats (#12333)

CyberROFL · web-flow · commit f3b158a4c245 · 2024-12-06T12:53:25.000+03:00
Test is flaky
diff --git a/ydb/core/protos/counters_replication.proto b/ydb/core/protos/counters_replication.proto
@@ -46,4 +46,5 @@ enum ETxTypes {
     TXTYPE_DESCRIBE_REPLICATION = 13 [(TxTypeOpts) = {Name: "TxDescribeReplication"}];
     TXTYPE_WORKER_ERROR = 14 [(TxTypeOpts) = {Name: "TxWorkerError"}];
     TXTYPE_ASSIGN_TX_ID = 15 [(TxTypeOpts) = {Name: "TxAssignTxId"}];
+    TXTYPE_HEARTBEAT = 16 [(TxTypeOpts) = {Name: "TxHeartbeat"}];
 }
diff --git a/ydb/core/tx/replication/controller/controller.cpp b/ydb/core/tx/replication/controller/controller.cpp
@@ -80,6 +80,7 @@ STFUNC(TController::StateWork) {
         HFunc(TEvService::TEvRunWorker, Handle);
         HFunc(TEvService::TEvWorkerDataEnd, Handle);
         HFunc(TEvService::TEvGetTxId, Handle);
+        HFunc(TEvService::TEvHeartbeat, Handle);
         HFunc(TEvTxAllocatorClient::TEvAllocateResult, Handle);
         HFunc(TEvInterconnect::TEvNodeDisconnected, Handle);
     default:
@@ -776,6 +777,22 @@ void TController::Handle(TEvService::TEvGetTxId::TPtr& ev, const TActorContext&
     RunTxAssignTxId(ctx);
 }
 
+void TController::Handle(TEvService::TEvHeartbeat::TPtr& ev, const TActorContext& ctx) {
+    CLOG_T(ctx, "Handle " << ev->Get()->ToString());
+
+    const auto nodeId = ev->Sender.NodeId();
+    if (!Sessions.contains(nodeId)) {
+        return;
+    }
+
+    const auto& record = ev->Get()->Record;
+    const auto id = TWorkerId::Parse(record.GetWorker());
+    const auto version = TRowVersion::Parse(record.GetVersion());
+    PendingHeartbeats[id] = version;
+
+    RunTxHeartbeat(ctx);
+}
+
 void TController::Handle(TEvInterconnect::TEvNodeDisconnected::TPtr& ev, const TActorContext& ctx) {
     const ui32 nodeId = ev->Get()->NodeId;
 
diff --git a/ydb/core/tx/replication/controller/controller_impl.h b/ydb/core/tx/replication/controller/controller_impl.h
@@ -97,6 +97,7 @@ class TController
     void Handle(TEvService::TEvRunWorker::TPtr& ev, const TActorContext& ctx);
     void Handle(TEvService::TEvWorkerDataEnd::TPtr& ev, const TActorContext& ctx);
     void Handle(TEvService::TEvGetTxId::TPtr& ev, const TActorContext& ctx);
+    void Handle(TEvService::TEvHeartbeat::TPtr& ev, const TActorContext& ctx);
     void Handle(TEvTxAllocatorClient::TEvAllocateResult::TPtr& ev, const TActorContext& ctx);
     void Handle(TEvInterconnect::TEvNodeDisconnected::TPtr& ev, const TActorContext& ctx);
 
@@ -133,6 +134,7 @@ class TController
     class TTxResolveSecretResult;
     class TTxWorkerError;
     class TTxAssignTxId;
+    class TTxHeartbeat;
 
     // tx runners
     void RunTxInitSchema(const TActorContext& ctx);
@@ -153,6 +155,7 @@ class TController
     void RunTxResolveSecretResult(TEvPrivate::TEvResolveSecretResult::TPtr& ev, const TActorContext& ctx);
     void RunTxWorkerError(const TWorkerId& id, const TString& error, const TActorContext& ctx);
     void RunTxAssignTxId(const TActorContext& ctx);
+    void RunTxHeartbeat(const TActorContext& ctx);
 
     // other
     template <typename T>
@@ -208,6 +211,11 @@ class TController
     TMap<TRowVersion, THashSet<ui32>> PendingTxId;
     bool AssignTxIdInFlight = false;
 
+    THashSet<TWorkerId> WorkersWithHeartbeat;
+    TMap<TRowVersion, THashSet<TWorkerId>> WorkersByHeartbeat;
+    THashMap<TWorkerId, TRowVersion> PendingHeartbeats;
+    bool ProcessHeartbeatsInFlight = false;
+
 }; // TController
 
 }
diff --git a/ydb/core/tx/replication/controller/schema.h b/ydb/core/tx/replication/controller/schema.h
@@ -70,12 +70,24 @@ struct TControllerSchema: NIceDb::Schema {
         using TColumns = TableColumns<VersionStep, VersionTxId, WriteTxId>;
     };
 
+    struct Workers: Table<6> {
+        struct ReplicationId: Column<1, NScheme::NTypeIds::Uint64> {};
+        struct TargetId: Column<2, NScheme::NTypeIds::Uint64> {};
+        struct WorkerId: Column<3, NScheme::NTypeIds::Uint64> {};
+        struct HeartbeatVersionStep: Column<4, NScheme::NTypeIds::Uint64> {};
+        struct HeartbeatVersionTxId: Column<5, NScheme::NTypeIds::Uint64> {};
+
+        using TKey = TableKey<ReplicationId, TargetId, WorkerId>;
+        using TColumns = TableColumns<ReplicationId, TargetId, WorkerId, HeartbeatVersionStep, HeartbeatVersionTxId>;
+    };
+
     using TTables = SchemaTables<
         SysParams,
         Replications,
         Targets,
         SrcStreams,
-        TxIds
+        TxIds,
+        Workers
     >;
 
 }; // TControllerSchema
diff --git a/ydb/core/tx/replication/controller/session_info.cpp b/ydb/core/tx/replication/controller/session_info.cpp
@@ -83,4 +83,17 @@ void TWorkerInfo::SetDataEnded(bool value) {
     DataEnded = value;
 }
 
+void TWorkerInfo::SetHeartbeat(const TRowVersion& value) {
+    Heartbeat = value;
+}
+
+bool TWorkerInfo::HasHeartbeat() const {
+    return Heartbeat.Defined();
+}
+
+const TRowVersion& TWorkerInfo::GetHeartbeat() const {
+    Y_ABORT_UNLESS(Heartbeat.Defined());
+    return *Heartbeat;
+}
+
 }
diff --git a/ydb/core/tx/replication/controller/session_info.h b/ydb/core/tx/replication/controller/session_info.h
@@ -1,5 +1,6 @@
 #pragma once
 
+#include <ydb/core/base/row_version.h>
 #include <ydb/core/tx/replication/common/worker_id.h>
 
 #include <util/generic/hash_set.h>
@@ -45,10 +46,15 @@ class TWorkerInfo {
     bool IsDataEnded() const;
     void SetDataEnded(bool value);
 
+    void SetHeartbeat(const TRowVersion& value);
+    bool HasHeartbeat() const;
+    const TRowVersion& GetHeartbeat() const;
+
 private:
     THolder<NKikimrReplication::TRunWorkerCommand> Command;
     TMaybe<ui32> Session;
     bool DataEnded = false;
+    TMaybe<TRowVersion> Heartbeat;
 };
 
 }
diff --git a/ydb/core/tx/replication/controller/tx_heartbeat.cpp b/ydb/core/tx/replication/controller/tx_heartbeat.cpp
@@ -0,0 +1,96 @@
+#include "controller_impl.h"
+
+namespace NKikimr::NReplication::NController {
+
+class TController::TTxHeartbeat: public TTxBase {
+public:
+    explicit TTxHeartbeat(TController* self)
+        : TTxBase("TxHeartbeat", self)
+    {
+    }
+
+    TTxType GetTxType() const override {
+        return TXTYPE_HEARTBEAT;
+    }
+
+    bool Execute(TTransactionContext& txc, const TActorContext& ctx) override {
+        CLOG_D(ctx, "Execute"
+            << ": pending# " << Self->PendingHeartbeats.size());
+
+        if (Self->Workers.empty()) {
+            CLOG_W(ctx, "There are no workers");
+            return true;
+        }
+
+        const auto prevMinVersion = !Self->WorkersByHeartbeat.empty()
+            ? std::make_optional<TRowVersion>(Self->WorkersByHeartbeat.begin()->first)
+            : std::nullopt; 
+
+        NIceDb::TNiceDb db(txc.DB);
+
+        for (const auto& [id, version] : Self->PendingHeartbeats) {
+            if (!Self->Workers.contains(id)) {
+                continue;
+            }
+
+            auto& worker = Self->Workers[id];
+            if (worker.HasHeartbeat()) {
+                auto it = Self->WorkersByHeartbeat.find(worker.GetHeartbeat());
+                if (it != Self->WorkersByHeartbeat.end()) {
+                    it->second.erase(id);
+                    if (it->second.empty()) {
+                        Self->WorkersByHeartbeat.erase(it);
+                    }
+                }
+            }
+
+            worker.SetHeartbeat(version);
+            Self->WorkersWithHeartbeat.insert(id);
+            Self->WorkersByHeartbeat[version].insert(id);
+
+            db.Table<Schema::Workers>().Key(id.ReplicationId(), id.TargetId(), id.WorkerId()).Update(
+                NIceDb::TUpdate<Schema::Workers::HeartbeatVersionStep>(version.Step),
+                NIceDb::TUpdate<Schema::Workers::HeartbeatVersionTxId>(version.TxId)
+            );
+        }
+
+        if (Self->Workers.size() != Self->WorkersWithHeartbeat.size()) {
+            return true;
+        }
+
+        Y_ABORT_UNLESS(!Self->WorkersByHeartbeat.empty());
+        const auto newMinVersion = Self->WorkersByHeartbeat.begin()->first;
+
+        if (newMinVersion <= prevMinVersion.value_or(TRowVersion::Min())) {
+            return true;
+        }
+
+        CLOG_N(ctx, "Min version has been changed"
+            << ": prev# " << prevMinVersion.value_or(TRowVersion::Min())
+            << ", new# " << newMinVersion);
+
+        // TODO: run commit
+        return true;
+    }
+
+    void Complete(const TActorContext& ctx) override {
+        CLOG_D(ctx, "Complete"
+            << ": pending# " << Self->PendingHeartbeats.size());
+
+        if (Self->PendingHeartbeats) {
+            Self->Execute(new TTxHeartbeat(Self), ctx);
+        } else {
+            Self->ProcessHeartbeatsInFlight = false;
+        }
+    }
+
+}; // TTxHeartbeat
+
+void TController::RunTxHeartbeat(const TActorContext& ctx) {
+    if (!ProcessHeartbeatsInFlight) {
+        ProcessHeartbeatsInFlight = true;
+        Execute(new TTxHeartbeat(this), ctx);
+    }
+}
+
+}
diff --git a/ydb/core/tx/replication/controller/tx_init.cpp b/ydb/core/tx/replication/controller/tx_init.cpp
@@ -157,13 +157,44 @@ class TController::TTxInit: public TTxBase {
         return true;
     }
 
+    bool LoadWorkers(NIceDb::TNiceDb& db) {
+        auto rowset = db.Table<Schema::Workers>().Select();
+        if (!rowset.IsReady()) {
+            return false;
+        }
+
+        while (!rowset.EndOfSet()) {
+            const auto id = TWorkerId(
+                rowset.GetValue<Schema::Workers::ReplicationId>(),
+                rowset.GetValue<Schema::Workers::TargetId>(),
+                rowset.GetValue<Schema::Workers::WorkerId>()
+            );
+            const auto version = TRowVersion(
+                rowset.GetValue<Schema::Workers::HeartbeatVersionStep>(),
+                rowset.GetValue<Schema::Workers::HeartbeatVersionTxId>()
+            );
+
+            auto* worker = Self->GetOrCreateWorker(id);
+            worker->SetHeartbeat(version);
+            Self->WorkersWithHeartbeat.insert(id);
+            Self->WorkersByHeartbeat[version].insert(id);
+
+            if (!rowset.Next()) {
+                return false;
+            }
+        }
+
+        return true;
+    }
+
     inline bool Load(NIceDb::TNiceDb& db) {
         Self->Reset();
         return LoadSysParams(db)
             && LoadReplications(db)
             && LoadTargets(db)
             && LoadSrcStreams(db)
-            && LoadTxIds(db);
+            && LoadTxIds(db)
+            && LoadWorkers(db);
     }
 
     inline bool Load(NTable::TDatabase& toughDb) {
diff --git a/ydb/core/tx/replication/controller/ya.make b/ydb/core/tx/replication/controller/ya.make
@@ -50,6 +50,7 @@ SRCS(
     tx_drop_dst_result.cpp
     tx_drop_replication.cpp
     tx_drop_stream_result.cpp
+    tx_heartbeat.cpp
     tx_init.cpp
     tx_init_schema.cpp
     tx_resolve_secret_result.cpp

Original file line number	Diff line number	Diff line change
`@@ -46,4 +46,5 @@ enum ETxTypes {`
`46`	`46`	`TXTYPE_DESCRIBE_REPLICATION = 13 [(TxTypeOpts) = {Name: "TxDescribeReplication"}];`
`47`	`47`	`TXTYPE_WORKER_ERROR = 14 [(TxTypeOpts) = {Name: "TxWorkerError"}];`
`48`	`48`	`TXTYPE_ASSIGN_TX_ID = 15 [(TxTypeOpts) = {Name: "TxAssignTxId"}];`
	`49`	`+ TXTYPE_HEARTBEAT = 16 [(TxTypeOpts) = {Name: "TxHeartbeat"}];`
`49`	`50`	`}`