YT-20637: Sample async spans only if parent is sampled

5 files changed, 33 insertions, 12 deletions

diff --git a/yt/yt/client/transaction_client/batching_timestamp_provider.cpp b/yt/yt/client/transaction_client/batching_timestamp_provider.cpp
index e6a5e299ab..ad4919bad3 100644
--- a/yt/yt/client/transaction_client/batching_timestamp_provider.cpp
+++ b/yt/yt/client/transaction_client/batching_timestamp_provider.cpp
@@ -110,7 +110,7 @@ public:
         std::vector<TRequest> requests;
         requests.swap(PendingRequests_);
 
-        auto [traceContext, sampled] = BatchTrace_.StartSpan("BatchingTimestampProvider:SendGenerateRequest");
+        auto [traceContext, _] = BatchTrace_.StartSpan("BatchingTimestampProvider:SendGenerateRequest");
 
         guard.Release();
 
diff --git a/yt/yt/library/tracing/async_queue_trace.cpp b/yt/yt/library/tracing/async_queue_trace.cpp
index 9af19f9188..87126a4fac 100644
--- a/yt/yt/library/tracing/async_queue_trace.cpp
+++ b/yt/yt/library/tracing/async_queue_trace.cpp
@@ -49,8 +49,8 @@ std::pair<TTraceContextPtr, bool> TAsyncQueueTrace::StartSpan(i64 startIndex, co
 
     for (const auto& [queueIndex, client] : Blocked_) {
         // If startIndex > queueIndex, client is not blocked by this span.
-        if (queueIndex >= startIndex && client->AddAsyncChild(traceContext->GetTraceId())) {
-            sampled = true;
+        if (queueIndex >= startIndex && client->AddAsyncChild(traceContext->GetTraceId()) && !sampled) {
+            sampled = client->IsSampled();
         }
     }
 
diff --git a/yt/yt/library/tracing/batch_trace.cpp b/yt/yt/library/tracing/batch_trace.cpp
index 4f81be4e12..f08ad38770 100644
--- a/yt/yt/library/tracing/batch_trace.cpp
+++ b/yt/yt/library/tracing/batch_trace.cpp
@@ -24,19 +24,19 @@ std::pair<TTraceContextPtr, bool> TBatchTrace::StartSpan(const TString& spanName
 {
     auto traceContext = TTraceContext::NewRoot(spanName);
 
-    bool hasBlockedClient = false;
+    bool sampled = false;
     for (const auto& client : Clients_) {
-        if (client->AddAsyncChild(traceContext->GetTraceId())) {
-            hasBlockedClient = true;
+        if (client->AddAsyncChild(traceContext->GetTraceId()) && !sampled) {
+            sampled = client->IsSampled();
         }
     }
     Clients_.clear();
 
-    if (hasBlockedClient) {
+    if (sampled) {
         traceContext->SetSampled();
     }
 
-    return {traceContext, hasBlockedClient};
+    return {traceContext, sampled};
 }
 
 ////////////////////////////////////////////////////////////////////////////////
diff --git a/yt/yt/library/tracing/example/main.cpp b/yt/yt/library/tracing/example/main.cpp
index 2f8b13d523..d2bbcb1972 100644
--- a/yt/yt/library/tracing/example/main.cpp
+++ b/yt/yt/library/tracing/example/main.cpp
@@ -1,13 +1,15 @@
-#include <random>
+#include <yt/yt/library/tracing/jaeger/tracer.h>
+
+#include <yt/yt/library/tracing/batch_trace.h>
 
 #include <yt/yt/core/tracing/trace_context.h>
 
 #include <util/generic/yexception.h>
 
-#include <yt/yt/library/tracing/jaeger/tracer.h>
-
 #include <util/system/env.h>
 
+#include <random>
+
 using namespace NYT;
 using namespace NYT::NTracing;
 
@@ -29,9 +31,18 @@ void SubrequestExample(std::optional<TString> endpoint)
 
     Sleep(TDuration::MilliSeconds(2));
     traceContext->AddLogEntry(GetCpuInstant(), "Request finished");
+
+    TBatchTrace batchTrace;
+    batchTrace.Join(traceContext);
+
+    auto [asyncChildTraceContext, asyncChildSampled] = batchTrace.StartSpan("TestBatchTrace");
+    YT_VERIFY(asyncChildSampled);
+
     traceContext->Finish();
 
-    Cout << ToString(traceContext->GetTraceId()) << Endl;
+    asyncChildTraceContext->Finish();
+
+    Cout << ToString(traceContext->GetTraceId()) << '\t' << ToString(asyncChildTraceContext->GetTraceId()) << Endl;
 }
 
 void DelayedSamplingExample(std::optional<TString> endpoint)
@@ -48,8 +59,16 @@ void DelayedSamplingExample(std::optional<TString> endpoint)
 
     auto slowRequestContext = startContext->CreateChild("SlowRequest");
 
+    TBatchTrace batchTrace;
+    batchTrace.Join(slowRequestContext);
+    auto [asyncChildTraceContext, asyncChildSampled] = batchTrace.StartSpan("TestBatchTrace");
+
+    YT_VERIFY(!asyncChildSampled);
+    YT_VERIFY(!slowRequestContext->IsSampled());
+
     traceContext->SetSampled();
     YT_VERIFY(slowRequestContext->IsSampled());
+    YT_VERIFY(!asyncChildTraceContext->IsSampled());
 
     slowRequestContext->Finish();
     traceContext->Finish();
diff --git a/yt/yt/library/tracing/public.h b/yt/yt/library/tracing/public.h
index 74635df631..ae473bb33b 100644
--- a/yt/yt/library/tracing/public.h
+++ b/yt/yt/library/tracing/public.h
@@ -8,6 +8,8 @@ namespace NYT::NTracing {
 
 DECLARE_REFCOUNTED_STRUCT(ITracer)
 
+class TBatchTrace;
+
 class TAsyncQueueTrace;
 class TAsyncQueueTraceGuard;