Support put tracing with sampling in TestShard (#6123)

Author: alexvru

ydb/core/protos/msgbus.proto

@@ -717,6 +717,8 @@ message TTestShardControlRequest {
         repeated TTimeInterval WritePeriods = 7; // time between two events
         repeated TTimeInterval RestartPeriods = 8; // time between automatic restarts
         optional uint32 PatchRequestsFractionPPM = 12;
+        optional uint32 PutTraceFractionPPM = 13;
+        optional uint32 PutTraceVerbosity = 14 [default = 15];
     }
 
     optional uint64 TabletId = 1;

ydb/core/test_tablet/load_actor_impl.h

@@ -22,6 +22,7 @@ namespace NKikimr::NTestShard {
             ::NTestShard::TStateServer::EEntityState ConfirmedState = ::NTestShard::TStateServer::ABSENT;
             ::NTestShard::TStateServer::EEntityState PendingState = ::NTestShard::TStateServer::ABSENT;
             std::unique_ptr<TEvKeyValue::TEvRequest> Request;
+            NWilson::TTraceId TraceId;
             size_t ConfirmedKeyIndex = Max<size_t>();
 
             TKeyInfo(ui32 len)
@@ -174,7 +175,8 @@ namespace NKikimr::NTestShard {
         std::deque<TKey*> TransitionInFlight;
 
         void RegisterTransition(TKey& key, ::NTestShard::TStateServer::EEntityState from,
-            ::NTestShard::TStateServer::EEntityState to, std::unique_ptr<TEvKeyValue::TEvRequest> ev = nullptr);
+            ::NTestShard::TStateServer::EEntityState to, std::unique_ptr<TEvKeyValue::TEvRequest> ev = nullptr,
+            NWilson::TTraceId traceId = {});
         void Handle(TEvStateServerWriteResult::TPtr ev);
 
         void MakeConfirmed(TKey& key);

ydb/core/test_tablet/load_actor_state.cpp

@@ -3,7 +3,8 @@
 namespace NKikimr::NTestShard {
 
     void TLoadActor::RegisterTransition(TKey& key, ::NTestShard::TStateServer::EEntityState from,
-            ::NTestShard::TStateServer::EEntityState to, std::unique_ptr<TEvKeyValue::TEvRequest> ev) {
+            ::NTestShard::TStateServer::EEntityState to, std::unique_ptr<TEvKeyValue::TEvRequest> ev,
+            NWilson::TTraceId traceId) {
         STLOG(PRI_DEBUG, TEST_SHARD, TS14, "RegisterTransition", (TabletId, TabletId), (Key, key.first), (From, from),
             (To, to));
 
@@ -38,7 +39,7 @@ namespace NKikimr::NTestShard {
                     : key.second.ConfirmedKeyIndex == Max<size_t>());
             }
             if (ev) {
-                Send(TabletActorId, ev.release());
+                Send(TabletActorId, ev.release(), 0, 0, std::move(traceId));
             }
             if (!DoSomeActionInFlight) {
                 TActivationContext::Send(new IEventHandle(EvDoSomeAction, 0, SelfId(), {}, nullptr, 0));
@@ -61,6 +62,7 @@ namespace NKikimr::NTestShard {
         // update local state
         key.second.PendingState = to;
         key.second.Request = std::move(ev);
+        key.second.TraceId = std::move(traceId);
         TransitionInFlight.push_back(&key);
     }
 
@@ -108,7 +110,7 @@ namespace NKikimr::NTestShard {
             if (const auto it = WritesInFlight.find(r->Record.GetCookie()); it != WritesInFlight.end()) {
                 StateServerWriteLatency.Add(TActivationContext::Monotonic(), TDuration::Seconds(it->second.Timer.PassedReset()));
             }
-            Send(TabletActorId, r.release());
+            Send(TabletActorId, r.release(), 0, 0, std::move(key.second.TraceId));
         }
         if (key.second.ConfirmedState == ::NTestShard::TStateServer::DELETED) {
             Y_ABORT_UNLESS(key.second.ConfirmedKeyIndex == Max<size_t>());

ydb/core/test_tablet/load_actor_write.cpp

@@ -33,13 +33,19 @@ namespace NKikimr::NTestShard {
 
         STLOG(PRI_INFO, TEST_SHARD, TS12, "writing data", (TabletId, TabletId), (Key, key), (Size, value.size()));
 
+        NWilson::TTraceId traceId;
+        if (RandomNumber(1000000u) < Settings.GetPutTraceFractionPPM()) {
+            traceId = NWilson::TTraceId::NewTraceId(Settings.GetPutTraceVerbosity(), Max<ui32>());
+        }
+
         auto [wifIt, wifInserted] = WritesInFlight.try_emplace(r.GetCookie(), key);
         Y_ABORT_UNLESS(wifInserted);
         Y_ABORT_UNLESS(wifIt->second.KeysInQuery.size() == 1);
 
         auto [it, inserted] = Keys.try_emplace(key, value.size());
         Y_ABORT_UNLESS(inserted);
-        RegisterTransition(*it, ::NTestShard::TStateServer::ABSENT, ::NTestShard::TStateServer::WRITE_PENDING, std::move(ev));
+        RegisterTransition(*it, ::NTestShard::TStateServer::ABSENT, ::NTestShard::TStateServer::WRITE_PENDING,
+            std::move(ev), std::move(traceId));
 
         ++KeysWritten;
         BytesProcessed += value.size();