Ensure jvmti agents don't share ThreadLocalStorage

Bug: 63665647
Test: ./test.py --host -j40

Change-Id: Iea33cca5b708f60390b8c79462ca991363ad33a2
diff --git a/runtime/openjdkjvmti/OpenjdkJvmTi.cc b/runtime/openjdkjvmti/OpenjdkJvmTi.cc
index 63892dd..3c1311b 100644
--- a/runtime/openjdkjvmti/OpenjdkJvmTi.cc
+++ b/runtime/openjdkjvmti/OpenjdkJvmTi.cc
@@ -1498,10 +1498,11 @@
 
   static jvmtiError DisposeEnvironment(jvmtiEnv* env) {
     ENSURE_VALID_ENV(env);
-    gEventHandler.RemoveArtJvmTiEnv(ArtJvmTiEnv::AsArtJvmTiEnv(env));
-    art::Runtime::Current()->RemoveSystemWeakHolder(
-        ArtJvmTiEnv::AsArtJvmTiEnv(env)->object_tag_table.get());
-    delete env;
+    ArtJvmTiEnv* tienv = ArtJvmTiEnv::AsArtJvmTiEnv(env);
+    gEventHandler.RemoveArtJvmTiEnv(tienv);
+    art::Runtime::Current()->RemoveSystemWeakHolder(tienv->object_tag_table.get());
+    ThreadUtil::RemoveEnvironment(tienv);
+    delete tienv;
     return OK;
   }
 
@@ -1671,6 +1672,7 @@
 }
 
 extern const jvmtiInterface_1 gJvmtiInterface;
+
 ArtJvmTiEnv::ArtJvmTiEnv(art::JavaVMExt* runtime, EventHandler* event_handler)
     : art_vm(runtime),
       local_data(nullptr),
diff --git a/runtime/openjdkjvmti/ti_thread.cc b/runtime/openjdkjvmti/ti_thread.cc
index 3d447dc..d1cee9a 100644
--- a/runtime/openjdkjvmti/ti_thread.cc
+++ b/runtime/openjdkjvmti/ti_thread.cc
@@ -159,6 +159,17 @@
   return ERR(NONE);
 }
 
+static art::Thread* GetNativeThreadLocked(jthread thread,
+                                          const art::ScopedObjectAccessAlreadyRunnable& soa)
+    REQUIRES_SHARED(art::Locks::mutator_lock_)
+    REQUIRES(art::Locks::thread_list_lock_) {
+  if (thread == nullptr) {
+    return art::Thread::Current();
+  }
+
+  return art::Thread::FromManagedThread(soa, thread);
+}
+
 // Get the native thread. The spec says a null object denotes the current thread.
 static art::Thread* GetNativeThread(jthread thread,
                                     const art::ScopedObjectAccessAlreadyRunnable& soa)
@@ -495,40 +506,82 @@
   return ERR(NONE);
 }
 
