8311556: GetThreadLocalStorage not working for vthreads mounted during JVMTI attach

Reviewed-by: lmesnik, cjplummer

src/hotspot/share/prims/jvmtiEnvBase.cpp

@@ -1602,13 +1602,8 @@ JvmtiEnvBase::is_in_thread_list(jint count, const jthread* list, oop jt_oop) {
 
 class VM_SetNotifyJvmtiEventsMode : public VM_Operation {
 private:
-  static bool _whitebox_used;
   bool _enable;
 
-  // This function is needed only for testing purposes to support multiple
-  // enable&disable notifyJvmti events. Otherwise, there can be only one call
-  // to enable_virtual_threads_notify_jvmti() for late binding agents. There
-  // have to be no JvmtiThreadState's and need to correct them in such a case.
   static void correct_jvmti_thread_state(JavaThread* jt) {
     oop  ct_oop = jt->threadObj();
     oop  vt_oop = jt->vthread();
@@ -1618,8 +1613,7 @@ private:
     bool virt = vt_oop != nullptr && java_lang_VirtualThread::is_instance(vt_oop);
 
     // Correct jt->jvmti_thread_state() and jt->jvmti_vthread().
-    // It was not maintained while notifyJvmti was disabled but there can be
-    // a leftover from previous cycle when notification were enabled.
+    // It was not maintained while notifyJvmti was disabled.
     if (virt) {
       jt->set_jvmti_thread_state(nullptr);  // reset jt->jvmti_thread_state()
       jt->set_jvmti_vthread(vt_oop);        // restore jt->jvmti_vthread()
@@ -1640,9 +1634,7 @@ private:
         count++;
         continue; // no need in JvmtiThreadState correction below if in transition
       }
-      if (_whitebox_used) {
-        correct_jvmti_thread_state(jt); // needed in testing environment only
-      }
+      correct_jvmti_thread_state(jt);
     }
     return count;
   }
@@ -1651,9 +1643,6 @@ public:
   VMOp_Type type() const { return VMOp_SetNotifyJvmtiEventsMode; }
   bool allow_nested_vm_operations() const { return false; }
   VM_SetNotifyJvmtiEventsMode(bool enable) : _enable(enable) {
-    if (!enable) {
-      _whitebox_used = true; // disabling is available via WhiteBox only
-    }
   }
 
   void doit() {
@@ -1664,8 +1653,6 @@ public:
   }
 };
 
-bool VM_SetNotifyJvmtiEventsMode::_whitebox_used = false;
-
 // This function is to support agents loaded into running VM.
 // Must be called in thread-in-native mode.
 bool

test/hotspot/jtreg/serviceability/jvmti/vthread/VThreadTLSTest/VThreadTLSTest.java

