feat: ensure tracecontext headers take precedence over datadog (AIT-10281) (#142)

zacharycmontoya · web-flow · commit dc5e762ee993 · 2024-08-02T13:06:38.000-07:00
Adds support for "W3C Phase 3" so that when multiple trace contexts are extracted (with the same trace-id) and the `tracecontext` format is used, the resulting trace context will have:
- `span.parent_id` set to the `tracecontext` parent-id value
- `span.tags["_dd.parent_id"]` set to either the value of `tracestate[dd][p]` or `hex(x-datadog-parent-id)`, in that order of precedence
diff --git a/src/datadog/extraction_util.cpp b/src/datadog/extraction_util.cpp
@@ -7,6 +7,7 @@
 #include <unordered_map>
 
 #include "extracted_data.h"
+#include "hex.h"
 #include "logger.h"
 #include "parse_util.h"
 #include "string_util.h"
@@ -250,51 +251,47 @@ void AuditedReader::visit(
   });
 }
 
-ExtractedData merge(const std::vector<ExtractedData>& contexts) {
+ExtractedData merge(
+    const PropagationStyle first_style,
+    const std::unordered_map<PropagationStyle, ExtractedData>& contexts) {
   ExtractedData result;
-
-  const auto found = std::find_if(
-      contexts.begin(), contexts.end(),
-      [](const ExtractedData& data) { return data.trace_id.has_value(); });
-
+  const auto found = contexts.find(first_style);
   if (found == contexts.end()) {
-    // Nothing extracted a trace ID. Return the first context that includes a
-    // parent ID, if any, or otherwise just return an empty `ExtractedData`.
-    // The purpose of looking for a parent ID is to allow for the error
-    // "extracted a parent ID without a trace ID," if that's what happened.
-    const auto other = std::find_if(
-        contexts.begin(), contexts.end(),
-        [](const ExtractedData& data) { return data.parent_id.has_value(); });
-    if (other != contexts.end()) {
-      result = *other;
-    }
     return result;
   }
 
   // `found` refers to the first extracted context that yielded a trace ID.
   // This will be our main context.
   //
-  // If the style of `found` is not W3C, then examine the remaining contexts
-  // for W3C-style tracestate that we might want to include in `result`.
-  result = *found;
-  if (result.style == PropagationStyle::W3C) {
-    return result;
-  }
+  // If the W3C style is present and its trace-id matches, we'll update the main
+  // context with tracestate information that we want to include in `result`. We
+  // may also need to use Datadog header information (only when the trace-id
+  // matches).
+  result = found->second;
 
-  const auto other =
-      std::find_if(found + 1, contexts.end(), [&](const ExtractedData& data) {
-        return data.style == PropagationStyle::W3C &&
-               data.trace_id == found->trace_id;
-      });
+  const auto w3c = contexts.find(PropagationStyle::W3C);
+  const auto dd = contexts.find(PropagationStyle::DATADOG);
 
-  if (other != contexts.end()) {
-    result.datadog_w3c_parent_id = other->datadog_w3c_parent_id;
-    result.additional_w3c_tracestate = other->additional_w3c_tracestate;
+  if (w3c != contexts.end() && w3c->second.trace_id == result.trace_id) {
+    result.additional_w3c_tracestate = w3c->second.additional_w3c_tracestate;
     result.additional_datadog_w3c_tracestate =
-        other->additional_datadog_w3c_tracestate;
+        w3c->second.additional_datadog_w3c_tracestate;
     result.headers_examined.insert(result.headers_examined.end(),
-                                   other->headers_examined.begin(),
-                                   other->headers_examined.end());
+                                   w3c->second.headers_examined.begin(),
+                                   w3c->second.headers_examined.end());
+
+    if (result.parent_id != w3c->second.parent_id) {
+      if (w3c->second.datadog_w3c_parent_id &&
+          w3c->second.datadog_w3c_parent_id != "0000000000000000") {
+        result.datadog_w3c_parent_id = w3c->second.datadog_w3c_parent_id;
+      } else if (dd != contexts.end() &&
+                 dd->second.trace_id == result.trace_id &&
+                 dd->second.parent_id.has_value()) {
+        result.datadog_w3c_parent_id = hex_padded(dd->second.parent_id.value());
+      }
+
+      result.parent_id = w3c->second.parent_id;
+    }
   }
 
   return result;
diff --git a/src/datadog/extraction_util.h b/src/datadog/extraction_util.h
@@ -69,9 +69,12 @@ struct AuditedReader : public DictReader {
 // Combine the specified trace `contexts`, each of which was extracted in a
 // particular propagation style, into one `ExtractedData` that includes fields
 // from compatible elements of `contexts`, and return the resulting
-// `ExtractedData`. The order of the elements of `contexts` must correspond to
-// the order of the configured extraction propagation styles.
-ExtractedData merge(const std::vector<ExtractedData>& contexts);
+// `ExtractedData`. The `first_style` specifies the first configured extraction
+// propagation style that has been extracted and the other contexts will be
+// merged with it, so long as the trace-ids match.
+ExtractedData merge(
+    const PropagationStyle first_style,
+    const std::unordered_map<PropagationStyle, ExtractedData>& contexts);
 
 }  // namespace tracing
 }  // namespace datadog
diff --git a/src/datadog/tracer.cpp b/src/datadog/tracer.cpp
@@ -150,7 +150,9 @@ Expected<Span> Tracer::extract_span(const DictReader& reader,
   AuditedReader audited_reader{reader};
 
   auto span_data = std::make_unique<SpanData>();
-  std::vector<ExtractedData> extracted_contexts;
+  Optional<PropagationStyle> first_style_with_trace_id;
+  Optional<PropagationStyle> first_style_with_parent_id;
+  std::unordered_map<PropagationStyle, ExtractedData> extracted_contexts;
 
   for (const auto style : extraction_styles_) {
     using Extractor = decltype(&extract_datadog);  // function pointer
@@ -175,11 +177,33 @@ Expected<Span> Tracer::extract_span(const DictReader& reader,
       return error->with_prefix(
           extraction_error_prefix(style, audited_reader.entries_found));
     }
-    extracted_contexts.push_back(std::move(*data));
-    extracted_contexts.back().headers_examined = audited_reader.entries_found;
+
+    if (!first_style_with_trace_id && data->trace_id.has_value()) {
+      first_style_with_trace_id = style;
+    }
+
+    if (!first_style_with_parent_id && data->parent_id.has_value()) {
+      first_style_with_parent_id = style;
+    }
+
+    data->headers_examined = audited_reader.entries_found;
+    extracted_contexts.emplace(style, std::move(*data));
   }
 
-  auto merged_context = merge(extracted_contexts);
+  ExtractedData merged_context;
+  if (!first_style_with_trace_id) {
+    // Nothing extracted a trace ID. Return the first context that includes a
+    // parent ID, if any, or otherwise just return an empty `ExtractedData`.
+    // The purpose of looking for a parent ID is to allow for the error
+    // "extracted a parent ID without a trace ID," if that's what happened.
+    if (first_style_with_parent_id) {
+      auto other = extracted_contexts.find(*first_style_with_parent_id);
+      assert(other != extracted_contexts.end());
+      merged_context = other->second;
+    }
+  } else {
+    merged_context = merge(*first_style_with_trace_id, extracted_contexts);
+  }
 
   // Some information might be missing.
   // Here are the combinations considered:
diff --git a/test/test_curl.cpp b/test/test_curl.cpp
@@ -359,7 +359,7 @@ TEST_CASE("setopt failures", "[curl]") {
 #define CASE(OPTION) \
   { OPTION, #OPTION }
 
-  const auto &[which_fails, name] = GENERATE(
+  auto test_case = GENERATE(
       values<TestCase>({CASE(CURLOPT_ERRORBUFFER), CASE(CURLOPT_HEADERDATA),
                         CASE(CURLOPT_HEADERFUNCTION), CASE(CURLOPT_HTTPHEADER),
                         CASE(CURLOPT_POST), CASE(CURLOPT_POSTFIELDS),
@@ -369,17 +369,17 @@ TEST_CASE("setopt failures", "[curl]") {
 
 #undef CASE
 
-  CAPTURE(name);
+  CAPTURE(test_case.name);
   MockCurlLibrary library;
-  library.fail = which_fails;
+  library.fail = test_case.which_fails;
 
   const auto clock = default_clock;
   const auto logger = std::make_shared<NullLogger>();
   const auto client = std::make_shared<Curl>(logger, clock, library);
 
   const auto ignore = [](auto &&...) {};
   HTTPClient::URL url;
-  if (which_fails == CURLOPT_UNIX_SOCKET_PATH) {
+  if (test_case.which_fails == CURLOPT_UNIX_SOCKET_PATH) {
     url.scheme = "unix";
     url.path = "/foo/bar.sock";
   } else {
diff --git a/test/test_span.cpp b/test/test_span.cpp
@@ -633,13 +633,14 @@ TEST_CASE("injecting W3C traceparent header") {
       int sampling_priority;
       std::string expected_flags;
     };
-    const auto& [sampling_priority, expected_flags] = GENERATE(
+    auto test_case = GENERATE(
         values<TestCase>({{-1, "00"}, {0, "00"}, {1, "01"}, {2, "01"}}));
 
-    CAPTURE(sampling_priority);
-    CAPTURE(expected_flags);
+    CAPTURE(test_case.sampling_priority);
+    CAPTURE(test_case.expected_flags);
 
-    span.trace_segment().override_sampling_priority(sampling_priority);
+    span.trace_segment().override_sampling_priority(
+        test_case.sampling_priority);
 
     MockDictWriter writer;
     span.inject(writer);
@@ -649,7 +650,7 @@ TEST_CASE("injecting W3C traceparent header") {
     // The "cafebabe"s come from `expected_id`.
     const std::string expected =
         "00-000000000000000000000000cafebabe-00000000cafebabe-" +
-        expected_flags;
+        test_case.expected_flags;
     REQUIRE(found->second == expected);
   }
 }
diff --git a/test/test_tracer.cpp b/test/test_tracer.cpp
@@ -1109,6 +1109,131 @@ TEST_CASE("span extraction") {
     REQUIRE(span_tags.empty());
   }
 
+  SECTION("W3C Phase 3 support - Preferring tracecontext") {
+    // Tests behavior from system-test
+    // test_headers_tracecontext.py::test_tracestate_w3c_p_extract_datadog_w3c
+    struct TestCase {
+      int line;
+      std::string name;
+      std::string traceparent;
+      Optional<std::string> tracestate;
+      Optional<std::string> dd_trace_id;
+      Optional<std::string> dd_parent_id;
+      Optional<std::string> dd_tags;
+      Optional<std::uint64_t> expected_parent_id;
+      Optional<std::string> expected_datadog_w3c_parent_id = {};
+    };
+
+    auto test_case = GENERATE(values<TestCase>({
+        {
+            __LINE__, "identical trace info",
+            "00-11111111111111110000000000000001-000000003ade68b1-"
+            "01",                               // traceparent
+            "dd=s:2;p:000000003ade68b1,foo=1",  // tracestate
+            "1",                                // x-datadog-trace-id
+            "987654321",                        // x-datadog-parent-id
+            "_dd.p.tid=1111111111111111",       // x-datadog-tags
+            987654321,                          // expected_parent_id
+            nullopt,  // expected_datadog_w3c_parent_id,
+        },
+
+        {
+            __LINE__, "trace ids do not match",
+            "00-11111111111111110000000000000002-000000003ade68b1-"
+            "01",                               // traceparent
+            "dd=s:2;p:000000000000000a,foo=1",  // tracestate
+            "2",                                // x-datadog-trace-id
+            "10",                               // x-datadog-parent-id
+            "_dd.p.tid=2222222222222222",       // x-datadog-tags
+            10,                                 // expected_parent_id
+            nullopt,  // expected_datadog_w3c_parent_id,
+        },
+
+        {
+            __LINE__, "same trace, non-matching parent ids",
+            "00-11111111111111110000000000000003-000000003ade68b1-"
+            "01",                               // traceparent
+            "dd=s:2;p:000000000000000a,foo=1",  // tracestate
+            "3",                                // x-datadog-trace-id
+            "10",                               // x-datadog-parent-id
+            "_dd.p.tid=1111111111111111",       // x-datadog-tags
+            987654321,                          // expected_parent_id
+            "000000000000000a",  // expected_datadog_w3c_parent_id,
+        },
+
+        {
+            __LINE__, "non-matching span, missing p value",
+            "00-00000000000000000000000000000004-000000003ade68b1-"
+            "01",                // traceparent
+            "dd=s:2,foo=1",      // tracestate
+            "4",                 // x-datadog-trace-id
+            "10",                // x-datadog-parent-id
+            nullopt,             // x-datadog-tags
+            987654321,           // expected_parent_id
+            "000000000000000a",  // expected_datadog_w3c_parent_id,
+        },
+
+        {
+            __LINE__, "non-matching span, non-matching p value",
+            "00-00000000000000000000000000000005-000000003ade68b1-"
+            "01",                               // traceparent
+            "dd=s:2;p:8fffffffffffffff,foo=1",  // tracestate
+            "5",                                // x-datadog-trace-id
+            "10",                               // x-datadog-parent-id
+            nullopt,                            // x-datadog-tags
+            987654321,                          // expected_parent_id
+            "8fffffffffffffff",  // expected_datadog_w3c_parent_id,
+        },
+    }));
+
+    CAPTURE(test_case.name);
+    CAPTURE(test_case.line);
+    CAPTURE(test_case.traceparent);
+    CAPTURE(test_case.tracestate);
+    CAPTURE(test_case.dd_trace_id);
+    CAPTURE(test_case.dd_parent_id);
+    CAPTURE(test_case.dd_tags);
+
+    const auto collector = std::make_shared<MockCollector>();
+    const auto logger = std::make_shared<MockLogger>();
+
+    TracerConfig config;
+    config.collector = collector;
+    config.logger = logger;
+    config.service = "service1";
+    config.delegate_trace_sampling = false;
+    std::vector<PropagationStyle> extraction_styles{
+        PropagationStyle::DATADOG, PropagationStyle::B3, PropagationStyle::W3C};
+    config.extraction_styles = extraction_styles;
+
+    auto valid_config = finalize_config(config);
+    REQUIRE(valid_config);
+    Tracer tracer{*valid_config};
+
+    std::unordered_map<std::string, std::string> headers;
+    headers["traceparent"] = test_case.traceparent;
+    if (test_case.tracestate) {
+      headers["tracestate"] = *test_case.tracestate;
+    }
+    if (test_case.dd_trace_id) {
+      headers["x-datadog-trace-id"] = *test_case.dd_trace_id;
+    }
+    if (test_case.dd_parent_id) {
+      headers["x-datadog-parent-id"] = *test_case.dd_parent_id;
+    }
+    if (test_case.dd_tags) {
+      headers["x-datadog-tags"] = *test_case.dd_tags;
+    }
+    MockDictReader reader{headers};
+
+    const auto span = tracer.extract_span(reader);
+    REQUIRE(span);
+
+    REQUIRE(span->parent_id() == test_case.expected_parent_id);
+    REQUIRE(span->lookup_tag(tags::internal::w3c_parent_id) ==
+            test_case.expected_datadog_w3c_parent_id);
+  }
+
   SECTION("_dd.parent_id") {
     auto finalized_config = finalize_config(config);
     REQUIRE(finalized_config);
