[ETW Exporter] - Add Trace flags in SpanContext (#1618)

Author: lalitb

exporters/etw/include/opentelemetry/exporters/etw/etw_tracer.h

@@ -64,9 +64,10 @@ class Span;
 template <class SpanType, class TracerType>
 SpanType *new_span(TracerType *objPtr,
                    nostd::string_view name,
-                   const opentelemetry::trace::StartSpanOptions &options)
+                   const opentelemetry::trace::StartSpanOptions &options,
+                   std::unique_ptr<opentelemetry::trace::SpanContext> spanContext)
 {
-  return new (std::nothrow) SpanType{*objPtr, name, options};
+  return new (std::nothrow) SpanType{*objPtr, name, options, std::move(spanContext)};
 }
 
 /**
@@ -374,30 +375,6 @@ class Tracer : public opentelemetry::trace::Tracer,
       const opentelemetry::trace::SpanContextKeyValueIterable &links,
       const opentelemetry::trace::StartSpanOptions &options = {}) noexcept override
   {
-    const auto &cfg = GetConfiguration(tracerProvider_);
-
-    // Parent Context:
-    // - either use current span
-    // - or attach to parent SpanContext specified in options
-    opentelemetry::trace::SpanContext parentContext = GetCurrentSpan()->GetContext();
-    if (nostd::holds_alternative<opentelemetry::trace::SpanContext>(options.parent))
-    {
-      auto span_context = nostd::get<opentelemetry::trace::SpanContext>(options.parent);
-      if (span_context.IsValid())
-      {
-        parentContext = span_context;
-      }
-    }
-    auto sampling_result =
-        GetSampler(tracerProvider_)
-            .ShouldSample(parentContext, traceId_, name, options.kind, attributes, links);
-    if (sampling_result.decision == sdk::trace::Decision::DROP)
-    {
-      static nostd::shared_ptr<trace::Span> noop_span(
-          new trace::NoopSpan{this->shared_from_this()});
-      return noop_span;
-    }
-
 #ifdef OPENTELEMETRY_RTTI_ENABLED
     common::KeyValueIterable &attribs = const_cast<common::KeyValueIterable &>(attributes);
     Properties *evt                   = dynamic_cast<Properties *>(&attribs);
@@ -433,12 +410,42 @@ class Tracer : public opentelemetry::trace::Tracer,
     opentelemetry::trace::SpanContext parentContext = GetCurrentSpan()->GetContext();
     if (nostd::holds_alternative<opentelemetry::trace::SpanContext>(options.parent))
     {
-      auto span_context = nostd::get<opentelemetry::trace::SpanContext>(options.parent);
-      if (span_context.IsValid())
+      auto spanContext = nostd::get<opentelemetry::trace::SpanContext>(options.parent);
+      if (spanContext.IsValid())
       {
-        parentContext = span_context;
+        parentContext = spanContext;
       }
     }
+    auto traceId = parentContext.IsValid() ? parentContext.trace_id() : traceId_;
+
+    // Sampling based on attributes is not supported for now, so passing empty below.
+    std::map<std::string, int> emptyAttributes = {{}};
+    opentelemetry::sdk::trace::SamplingResult sampling_result =
+        GetSampler(tracerProvider_)
+            .ShouldSample(parentContext, traceId, name, options.kind,
+                          opentelemetry::common::KeyValueIterableView<std::map<std::string, int>>(
+                              emptyAttributes),
+                          links);
+
+    opentelemetry::trace::TraceFlags traceFlags =
+        sampling_result.decision == opentelemetry::sdk::trace::Decision::DROP
+            ? opentelemetry::trace::TraceFlags{}
+            : opentelemetry::trace::TraceFlags{opentelemetry::trace::TraceFlags::kIsSampled};
+
+    auto spanContext =
+        std::unique_ptr<opentelemetry::trace::SpanContext>(new opentelemetry::trace::SpanContext(
+            traceId, GetIdGenerator(tracerProvider_).GenerateSpanId(), traceFlags, false,
+            sampling_result.trace_state
+                ? sampling_result.trace_state
+                : parentContext.IsValid() ? parentContext.trace_state()
+                                          : opentelemetry::trace::TraceState::GetDefault()));
+
+    if (sampling_result.decision == sdk::trace::Decision::DROP)
+    {
+      auto noopSpan = nostd::shared_ptr<trace::Span>{
+          new (std::nothrow) trace::NoopSpan(this->shared_from_this(), std::move(spanContext))};
+      return noopSpan;
+    }
 
     // Populate Etw.RelatedActivityId at envelope level if enabled
     GUID RelatedActivityId;
@@ -456,11 +463,9 @@ class Tracer : public opentelemetry::trace::Tracer,
 
     // This template pattern allows us to forward-declare the etw::Span,
     // create an instance of it, then assign it to tracer::Span result.
-    auto currentSpan = new_span<Span, Tracer>(this, name, options);
+    auto currentSpan = new_span<Span, Tracer>(this, name, options, std::move(spanContext));
     nostd::shared_ptr<opentelemetry::trace::Span> result = to_span_ptr<Span>(currentSpan);
 
-    auto spanContext = result->GetContext();
-
     // Decorate with additional standard fields
     std::string eventName = name.data();
 
@@ -475,13 +480,13 @@ class Tracer : public opentelemetry::trace::Tracer,
       {
         evt[ETW_FIELD_SPAN_PARENTID] = ToLowerBase16(parentContext.span_id());
       }
-      evt[ETW_FIELD_SPAN_ID] = ToLowerBase16(spanContext.span_id());
+      evt[ETW_FIELD_SPAN_ID] = ToLowerBase16(result.get()->GetContext().span_id());
     }
 
     // Populate Etw.Payload["TraceId"] attribute
     if (cfg.enableTraceId)
     {
-      evt[ETW_FIELD_TRACE_ID] = ToLowerBase16(spanContext.trace_id());
+      evt[ETW_FIELD_TRACE_ID] = ToLowerBase16(result.get()->GetContext().trace_id());
     }
 
     // Populate Etw.ActivityId at envelope level if enabled
@@ -705,7 +710,7 @@ class Span : public opentelemetry::trace::Span
    */
   Span *GetParent() const { return parent_; }
 
