[PROF-14883] Publish thread context attribute keys at process start

ddprof-lib/src/main/cpp/context_api.cpp

@@ -18,15 +18,10 @@
 #include "context.h"
 #include "guards.h"
 #include "otel_context.h"
-#include "otel_process_ctx.h"
 #include "profiler.h"
 #include "thread.h"
 #include <cstring>
 
-// Reserved attribute index for local root span ID in OTEL attrs_data.
-// Only used within this translation unit; not part of the public ContextApi header.
-static const uint8_t LOCAL_ROOT_SPAN_ATTR_INDEX = 0;
-
 /**
  * Initialize context TLS for the current thread on first use.
  * Must be called with signals blocked to prevent musl TLS deadlock:
@@ -71,55 +66,3 @@ Context ContextApi::snapshot() {
     size_t numAttrs = Profiler::instance()->numContextAttributes();
     return thrd->snapshotContext(numAttrs);
 }
-
-void ContextApi::registerAttributeKeys(const char** keys, int count) {
-    // Clip to DD_TAGS_CAPACITY: that is the actual sidecar slot limit and the
-    // maximum keyIndex accepted by ThreadContext.setContextAttribute.
-    int n = count < (int)DD_TAGS_CAPACITY ? count : (int)DD_TAGS_CAPACITY;
-
-    // Build NULL-terminated key array for the process context config.
-    // Index LOCAL_ROOT_SPAN_ATTR_INDEX (0) is reserved for local_root_span_id; user keys start at index 1.
-    // otel_process_ctx_publish copies all strings, so no strdup is needed.
-    const char* key_ptrs[DD_TAGS_CAPACITY + 2]; // +1 reserved, +1 null
-    key_ptrs[LOCAL_ROOT_SPAN_ATTR_INDEX] = "datadog.local_root_span_id";
-    for (int i = 0; i < n; i++) {
-        key_ptrs[i + 1] = keys[i];
-    }
-    key_ptrs[n + 1] = nullptr;
-
-    otel_thread_ctx_config_data config = {
-        .schema_version = "tlsdesc_v1_dev",
-        .attribute_key_map = key_ptrs,
-    };
-
-#ifndef OTEL_PROCESS_CTX_NO_READ
-    otel_process_ctx_read_result read_result = otel_process_ctx_read();
-    if (read_result.success) {
-        otel_process_ctx_data data = {
-            .deployment_environment_name = read_result.data.deployment_environment_name,
-            .service_instance_id = read_result.data.service_instance_id,
-            .service_name = read_result.data.service_name,
-            .service_version = read_result.data.service_version,
-            .telemetry_sdk_language = read_result.data.telemetry_sdk_language,
-            .telemetry_sdk_version = read_result.data.telemetry_sdk_version,
-            .telemetry_sdk_name = read_result.data.telemetry_sdk_name,
-            .resource_attributes = read_result.data.resource_attributes,
-            .extra_attributes = read_result.data.extra_attributes,
-            .thread_ctx_config = &config,
-        };
-
-        otel_process_ctx_publish(&data);
-        otel_process_ctx_read_drop(&read_result);
-    }
-#endif
-}
-
-void ContextApi::registerAttributeKeys(const std::vector<std::string>& keys) {
-    // Clip to DD_TAGS_CAPACITY before materializing C string pointers.
-    size_t n = keys.size() < DD_TAGS_CAPACITY ? keys.size() : DD_TAGS_CAPACITY;
-    const char* key_ptrs[DD_TAGS_CAPACITY];
-    for (size_t i = 0; i < n; i++) {
-        key_ptrs[i] = keys[i].c_str();
-    }
-    registerAttributeKeys(key_ptrs, (int)n);
-}

ddprof-lib/src/main/cpp/context_api.h

@@ -20,8 +20,6 @@
 #include "arch.h"
 #include "context.h"
 #include <cstdint>
-#include <string>
-#include <vector>
 
 class ProfiledThread;
 
@@ -69,23 +67,6 @@ class ContextApi {
      * @return A Context struct representing the current thread's context
      */
     static Context snapshot();
