fix: address JNI bugs

Author: kaahos

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

@@ -17,16 +17,36 @@
 #include <jni.h>
 
 
+using JniFunction = void (JNICALL*)();
+using IsVirtualThreadFunction = jboolean (JNICALL*)(JNIEnv*, jobject);
+
+static constexpr jint JNI_VERSION_21_VALUE = 0x00150000;
+static constexpr int IS_VIRTUAL_THREAD_INDEX = 234;
+
+static_assert(sizeof(JniFunction) == sizeof(void*), "JNI function table entries must be pointer-sized");
+
 volatile JVMSupport::JMethodIDLoadStats JVMSupport::jmethodID_load_state = JVMSupport::No_loaded;
 Mutex JVMSupport::_initialization_lock;
 
 bool JVMSupport::isPlatformThread(JNIEnv* jni, jthread thread) {
     if (jni == nullptr || thread == nullptr) {
         return false;
     }
-    int java_version = VM::java_version();
-    return java_version > 0 &&
-        (java_version < 21 || jni->IsVirtualThread(thread) == JNI_FALSE);
+
+    jint jni_version = jni->GetVersion();
+    if (jni_version <= 0) {
+        return false;
+    }
+    if (jni_version < JNI_VERSION_21_VALUE) {
+        return true;
+    }
+
+    // IsVirtualThread is slot 234 in the JNI 21 function table. Access the
+    // standardized slot directly so this file also compiles with older headers.
+    const JniFunction* functions = reinterpret_cast<const JniFunction*>(jni->functions);
+    IsVirtualThreadFunction is_virtual_thread =
+        reinterpret_cast<IsVirtualThreadFunction>(functions[IS_VIRTUAL_THREAD_INDEX]);
+    return is_virtual_thread != nullptr && is_virtual_thread(jni, thread) == JNI_FALSE;
 }
 
 bool JVMSupport::initialize() {

ddprof-lib/src/test/cpp/jvmSupport_ut.cpp

@@ -37,8 +37,87 @@ class JvmSupportGlobalSetup {
 };
 static JvmSupportGlobalSetup jvm_support_global_setup;
 
-TEST(JvmSupportThreadClassificationTest, NullInputsFailClosed) {
-    EXPECT_FALSE(JVMSupport::isPlatformThread(nullptr, nullptr));
+class JvmSupportThreadClassificationTest : public ::testing::Test {
+protected:
+    using JniFunction = void (JNICALL*)();
+
+    static constexpr int GET_VERSION_INDEX = 4;
+    static constexpr int IS_VIRTUAL_THREAD_INDEX = 234;
+    static constexpr int FUNCTION_TABLE_SIZE = IS_VIRTUAL_THREAD_INDEX + 1;
+
+    inline static jint jni_version;
+    inline static jboolean virtual_thread;
+    inline static int is_virtual_thread_calls;
+    inline static jobject last_thread;
+
+    JniFunction function_table[FUNCTION_TABLE_SIZE]{};
+    JNIEnv jni{};
+    _jobject thread_object;
+    jthread thread = &thread_object;
+
+    static jint JNICALL getVersion(JNIEnv*) {
+        return jni_version;
+    }
+
+    static jboolean JNICALL isVirtualThread(JNIEnv*, jobject candidate) {
+        is_virtual_thread_calls++;
+        last_thread = candidate;
+        return virtual_thread;
+    }
+
+    void SetUp() override {
+        jni_version = 0x00150000;
+        virtual_thread = JNI_FALSE;
+        is_virtual_thread_calls = 0;
+        last_thread = nullptr;
+
+        function_table[GET_VERSION_INDEX] = reinterpret_cast<JniFunction>(&getVersion);
+        function_table[IS_VIRTUAL_THREAD_INDEX] = reinterpret_cast<JniFunction>(&isVirtualThread);
+        jni.functions = reinterpret_cast<const JNINativeInterface_*>(function_table);
+    }
+};
+
+TEST_F(JvmSupportThreadClassificationTest, NullInputsFailClosed) {
+    EXPECT_FALSE(JVMSupport::isPlatformThread(nullptr, thread));
+    EXPECT_FALSE(JVMSupport::isPlatformThread(&jni, nullptr));
+}
+
+TEST_F(JvmSupportThreadClassificationTest, InvalidJniVersionFailsClosed) {
+    jni_version = 0;
+
+    EXPECT_FALSE(JVMSupport::isPlatformThread(&jni, thread));
+    EXPECT_EQ(0, is_virtual_thread_calls);
+}
+
+TEST_F(JvmSupportThreadClassificationTest, PreJni21ThreadIsPlatform) {
+    jni_version = 0x000a0000;
+    function_table[IS_VIRTUAL_THREAD_INDEX] = nullptr;
+
+    EXPECT_TRUE(JVMSupport::isPlatformThread(&jni, thread));
+    EXPECT_EQ(0, is_virtual_thread_calls);
+}
+
+TEST_F(JvmSupportThreadClassificationTest, Jni21PlatformThreadIsAccepted) {
+    virtual_thread = JNI_FALSE;
+
+    EXPECT_TRUE(JVMSupport::isPlatformThread(&jni, thread));
+    EXPECT_EQ(1, is_virtual_thread_calls);
+    EXPECT_EQ(thread, last_thread);
+}
+
+TEST_F(JvmSupportThreadClassificationTest, Jni21VirtualThreadIsRejected) {
+    virtual_thread = JNI_TRUE;
+
+    EXPECT_FALSE(JVMSupport::isPlatformThread(&jni, thread));
+    EXPECT_EQ(1, is_virtual_thread_calls);
+    EXPECT_EQ(thread, last_thread);
+}
+
+TEST_F(JvmSupportThreadClassificationTest, MissingJni21FunctionFailsClosed) {
+    function_table[IS_VIRTUAL_THREAD_INDEX] = nullptr;
+
+    EXPECT_FALSE(JVMSupport::isPlatformThread(&jni, thread));
+    EXPECT_EQ(0, is_virtual_thread_calls);
 }
 
 // ---------------------------------------------------------------------------

ddprof-lib/src/test/cpp/nativeSocketInterposer_ut.cpp

@@ -592,9 +592,10 @@ TEST_F(NativeSocketInterposerHookTest, CloseForwardsAndPreservesErrno) {
   int fds[2];
   ASSERT_EQ(0, pipe(fds));
   g_close_ret = 0;
+  g_close_errno = E2BIG;
 
   LibraryPatcher::_socket_active.store(true, std::memory_order_release);
-  errno = E2BIG;
+  errno = ERANGE;
   int ret = NativeSocketInterposer::close_hook(fds[0]);
 
   EXPECT_EQ(0, ret);

ddprof-test/src/test/java/com/datadoghq/profiler/AbstractProfilerTest.java

@@ -1,3 +1,8 @@
+/*
+ * Copyright 2026, Datadog, Inc.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
 package com.datadoghq.profiler;
 
 import java.nio.file.Files;
@@ -329,6 +334,7 @@ protected void runTests(Runnable... runnables) throws InterruptedException {
 
   public final void stopProfiler() {
     if (!stopped) {
+      profiler.clearTraceContext();
       profiler.stop();
       profiler.resetThreadContext();
       stopped = true;
@@ -504,4 +510,4 @@ public long getRecordedCounterValue(String counterName) {
     }
     return -1;
   }
-}
+}

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

@@ -63,6 +63,7 @@ public static void setup() throws IOException {
     @AfterEach
     public void cleanup() {
         if (profilerStarted) {
+            profiler.clearTraceContext();
             profiler.stop();
             profiler.resetThreadContext();
             profilerStarted = false;

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

@@ -45,6 +45,7 @@ public static void setup() throws IOException {
     @AfterEach
     public void cleanup() {
         if (profilerStarted) {
+            profiler.clearTraceContext();
             profiler.stop();
             profiler.resetThreadContext();
             profilerStarted = false;