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. Add JavaProfiler.resetThreadContext()
to flush the cached ThreadLocal so tests that switch storage modes between
profiler restarts get a fresh ThreadContext with the correct otelMode flag.

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/JavaProfiler.java

@@ -360,4 +360,13 @@ public ThreadContext getThreadContext() {
     public static native void testlog(String msg);
 
     public static native void dumpContext();
+
+    /**
+     * Resets the cached ThreadContext for the current thread.
+     * The next call to {@link #getThreadContext()} or {@link #setContext(long, long)}
+     * will re-create it, picking up the current storage mode.
+     */
+    public void resetThreadContext() {
+        tlsContextStorage.remove();
+    }
 }

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);
 
@@ -234,15 +248,4 @@ public boolean setContextAttribute(int keyIndex, String value) {
      * @return array with [spanId, rootSpanId], or null on error
      */
     private static native long[] getContext0();
-
-    /**
-     * Reads 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.
-     *
-     * @return array with [spanId, rootSpanId], or null on error
-     */
-    public static long[] readOtelContext() {
-        return getContext0();
-    }
 }

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

@@ -96,18 +96,16 @@ 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
+        // Reset cached ThreadContext so it picks up the current OTEL mode
+        profiler.resetThreadContext();
+
         long spanId = 0xAAAABBBBCCCCDDDDL;
         long rootSpanId = 0x1111222233334444L;
         profiler.setContext(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");
-        assertEquals(rootSpanId, ctx[1], "RootSpanId should match in OTEL mode");
+        ThreadContext ctx = profiler.getThreadContext();
+        assertEquals(spanId, ctx.getSpanId(), "SpanId should match in OTEL mode");
+        assertEquals(rootSpanId, ctx.getRootSpanId(), "RootSpanId should match in OTEL mode");
     }
 
     /**
@@ -120,6 +118,8 @@ public void testOtelModeStartsOnAnyPlatform() throws Exception {
         profiler.execute(String.format("start,cpu=1ms,ctxstorage=otel,jfr,file=%s", jfrFile.toAbsolutePath()));
         profilerStarted = true;
 
+        profiler.resetThreadContext();
+
         // Context operations should not crash
         profiler.setContext(0x123L, 0x456L);
     }
@@ -134,19 +134,14 @@ 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);
+        profiler.resetThreadContext();
 
-        // Clear context
+        profiler.setContext(0xDEADBEEFL, 0xCAFEBABEL);
         profiler.setContext(0, 0);
 
-        // Verify context is cleared (use direct JNI read, not cached ThreadContext)
-        long[] ctx = ThreadContext.readOtelContext();
-        // After clear, readContextDirect returns zeros (OTEL record is invalidated)
-        long clearedSpanId = (ctx != null) ? ctx[0] : 0;
-        long clearedRootSpanId = (ctx != null) ? ctx[1] : 0;
-        assertEquals(0, clearedSpanId, "SpanId should be zero after clear");
-        assertEquals(0, clearedRootSpanId, "RootSpanId should be zero after clear");
+        ThreadContext ctx = profiler.getThreadContext();
+        assertEquals(0, ctx.getSpanId(), "SpanId should be zero after clear");
+        assertEquals(0, ctx.getRootSpanId(), "RootSpanId should be zero after clear");
     }
 
     /**
@@ -164,24 +159,22 @@ public void testOtelModeCustomAttributes() throws Exception {
         // Register attribute keys
         OTelContext.getInstance().registerAttributeKeys("http.route", "db.system");
 
-        // Set context first (trace/span)
+        profiler.resetThreadContext();
+
         long spanId = 0xAAAABBBBCCCCDDDDL;
         long rootSpanId = 0x1111222233334444L;
         profiler.setContext(spanId, rootSpanId);
 
-        // Set custom attribute via new API
         ThreadContext ctx = profiler.getThreadContext();
         boolean result = ctx.setContextAttribute(0, "GET /api/users");
         assertTrue(result, "setContextAttribute should succeed");
 
         result = ctx.setContextAttribute(1, "postgresql");
         assertTrue(result, "setContextAttribute for second key should succeed");
 
-        // Verify trace context is still intact (use direct JNI read)
-        long[] context = ThreadContext.readOtelContext();
-        assertNotNull(context, "Context should be readable after setAttribute");
-        assertEquals(spanId, context[0], "SpanId should match after setAttribute");
-        assertEquals(rootSpanId, context[1], "RootSpanId should match after setAttribute");
+        // Verify trace context is still intact
+        assertEquals(spanId, ctx.getSpanId(), "SpanId should match after setAttribute");
+        assertEquals(rootSpanId, ctx.getRootSpanId(), "RootSpanId should match after setAttribute");
     }
 
     /**
@@ -214,6 +207,8 @@ public void testOtelModeAttributeOverflow() throws Exception {
 
         OTelContext.getInstance().registerAttributeKeys("k0", "k1", "k2", "k3", "k4");
 
+        profiler.resetThreadContext();
+
         profiler.setContext(0x1L, 0x2L);
 
         ThreadContext ctx = profiler.getThreadContext();