-  opentelemetry::trace::SpanContext context_;
+  std::unique_ptr<opentelemetry::trace::SpanContext> context_;
 
 public:
   /**
@@ -759,13 +764,13 @@ class Span : public opentelemetry::trace::Span
   Span(Tracer &owner,
        nostd::string_view name,
        const opentelemetry::trace::StartSpanOptions &options,
+       std::unique_ptr<opentelemetry::trace::SpanContext> spanContext,
        Span *parent = nullptr) noexcept
       : opentelemetry::trace::Span(),
         start_time_(std::chrono::system_clock::now()),
         owner_(owner),
-        parent_(parent),
-        context_{owner.traceId_, GetIdGenerator(owner.tracerProvider_).GenerateSpanId(),
-                 opentelemetry::trace::TraceFlags{0}, false}
+        context_(std::move(spanContext)),
+        parent_(parent)
   {
     name_ = name;
     UNREFERENCED_PARAMETER(options);
@@ -883,7 +888,7 @@ class Span : public opentelemetry::trace::Span
    * @brief Obtain SpanContext
    * @return
    */
-  opentelemetry::trace::SpanContext GetContext() const noexcept override { return context_; }
+  opentelemetry::trace::SpanContext GetContext() const noexcept override { return *context_.get(); }
 
   /**
    * @brief Check if Span is recording data.

exporters/etw/test/etw_tracer_test.cc

@@ -7,13 +7,15 @@
 #  include <map>
 #  include <string>
 
+#  include "opentelemetry//sdk/trace/sampler.h"
 #  include "opentelemetry/exporters/etw/etw_tracer_exporter.h"
 #  include "opentelemetry/sdk/trace/samplers/always_off.h"
 #  include "opentelemetry/sdk/trace/simple_processor.h"
 
 using namespace OPENTELEMETRY_NAMESPACE;
 
 using namespace opentelemetry::exporter::etw;
+using namespace opentelemetry::sdk::trace;
 
 const char *kGlobalProviderName = "OpenTelemetry-ETW-TLD";
 
@@ -40,6 +42,42 @@ class MockIdGenerator : public sdk::trace::IdGenerator
   uint8_t buf_trace[16] = {1, 2, 3, 4, 5, 6, 7, 8, 8, 7, 6, 5, 4, 3, 2, 1};
 };
 
+/* A Custom Sampler, implementing parent based sampler*/
+class MockSampler : public sdk::trace::Sampler
+{
+public:
+  MockSampler(std::shared_ptr<Sampler> delegate_sampler) noexcept
+      : delegate_sampler_(delegate_sampler)
+  {}
+  sdk::trace::SamplingResult ShouldSample(
+      const trace_api::SpanContext &parent_context,
+      trace_api::TraceId trace_id,
+      nostd::string_view name,
+      trace_api::SpanKind span_kind,
+      const opentelemetry::common::KeyValueIterable &attributes,
+      const trace_api::SpanContextKeyValueIterable &links) noexcept
+  {
+    if (!parent_context.IsValid())
+    {
+      // If no parent (root span) exists returns the result of the delegateSampler
+      return delegate_sampler_->ShouldSample(parent_context, trace_id, name, span_kind, attributes,
+                                             links);
+    }
+
+    // If parent exists:
+    if (parent_context.IsSampled())
+    {
+      return {Decision::RECORD_AND_SAMPLE, nullptr, parent_context.trace_state()};
+    }
+    return {Decision::DROP, nullptr, parent_context.trace_state()};
+  }
+
+  nostd::string_view GetDescription() const noexcept { return "Custom Sampler"; }
+
+private:
+  std::shared_ptr<Sampler> delegate_sampler_;
+};
+
 /* clang-format off */
 TEST(ETWTracer, TracerCheck)
 {
@@ -417,7 +455,8 @@ TEST(ETWTracer, AlwayOffSampler)
      std::move(always_off));
   auto tracer = tp.GetTracer(providerName);
   auto span = tracer->StartSpan("span_off");
