feat: add data in tls for trace to profile correlation

* BUILD.bazel
* include/datadog/tls_storage.h
* include/datadog/trace_segment.h
* include/datadog/tracer.h
* src/datadog/span.cpp
* src/datadog/trace_segment.cpp
* src/datadog/tracer.cpp

Author: elsa

BUILD.bazel

@@ -113,6 +113,7 @@ cc_library(
       "include/datadog/span_matcher.h",
       "include/datadog/span_sampler_config.h",
       "include/datadog/string_view.h",
+      "include/datadog/tls_storage.h",
       "include/datadog/tracer.h",
       "include/datadog/tracer_config.h",
       "include/datadog/tracer_signature.h",

include/datadog/tls_storage.h

@@ -0,0 +1,25 @@
+#pragma once
+
+#include <datadog/trace_id.h>
+
+#include <array>
+#include <cstdint>
+
+// Global struct used to exposed thread-specific information.
+// https://github.com/elastic/apm/blob/149cd3e39a77a58002344270ed2ad35357bdd02d/specs/agents/universal-profiling-integration.md#thread-local-storage-layout
+
+namespace datadog {
+namespace tracing {
+struct __attribute__((packed)) TLSStorage {
+  uint16_t layout_minor_version;
+  uint8_t valid;
+  uint8_t trace_present;
+  uint8_t trace_flags;
+  uint64_t trace_id_low;
+  uint64_t trace_id_high;
+  uint64_t span_id;
+  uint64_t transaction_id;
+};
+
+}  // namespace tracing
+}  // namespace datadog

include/datadog/trace_segment.h

@@ -102,6 +102,8 @@ class TraceSegment {
   const Optional<std::string>& origin() const;
   Optional<SamplingDecision> sampling_decision() const;
 
+  uint64_t local_root_id() const;
+
   Logger& logger() const;
 
   // Inject trace context for the specified `span` into the specified `writer`.

include/datadog/tracer.h

@@ -10,6 +10,10 @@
 // obtained from a `TracerConfig` via the `finalize_config` function.  See
 // `tracer_config.h`.
 
+#ifdef __linux__
+#include <datadog/tls_storage.h>
+#endif
+
 #include <cstddef>
 #include <memory>
 
@@ -25,6 +29,8 @@
 
 #ifdef __linux__
 extern const void* elastic_apm_profiling_correlation_process_storage_v1;
+extern thread_local struct datadog::tracing::TLSStorage*
+    elastic_apm_profiling_correlation_tls_v1;
 #endif
 
 namespace datadog {
@@ -109,6 +115,9 @@ class Tracer {
   std::string config() const;
 
  private:
+#ifdef __linux__
+  void correlate(const Span& span);
+#endif
   void store_config();
 };
 

src/datadog/span.cpp

@@ -4,6 +4,7 @@
 #include <datadog/span_config.h>
 #include <datadog/string_view.h>
 #include <datadog/trace_segment.h>
+#include <datadog/tracer.h>
 
 #include <cassert>
 #include <string>
@@ -40,6 +41,16 @@ Span::~Span() {
     data_->duration = now - data_->start;
   }
 
+#ifdef __linux__
+  // When a span is finished, we must update the span_id to its parent's.
+  if (elastic_apm_profiling_correlation_process_storage_v1 != nullptr &&
+      parent_id().has_value()) {
+    elastic_apm_profiling_correlation_tls_v1->valid = 0;
+    elastic_apm_profiling_correlation_tls_v1->span_id = parent_id().value();
+    elastic_apm_profiling_correlation_tls_v1->valid = 1;
+  }
+#endif
+
   trace_segment_->span_finished();
 }
 

src/datadog/trace_segment.cpp

@@ -9,6 +9,7 @@
 #include <datadog/telemetry/metrics.h>
 #include <datadog/telemetry/telemetry.h>
 #include <datadog/trace_segment.h>
+#include <datadog/tracer.h>
 
 #include <cassert>
 #include <string>
@@ -137,6 +138,8 @@ Optional<SamplingDecision> TraceSegment::sampling_decision() const {
   return sampling_decision_;
 }
 
+uint64_t TraceSegment::local_root_id() const { return spans_.front()->span_id; }
+
 Logger& TraceSegment::logger() const { return *logger_; }
 
 void TraceSegment::register_span(std::unique_ptr<SpanData> span) {
@@ -255,6 +258,11 @@ void TraceSegment::span_finished() {
   }
 
   telemetry::counter::increment(metrics::tracer::trace_segments_closed);
+
+#ifdef __linux__
+  // When all spans are finished, so is the current trace.
+  elastic_apm_profiling_correlation_tls_v1->trace_present = 0;
+#endif
 }
 
 void TraceSegment::override_sampling_priority(SamplingPriority priority) {

src/datadog/tracer.cpp

@@ -35,6 +35,10 @@
 
 #ifdef __linux__
 const void* elastic_apm_profiling_correlation_process_storage_v1 = nullptr;
+thread_local struct datadog::tracing::TLSStorage*
+    elastic_apm_profiling_correlation_tls_v1 = nullptr;
+thread_local std::unique_ptr<datadog::tracing::TLSStorage> tls_info_holder =
+    nullptr;
 #endif
 
 namespace datadog {
@@ -115,6 +119,33 @@ Tracer::Tracer(const FinalizedTracerConfig& config,
   store_config();
 }
 
+#ifdef __linux__
+void Tracer::correlate(const Span& span) {
+  // See Layout:
+  // https://github.com/elastic/apm/blob/149cd3e39a77a58002344270ed2ad35357bdd02d/specs/agents/universal-profiling-integration.md#thread-local-storage-layout
+  tls_info_holder = std::make_unique<datadog::tracing::TLSStorage>();
+  elastic_apm_profiling_correlation_tls_v1 = tls_info_holder.get();
+
+  struct TLSStorage* tls_data = elastic_apm_profiling_correlation_tls_v1;
+  tls_data->valid = 0;
+
+  tls_data->layout_minor_version = 1;
+  tls_data->trace_present = 1;  // We are in a span so no errors
+  tls_data->trace_flags =
+      span.trace_segment().sampling_decision().has_value() &&
+              (span.trace_segment().sampling_decision().value().priority > 0)
+          ? 1
+          : 0;
+  auto trace_id = span.trace_id();
+  tls_data->trace_id_low = trace_id.low;
+  tls_data->trace_id_high = trace_id.high;
+  tls_data->span_id = span.id();
+  tls_data->transaction_id = span.trace_segment().local_root_id();
+
+  tls_data->valid = 1;
+}
+#endif
+
 std::string Tracer::config() const {
   // clang-format off
   auto config = nlohmann::json::object({
@@ -206,6 +237,11 @@ Span Tracer::create_span(const SpanConfig& config) {
   Span span{span_data_ptr, segment,
             [generator = generator_]() { return generator->span_id(); },
             clock_};
+
+#ifdef __linux__
+  correlate(span);
+#endif
+
   return span;
 }
 
@@ -417,6 +453,11 @@ Expected<Span> Tracer::extract_span(const DictReader& reader,
   Span span{span_data_ptr, segment,
             [generator = generator_]() { return generator->span_id(); },
             clock_};
+
+#ifdef __linux__
+  correlate(span);
+#endif
+
   return span;
 }