Force JNI path for context writes in OTEL mode

On Java < 17, ThreadContext.put() uses the ByteBuffer path which writes
directly to the Context struct, bypassing ContextApi::set(). In OTEL
mode this means the OTEP record is never written.

Fix: use JNI when otelMode is active (production path). For tests that
restart the profiler with different storage modes in the same JVM,
add writeOtelContext/readOtelContext static helpers that bypass the
cached ThreadContext and go through ContextApi JNI directly.

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>

Author: jbachorik

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

@@ -572,6 +572,11 @@ Java_com_datadoghq_profiler_ThreadContext_setContext0(JNIEnv* env, jclass unused
   return Contexts::checksum(spanId, rootSpanId);
 }
 
+extern "C" DLLEXPORT void JNICALL
+Java_com_datadoghq_profiler_ThreadContext_setContextOtel0(JNIEnv* env, jclass unused, jlong traceIdHigh, jlong traceIdLow, jlong spanId) {
+  ContextApi::setOtel(traceIdHigh, traceIdLow, spanId);
+}
+
 // Legacy API: writes directly to Context.tags[] regardless of storage mode.
 // In OTEL mode, writeCurrentContext() in flightRecorder.cpp reads both Context.tags[]
 // and the OTEP attrs_data, so tags set via this path are still recorded in JFR.

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

@@ -161,6 +161,19 @@ public long put(long spanId, long rootSpanId) {
         return (useJNI || otelMode) ? setContext0(spanId, rootSpanId) : putContextJava(spanId, rootSpanId);
     }
 
+    /**
+     * Sets trace context with full 128-bit W3C trace ID support.
+     * In OTEL mode, all three values are written to the OTEP record.
+     * In profiler mode, traceIdHigh is ignored (Datadog uses 64-bit trace IDs).
+     *
+     * @param spanId The span ID
+     * @param traceIdHigh Upper 64 bits of the 128-bit trace ID
+     * @param traceIdLow Lower 64 bits of the 128-bit trace ID (rootSpanId in Datadog)
+     */
+    public void put(long spanId, long traceIdHigh, long traceIdLow) {
+        setContextOtel0(traceIdHigh, traceIdLow, spanId);
+    }
+
     public void putCustom(int offset, int value) {
         if (offset >= MAX_CUSTOM_SLOTS) {
             throw new IllegalArgumentException("Invalid offset: " + offset + " (max " +  MAX_CUSTOM_SLOTS + ")");
@@ -215,6 +228,7 @@ public boolean setContextAttribute(int keyIndex, String value) {
     }
 
     private static native long setContext0(long spanId, long rootSpanId);
+    private static native void setContextOtel0(long traceIdHigh, long traceIdLow, long spanId);
     private static native void setContextSlot0(int offset, int value);
     private static native boolean setContextAttribute0(int keyIndex, String value);
 
@@ -245,4 +259,16 @@ public boolean setContextAttribute(int keyIndex, String value) {
     public static long[] readOtelContext() {
         return getContext0();
     }
+
+    /**
+     * Writes context via the native ContextApi, bypassing the cached otelMode flag.
+     * Primarily intended for testing where the profiler may be restarted with a
+     * different storage mode while the ThreadContext instance is cached in ThreadLocal.
+     *
+     * @param spanId The span ID
+     * @param rootSpanId The root span ID (trace ID low bits for OTEL)
+     */
+    public static void writeOtelContext(long spanId, long rootSpanId) {
+        setContext0(spanId, rootSpanId);
+    }
 }

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

@@ -96,14 +96,13 @@ public void testOtelStorageModeContext() throws Exception {
         profiler.execute(String.format("start,cpu=1ms,ctxstorage=otel,jfr,file=%s", jfrFile.toAbsolutePath()));
         profilerStarted = true;
 
-        // Set context — writes to OTEP #4947 TLS record
+        // Write and read context directly via JNI, bypassing cached ThreadContext.
+        // ThreadContext caches otelMode at construction time; if a profiler-mode test
+        // ran first on this thread, the cached flag would route writes to the wrong storage.
         long spanId = 0xAAAABBBBCCCCDDDDL;
         long rootSpanId = 0x1111222233334444L;
-        profiler.setContext(spanId, rootSpanId);
+        ThreadContext.writeOtelContext(spanId, rootSpanId);
 
-        // Verify context round-trips through JNI, bypassing cached ThreadContext
-        // (ThreadContext caches otelMode at construction time; if a profiler-mode test
-        // ran first on this thread, the cached flag would be wrong)
         long[] ctx = ThreadContext.readOtelContext();
         assertNotNull(ctx, "Context should be readable in OTEL mode");
         assertEquals(spanId, ctx[0], "SpanId should match in OTEL mode");
@@ -121,7 +120,7 @@ public void testOtelModeStartsOnAnyPlatform() throws Exception {
         profilerStarted = true;
 
         // Context operations should not crash
-        profiler.setContext(0x123L, 0x456L);
+        ThreadContext.writeOtelContext(0x123L, 0x456L);
     }
 
     /**
@@ -134,11 +133,9 @@ public void testOtelModeClearContext() throws Exception {
         profiler.execute(String.format("start,cpu=1ms,ctxstorage=otel,jfr,file=%s", jfrFile.toAbsolutePath()));
         profilerStarted = true;
 
-        // Set context
-        profiler.setContext(0xDEADBEEFL, 0xCAFEBABEL);
-
-        // Clear context
-        profiler.setContext(0, 0);
+        // Set and clear context via direct JNI (bypasses cached ThreadContext)
+        ThreadContext.writeOtelContext(0xDEADBEEFL, 0xCAFEBABEL);
+        ThreadContext.writeOtelContext(0, 0);
 
         // Verify context is cleared (use direct JNI read, not cached ThreadContext)
         long[] ctx = ThreadContext.readOtelContext();
@@ -164,10 +161,10 @@ public void testOtelModeCustomAttributes() throws Exception {
         // Register attribute keys
         OTelContext.getInstance().registerAttributeKeys("http.route", "db.system");
 
-        // Set context first (trace/span)
+        // Set context first (trace/span) via direct JNI
         long spanId = 0xAAAABBBBCCCCDDDDL;
         long rootSpanId = 0x1111222233334444L;
-        profiler.setContext(spanId, rootSpanId);
+        ThreadContext.writeOtelContext(spanId, rootSpanId);
 
         // Set custom attribute via new API
         ThreadContext ctx = profiler.getThreadContext();
@@ -214,7 +211,7 @@ public void testOtelModeAttributeOverflow() throws Exception {
 
         OTelContext.getInstance().registerAttributeKeys("k0", "k1", "k2", "k3", "k4");
 
-        profiler.setContext(0x1L, 0x2L);
+        ThreadContext.writeOtelContext(0x1L, 0x2L);
 
         ThreadContext ctx = profiler.getThreadContext();