-  EXPECT_EQ(span->GetContext().IsValid(), false);
+  EXPECT_EQ(span->GetContext().IsValid(), true);
+  EXPECT_EQ(span->GetContext().IsSampled(), false);
 }
 
 TEST(ETWTracer, CustomIdGenerator)
@@ -440,6 +479,35 @@ TEST(ETWTracer, CustomIdGenerator)
   EXPECT_EQ(span->GetContext().trace_id(), id_generator->GenerateTraceId());
 }
 
+TEST(ETWTracer, CustomSampler)
+{
+  std::string providerName = kGlobalProviderName; // supply unique instrumentation name here
+  auto parent_off = std::unique_ptr<Sampler>(new MockSampler(std::make_shared<AlwaysOnSampler>()));
+  exporter::etw::TracerProvider tp
+    ({
+      {"enableTraceId", true},
+      {"enableSpanId", true},
+      {"enableActivityId", true},
+      {"enableRelatedActivityId", true},
+      {"enableAutoParent", true}
+     },
+     std::move(parent_off));
+  auto tracer = tp.GetTracer(providerName);
+  {
+    auto span = tracer->StartSpan("span_off");
+    EXPECT_EQ(span->GetContext().IsValid(), true);
+    EXPECT_EQ(span->GetContext().IsSampled(), true);
+    auto scope = tracer->WithActiveSpan(span);
+    auto trace_id = span->GetContext().trace_id();
+    {
+      auto child_span = tracer->StartSpan("span on");
+      EXPECT_EQ(child_span->GetContext().IsValid(), true);
+      EXPECT_EQ(child_span->GetContext().IsSampled(), true);
+      EXPECT_EQ(child_span->GetContext().trace_id(), trace_id);
+    }
+  }
+}
+
 /* clang-format on */
 
 #endif