-jvmtiError ThreadUtil::SetThreadLocalStorage(jvmtiEnv* env ATTRIBUTE_UNUSED,
-                                             jthread thread,
-                                             const void* data) {
-  art::ScopedObjectAccess soa(art::Thread::Current());
-  art::Thread* self = GetNativeThread(thread, soa);
-  if (self == nullptr && thread == nullptr) {
+// The struct that we store in the art::Thread::custom_tls_ that maps the jvmtiEnvs to the data
+// stored with that thread. This is needed since different jvmtiEnvs are not supposed to share TLS
+// data but we only have a single slot in Thread objects to store data.
+struct JvmtiGlobalTLSData {
+  std::unordered_map<jvmtiEnv*, const void*> data GUARDED_BY(art::Locks::thread_list_lock_);
+};
+
+static void RemoveTLSData(art::Thread* target, void* ctx) REQUIRES(art::Locks::thread_list_lock_) {
+  jvmtiEnv* env = reinterpret_cast<jvmtiEnv*>(ctx);
+  art::Locks::thread_list_lock_->AssertHeld(art::Thread::Current());
+  JvmtiGlobalTLSData* global_tls = reinterpret_cast<JvmtiGlobalTLSData*>(target->GetCustomTLS());
+  if (global_tls != nullptr) {
+    global_tls->data.erase(env);
+  }
+}
+
+void ThreadUtil::RemoveEnvironment(jvmtiEnv* env) {
+  art::Thread* self = art::Thread::Current();
+  art::MutexLock mu(self, *art::Locks::thread_list_lock_);
+  art::ThreadList* list = art::Runtime::Current()->GetThreadList();
+  list->ForEach(RemoveTLSData, env);
+}
+
+jvmtiError ThreadUtil::SetThreadLocalStorage(jvmtiEnv* env, jthread thread, const void* data) {
+  art::Thread* self = art::Thread::Current();
+  art::ScopedObjectAccess soa(self);
+  art::MutexLock mu(self, *art::Locks::thread_list_lock_);
+  art::Thread* target = GetNativeThreadLocked(thread, soa);
+  if (target == nullptr && thread == nullptr) {
     return ERR(INVALID_THREAD);
   }
-  if (self == nullptr) {
+  if (target == nullptr) {
     return ERR(THREAD_NOT_ALIVE);
   }
 
-  self->SetCustomTLS(data);
+  JvmtiGlobalTLSData* global_tls = reinterpret_cast<JvmtiGlobalTLSData*>(target->GetCustomTLS());
+  if (global_tls == nullptr) {
+    target->SetCustomTLS(new JvmtiGlobalTLSData);
+    global_tls = reinterpret_cast<JvmtiGlobalTLSData*>(target->GetCustomTLS());
+  }
+
+  global_tls->data[env] = data;
 
   return ERR(NONE);
 }
 
-jvmtiError ThreadUtil::GetThreadLocalStorage(jvmtiEnv* env ATTRIBUTE_UNUSED,
+jvmtiError ThreadUtil::GetThreadLocalStorage(jvmtiEnv* env,
                                              jthread thread,
                                              void** data_ptr) {
   if (data_ptr == nullptr) {
     return ERR(NULL_POINTER);
   }
 
-  art::ScopedObjectAccess soa(art::Thread::Current());
-  art::Thread* self = GetNativeThread(thread, soa);
-  if (self == nullptr && thread == nullptr) {
+  art::Thread* self = art::Thread::Current();
+  art::ScopedObjectAccess soa(self);
+  art::MutexLock mu(self, *art::Locks::thread_list_lock_);
+  art::Thread* target = GetNativeThreadLocked(thread, soa);
+  if (target == nullptr && thread == nullptr) {
     return ERR(INVALID_THREAD);
   }
-  if (self == nullptr) {
+  if (target == nullptr) {
     return ERR(THREAD_NOT_ALIVE);
   }
 
-  *data_ptr = const_cast<void*>(self->GetCustomTLS());
+  JvmtiGlobalTLSData* global_tls = reinterpret_cast<JvmtiGlobalTLSData*>(target->GetCustomTLS());
+  if (global_tls == nullptr) {
+    *data_ptr = nullptr;
+    return OK;
+  }
+  auto it = global_tls->data.find(env);
+  if (it != global_tls->data.end()) {
+    *data_ptr = const_cast<void*>(it->second);
+  } else {
+    *data_ptr = nullptr;
+  }
+
   return ERR(NONE);
 }
 
diff --git a/runtime/openjdkjvmti/ti_thread.h b/runtime/openjdkjvmti/ti_thread.h
index 57967eb..d07dc06 100644
--- a/runtime/openjdkjvmti/ti_thread.h
+++ b/runtime/openjdkjvmti/ti_thread.h
@@ -54,6 +54,9 @@
   // To be called when it is safe to cache data.
   static void CacheData();
 
+  // Handle a jvmtiEnv going away.
+  static void RemoveEnvironment(jvmtiEnv* env);
+
   static jvmtiError GetAllThreads(jvmtiEnv* env, jint* threads_count_ptr, jthread** threads_ptr);
 
   static jvmtiError GetCurrentThread(jvmtiEnv* env, jthread* thread_ptr);
diff --git a/runtime/thread.h b/runtime/thread.h
index ad35ef2..e1102ed 100644
--- a/runtime/thread.h
+++ b/runtime/thread.h
@@ -1182,11 +1182,11 @@
     return debug_disallow_read_barrier_;
   }
 
-  const void* GetCustomTLS() const {
+  void* GetCustomTLS() const REQUIRES(Locks::thread_list_lock_) {
     return custom_tls_;
   }
 
-  void SetCustomTLS(const void* data) {
+  void SetCustomTLS(void* data) REQUIRES(Locks::thread_list_lock_) {
     custom_tls_ = data;
   }
 
@@ -1683,7 +1683,7 @@
 
   // Custom TLS field that can be used by plugins.
   // TODO: Generalize once we have more plugins.
-  const void* custom_tls_;
+  void* custom_tls_;
 
   // True if the thread is allowed to call back into java (for e.g. during class resolution).
   // By default this is true.
diff --git a/test/1909-per-agent-tls/agent_tls.cc b/test/1909-per-agent-tls/agent_tls.cc
new file mode 100644
index 0000000..14c82e3
--- /dev/null
+++ b/test/1909-per-agent-tls/agent_tls.cc
@@ -0,0 +1,75 @@
+/*
+ * Copyright (C) 2013 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "jvmti.h"
+
+// Test infrastructure
+#include "jvmti_helper.h"
+#include "scoped_local_ref.h"
+#include "test_env.h"
+
+namespace art {
+namespace Test1909AgentTLS {
+
+extern "C" JNIEXPORT void JNICALL Java_art_Test1909_setTLS(JNIEnv* env,
+                                                           jclass,
+                                                           jlong jvmti_env_ptr,
+                                                           jthread thr,
+                                                           jlong data) {
+  JvmtiErrorToException(env,
+                        reinterpret_cast<jvmtiEnv*>(jvmti_env_ptr),
+                        reinterpret_cast<jvmtiEnv*>(jvmti_env_ptr)->SetThreadLocalStorage(
+                            thr, reinterpret_cast<const void*>(static_cast<intptr_t>(data))));
+}
+
+extern "C" JNIEXPORT jlong JNICALL Java_art_Test1909_getTLS(JNIEnv* env,
+                                                            jclass,
+                                                            jlong jvmti_env_ptr,
+                                                            jthread thr) {
+  void* res = nullptr;
+  JvmtiErrorToException(
+      env,
+      reinterpret_cast<jvmtiEnv*>(jvmti_env_ptr),
+      reinterpret_cast<jvmtiEnv*>(jvmti_env_ptr)->GetThreadLocalStorage(thr, &res));
+  return static_cast<jlong>(reinterpret_cast<intptr_t>(res));
+}
+
+extern "C" JNIEXPORT void Java_art_Test1909_destroyJvmtiEnv(JNIEnv* env,
+                                                            jclass,
+                                                            jlong jvmti_env_ptr) {
+  JvmtiErrorToException(env,
+                        jvmti_env,
+                        reinterpret_cast<jvmtiEnv*>(jvmti_env_ptr)->DisposeEnvironment());
+}
+
+extern "C" JNIEXPORT jlong Java_art_Test1909_newJvmtiEnv(JNIEnv* env, jclass) {
+  JavaVM* vm = nullptr;
+  if (env->GetJavaVM(&vm) != 0) {
+    ScopedLocalRef<jclass> rt_exception(env, env->FindClass("java/lang/RuntimeException"));
+    env->ThrowNew(rt_exception.get(), "Unable to get JavaVM");
+    return -1;
+  }
+  jvmtiEnv* new_env = nullptr;
+  if (vm->GetEnv(reinterpret_cast<void**>(&new_env), JVMTI_VERSION_1_0) != 0) {
+    ScopedLocalRef<jclass> rt_exception(env, env->FindClass("java/lang/RuntimeException"));
+    env->ThrowNew(rt_exception.get(), "Unable to create new jvmtiEnv");
+    return -1;
+  }
+  return static_cast<jlong>(reinterpret_cast<intptr_t>(new_env));
+}
+
+}  // namespace Test1909AgentTLS
+}  // namespace art
diff --git a/test/1909-per-agent-tls/expected.txt b/test/1909-per-agent-tls/expected.txt
new file mode 100644
index 0000000..386f3d2
--- /dev/null
+++ b/test/1909-per-agent-tls/expected.txt
@@ -0,0 +1 @@
+Test passed
diff --git a/test/1909-per-agent-tls/info.txt b/test/1909-per-agent-tls/info.txt
new file mode 100644
index 0000000..00acefd
--- /dev/null
+++ b/test/1909-per-agent-tls/info.txt
@@ -0,0 +1 @@
+Tests jvmti behavior of GetThreadLocalStorage with multiple threads.
diff --git a/test/1909-per-agent-tls/run b/test/1909-per-agent-tls/run
new file mode 100755
index 0000000..c6e62ae
--- /dev/null
+++ b/test/1909-per-agent-tls/run
@@ -0,0 +1,17 @@
+#!/bin/bash
+#
+# Copyright 2016 The Android Open Source Project
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+./default-run "$@" --jvmti
diff --git a/test/1909-per-agent-tls/src/Main.java b/test/1909-per-agent-tls/src/Main.java
new file mode 100644
index 0000000..befebea
--- /dev/null
+++ b/test/1909-per-agent-tls/src/Main.java
@@ -0,0 +1,21 @@
+/*
+ * Copyright (C) 2016 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+public class Main {
+  public static void main(String[] args) throws Exception {
+    art.Test1909.run();
+  }
+}
diff --git a/test/1909-per-agent-tls/src/art/Main.java b/test/1909-per-agent-tls/src/art/Main.java
new file mode 100644
index 0000000..aa5498b
--- /dev/null
+++ b/test/1909-per-agent-tls/src/art/Main.java
@@ -0,0 +1,32 @@
+/*
+ * Copyright (C) 2017 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package art;
+
+// Binder class so the agent's C code has something that can be bound and exposed to tests.
+// In a package to separate cleanly and work around CTS reference issues (though this class
+// should be replaced in the CTS version).
+public class Main {
+  // Load the given class with the given classloader, and bind all native methods to corresponding
+  // C methods in the agent. Will abort if any of the steps fail.
+  public static native void bindAgentJNI(String className, ClassLoader classLoader);
+  // Same as above, giving the class directly.
+  public static native void bindAgentJNIForClass(Class<?> klass);
+
+  // Common infrastructure.
+  public static native void setTag(Object o, long tag);
+  public static native long getTag(Object o);
+}
diff --git a/test/1909-per-agent-tls/src/art/Test1909.java b/test/1909-per-agent-tls/src/art/Test1909.java
new file mode 100644
index 0000000..245397d
--- /dev/null
+++ b/test/1909-per-agent-tls/src/art/Test1909.java
@@ -0,0 +1,176 @@
+/*
+ * Copyright (C) 2016 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package art;
+
+import java.lang.ref.WeakReference;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.Semaphore;
+
+public class Test1909 {
+
+  public static class ThreadHolder {
+    public Thread thr;
+    public ThreadHolder(Thread thr) {
+      this.thr = thr;
+    }
+
+    public long getTLS(long jvmtienv) {
+      return Test1909.getTLS(jvmtienv, this.thr);
+    }
+    public void setTLS(long jvmtienv, long ptr) {
+      Test1909.setTLS(jvmtienv, this.thr, ptr);
+    }
+  }
+
+  public static class ThreadWaiter {
+    public boolean exit;
+    public Thread thr;
+    public final Object lock;
+
+    public ThreadWaiter() {
+      this.exit = false;
+      this.lock = new Object();
+      this.thr = new Thread(() -> {
+        try {
+          synchronized (lock) {
+            while (!this.exit) {
+              this.lock.wait();
+            }
+          }
+        } catch (Exception e) {
+          e.printStackTrace();
+        }
+      });
+      // Kill threads if we exit.
+      thr.setDaemon(true);
+      thr.start();
+    }
+
+    public void cleanup() throws Exception {
+      synchronized (lock) {
+        exit = true;
+        lock.notifyAll();
+      }
+      thr.join();
+    }
+    public long getTLS(long jvmtienv) {
+      return Test1909.getTLS(jvmtienv, this.thr);
+    }
+    public void setTLS(long jvmtienv, long ptr) {
+      Test1909.setTLS(jvmtienv, this.thr, ptr);
+    }
+  }
+
+  public static void checkEq(long a, long b) {
+    if (a != b) {
+      throw new Error("Expected: " + a + " got: " + b);
+    }
+  }
+
+  public static void run() throws Exception {
+    ThreadHolder tc = new ThreadHolder(Thread.currentThread());
+    ThreadWaiter t1 = new ThreadWaiter();
+    long e1 = newJvmtiEnv();
+    long e2 = newJvmtiEnv();
+
+    // Everything should be 0
+    checkEq(0, tc.getTLS(e1));
+    checkEq(0, t1.getTLS(e1));
+    checkEq(0, tc.getTLS(e2));
+    checkEq(0, t1.getTLS(e2));
+
+    // Set in one pair.
+    tc.setTLS(e1, 1);
+    checkEq(1, tc.getTLS(e1));
+    checkEq(0, t1.getTLS(e1));
+    checkEq(0, tc.getTLS(e2));
+    checkEq(0, t1.getTLS(e2));
+
+    // Set in another pair.
+    t1.setTLS(e1, 2);
+    checkEq(1, tc.getTLS(e1));
+    checkEq(2, t1.getTLS(e1));
+    checkEq(0, tc.getTLS(e2));
+    checkEq(0, t1.getTLS(e2));
+
+    // Set in third pair.
+    tc.setTLS(e2, 3);
+    checkEq(1, tc.getTLS(e1));
+    checkEq(2, t1.getTLS(e1));
+    checkEq(3, tc.getTLS(e2));
+    checkEq(0, t1.getTLS(e2));
+
+    // Set in fourth pair.
+    t1.setTLS(e2, 4);
+    checkEq(1, tc.getTLS(e1));
+    checkEq(2, t1.getTLS(e1));
+    checkEq(3, tc.getTLS(e2));
+    checkEq(4, t1.getTLS(e2));
+
+    // Create a new thread and make sure everything is 0.
+    ThreadWaiter t2 = new ThreadWaiter();
+    checkEq(1, tc.getTLS(e1));
+    checkEq(2, t1.getTLS(e1));
+    checkEq(0, t2.getTLS(e1));
+    checkEq(3, tc.getTLS(e2));
+    checkEq(4, t1.getTLS(e2));
+    checkEq(0, t2.getTLS(e2));
+
+    // Create a new jvmtienv and make sure everything is 0.
+    long e3 = newJvmtiEnv();
+    checkEq(1, tc.getTLS(e1));
+    checkEq(2, t1.getTLS(e1));
+    checkEq(0, t2.getTLS(e1));
+    checkEq(3, tc.getTLS(e2));
+    checkEq(4, t1.getTLS(e2));
+    checkEq(0, t2.getTLS(e2));
+    checkEq(0, tc.getTLS(e3));
+    checkEq(0, t1.getTLS(e3));
+    checkEq(0, t2.getTLS(e3));
+
+    // Delete an env without data and make sure everything is still there.
+    destroyJvmtiEnv(e3);
+    checkEq(1, tc.getTLS(e1));
+    checkEq(2, t1.getTLS(e1));
+    checkEq(0, t2.getTLS(e1));
+    checkEq(3, tc.getTLS(e2));
+    checkEq(4, t1.getTLS(e2));
+    checkEq(0, t2.getTLS(e2));
+
+    // Delete an env with data and make sure everything is still there.
+    destroyJvmtiEnv(e2);
+    checkEq(1, tc.getTLS(e1));
+    checkEq(2, t1.getTLS(e1));
+    checkEq(0, t2.getTLS(e1));
+
+    // Delete a thread. Make sure other thread still has data.
+    t1.cleanup();
+    checkEq(1, tc.getTLS(e1));
+    checkEq(0, t2.getTLS(e1));
+
+    t2.cleanup();
+
+    System.out.println("Test passed");
+  }
+
+  public static native long getTLS(long jvmtienv, Thread thr);
+  public static native void setTLS(long jvmtienv, Thread thr, long ptr);
+  public static native long newJvmtiEnv();
+  public static native void destroyJvmtiEnv(long jvmtienv);
+}
diff --git a/test/Android.bp b/test/Android.bp
index 982d051..44cb4f6 100644
--- a/test/Android.bp
+++ b/test/Android.bp
@@ -291,6 +291,7 @@
         "1901-get-bytecodes/bytecodes.cc",
         "1905-suspend-native/native_suspend.cc",
         "1908-suspend-native-resume-self/native_suspend_resume.cc",
+        "1909-per-agent-tls/agent_tls.cc",
     ],
     shared_libs: [
         "libbase",