From 18ed18eada557111bb9c435ff211a95109e4a5dc Mon Sep 17 00:00:00 2001
From: Alexander Rutkovsky <alexvru@ydb.tech>
Date: Tue, 24 Dec 2024 08:05:58 +0300
Subject: [PATCH] Allow some quiet period before starting TestShard load
 (#12887)

---
 ydb/core/protos/msgbus.proto             |  1 +
 ydb/core/test_tablet/load_actor_impl.cpp | 22 ++++++++++++++++------
 ydb/core/test_tablet/load_actor_impl.h   |  6 ++++--
 3 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/ydb/core/protos/msgbus.proto b/ydb/core/protos/msgbus.proto
index ddfb74a68fe8..799c94c7c5ba 100644
--- a/ydb/core/protos/msgbus.proto
+++ b/ydb/core/protos/msgbus.proto
@@ -704,6 +704,7 @@ message TTestShardControlRequest {
         optional uint32 PatchRequestsFractionPPM = 12;
         optional uint32 PutTraceFractionPPM = 13;
         optional uint32 PutTraceVerbosity = 14 [default = 15];
+        optional uint32 SecondsBeforeLoadStart = 15; // number of seconds to wait before starting load
     }
 
     optional uint64 TabletId = 1;
diff --git a/ydb/core/test_tablet/load_actor_impl.cpp b/ydb/core/test_tablet/load_actor_impl.cpp
index 866d62a5771b..ae4386568729 100644
--- a/ydb/core/test_tablet/load_actor_impl.cpp
+++ b/ydb/core/test_tablet/load_actor_impl.cpp
@@ -4,7 +4,8 @@ namespace NKikimr::NTestShard {
 
     TLoadActor::TLoadActor(ui64 tabletId, ui32 generation, TActorId tablet,
             const NKikimrClient::TTestShardControlRequest::TCmdInitialize& settings)
-        : TabletId(tabletId)
+        : TActor(&TThis::StateFunc)
+        , TabletId(tabletId)
         , Generation(generation)
         , Tablet(tablet)
         , Settings(settings)
@@ -14,6 +15,17 @@ namespace NKikimr::NTestShard {
         ClearKeys();
     }
 
+    void TLoadActor::Registered(TActorSystem *sys, const TActorId& owner) {
+        TActor::Registered(sys, owner);
+        TabletActorId = owner;
+        auto ev = std::make_unique<IEventHandle>(TEvents::TSystem::Bootstrap, 0, SelfId(), owner, nullptr, 0);
+        if (Settings.HasSecondsBeforeLoadStart()) {
+            sys->Schedule(TDuration::Seconds(Settings.GetSecondsBeforeLoadStart()), ev.release());
+        } else {
+            sys->Send(ev.release());
+        }
+    }
+
     void TLoadActor::ClearKeys() {
         for (auto& [key, info] : Keys) {
             Y_ABORT_UNLESS(info.ConfirmedState == ::NTestShard::TStateServer::CONFIRMED
@@ -25,18 +37,16 @@ namespace NKikimr::NTestShard {
         ConfirmedKeys.clear();
     }
 
-    void TLoadActor::Bootstrap(const TActorId& parentId) {
+    void TLoadActor::Bootstrap() {
         STLOG(PRI_DEBUG, TEST_SHARD, TS31, "TLoadActor::Bootstrap", (TabletId, TabletId));
-        TabletActorId = parentId;
         if (Settings.HasStorageServerHost()) {
             Send(MakeStateServerInterfaceActorId(), new TEvStateServerConnect(Settings.GetStorageServerHost(),
                 Settings.GetStorageServerPort()));
-            Send(parentId, new TTestShard::TEvSwitchMode(TTestShard::EMode::STATE_SERVER_CONNECT));
+            Send(TabletActorId, new TTestShard::TEvSwitchMode(TTestShard::EMode::STATE_SERVER_CONNECT));
         } else {
             RunValidation(true);
         }
         NextWriteTimestamp = TActivationContext::Monotonic();
-        Become(&TThis::StateFunc);
     }
 
     void TLoadActor::PassAway() {
@@ -46,7 +56,7 @@ namespace NKikimr::NTestShard {
         if (ValidationActorId) {
             TActivationContext::Send(new IEventHandle(TEvents::TSystem::Poison, 0, ValidationActorId, SelfId(), nullptr, 0));
         }
-        TActorBootstrapped::PassAway();
+        TActor::PassAway();
     }
 
     void TLoadActor::HandleWakeup() {
diff --git a/ydb/core/test_tablet/load_actor_impl.h b/ydb/core/test_tablet/load_actor_impl.h
index 9fb2c2989b1a..1e8c5d9a065b 100644
--- a/ydb/core/test_tablet/load_actor_impl.h
+++ b/ydb/core/test_tablet/load_actor_impl.h
@@ -8,7 +8,7 @@
 
 namespace NKikimr::NTestShard {
 
-    class TLoadActor : public TActorBootstrapped<TLoadActor> {
+    class TLoadActor : public TActor<TLoadActor> {
         const ui64 TabletId;
         const ui32 Generation;
         const TActorId Tablet;
@@ -58,14 +58,16 @@ namespace NKikimr::NTestShard {
         TLoadActor(ui64 tabletId, ui32 generation, const TActorId tablet,
             const NKikimrClient::TTestShardControlRequest::TCmdInitialize& settings);
         ~TLoadActor();
+        void Registered(TActorSystem *sys, const TActorId& owner) override;
         void ClearKeys();
-        void Bootstrap(const TActorId& parentId);
+        void Bootstrap();
         void PassAway() override;
         void HandleWakeup();
         void Action();
         void Handle(TEvStateServerStatus::TPtr ev);
 
         STRICT_STFUNC(StateFunc,
+            cFunc(TEvents::TSystem::Bootstrap, Bootstrap);
             hFunc(TEvKeyValue::TEvResponse, Handle);
             hFunc(NMon::TEvRemoteHttpInfo, Handle);
             hFunc(TEvStateServerStatus, Handle);