KeyValue tracing sampling (#1047)

15 files changed, 346 insertions, 15 deletions

diff --git a/ydb/core/cms/json_proxy_proto.h b/ydb/core/cms/json_proxy_proto.h
index a422cbe8292..4844919de31 100644
--- a/ydb/core/cms/json_proxy_proto.h
+++ b/ydb/core/cms/json_proxy_proto.h
@@ -76,6 +76,16 @@ protected:
                 return ReplyWithTypeDescription(*NKikimrConfig::TImmediateControlsConfig::TSchemeShardControls::descriptor(), ctx);
             else if (name == ".NKikimrConfig.TImmediateControlsConfig.TTCMallocControls")
                 return ReplyWithTypeDescription(*NKikimrConfig::TImmediateControlsConfig::TTCMallocControls::descriptor(), ctx);
+            else if (name == ".NKikimrConfig.TImmediateControlsConfig.TTracingControls")
+                return ReplyWithTypeDescription(*NKikimrConfig::TImmediateControlsConfig::TTracingControls::descriptor(), ctx);
+            else if (name == ".NKikimrConfig.TImmediateControlsConfig.TTracingControls.TSamplingThrottlingOptions")
+                return ReplyWithTypeDescription(*NKikimrConfig::TImmediateControlsConfig::TTracingControls::TSamplingThrottlingOptions::descriptor(), ctx);
+            else if (name == ".NKikimrConfig.TImmediateControlsConfig.TTracingControls.TSamplingThrottlingOptions.TThrottlingOptions")
+                return ReplyWithTypeDescription(*NKikimrConfig::TImmediateControlsConfig::TTracingControls::TSamplingThrottlingOptions::TThrottlingOptions::descriptor(), ctx);
+            else if (name == ".NKikimrConfig.TImmediateControlsConfig.TTracingControls.TSamplingThrottlingOptions.TSamplingOptions")
+                return ReplyWithTypeDescription(*NKikimrConfig::TImmediateControlsConfig::TTracingControls::TSamplingThrottlingOptions::TSamplingOptions::descriptor(), ctx);
+            else if (name == ".NKikimrConfig.TImmediateControlsConfig.TTracingControls.TKeyValue")
+                return ReplyWithTypeDescription(*NKikimrConfig::TImmediateControlsConfig::TTracingControls::TKeyValue::descriptor(), ctx);
         }
 
         ctx.Send(RequestEvent->Sender,
diff --git a/ydb/core/control/common_controls/tracing_control.cpp b/ydb/core/control/common_controls/tracing_control.cpp
new file mode 100644
index 00000000000..536738dd289
--- /dev/null
+++ b/ydb/core/control/common_controls/tracing_control.cpp
@@ -0,0 +1,62 @@
+#include "tracing_control.h"
+
+#include <ydb/core/control/immediate_control_board_impl.h>
+#include <ydb/core/protos/config.pb.h>
+#include <library/cpp/random_provider/random_provider.h>
+
+namespace NKikimr {
+
+namespace {
+
+const NKikimrConfig::TImmediateControlOptions& GetImmediateControlOptionsForField(
+        const google::protobuf::Descriptor& descriptor, TString fieldName) {
+    auto field = descriptor.FindFieldByName(fieldName);
+    Y_ABORT_UNLESS(field);
+    auto& fieldOptions = field->options();
+    return fieldOptions.GetExtension(NKikimrConfig::ControlOptions);
+}
+
+TThrottler CreateThrottler(TIntrusivePtr<TControlBoard>& icb, TIntrusivePtr<ITimeProvider> timeProvider, TString domain) {
+    TControlWrapper maxRatePerMinute;
+    TControlWrapper maxBurst;
+
+    const std::array<std::pair<TControlWrapper&, TStringBuf>, 2> controls = {{
+        {maxRatePerMinute, "MaxRatePerMinute"},
+        {maxBurst, "MaxBurst"},
+    }};
+    const auto& throttlingOptions = *NKikimrConfig::TImmediateControlsConfig::TTracingControls::TSamplingThrottlingOptions::TThrottlingOptions::descriptor();
+    for (auto& [control, fieldName] : controls) {
+        const auto& controlOptions = GetImmediateControlOptionsForField(throttlingOptions, TString(fieldName));
+
+        control.Reset(controlOptions.GetDefaultValue(), controlOptions.GetMinValue(), controlOptions.GetMaxValue());
+        icb->RegisterSharedControl(control, domain + "." + fieldName);
+    }
+
+    return TThrottler(std::move(maxRatePerMinute), std::move(maxBurst), std::move(timeProvider));
+}
+
+}
+
+TTracingControl::TTracingControl(TIntrusivePtr<TControlBoard>& icb, TIntrusivePtr<ITimeProvider> timeProvider,
+        TIntrusivePtr<IRandomProvider>& randomProvider, TString controlDomain)
+{
+    SampledThrottler = CreateThrottler(icb, timeProvider, controlDomain + ".SampledThrottling");
+    ExternalThrottler = CreateThrottler(icb, timeProvider, controlDomain + ".ExternalThrottling");
+
+    TControlWrapper samplingPPM;
+    const std::array<std::pair<TControlWrapper&, TStringBuf>, 2> controls = {{
+        {samplingPPM, "PPM"},
+        {SampledLevel, "Level"},
+    }};
+
+    const auto& samplingOptions = *NKikimrConfig::TImmediateControlsConfig::TTracingControls::TSamplingThrottlingOptions::TSamplingOptions::descriptor();
+    for (auto [control, name] : controls) {
+        const auto& controlOptions = GetImmediateControlOptionsForField(samplingOptions, TString(name));
+        control.Reset(controlOptions.GetDefaultValue(), controlOptions.GetMinValue(), controlOptions.GetMaxValue());
+        icb->RegisterSharedControl(control, controlDomain + ".Sampling." + name);
+    }
+
+    Sampler = TSampler(std::move(samplingPPM), randomProvider->GenRand64());
+}
+
+} // namespace NKikimr
diff --git a/ydb/core/control/common_controls/tracing_control.h b/ydb/core/control/common_controls/tracing_control.h
new file mode 100644
index 00000000000..56b7f45966d
--- /dev/null
+++ b/ydb/core/control/common_controls/tracing_control.h
@@ -0,0 +1,33 @@
+#pragma once
+
+#include <ydb/core/base/appdata_fwd.h>
+#include <ydb/core/control/immediate_control_board_sampler.h>
+#include <ydb/core/control/immediate_control_board_throttler.h>
+
+namespace NKikimr {
+
+class TTracingControl {
+public:
+    TTracingControl(TIntrusivePtr<TControlBoard>& icb, TIntrusivePtr<ITimeProvider> timeProvider,
+            TIntrusivePtr<IRandomProvider>& randomProvider, TString controlDomain);
+
+    bool SampleThrottle() {
+        return Sampler.Sample() && !SampledThrottler.Throttle();
+    }
+
+    bool ThrottleExternal() {
+        return ExternalThrottler.Throttle();
+    }
+
+    ui8 SampledVerbosity() const {
+        return SampledLevel;
+    }
+
+private:
+    TSampler Sampler;
+    TThrottler SampledThrottler;
+    TThrottler ExternalThrottler;
+    TControlWrapper SampledLevel;
+};
+
+} // namespace NKikimr
diff --git a/ydb/core/control/common_controls/ya.make b/ydb/core/control/common_controls/ya.make
new file mode 100644
index 00000000000..afc6df1f79d
--- /dev/null
+++ b/ydb/core/control/common_controls/ya.make
@@ -0,0 +1,13 @@
+LIBRARY()
+
+PEERDIR(
+    ydb/library/actors/wilson
+    ydb/core/protos
+)
+
+SRCS(
+    tracing_control.h
+    tracing_control.cpp
+)
+
+END()
diff --git a/ydb/core/control/immediate_control_board_sampler.h b/ydb/core/control/immediate_control_board_sampler.h
new file mode 100644
index 00000000000..e6d6784540b
--- /dev/null
+++ b/ydb/core/control/immediate_control_board_sampler.h
@@ -0,0 +1,25 @@
+#pragma once
+
+#include <ydb/core/control/immediate_control_board_wrapper.h>
+
+namespace NKikimr {
+
+class TSampler {
+public:
+    TSampler() : Rng(0) {}
+
+    TSampler(TControlWrapper samplingPPM, ui64 seed)
+        : SamplingPPM(std::move(samplingPPM))
+        , Rng(seed)
+    {}
+
+    bool Sample() {
+        return Rng() % 1'000'000 < SamplingPPM;
+    }
+
+private:
+    TControlWrapper SamplingPPM;
+    TReallyFastRng32 Rng;
+};
+
+} // namespace NKikimr
diff --git a/ydb/core/control/immediate_control_board_throttler.h b/ydb/core/control/immediate_control_board_throttler.h
new file mode 100644
index 00000000000..c3ee83bf0f4
--- /dev/null
+++ b/ydb/core/control/immediate_control_board_throttler.h
@@ -0,0 +1,66 @@
+#pragma once
+
+#include <ydb/core/control/immediate_control_board_wrapper.h>
+#include <library/cpp/time_provider/time_provider.h>
+
+namespace NKikimr {
+
+class TThrottler {
+public:
+    TThrottler() = default;
+
+    TThrottler(TControlWrapper maxRatePerMinute, TControlWrapper maxBurst,
+            TIntrusivePtr<ITimeProvider> timeProvider)
+        : TimeProvider(std::move(timeProvider))
+        , MaxRatePerMinute(std::move(maxRatePerMinute))
+        , MaxBurst(std::move(maxBurst))
+        , LastUpdate(TimeProvider->Now())
+    {}
+
+    bool Throttle() {
+        auto maxRatePerMinute = static_cast<i64>(MaxRatePerMinute);
+        auto maxBurst = static_cast<i64>(MaxBurst);
+        auto maxTotal = maxBurst + 1;
+        CurrentBurst = std::min(CurrentBurst, maxTotal);
+        if (maxRatePerMinute == 0) {
+            return true;
+        }
+
+        auto now = TimeProvider->Now();
+        if (now < LastUpdate) {
+            return true;
+        }
+
+        const auto deltaBetweenSends = TDuration::Minutes(1) / maxRatePerMinute;
+        UpdateStats(now, deltaBetweenSends);
+
+        if (CurrentBurst < maxTotal) {
+            CurrentBurst += 1;
+            return false;
+        }
+
+        return true;
+    }
+
+private:
+    void UpdateStats(TInstant now, TDuration deltaBetweenSends) {
+        i64 decrease = (now - LastUpdate) / deltaBetweenSends;
+        decrease = std::min(decrease, CurrentBurst);
+        Y_ABORT_UNLESS(decrease >= 0);
+        CurrentBurst -= decrease;
+        LastUpdate += decrease * deltaBetweenSends;
+        if (CurrentBurst == 0) {
+            LastUpdate = now;
+        }
+    }
+
+    TIntrusivePtr<ITimeProvider> TimeProvider;
+
+    TControlWrapper MaxRatePerMinute;
+    TControlWrapper MaxBurst;
+
+    TInstant LastUpdate = TInstant::Zero();
+    i64 CurrentBurst = 0;
+};
+
+} // namespace NKikimr
diff --git a/ydb/core/control/ya.make b/ydb/core/control/ya.make
index 8c1c83bff96..4faca2369cf 100644
--- a/ydb/core/control/ya.make
+++ b/ydb/core/control/ya.make
@@ -19,10 +19,16 @@ SRCS(
     immediate_control_board_impl.cpp
     immediate_control_board_impl.h
     immediate_control_board_wrapper.h
+    immediate_control_board_throttler.h
+    immediate_control_board_sampler.h
 )
 
 END()
 
+RECURSE(
+    common_controls
+)
+
 RECURSE_FOR_TESTS(
     ut
 )
diff --git a/ydb/core/driver_lib/run/kikimr_services_initializers.cpp b/ydb/core/driver_lib/run/kikimr_services_initializers.cpp
index f327d04ad7c..38da1001665 100644
--- a/ydb/core/driver_lib/run/kikimr_services_initializers.cpp
+++ b/ydb/core/driver_lib/run/kikimr_services_initializers.cpp
@@ -1619,7 +1619,7 @@ void TGRpcServicesInitializer::InitializeServices(NActors::TActorSystemSetup* se
         for (size_t i = 0; i < proxyCount; ++i) {
             auto grpcReqProxy = Config.HasGRpcConfig() && Config.GetGRpcConfig().GetSkipSchemeCheck()
                 ? NGRpcService::CreateGRpcRequestProxySimple(Config)
-                : NGRpcService::CreateGRpcRequestProxy(Config);
+                : NGRpcService::CreateGRpcRequestProxy(Config, appData->Icb);
             setup->LocalServices.push_back(std::pair<TActorId,
                                            TActorSetupCmd>(NGRpcService::CreateGRpcRequestProxyId(i),
                                                            TActorSetupCmd(grpcReqProxy, TMailboxType::ReadAsFilled,
diff --git a/ydb/core/grpc_services/base/base.h b/ydb/core/grpc_services/base/base.h
index bd5ae1bc00b..12272ce9af6 100644
--- a/ydb/core/grpc_services/base/base.h
+++ b/ydb/core/grpc_services/base/base.h
@@ -367,6 +367,9 @@ public:
     virtual void StartTracing(NWilson::TSpan&& span) = 0;
     virtual void LegacyFinishSpan() = 0;
 
+    // Used for per-type sampling
+    virtual const TString& GetInternalRequestType() const = 0;
+
     // validation
     virtual bool Validate(TString& error) = 0;
 
@@ -485,6 +488,10 @@ public:
 
     void StartTracing(NWilson::TSpan&& /*span*/) override {}
     void LegacyFinishSpan() override {}
+    const TString& GetInternalRequestType() const final {
+        static const TString empty = "";
+        return empty;
+    }
 
     void UpdateAuthState(NYdbGrpc::TAuthState::EAuthState state) override {
         State_.State = state;
@@ -890,6 +897,10 @@ public:
         Span_.End();
     }
 
+    const TString& GetInternalRequestType() const final {
+        return TRequest::descriptor()->full_name();
+    }
+
     // IRequestCtxBase
     //
     void AddAuditLogPart(const TStringBuf&, const TString&) override {
@@ -1302,6 +1313,10 @@ public:
 
     void LegacyFinishSpan() override {}
 
+    const TString& GetInternalRequestType() const final {
+        return TRequest::descriptor()->full_name();
+    }
+
     void ReplyGrpcError(grpc::StatusCode code, const TString& msg, const TString& details = "") {
         Ctx_->ReplyError(code, msg, details);
     }
diff --git a/ydb/core/grpc_services/grpc_request_proxy.cpp b/ydb/core/grpc_services/grpc_request_proxy.cpp
index 3ef705d9954..3e44414f5de 100644
--- a/ydb/core/grpc_services/grpc_request_proxy.cpp
+++ b/ydb/core/grpc_services/grpc_request_proxy.cpp
@@ -8,6 +8,7 @@
 #include <ydb/core/base/nameservice.h>
 #include <ydb/core/cms/console/configs_dispatcher.h>
 #include <ydb/core/cms/console/console.h>
+#include <ydb/core/control/common_controls/tracing_control.h>
 #include <ydb/core/grpc_services/counters/proxy_counters.h>
 #include <ydb/core/tx/tx_proxy/proxy.h>
 #include <ydb/core/tx/scheme_board/scheme_board.h>
@@ -60,8 +61,9 @@ class TGRpcRequestProxyImpl
 {
     using TBase = TActorBootstrapped<TGRpcRequestProxyImpl>;
 public:
-    explicit TGRpcRequestProxyImpl(const NKikimrConfig::TAppConfig& appConfig)
+    explicit TGRpcRequestProxyImpl(const NKikimrConfig::TAppConfig& appConfig, TIntrusivePtr<TControlBoard> icb)
         : ChannelBufferSize(appConfig.GetTableServiceConfig().GetResourceManager().GetChannelBufferSize())
+        , Icb(std::move(icb))
     { }
 
     void Bootstrap(const TActorContext& ctx);
@@ -80,7 +82,9 @@ private:
     void HandleSchemeBoard(TSchemeBoardEvents::TEvNotifyDelete::TPtr& ev);
     void ReplayEvents(const TString& databaseName, const TActorContext& ctx);
 
-    void StartTracing(IRequestProxyCtx& ctx);
+    static TString InternalRequestTypeToControlDomain(const TString& type);
+    TTracingControl& GetTracingControl(const TString& type);
+    void MaybeStartTracing(IRequestProxyCtx& ctx);
 
     static bool IsAuthStateOK(const IRequestProxyCtx& ctx);
 
@@ -151,7 +155,7 @@ private:
         }
 
 
-        StartTracing(*requestBaseCtx);
+        MaybeStartTracing(*requestBaseCtx);
 
         if (IsAuthStateOK(*requestBaseCtx)) {
             Handle(event, ctx);
@@ -311,6 +315,8 @@ private:
     bool DynamicNode = false;
     TString RootDatabase;
     IGRpcProxyCounters::TPtr Counters;
+    THashMap<TString, TTracingControl> TracingControls;
+    TIntrusivePtr<TControlBoard> Icb;
 };
 
 void TGRpcRequestProxyImpl::Bootstrap(const TActorContext& ctx) {
@@ -409,12 +415,52 @@ bool TGRpcRequestProxyImpl::IsAuthStateOK(const IRequestProxyCtx& ctx) {
            state.NeedAuth == false && !ctx.GetYdbToken();
 }
 
-void TGRpcRequestProxyImpl::StartTracing(IRequestProxyCtx& ctx) {
+TString TGRpcRequestProxyImpl::InternalRequestTypeToControlDomain(const TString& type) {
+    static constexpr TStringBuf ydbNamespacePrefix = "Ydb.";
+    static constexpr TStringBuf requestSuffix = "Request";
+
+    TString controlDomain = type;
+    if (controlDomain.StartsWith(ydbNamespacePrefix)) {
+        controlDomain.erase(0, ydbNamespacePrefix.size());
+    }
+    if (controlDomain.EndsWith(requestSuffix)) {
+        controlDomain.erase(controlDomain.size() - requestSuffix.size());
+    }
+
+    return controlDomain;
+}
+
+TTracingControl& TGRpcRequestProxyImpl::GetTracingControl(const TString& type) {
+    if (auto it = TracingControls.find(type); it != TracingControls.end()) {
+        return it->second;
+    }
+    auto tracingControlsDomain = InternalRequestTypeToControlDomain(type);
+    auto domain = TString::Join("TracingControls.", tracingControlsDomain);
+    TTracingControl control(Icb, TAppData::TimeProvider, TAppData::RandomProvider, std::move(domain));
+    return TracingControls.emplace(type, std::move(control)).first->second;
+}
+
+void TGRpcRequestProxyImpl::MaybeStartTracing(IRequestProxyCtx& ctx) {
+    auto requestType = ctx.GetInternalRequestType();
+    if (requestType.empty()) {
+        return;
+    }
+    NWilson::TTraceId traceId;
     if (const auto otelHeader = ctx.GetPeerMetaValues(NYdb::OTEL_TRACE_HEADER)) {
-        if (auto traceId = NWilson::TTraceId::FromTraceparentHeader(otelHeader.GetRef())) {
-            NWilson::TSpan grpcRequestProxySpan(TWilsonGrpc::RequestProxy, std::move(traceId), "GrpcRequestProxy");
-            ctx.StartTracing(std::move(grpcRequestProxySpan));
-        }
+        traceId = NWilson::TTraceId::FromTraceparentHeader(otelHeader.GetRef());
+    }
+    auto& control = GetTracingControl(requestType);
+    if (traceId && control.ThrottleExternal()) {
+        LOG_DEBUG_S(*TlsActivationContext, NKikimrServices::GRPC_SERVER, "Dropping external traceId " << traceId.GetHexTraceId() << " for request type " << requestType);
+        traceId = {};
+    }
+    if (!traceId && control.SampleThrottle()) {
+        traceId = NWilson::TTraceId::NewTraceId(control.SampledVerbosity(), 4095);
+        LOG_DEBUG_S(*TlsActivationContext, NKikimrServices::GRPC_SERVER, "Created new traceId " << traceId.GetHexTraceId() << " for request type " << requestType);
+    }
+    if (traceId) {
+        NWilson::TSpan grpcRequestProxySpan(TWilsonGrpc::RequestProxy, std::move(traceId), "GrpcRequestProxy");
+        ctx.StartTracing(std::move(grpcRequestProxySpan));
     }
 }
 
@@ -583,8 +629,8 @@ void TGRpcRequestProxyImpl::StateFunc(TAutoPtr<IEventHandle>& ev) {
     }
 }
 
-IActor* CreateGRpcRequestProxy(const NKikimrConfig::TAppConfig& appConfig) {
-    return new TGRpcRequestProxyImpl(appConfig);
+IActor* CreateGRpcRequestProxy(const NKikimrConfig::TAppConfig& appConfig, TIntrusivePtr<TControlBoard> icb) {
+    return new TGRpcRequestProxyImpl(appConfig, std::move(icb));
 }
 
 } // namespace NGRpcService
diff --git a/ydb/core/grpc_services/grpc_request_proxy.h b/ydb/core/grpc_services/grpc_request_proxy.h
index d2eb347c979..b4eedb51c5f 100644
--- a/ydb/core/grpc_services/grpc_request_proxy.h
+++ b/ydb/core/grpc_services/grpc_request_proxy.h
@@ -5,6 +5,8 @@
 
 #include "grpc_request_proxy_handle_methods.h"
 
+#include <ydb/core/base/appdata_fwd.h>
+
 #include <ydb/library/actors/core/actor.h>
 
 #include <util/generic/ptr.h>
@@ -21,7 +23,7 @@ struct TAppData;
 namespace NGRpcService {
 
 TString DatabaseFromDomain(const TAppData* appdata);
-IActor* CreateGRpcRequestProxy(const NKikimrConfig::TAppConfig& appConfig);
+IActor* CreateGRpcRequestProxy(const NKikimrConfig::TAppConfig& appConfig, TIntrusivePtr<TControlBoard> icb);
 IActor* CreateGRpcRequestProxySimple(const NKikimrConfig::TAppConfig& appConfig);
 
 class TGRpcRequestProxy : public TGRpcRequestProxyHandleMethods, public IFacilityProvider {
diff --git a/ydb/core/keyvalue/keyvalue_flat_impl.h b/ydb/core/keyvalue/keyvalue_flat_impl.h
index 0e200ad0841..2616d4229d0 100644
--- a/ydb/core/keyvalue/keyvalue_flat_impl.h
+++ b/ydb/core/keyvalue/keyvalue_flat_impl.h
@@ -7,6 +7,7 @@
 #include "keyvalue_simple_db.h"
 #include "keyvalue_simple_db_flat.h"
 #include "keyvalue_state.h"
+
 #include <ydb/core/tablet_flat/tablet_flat_executed.h>
 #include <ydb/core/tablet_flat/flat_database.h>
 #include <ydb/core/engine/minikql/flat_local_tx_factory.h>
diff --git a/ydb/core/keyvalue/ya.make b/ydb/core/keyvalue/ya.make
index 3a014e31316..740eae043d0 100644
--- a/ydb/core/keyvalue/ya.make
+++ b/ydb/core/keyvalue/ya.make
@@ -45,6 +45,7 @@ PEERDIR(
     ydb/library/actors/protos
     ydb/core/base
     ydb/core/blobstorage/base
+    ydb/core/control/common_controls
     ydb/core/engine/minikql
     ydb/core/keyvalue/protos
     ydb/core/protos
diff --git a/ydb/core/protos/config.proto b/ydb/core/protos/config.proto
index d9d1c9d377a..46aa0dacfc3 100644
--- a/ydb/core/protos/config.proto
+++ b/ydb/core/protos/config.proto
@@ -1245,11 +1245,61 @@ message TImmediateControlsConfig {
             DefaultValue: 8388608 }];
     }
 
+    message TTracingControls {
+        message TSamplingThrottlingOptions {
+            message TThrottlingOptions {
+                optional uint64 MaxRatePerMinute = 1 [(ControlOptions) = {
+                    Description: "Maximum amount of traced requests per minute",
+                    MinValue: 0,
+                    MaxValue: 300,
+                    DefaultValue: 0,
+                }];
+                optional uint64 MaxBurst = 2 [(ControlOptions) = {
+                    Description: "Maximum burst of traced events",
+                    MinValue: 0,
+                    MaxValue: 300, 
+                    DefaultValue: 0,
+                }];
+            }
+
+            message TSamplingOptions {
+                optional uint64 PPM = 1 [(ControlOptions) = {
+                    Description: "Average amount of sampled requests per one million",
+                    MinValue: 0,
+                    MaxValue: 1000000,
+                    DefaultValue: 0,
+                }];
+                optional uint64 Level = 2 [(ControlOptions) = {
+                    Description: "Tracing level of sampled requests",
+                    MinValue: 1,
+                    MaxValue: 15,
+                    DefaultValue: 15,
+                }];
+            }
+
+            optional TSamplingOptions Sampling = 1;
+            optional TThrottlingOptions SampledThrottling = 2;
+            optional TThrottlingOptions ExternalThrottling = 3;
+        }
+
+        message TKeyValue {
+            optional TSamplingThrottlingOptions AcquireLock = 1;
+            optional TSamplingThrottlingOptions ExecuteTransaction = 2;
+            optional TSamplingThrottlingOptions Read = 3;
+            optional TSamplingThrottlingOptions ReadRange = 4;
+            optional TSamplingThrottlingOptions ListRange = 5;
+            optional TSamplingThrottlingOptions GetStorageChannelStatus = 6;
+        }
+
+        optional TKeyValue KeyValue = 1;
+    }
+
     optional TDataShardControls DataShardControls = 1;
     optional TTxLimitControls TxLimitControls = 2;
     optional TCoordinatorControls CoordinatorControls = 3;
     optional TSchemeShardControls SchemeShardControls = 4;
     optional TTCMallocControls TCMallocControls = 5;
+    optional TTracingControls TracingControls = 6;
 };
 
 message TMeteringConfig {
diff --git a/ydb/core/testlib/test_client.cpp b/ydb/core/testlib/test_client.cpp
index 6ef5ad171ce..58b5e29299b 100644
--- a/ydb/core/testlib/test_client.cpp
+++ b/ydb/core/testlib/test_client.cpp
@@ -307,8 +307,11 @@ namespace Tests {
         const size_t proxyCount = Max(ui32{1}, Settings->AppConfig->GetGRpcConfig().GetGRpcProxyCount());
         TVector<TActorId> grpcRequestProxies;
         grpcRequestProxies.reserve(proxyCount);
+
+        auto& appData = Runtime->GetAppData();
+
         for (size_t i = 0; i < proxyCount; ++i) {
-            auto grpcRequestProxy = NGRpcService::CreateGRpcRequestProxy(*Settings->AppConfig);
+            auto grpcRequestProxy = NGRpcService::CreateGRpcRequestProxy(*Settings->AppConfig, appData.Icb);
             auto grpcRequestProxyId = system->Register(grpcRequestProxy, TMailboxType::ReadAsFilled);
             system->RegisterLocalService(NGRpcService::CreateGRpcRequestProxyId(), grpcRequestProxyId);
             grpcRequestProxies.push_back(grpcRequestProxyId);
@@ -320,8 +323,6 @@ namespace Tests {
         GRpcServerRootCounters = MakeIntrusive<::NMonitoring::TDynamicCounters>();
         auto& counters = GRpcServerRootCounters;
 
-        auto& appData = Runtime->GetAppData();
-
         // Setup discovery for typically used services on the node
         {
             TIntrusivePtr<NGRpcService::TGrpcEndpointDescription> desc = new NGRpcService::TGrpcEndpointDescription();