Use custom labels library for trace correlation with OTel eBPF profiler

BUILD.bazel

@@ -1,3 +1,8 @@
+config_setting(
+    name              = "is_linux",
+    constraint_values = ["@platforms//os:linux"],
+)
+
 cc_library(
     name = "dd_trace_cpp",
     srcs = [
@@ -67,6 +72,7 @@ cc_library(
       "src/datadog/limiter.h",
       "src/datadog/msgpack.h",
       "src/datadog/null_logger.h",
+      "src/datadog/otel_identifiers.h",
       "src/datadog/parse_util.h",
       "src/datadog/platform_util.h",
       "src/datadog/random.h",
@@ -134,5 +140,8 @@ cc_library(
     deps = [
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:optional",
-    ],
+    ] + select({
+        ":is_linux": ["@customlabels//:customlabels"],
+        "//conditions:default": [],
+    }),
 )

CMakeLists.txt

@@ -24,6 +24,11 @@ endif ()
 
 set(DD_TRACE_TRANSPORT "curl" CACHE STRING "HTTP transport that dd-trace-cpp uses to communicate with the Datadog Agent, can be either 'none' or 'curl'")
 
+if (CMAKE_SYSTEM_NAME STREQUAL "Linux")
+  include(cmake/deps/customlabels.cmake)
+  message(STATUS "customlabels Library is available")
+endif()
+
 if(DD_TRACE_TRANSPORT STREQUAL "curl")
   include(cmake/deps/curl.cmake)
   message(STATUS "DD_TRACE_TRANSPORT is set to 'curl', including curl")
@@ -207,6 +212,13 @@ if (BUILD_SHARED_LIBS)
       dd_trace::specs
   )
 
+  if (CMAKE_SYSTEM_NAME STREQUAL "Linux")
+   target_link_libraries(dd_trace_cpp-shared
+     PUBLIC
+       customlabels
+   )
+  endif ()
+
   install(TARGETS dd_trace_cpp-objects dd_trace_cpp-shared
     EXPORT dd_trace_cpp-export
     FILE_SET public_headers DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}
@@ -249,6 +261,13 @@ if (BUILD_STATIC_LIBS)
       dd_trace::specs
   )
 
+  if (CMAKE_SYSTEM_NAME STREQUAL "Linux")
+   target_link_libraries(dd_trace_cpp-static
+     PUBLIC
+       customlabels
+   )
+  endif ()
+
   install(TARGETS dd_trace_cpp-objects dd_trace_cpp-static
     EXPORT dd_trace_cpp-export
     FILE_SET public_headers DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}
@@ -257,3 +276,10 @@ if (BUILD_STATIC_LIBS)
     RUNTIME DESTINATION ${CMAKE_INSTALL_LIBDIR}
   )
 endif ()
+
+if (CMAKE_SYSTEM_NAME STREQUAL "Linux")
+  target_include_directories(dd_trace_cpp-objects
+  PRIVATE
+    ${customlabels_SOURCE_DIR}/src
+  )
+endif()

MODULE.bazel

@@ -11,10 +11,24 @@ bazel_dep(
     name = "rules_cc",
     version = "0.0.9",
 )
+bazel_dep(
+    name = "customlabels",
+    version = "1.0.0",
+)
+bazel_dep(
+    name = "platforms",
+    version = "0.0.11",
+)
 # -- bazel_dep definitions -- #
 
 non_module_dependencies = use_extension("//:extensions.bzl", "non_module_dependencies")
 use_repo(
     non_module_dependencies,
     "com_google_absl",
 )
+
+git_override(
+    module_name = "customlabels",
+    remote = "https://github.com/DataDog/custom-labels.git",
+    commit = "1bc74232c0e414f2c63139f7cf7ac0bdce57d258",
+)

cmake/deps/customlabels.cmake

