GetThreadLocalStorage not working for vthreads mounted during JVMTI attach

ADDITIONAL SYSTEM INFORMATION :
Ubunu 22.10 x86_64
openjdk version "21-ea" 2023-09-19
OpenJDK Runtime Environment (build 21-ea+28-2377)
OpenJDK 64-Bit Server VM (build 21-ea+28-2377, mixed mode, sharing)

or

openjdk version "22-ea" 2024-03-19
OpenJDK Runtime Environment (build 22-ea+3-184)
OpenJDK 64-Bit Server VM (build 22-ea+3-184, mixed mode, sharing)

A DESCRIPTION OF THE PROBLEM :
GetThreadLocalStorage/SetThreadLocalStorage for virtual threads that are mounted while an agent is loaded into a running JVM does not work as expected. A value set with SetThreadLocalStorage on such a thread cannot be retrieved with GetThreadLocalStorage after the thread is unmounted and remounted.

STEPS TO FOLLOW TO REPRODUCE THE PROBLEM :
With the attached sample and JAVA_HOME set to a JDK 21 or 22:

g++ -shared -I$JAVA_HOME/include -I$JAVA_HOME/include/linux -fPIC VThreadTls.cpp -o libVThreadTls.so
$JAVA_HOME/bin/javac VThreadTls.java

LD_LIBRARY_PATH=. $JAVA_HOME/bin/java -Djdk.attach.allowAttachSelf=true -XX:+EnableDynamicAgentLoading VThreadTls attach

With JVMTI initialized at launch the sample works as expected:

LD_LIBRARY_PATH=. $JAVA_HOME/bin/java -agentlib:VThreadTls VThreadTls


EXPECTED VERSUS ACTUAL BEHAVIOR :
EXPECTED -
Process exits with 0.
ACTUAL -
Process exits with 1 and prints errors about non-matching values

---------- BEGIN SOURCE ----------
-- VThreadTls.java -----------------------------------------------------------
import com.sun.tools.attach.VirtualMachine;

import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

public class VThreadTls {
    public static native long getTls();
    public static native void setTls(long value);
    
    public static volatile boolean attached;
    public static volatile boolean error;

    public static void main(String[] args) throws Exception {
        try (ExecutorService executorService = Executors.newVirtualThreadPerTaskExecutor()) {
            for (int threadCount = 0; threadCount < 100; threadCount++) {
                executorService.execute(() -> {
                    try {
                        while (!attached) {
                            // keep mounted
                        }
                        long threadId = Thread.currentThread().threadId();
                        setTls(threadId);
                        long mountedValue = getTls();
                        if (mountedValue != threadId) {
                            System.out.println("wrong tls value while still mounted: " + threadId + ", " + mountedValue);
                            error = true;
                            return;
                        }

                        for (int repetion = 0; repetion < 1000; repetion++) {
                            Thread.sleep(1);
                            long tlsValue = getTls();
                            if (tlsValue != threadId) {
                                System.out.println("wrong tls value after yield: " + threadId + ", " + tlsValue);
                                error = true;
                                return;
                            }
                        }
                    } catch (InterruptedException e) {
                        throw new RuntimeException(e);
                    }
                });
            }
            if (args.length == 1 && args[0].equals("attach")) {
                VirtualMachine vm = VirtualMachine.attach(String.valueOf(ProcessHandle.current().pid()));
                vm.loadAgentLibrary("VThreadTls");
            }
            Thread.sleep(500);
            attached = true;
        }
        if (error) {
            System.exit(1);
        }
    }
}
------------------------------------------------------------------------------
-- VThreadTls.cpp ------------------------------------------------------------
#include <jvmti.h>
#include <cstdlib>
#include <cstring>

namespace {
    jvmtiEnv *jvmti = nullptr;

    void checkJvmti(int code, const char* message) {
        if (code != JVMTI_ERROR_NONE) {
            printf("Error %s: %d\n", message, code);
            abort();
        }
    }
}

extern "C" JNIEXPORT jlong JNICALL Java_VThreadTls_getTls(JNIEnv* jni_env, jclass clazz) {
    void* data;
    checkJvmti(jvmti->GetThreadLocalStorage(nullptr, &data), "get tls");
    return (jlong)data;
}

extern "C" JNIEXPORT void JNICALL Java_VThreadTls_setTls(JNIEnv* jni_env, jclass clazz, jlong value) {
    checkJvmti(jvmti->SetThreadLocalStorage(nullptr, (void*)value), "get tls");
}

extern "C" JNIEXPORT jint JNICALL Agent_OnLoad(JavaVM *vm, char *options, void *reserved) {
    if (vm->GetEnv(reinterpret_cast<void **>(&jvmti), JVMTI_VERSION) != JNI_OK || !jvmti) {
        printf("Could not initialize JVMTI\n");
        abort();
    }
    jvmtiCapabilities capabilities;
    memset(&capabilities, 0, sizeof(capabilities));
    capabilities.can_support_virtual_threads = 1;
    checkJvmti(jvmti->AddCapabilities(&capabilities), "adding capabilities");
    return JVMTI_ERROR_NONE;
}

extern "C" JNIEXPORT jint JNICALL Agent_OnAttach(JavaVM *vm, char *options, void *reserved) {
    printf("attached\n");
    fflush(stdout);
    return Agent_OnLoad(vm, options, reserved);
}
------------------------------------------------------------------------------
---------- END SOURCE ----------

CUSTOMER SUBMITTED WORKAROUND :
don't use GetThreadLocalStorage/SetThreadLocalStorage for virtual threads when the agent is loaded into a running JVM.

FREQUENCY : always