@@ -0,0 +1,102 @@
+/*
+ * Copyright (c) 2023, Oracle and/or its affiliates. All rights reserved.
+ * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
+ *
+ * This code is free software; you can redistribute it and/or modify it
+ * under the terms of the GNU General Public License version 2 only, as
+ * published by the Free Software Foundation.
+ *
+ * This code is distributed in the hope that it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
+ * version 2 for more details (a copy is included in the LICENSE file that
+ * accompanied this code).
+ *
+ * You should have received a copy of the GNU General Public License version
+ * 2 along with this work; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
+ *
+ * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
+ * or visit www.oracle.com if you need additional information or have any
+ * questions.
+ */
+
+/**
+ * @test
+ * @summary Verifies JVMTI GetLocalStorage/SetLocalStorage
+ * @requires vm.continuations
+ * @requires vm.jvmti
+ * @run main/othervm/native -agentlib:VThreadTLSTest VThreadTLSTest
+ */
+
+/**
+ * @test
+ * @bug 8311556
+ * @summary Verifies JVMTI GetLocalStorage/SetLocalStorage
+ * @requires vm.continuations
+ * @requires vm.jvmti
+ * @run main/othervm/native -Djdk.attach.allowAttachSelf=true -XX:+EnableDynamicAgentLoading VThreadTLSTest attach
+ */
+
+import com.sun.tools.attach.VirtualMachine;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+
+public class VThreadTLSTest {
+    static final String AGENT_LIB = "VThreadTLSTest";
+    static volatile boolean attached;
+    static volatile boolean failed;
+
+    static void log(String msg) { System.out.println(msg); }
+    static native long getTLS();
+    static native void setTLS(long value);
+
+    static void test() {
+        try {
+            while (!attached) {
+                // keep mounted
+            }
+            long threadId = Thread.currentThread().threadId();
+            setTLS(threadId);
+            long mountedValue = getTLS();
+
+            if (mountedValue != threadId) {
+                log("Error: wrong TLS value while mounted: " + threadId + ", " + mountedValue);
+                failed = true;
+                return;
+            }
+            for (int count = 0; count < 10; count++) {
+                Thread.sleep(1);
+                long tlsValue = getTLS();
+                if (tlsValue != threadId) {
+                    log("Error: wrong TLS value after yield: expected: " + threadId + " got: " + tlsValue);
+                    failed = true;
+                    return;
+                }
+            }
+        } catch (InterruptedException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    public static void main(String[] args) throws Exception {
+        try (ExecutorService execService = Executors.newVirtualThreadPerTaskExecutor()) {
+            for (int threadCount = 0; threadCount < 20; threadCount++) {
+                execService.execute(() -> test());
+            }
+            if (args.length == 1 && args[0].equals("attach")) {
+                log("loading " + AGENT_LIB + " lib");
+                VirtualMachine vm = VirtualMachine.attach(String.valueOf(ProcessHandle.current().pid()));
+                vm.loadAgentLibrary(AGENT_LIB);
+            }
+            Thread.sleep(10);
+            attached = true;
+        }
+        if (failed) {
+            throw new RuntimeException("Test FAILED: errors encountered");
+        } else {
+            log("Test passed");
+        }
+    }
+}
+

test/hotspot/jtreg/serviceability/jvmti/vthread/VThreadTLSTest/libVThreadTLSTest.cpp

@@ -0,0 +1,79 @@
+/*
+ * Copyright (c) 2023, Oracle and/or its affiliates. All rights reserved.
+ * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
+ *
+ * This code is free software; you can redistribute it and/or modify it
+ * under the terms of the GNU General Public License version 2 only, as
+ * published by the Free Software Foundation.
+ *
+ * This code is distributed in the hope that it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
+ * version 2 for more details (a copy is included in the LICENSE file that
+ * accompanied this code).
+ *
+ * You should have received a copy of the GNU General Public License version
+ * 2 along with this work; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
+ *
+ * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
+ * or visit www.oracle.com if you need additional information or have any
+ * questions.
+ */
+
+#include <cstdlib>
+#include <cstring>
+#include <jvmti.h>
+#include "jvmti_common.h"
+
+extern "C" {
+
+static jvmtiEnv *jvmti;
+
+JNIEXPORT jlong JNICALL
+Java_VThreadTLSTest_getTLS(JNIEnv* jni, jclass clazz) {
+  void* data;
+  jvmtiError err = jvmti->GetThreadLocalStorage(nullptr, &data);
+  check_jvmti_status(jni, err, "getTLS: Failed in JVMTI GetThreadLocalStorage");
+  return (jlong)data;
+}
+
+JNIEXPORT void JNICALL
+Java_VThreadTLSTest_setTLS(JNIEnv* jni, jclass clazz, jlong value) {
+  jvmtiError err = jvmti->SetThreadLocalStorage(nullptr, (void*)value);
+  check_jvmti_status(jni, err, "setTLS: Failed in JVMTI SetThreadLocalStorage");
+}
+
+jint agent_init(JavaVM *jvm, char *options, void *reserved) {
+  jvmtiCapabilities caps;
+  jvmtiError err;
+
+  if (jvm->GetEnv((void **) (&jvmti), JVMTI_VERSION) != JNI_OK) {
+    LOG("agent_init: could not initialize JVMTI\n");
+    return JNI_ERR;
+  }
+  memset(&caps, 0, sizeof(caps));
+  caps.can_support_virtual_threads = 1;
+
+  err = jvmti->AddCapabilities(&caps);
+  if (err != JVMTI_ERROR_NONE) {
+    LOG("agent_init: error in JVMTI AddCapabilities: %s (%d)\n", TranslateError(err), err);
+    return JNI_ERR;
+  }
+  return JNI_OK;
+}
+
+JNIEXPORT jint JNICALL
+Agent_OnLoad(JavaVM *jvm, char *options, void *reserved) {
+  LOG("Agent_OnLoad\n");
+  return agent_init(jvm, options, reserved);
+}
+
+JNIEXPORT jint JNICALL
+Agent_OnAttach(JavaVM *jvm, char *options, void *reserved) {
+  LOG("Agent_OnAttach\n");
+  return agent_init(jvm, options, reserved);
+}
+
+} // extern "C"
+