@@ -0,0 +1,38 @@
+include(FetchContent)
+
+FetchContent_Declare(customlabels
+  GIT_REPOSITORY https://github.com/DataDog/custom-labels.git
+  GIT_TAG  elsa/add-process-storage
+)
+
+FetchContent_MakeAvailable(customlabels)
+
+add_library(customlabels SHARED
+  ${customlabels_SOURCE_DIR}/src/customlabels.c
+)
+
+if (CMAKE_SYSTEM_NAME STREQUAL "Linux")
+  if (CMAKE_SYSTEM_PROCESSOR STREQUAL "aarch64")
+    set(TLS_DIALECT desc)
+  elseif (CMAKE_SYSTEM_PROCESSOR STREQUAL "x86_64" OR CMAKE_SYSTEM_PROCESSOR STREQUAL "AMD64")
+    set(TLS_DIALECT gnu2)
+  else()
+    message(FATAL_ERROR "Only aarch64 and x86-64 are supported (found: ${CMAKE_SYSTEM_PROCESSOR})")
+  endif()
+
+  include(CheckCompilerFlag)
+  check_compiler_flag(CXX "-ftls-dialect=${TLS_DIALECT}" TLS_DIALECT_OK)
+  if (TLS_DIALECT_OK)
+      target_compile_options(customlabels PRIVATE
+       -g
+       -fPIC
+       -ftls-model=global-dynamic
+       -mtls-dialect=${TLS_DIALECT}
+      )
+  endif()
+endif()
+
+target_include_directories(customlabels
+PUBLIC
+  ${customlabels_SOURCE_DIR}/src
+)

include/datadog/environment.h

