Move TxIds & PendingTxId into session (#13039)

Author: CyberROFL

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

@@ -113,7 +113,9 @@ class TSessionInfo {
         ops->Send(ActorId, ev.Release());
     }
 
-    void SendWorkerDataEnd(IActorOps* ops, const TWorkerId& id, ui64 partitionId, const TVector<ui64>&& adjacentPartitionsIds, const TVector<ui64>&& childPartitionsIds) {
+    void SendWorkerDataEnd(IActorOps* ops, const TWorkerId& id, ui64 partitionId,
+            const TVector<ui64>&& adjacentPartitionsIds, const TVector<ui64>&& childPartitionsIds)
+    {
         auto ev = MakeHolder<TEvService::TEvWorkerDataEnd>();
         auto& record = ev->Record;
 
@@ -129,18 +131,87 @@ class TSessionInfo {
         ops->Send(ActorId, ev.Release());
     }
 
+    void Handle(IActorOps* ops, TEvService::TEvGetTxId::TPtr& ev) {
+        TMap<TRowVersion, ui64> result;
+        TVector<TRowVersion> versionsWithoutTxId;
+
+        for (const auto& v : ev->Get()->Record.GetVersions()) {
+            const auto version = TRowVersion::FromProto(v);
+            if (auto it = TxIds.upper_bound(version); it != TxIds.end()) {
+                result[it->first] = it->second;
+            } else {
+                versionsWithoutTxId.push_back(version);
+                PendingTxId[version].insert(ev->Sender);
+            }
+        }
+
+        if (versionsWithoutTxId) {
+            ops->Send(ActorId, new TEvService::TEvGetTxId(versionsWithoutTxId));
+        }
+
+        if (result) {
+            SendTxIdResult(ops, ev->Sender, result);
+        }
+    }
+
+    void Handle(IActorOps* ops, TEvService::TEvTxIdResult::TPtr& ev) {
+        THashMap<TActorId, TMap<TRowVersion, ui64>> results;
+
+        for (const auto& kv : ev->Get()->Record.GetVersionTxIds()) {
+            const auto version = TRowVersion::FromProto(kv.GetVersion());
+            TxIds.emplace(version, kv.GetTxId());
+
+            for (auto it = PendingTxId.begin(); it != PendingTxId.end();) {
+                if (it->first >= version) {
+                    break;
+                }
+
+                for (const auto& actorId : it->second) {
+                    results[actorId].emplace(version, kv.GetTxId());
+                }
+
+                PendingTxId.erase(it++);
+            }
+        }
+
+        for (const auto& [actorId, result] : results) {
+            SendTxIdResult(ops, actorId, result);
+        }
+    }
+
+    void Handle(IActorOps* ops, TEvService::TEvHeartbeat::TPtr& ev) {
+        GetWorkerId(ev->Sender).Serialize(*ev->Get()->Record.MutableWorker());
+        ops->Send(ActorId, ev->ReleaseBase().Release(), ev->Flags, ev->Cookie);
+    }
+
     void Shutdown(IActorOps* ops) const {
         for (const auto& [_, actorId] : Workers) {
             ops->Send(actorId, new TEvents::TEvPoison());
         }
     }
 
+private:
+    static void SendTxIdResult(IActorOps* ops, const TActorId& recipient, const TMap<TRowVersion, ui64>& result) {
+        auto ev = MakeHolder<TEvService::TEvTxIdResult>();
+
+        for (const auto& [version, txId] : result) {
+            auto& item = *ev->Record.AddVersionTxIds();
+            version.ToProto(item.MutableVersion());
+            item.SetTxId(txId);
+        }
+
+        ops->Send(recipient, ev.Release());
+    }
+
 private:
     TActorId ActorId;
     ui64 Generation;
     THashMap<TWorkerId, TActorId> Workers;
     THashMap<TActorId, TWorkerId> ActorIdToWorkerId;
 
+    TMap<TRowVersion, ui64> TxIds;
+    TMap<TRowVersion, THashSet<TActorId>> PendingTxId;
+
 }; // TSessionInfo
 
 struct TConnectionParams: std::tuple<TString, TString, bool, TString> {
@@ -365,22 +436,10 @@ class TReplicationService: public TActorBootstrapped<TReplicationService> {
         session.StopWorker(this, id);
     }
 
-    void SendTxIdResult(const TActorId& recipient, const TMap<TRowVersion, ui64>& result) {
-        auto ev = MakeHolder<TEvService::TEvTxIdResult>();
-
-        for (const auto& [version, txId] : result) {
-            auto& item = *ev->Record.AddVersionTxIds();
-            version.ToProto(item.MutableVersion());
-            item.SetTxId(txId);
-        }
-
-        Send(recipient, ev.Release());
-    }
-
     void Handle(TEvService::TEvGetTxId::TPtr& ev) {
         LOG_D("Handle " << ev->Get()->ToString());
 
-        const auto* session = SessionFromWorker(ev->Sender);
+        auto* session = SessionFromWorker(ev->Sender);
         if (!session) {
             return;
         }
@@ -391,26 +450,7 @@ class TReplicationService: public TActorBootstrapped<TReplicationService> {
             return;
         }
 
-        TMap<TRowVersion, ui64> result;
-        TVector<TRowVersion> versionsWithoutTxId;
-
-        for (const auto& v : ev->Get()->Record.GetVersions()) {
-            const auto version = TRowVersion::FromProto(v);
-            if (auto it = TxIds.upper_bound(version); it != TxIds.end()) {
-                result[it->first] = it->second;
-            } else {
-                versionsWithoutTxId.push_back(version);
-                PendingTxId[version].insert(ev->Sender);
-            }
-        }
-
-        if (versionsWithoutTxId) {
-            Send(*session, new TEvService::TEvGetTxId(versionsWithoutTxId));
-        }
-
-        if (result) {
-            SendTxIdResult(ev->Sender, result);
-        }
+        session->Handle(this, ev);
     }
 
     void Handle(TEvService::TEvTxIdResult::TPtr& ev) {
@@ -436,34 +476,13 @@ class TReplicationService: public TActorBootstrapped<TReplicationService> {
             return;
         }
 
-        THashMap<TActorId, TMap<TRowVersion, ui64>> results;
-
-        for (const auto& kv : record.GetVersionTxIds()) {
-            const auto version = TRowVersion::FromProto(kv.GetVersion());
-            TxIds.emplace(version, kv.GetTxId());
-
-            for (auto it = PendingTxId.begin(); it != PendingTxId.end();) {
-                if (it->first >= version) {
-                    break;
-                }
-
-                for (const auto& actorId : it->second) {
-                    results[actorId].emplace(version, kv.GetTxId());
-                }
-
-                PendingTxId.erase(it++);
-            }
-        }
-
-        for (const auto& [actorId, result] : results) {
-            SendTxIdResult(actorId, result);
-        }
+        session.Handle(this, ev);
     }
 
     void Handle(TEvService::TEvHeartbeat::TPtr& ev) {
         LOG_D("Handle " << ev->Get()->ToString());
 
-        const auto* session = SessionFromWorker(ev->Sender);
+        auto* session = SessionFromWorker(ev->Sender);
         if (!session) {
             return;
         }
@@ -474,14 +493,10 @@ class TReplicationService: public TActorBootstrapped<TReplicationService> {
             return;
         }
 
-        auto& record = ev->Get()->Record;
-
         LOG_I("Heartbeat"
             << ": worker# " << ev->Sender
-            << ", version# " << TRowVersion::FromProto(record.GetVersion()));
-
-        session->GetWorkerId(ev->Sender).Serialize(*record.MutableWorker());
-        Send(ev->Forward(*session));
+            << ", version# " << TRowVersion::FromProto(ev->Get()->Record.GetVersion()));
+        session->Handle(this, ev);
     }
 
     void Handle(TEvWorker::TEvDataEnd::TPtr& ev) {
@@ -608,8 +623,6 @@ class TReplicationService: public TActorBootstrapped<TReplicationService> {
     THashMap<ui64, TSessionInfo> Sessions;
     THashMap<TConnectionParams, TActorId> YdbProxies;
     THashMap<TActorId, ui64> WorkerActorIdToSession;
-    TMap<TRowVersion, ui64> TxIds;
-    TMap<TRowVersion, THashSet<TActorId>> PendingTxId;
 
 }; // TReplicationService