fix for sampling root span (#784)

open-telemetry · May 24, 2021 · 57d80f7 · 57d80f7
1 parent 5a1b189
commit 57d80f7
Show file tree

Hide file tree

Showing 4 changed files with 66 additions and 43 deletions.
diff --git a/sdk/src/trace/span.cc b/sdk/src/trace/span.cc
@@ -47,11 +47,11 @@ Span::Span(std::shared_ptr<Tracer> &&tracer,
            const trace_api::SpanContextKeyValueIterable &links,
            const trace_api::StartSpanOptions &options,
            const trace_api::SpanContext &parent_span_context,
-           const nostd::shared_ptr<opentelemetry::trace::TraceState> trace_state,
-           const bool sampled) noexcept
+           std::unique_ptr<trace_api::SpanContext> span_context) noexcept
     : tracer_{std::move(tracer)},
       recordable_{tracer_->GetProcessor().MakeRecordable()},
       start_steady_time{options.start_steady_time},
+      span_context_(std::move(span_context)),
       has_ended_{false}
 {
   if (recordable_ == nullptr)
@@ -60,32 +60,9 @@ Span::Span(std::shared_ptr<Tracer> &&tracer,
   }
   recordable_->SetName(name);
   recordable_->SetInstrumentationLibrary(tracer_->GetInstrumentationLibrary());
-
-  trace_api::TraceId trace_id;
-  trace_api::SpanId span_id = tracer_->GetIdGenerator().GenerateSpanId();
-  trace_api::SpanId parent_span_id;
-  bool is_parent_span_valid = false;
-
-  if (parent_span_context.IsValid())
-  {
-    trace_id             = parent_span_context.trace_id();
-    parent_span_id       = parent_span_context.span_id();
-    is_parent_span_valid = true;
-  }
-  else
-  {
-    trace_id = tracer_->GetIdGenerator().GenerateTraceId();
-  }
-
-  span_context_ = std::unique_ptr<trace_api::SpanContext>(new trace_api::SpanContext(
-      trace_id, span_id,
-      sampled ? trace_api::TraceFlags{trace_api::TraceFlags::kIsSampled} : trace_api::TraceFlags{},
-      false,
-      trace_state ? trace_state
-                  : is_parent_span_valid ? parent_span_context.trace_state()
-                                         : trace_api::TraceState::GetDefault()));
-
-  recordable_->SetIdentity(*span_context_, parent_span_id);
+  recordable_->SetIdentity(*span_context_, parent_span_context.IsValid()
+                                               ? parent_span_context.span_id()
+                                               : trace_api::SpanId());
 
   attributes.ForEachKeyValue(
       [&](nostd::string_view key, opentelemetry::common::AttributeValue value) noexcept {

diff --git a/sdk/src/trace/span.h b/sdk/src/trace/span.h
@@ -21,9 +21,7 @@ class Span final : public trace_api::Span
        const trace_api::SpanContextKeyValueIterable &links,
        const trace_api::StartSpanOptions &options,
        const trace_api::SpanContext &parent_span_context,
-       const nostd::shared_ptr<opentelemetry::trace::TraceState> trace_state =
-           trace_api::TraceState::GetDefault(),
-       const bool sampled = false) noexcept;
+       std::unique_ptr<trace_api::SpanContext> span_context) noexcept;
 
   ~Span() override;
 

diff --git a/sdk/src/trace/tracer.cc b/sdk/src/trace/tracer.cc
@@ -5,6 +5,8 @@
 #include "opentelemetry/version.h"
 #include "src/trace/span.h"
 
+#include <memory>
+
 OPENTELEMETRY_BEGIN_NAMESPACE
 namespace sdk
 {
@@ -22,11 +24,26 @@ nostd::shared_ptr<trace_api::Span> Tracer::StartSpan(
     const trace_api::SpanContextKeyValueIterable &links,
     const trace_api::StartSpanOptions &options) noexcept
 {
-  trace_api::SpanContext parent =
+  trace_api::SpanContext parent_context =
       options.parent.IsValid() ? options.parent : GetCurrentSpan()->GetContext();
 
-  auto sampling_result = context_->GetSampler().ShouldSample(parent, parent.trace_id(), name,
+  trace_api::TraceId trace_id;
+  trace_api::SpanId span_id = GetIdGenerator().GenerateSpanId();
+  bool is_parent_span_valid = false;
+
+  if (parent_context.IsValid())
+  {
+    trace_id             = parent_context.trace_id();
+    is_parent_span_valid = true;
+  }
+  else
+  {
+    trace_id = GetIdGenerator().GenerateTraceId();
+  }
+
+  auto sampling_result = context_->GetSampler().ShouldSample(parent_context, trace_id, name,
                                                              options.kind, attributes, links);
+
   if (sampling_result.decision == Decision::DROP)
   {
     // Don't allocate a no-op span for every DROP decision, but use a static
@@ -38,9 +55,16 @@ nostd::shared_ptr<trace_api::Span> Tracer::StartSpan(
   }
   else
   {
+
+    auto span_context = std::unique_ptr<trace_api::SpanContext>(new trace_api::SpanContext(
+        trace_id, span_id, trace_api::TraceFlags{trace_api::TraceFlags::kIsSampled}, false,
+        sampling_result.trace_state ? sampling_result.trace_state
+                                    : is_parent_span_valid ? parent_context.trace_state()
+                                                           : trace_api::TraceState::GetDefault()));
+
     auto span = nostd::shared_ptr<trace_api::Span>{
-        new (std::nothrow) Span{this->shared_from_this(), name, attributes, links, options, parent,
-                                sampling_result.trace_state, true}};
+        new (std::nothrow) Span{this->shared_from_this(), name, attributes, links, options,
+                                parent_context, std::move(span_context)}};
 
     // if the attributes is not nullptr, add attributes to the span.
     if (sampling_result.attributes)

diff --git a/sdk/test/trace/tracer_test.cc b/sdk/test/trace/tracer_test.cc
@@ -21,25 +21,37 @@ using opentelemetry::exporter::memory::InMemorySpanExporter;
 using opentelemetry::trace::SpanContext;
 
 /**
- * A mock sampler that returns non-empty sampling results attributes.
+ * A mock sampler with ShouldSample returning:
+ *  Decision::RECORD_AND_SAMPLE if trace_id is valid
+ *  Decision::DROP otherwise.
  */
 class MockSampler final : public Sampler
 {
 public:
   SamplingResult ShouldSample(
       const SpanContext & /*parent_context*/,
-      trace_api::TraceId /*trace_id*/,
+      trace_api::TraceId trace_id,
       nostd::string_view /*name*/,
       trace_api::SpanKind /*span_kind*/,
       const opentelemetry::common::KeyValueIterable & /*attributes*/,
       const opentelemetry::trace::SpanContextKeyValueIterable & /*links*/) noexcept override
   {
-    // Return two pairs of attributes. These attributes should be added to the
-    // span attributes
-    return {Decision::RECORD_AND_SAMPLE,
-            nostd::unique_ptr<const std::map<std::string, opentelemetry::common::AttributeValue>>(
-                new const std::map<std::string, opentelemetry::common::AttributeValue>(
-                    {{"sampling_attr1", 123}, {"sampling_attr2", "string"}}))};
+    // Sample only if valid trace_id ( This is to test Sampler get's valid trace id)
+    if (trace_id.IsValid())
+    {
+      // Return two pairs of attributes. These attributes should be added to the
+      // span attributes
+      return {Decision::RECORD_AND_SAMPLE,
+              nostd::unique_ptr<const std::map<std::string, opentelemetry::common::AttributeValue>>(
+                  new const std::map<std::string, opentelemetry::common::AttributeValue>(
+                      {{"sampling_attr1", 123}, {"sampling_attr2", "string"}}))};
+    }
+    else
+    {
+      // we should never reach here
+      assert(false);
+      return {Decision::DROP};
+    }
   }
 
   nostd::string_view GetDescription() const noexcept override { return "MockSampler"; }
@@ -599,3 +611,15 @@ TEST(Tracer, ExpectParent)
   EXPECT_EQ(spandata_first->GetSpanId(), spandata_second->GetParentSpanId());
   EXPECT_EQ(spandata_second->GetSpanId(), spandata_third->GetParentSpanId());
 }
+
+TEST(Tracer, ValidTraceIdToSampler)
+{
+  std::unique_ptr<InMemorySpanExporter> exporter(new InMemorySpanExporter());
+  std::shared_ptr<InMemorySpanData> span_data = exporter->GetData();
+  auto tracer                                 = initTracer(std::move(exporter), new MockSampler());
+
+  auto span = tracer->StartSpan("span 1");
+  // sampler was fed with valid trace_id, so span shouldn't be NoOp Span.
+  EXPECT_TRUE(span->IsRecording());
+  EXPECT_TRUE(span->GetContext().IsValid());
+}