@@ -50,6 +50,7 @@ namespace environment {
   MACRO(DD_TRACE_TAGS_PROPAGATION_MAX_LENGTH)        \
   MACRO(DD_VERSION)                                  \
   MACRO(DD_TRACE_128_BIT_TRACEID_GENERATION_ENABLED) \
+  MACRO(DD_TRACE_CORRELATE_FULL_HOST_PROFILES)       \
   MACRO(DD_TELEMETRY_HEARTBEAT_INTERVAL)             \
   MACRO(DD_TELEMETRY_METRICS_ENABLED)                \
   MACRO(DD_TELEMETRY_METRICS_INTERVAL_SECONDS)       \

include/datadog/tracer.h

@@ -54,6 +54,7 @@ class Tracer {
   Baggage::Options baggage_opts_;
   bool baggage_injection_enabled_;
   bool baggage_extraction_enabled_;
+  bool correlate_full_host_profiles_;
 
  public:
   // Create a tracer configured using the specified `config`, and optionally:
@@ -105,6 +106,9 @@ class Tracer {
   std::string config() const;
 
  private:
+#ifdef __linux__
+  void correlate(const Span& span);
+#endif
   void store_config();
 };
 

include/datadog/tracer_config.h

@@ -133,6 +133,14 @@ struct TracerConfig {
   // the `DD_TRACE_128_BIT_TRACEID_GENERATION_ENABLED` environment variable.
   Optional<bool> generate_128bit_trace_ids;
 
+  // `correlate_full_host_profiles` indicates whether we want to correlate
+  // traces and spans with profiles generated by the eBPF full host profiler.
+  // This correlation only works on linux, due to the eBPF-based nature of
+  // the profiler. It implies writing some process-level and thread-level
+  // data in variables which the profiler will then read from the process's
+  // memory.
+  Optional<bool> correlate_full_host_profiles;
+
   // `runtime_id` denotes the current run of the application in which the tracer
   // is embedded. If `runtime_id` is not specified, then it defaults to a
   // pseudo-randomly generated value. A server that contains multiple tracers,
@@ -197,6 +205,7 @@ class FinalizedTracerConfig final {
   std::shared_ptr<Logger> logger;
   bool log_on_startup;
   bool generate_128bit_trace_ids;
+  bool correlate_full_host_profiles;
   Optional<RuntimeID> runtime_id;
   Clock clock;
   std::string integration_name;

include/datadog/tracer_signature.h

@@ -19,6 +19,14 @@
 //    polling the Datadog Agent. See
 //    `RemoteConfigurationManager::process_response` in `remote_config.h`.
 
+#ifdef __linux__
+#include <cstring>
+#include <vector>
+extern "C" {
+#include <customlabels.h>
+}
+#endif
+
 #include <string>
 
 #include "runtime_id.h"
@@ -31,6 +39,17 @@
 namespace datadog {
 namespace tracing {
 
+#ifdef __linux__
+namespace {
+void write_utf8_string(std::vector<uint8_t>& buffer, const std::string& str) {
+  uint32_t length = str.length();
+  buffer.insert(buffer.end(), reinterpret_cast<uint8_t*>(&length),
+                reinterpret_cast<uint8_t*>(&length) + sizeof(length));
+  buffer.insert(buffer.end(), str.begin(), str.end());
+}
+}  // namespace
+#endif
+
 struct TracerSignature {
   RuntimeID runtime_id;
   std::string default_service;
@@ -47,6 +66,30 @@ struct TracerSignature {
         library_version(tracer_version),
         library_language("cpp"),
         library_language_version(DD_TRACE_STRINGIFY(__cplusplus), 6) {}
+
+#ifdef __linux__
+  // The process correlation storage contains information needed to
+  // correlate traces to profiles generated by dd-otel-host-profiler.
+  void generate_process_correlation_storage() {
+    std::vector<uint8_t> buffer;
+
+    // Currently, layout minor version is 2 to differ from Elastic's
+    // version which includes a socket path.
+    // Layout:
+    // https://github.com/elastic/apm/blob/149cd3e39a77a58002344270ed2ad35357bdd02d/specs/agents/universal-profiling-integration.md#process-storage-layout
+    uint16_t layout_minor_version = 2;
+    buffer.insert(buffer.end(),
+                  reinterpret_cast<uint8_t*>(&layout_minor_version),
+                  reinterpret_cast<uint8_t*>(&layout_minor_version) +
+                      sizeof(layout_minor_version));
+
+    write_utf8_string(buffer, default_service);
+    write_utf8_string(buffer, default_environment);
+    write_utf8_string(buffer, runtime_id.string());
+
+    proc_storage_set(buffer.data(), buffer.size());
+  }
+#endif
 };
 
 }  // namespace tracing

src/datadog/datadog_agent.cpp

@@ -185,8 +185,11 @@ DatadogAgent::~DatadogAgent() {
 Expected<void> DatadogAgent::send(
     std::vector<std::unique_ptr<SpanData>>&& spans,
     const std::shared_ptr<TraceSampler>& response_handler) {
-  std::lock_guard<std::mutex> lock(mutex_);
-  trace_chunks_.push_back(TraceChunk{std::move(spans), response_handler});
+  {
+    std::lock_guard<std::mutex> lock(mutex_);
+    trace_chunks_.push_back(TraceChunk{std::move(spans), response_handler});
+  }
+  flush();
   return nullopt;
 }
 

src/datadog/otel_identifiers.h

@@ -0,0 +1,12 @@
+#pragma once
+
+namespace datadog {
+namespace tracing {
+
+// Based on the OTel Trace API
+// https://opentelemetry.io/docs/specs/otel/trace/api
+#define OTEL_TRACE_ID_IDENTIFIER "TraceId"
+#define OTEL_SPAN_ID_IDENTIFIER "SpanId"
+
+}  // namespace tracing
+}  // namespace datadog

src/datadog/span.cpp

@@ -4,10 +4,12 @@
 #include <datadog/span_config.h>
 #include <datadog/string_view.h>
 #include <datadog/trace_segment.h>
+#include <datadog/tracer.h>
 
 #include <cassert>
 #include <string>
 
+#include "otel_identifiers.h"
 #include "span_data.h"
 #include "tags.h"
 
@@ -40,6 +42,19 @@ 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 (process_storage != nullptr && parent_id().has_value()) {
+    auto span_id = std::to_string(parent_id().value());
+    custom_labels_labelset_set(
+        custom_labels_current_set,
+        {sizeof(OTEL_SPAN_ID_IDENTIFIER),
+         reinterpret_cast<const unsigned char*>(OTEL_SPAN_ID_IDENTIFIER)},
+        {span_id.size(),
+         reinterpret_cast<const unsigned char*>(span_id.c_str())});
+  }
+#endif
+
   trace_segment_->span_finished();
 }
 

src/datadog/trace_segment.cpp

@@ -1,3 +1,9 @@
+#ifdef __linux__
+extern "C" {
+#include <customlabels.h>
+}
+#endif
+
 #include <datadog/collector.h>
 #include <datadog/dict_reader.h>
 #include <datadog/dict_writer.h>
@@ -254,6 +260,15 @@ void TraceSegment::span_finished() {
     }
   }
 
+#ifdef __linux__
+  // When all spans are finished, so is the current trace.
+  if (process_storage != nullptr) {
+    auto saved_set = custom_labels_current_set;
+    custom_labels_current_set = nullptr;
+    custom_labels_labelset_free(saved_set);
+  }
+#endif
+
   telemetry::counter::increment(metrics::tracer::trace_segments_closed);
 }