-
-    /**
-     * Register attribute key names and publish them in the process context.
-     * Must be called before setAttribute().
-     * Keys beyond DD_TAGS_CAPACITY are silently clipped.
-     *
-     * @param keys Array of key name strings
-     * @param count Number of keys (clipped to DD_TAGS_CAPACITY)
-     */
-    static void registerAttributeKeys(const char** keys, int count);
-
-    /**
-     * std::vector overload of registerAttributeKeys, used from the C++ start
-     * path so that the attributes=... CLI argument auto-publishes the OTEP
-     * attribute_key_map without requiring an explicit Java-side call.
-     */
-    static void registerAttributeKeys(const std::vector<std::string>& keys);
 };
 
 #endif /* _CONTEXT_API_H */

ddprof-lib/src/main/cpp/javaApi.cpp

@@ -438,7 +438,8 @@ Java_com_datadoghq_profiler_OTelContext_setProcessCtx0(JNIEnv *env,
                                                          jstring runtime_id,
                                                          jstring service,
                                                          jstring version,
-                                                         jstring tracer_version
+                                                         jstring tracer_version,
+                                                         jobjectArray attribute_keys
                                                         ) {
   JniString env_str(env, env_data);
   JniString hostname_str(env, hostname);
@@ -449,6 +450,47 @@ Java_com_datadoghq_profiler_OTelContext_setProcessCtx0(JNIEnv *env,
 
   const char *host_name_attrs[] = {"host.name", hostname_str.c_str(), NULL};
 
+  // Build the thread context attribute_key_map published alongside the process
+  // context: index 0 is the reserved datadog.local_root_span_id slot, followed by
+  // the caller-provided keys (clipped to DD_TAGS_CAPACITY)
+  int count = (attribute_keys != nullptr) ? env->GetArrayLength(attribute_keys) : 0;
+  int n = count < (int)DD_TAGS_CAPACITY ? count : (int)DD_TAGS_CAPACITY;
+  if (count > n) {
+    Log::warn("setProcessContext: %d attribute keys requested but capacity is %d; extra keys will be ignored",
+              count, (int)DD_TAGS_CAPACITY);
+  }
+
+  const char *key_ptrs[DD_TAGS_CAPACITY + 2]; // +1 reserved slot, +1 NULL terminator
+  JniString *jni_keys[DD_TAGS_CAPACITY];
+  int built = 0;
+  key_ptrs[0] = "datadog.local_root_span_id";
+  for (int i = 0; i < n; i++) {
+    jstring jstr = (jstring)env->GetObjectArrayElement(attribute_keys, i);
+    if (jstr == nullptr) {
+      // A null key would corrupt the index mapping; abort the publish.
+      for (int j = 0; j < built; j++) delete jni_keys[j];
+      Log::warn("setProcessContext: null attribute key at index %d; skipping publish", i);
+      return;
+    }
+    jni_keys[built] = new JniString(env, jstr);
+    if (jni_keys[built]->c_str() == nullptr) {
+      // GetStringUTFChars failed (e.g. OOM); a NULL key pointer would truncate
+      // the published map mid-array, so abort the publish.
+      delete jni_keys[built];
+      for (int j = 0; j < built; j++) delete jni_keys[j];
+      Log::warn("setProcessContext: failed to read attribute key at index %d; skipping publish", i);
+      return;
+    }
+    key_ptrs[i + 1] = jni_keys[built]->c_str();
+    built++;
+  }
+  key_ptrs[n + 1] = nullptr;
+
+  otel_thread_ctx_config_data thread_ctx_config = {
+    .schema_version = "tlsdesc_v1_dev",
+    .attribute_key_map = key_ptrs,
+  };
+
   otel_process_ctx_data data = {
     .deployment_environment_name = env_str.c_str(),
     .service_instance_id = runtime_id_str.c_str(),
@@ -459,13 +501,15 @@ Java_com_datadoghq_profiler_OTelContext_setProcessCtx0(JNIEnv *env,
     .telemetry_sdk_name = "dd-trace-java",
     .resource_attributes = host_name_attrs,
     .extra_attributes = NULL,
-    .thread_ctx_config = NULL  // Set later by ContextApi::registerAttributeKeys() when keys are known
+    .thread_ctx_config = &thread_ctx_config
   };
 
   otel_process_ctx_result result = otel_process_ctx_publish(&data);
   if (!result.success) {
     Log::warn("Failed to publish process context: %s", result.error_message);
   }
+
+  for (int i = 0; i < built; i++) delete jni_keys[i];
 }
 
 extern "C" DLLEXPORT jobject JNICALL
@@ -599,35 +643,6 @@ Java_com_datadoghq_profiler_ThreadContext_registerConstant0(JNIEnv* env, jclass
   return encoding == 0 ? -1 : (jint)encoding;
 }
 
-extern "C" DLLEXPORT void JNICALL
-Java_com_datadoghq_profiler_OTelContext_registerAttributeKeys0(JNIEnv* env, jclass unused, jobjectArray keys) {
-  int count = (keys != nullptr) ? env->GetArrayLength(keys) : 0;
-  int n = count < (int)DD_TAGS_CAPACITY ? count : (int)DD_TAGS_CAPACITY;
-  if (count > n) {
-    LOG_WARN("registerAttributeKeys: %d keys requested but capacity is %d; extra keys will be ignored",
-             count, (int)DD_TAGS_CAPACITY);
-  }
-
-  const char* key_ptrs[DD_TAGS_CAPACITY];
-  JniString* jni_strings[DD_TAGS_CAPACITY];
-
-  for (int i = 0; i < n; i++) {
-    jstring jstr = (jstring)env->GetObjectArrayElement(keys, i);
-    if (jstr == nullptr) {
-      for (int j = 0; j < i; j++) delete jni_strings[j];
-      return;
-    }
-    jni_strings[i] = new JniString(env, jstr);
-    key_ptrs[i] = jni_strings[i]->c_str();
-  }
-
-  // Always call registerAttributeKeys even with n==0 so the reserved
-  // datadog.local_root_span_id key (index 0) is published in the process context.
-  ContextApi::registerAttributeKeys(key_ptrs, n);
-
-  for (int i = 0; i < n; i++) delete jni_strings[i];
-}
-
 // ---- test and debug utilities
 extern "C" DLLEXPORT void JNICALL
 Java_com_datadoghq_profiler_JavaProfiler_testlog(JNIEnv* env, jclass unused, jstring msg) {

ddprof-lib/src/main/cpp/profiler.cpp

@@ -9,7 +9,6 @@
 #include "asyncSampleMutex.h"
 #include "mallocTracer.h"
 #include "context.h"
-#include "context_api.h"
 #include "guards.h"
 #include "common.h"
 #include "counters.h"
@@ -1355,9 +1354,6 @@ Error Profiler::start(Arguments &args, bool reset) {
   JfrMetadata::reset();
   JfrMetadata::initialize(args._context_attributes);
   _num_context_attributes = args._context_attributes.size();
-  // Initialize the OTel thread context so external profilers can decode
-  // the per-thread context, including custom attributes
-  ContextApi::registerAttributeKeys(args._context_attributes);
   error = _jfr.start(args, reset);
   if (error) {
     disableEngines();

ddprof-lib/src/main/java/com/datadoghq/profiler/OTelContext.java

@@ -1,5 +1,6 @@
 package com.datadoghq.profiler;
 
+import java.util.Objects;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
 import java.util.function.Consumer;
 
@@ -145,7 +146,7 @@ public OTelContext(String libLocation, String scratchDir, Consumer<Throwable> er
      * Reads the currently published OpenTelemetry process context, if any.
      * 
      * <p>This method attempts to read back the process context that was previously
-     * published via {@link #setProcessContext(String, String, String, String, String, String)}. This is
+     * published via {@link #setProcessContext(String, String, String, String, String, String, String[])}. This is
      * primarily useful for debugging and testing purposes.
      * 
      * <p><b>Platform Support:</b> Currently only supported on Linux. On other
@@ -193,7 +194,8 @@ public ProcessContext readProcessContext() {
      *     "instance-12345",    // runtime-id
      *     "my-service",        // service
      *     "1.0.0",             // version
-     *     "3.5.0"              // tracer-version
+     *     "3.5.0",             // tracer-version
+     *     new String[] {"http.route", "db.system"}  // thread-context attribute keys
      * );
      * }</pre>
      *
@@ -215,58 +217,35 @@ public ProcessContext readProcessContext() {
      * @param tracerVersion the version of the tracer as defined by OpenTelemetry
      *                      semantic conventions (telemetry.sdk.version). Must not be null.
      *                      Examples: "3.5.0", "4.2.0"
-     *     *
+     * @param attributeKeys the thread-context attribute key names whose per-thread
+     *                      values are recorded in the OTEP thread-local record (e.g.
+     *                      "http.route", "db.system"). Published in the process context's
+     *                      thread_ctx_config as the attribute_key_map, preceded by the
+     *                      reserved datadog.local_root_span_id slot. Must not be null
+     *                      (may be empty); keys beyond capacity are clipped. If any
+     *                      element is null, the entire process context publish is
+     *                      skipped - no context is published (a warning is logged) and
+     *                      no exception is thrown. Order must match the indices used
+     *                      with {@link ThreadContext#setContextAttribute(int, String)}.
+     *
+     * @throws NullPointerException if {@code attributeKeys} is null
+     *
      * @see <a href="https://opentelemetry.io/docs/specs/semconv/registry/attributes/service/">OpenTelemetry Service Attributes</a>
      * @see <a href="https://opentelemetry.io/docs/specs/semconv/registry/attributes/deployment/">OpenTelemetry Deployment Attributes</a>
      */
-    public void setProcessContext(String env, String hostname, String runtimeId, String service, String version, String tracerVersion) {
+    public void setProcessContext(String env, String hostname, String runtimeId, String service, String version, String tracerVersion, String[] attributeKeys) {
+        Objects.requireNonNull(attributeKeys, "attributeKeys");
         if (!libraryLoadResult.succeeded) {
             return;
         }
         try {
             lock.writeLock().lock();
-            setProcessCtx0(env, hostname, runtimeId, service, version, tracerVersion);
+            setProcessCtx0(env, hostname, runtimeId, service, version, tracerVersion, attributeKeys);
         } finally {
             lock.writeLock().unlock();
         }
     }
 
-    /**
-     * Registers attribute key names for the thread context.
-     *
-     * <p>These keys define the attribute_key_map in the process context's
-     * thread_ctx_config. In OTEL mode, attribute values set via
-     * {@link ThreadContext#setContextAttribute(int, String)} are encoded
-     * with the key index corresponding to position in this array.
-     *
-     * <p>Calling this method explicitly is <b>optional</b> when the profiler
-     * is started with the {@code attributes=...} argument (e.g.
-     * {@code execute("start,attributes=http.route;db.system,...")}); the
-     * native start path auto-registers those keys.
-     *
-     * <p>This method reads the currently published process context and
-     * republishes it with thread_ctx_config attached. It must therefore be
-     * called <b>after</b>
-     * {@link #setProcessContext(String, String, String, String, String, String)};
-     * if no process context has been published yet, the registration is a
-     * no-op for the process-level attribute_key_map (per-thread
-     * setContextAttribute writes still work). Conversely, calling
-     * setProcessContext after this method drops the previously published
-     * thread_ctx_config — re-register the keys after each setProcessContext
-     * if you need them to persist.
-     *
-     * <p>Must be called before any calls to setContextAttribute.
-     *
-     * @param keys Attribute key names (e.g. "http.route", "db.system")
-     */
-    public void registerAttributeKeys(String... keys) {
-        if (!libraryLoadResult.succeeded) {
-            return;
-        }
-        registerAttributeKeys0(keys);
-    }
-
-    private static native void setProcessCtx0(String env, String hostname, String runtimeId, String service, String version, String tracerVersion);
+    private static native void setProcessCtx0(String env, String hostname, String runtimeId, String service, String version, String tracerVersion, String[] attributeKeys);
     private static native ProcessContext readProcessCtx0();
-    private static native void registerAttributeKeys0(String[] keys);
 }

ddprof-stresstest/src/jmh/java/com/datadoghq/profiler/stresstest/scenarios/throughput/ThreadContextBenchmark.java

@@ -16,7 +16,6 @@
 package com.datadoghq.profiler.stresstest.scenarios.throughput;
 
 import com.datadoghq.profiler.JavaProfiler;
-import com.datadoghq.profiler.OTelContext;
 import com.datadoghq.profiler.ThreadContext;
 import java.nio.file.Files;
 import java.nio.file.Path;
@@ -64,7 +63,6 @@ public void setup() throws Exception {
             profiler = JavaProfiler.getInstance();
             Path jfr = Files.createTempFile("bench", ".jfr");
             profiler.execute("start,cpu=10ms,attributes=http.route,jfr,file=" + jfr.toAbsolutePath());
-            OTelContext.getInstance().registerAttributeKeys("http.route");
         }
     }
 

ddprof-test/src/test/java/com/datadoghq/profiler/context/OtelContextStorageModeTest.java

@@ -16,7 +16,6 @@
 package com.datadoghq.profiler.context;
 
 import com.datadoghq.profiler.JavaProfiler;
-import com.datadoghq.profiler.OTelContext;
 import com.datadoghq.profiler.ThreadContext;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeAll;
@@ -88,9 +87,6 @@ public void testOtelModeCustomAttributes() throws Exception {
         profiler.execute(String.format("start,cpu=1ms,attributes=http.route;db.system,jfr,file=%s", jfrFile.toAbsolutePath()));
         profilerStarted = true;
 
-        // Register attribute keys
-        OTelContext.getInstance().registerAttributeKeys("http.route", "db.system");
-
         long localRootSpanId = 0x1111222233334444L;
         long spanId = 0xAAAABBBBCCCCDDDDL;
         profiler.setContext(localRootSpanId, spanId, 0L, 0x9999L);
@@ -121,8 +117,6 @@ public void testOtelModeAttributeOverflow() throws Exception {
         profiler.execute(String.format("start,cpu=1ms,attributes=k0;k1;k2;k3;k4,jfr,file=%s", jfrFile.toAbsolutePath()));
         profilerStarted = true;
 
-        OTelContext.getInstance().registerAttributeKeys("k0", "k1", "k2", "k3", "k4");
-
         profiler.setContext(0x2L, 0x1L, 0L, 0x3L);
 
         ThreadContext ctx = profiler.getThreadContext();
@@ -212,8 +206,6 @@ public void testSpanTransitionClearsAttributes() throws Exception {
         profiler.execute(String.format("start,cpu=1ms,attributes=http.route,jfr,file=%s", jfrFile.toAbsolutePath()));
         profilerStarted = true;
 
-        OTelContext.getInstance().registerAttributeKeys("http.route");
-
         // Span A: set a custom attribute
         profiler.setContext(0x1L, 0x1L, 0L, 0x1L);
         ThreadContext ctx = profiler.getThreadContext();
@@ -256,7 +248,6 @@ public void testAttributeCacheIsolation() throws Exception {
         Path jfrFile = Files.createTempFile("otel-attr-cache-iso", ".jfr");
         profiler.execute(String.format("start,cpu=1ms,attributes=attr0,jfr,file=%s", jfrFile.toAbsolutePath()));
         profilerStarted = true;
-        OTelContext.getInstance().registerAttributeKeys("attr0");
 
         final String valueA = "FB"; // hashCode = 2236, slot 188
         final String valueB = "Ea"; // hashCode = 2236, same slot