Cleanup and consolidate JVMTI event code.

Various cleanups around JVMTI event code.

Ensure that we always store and restore exceptions.

Ensure we always give agents a local frame.

Ensure that we have static_asserts to verify that we are calling
events with appropriate types.

Various other improvements.

Test: ./test.py --host -j50

Change-Id: I71937d1575efca5096c9d5218203dc8201e3bb79

6 files changed, 335 insertions, 234 deletions

diff --git a/openjdkjvmti/events-inl.h b/openjdkjvmti/events-inl.h
index 31edc73ac6..ab8e6def2d 100644
--- a/openjdkjvmti/events-inl.h
+++ b/openjdkjvmti/events-inl.h
@@ -18,10 +18,13 @@
 #define ART_OPENJDKJVMTI_EVENTS_INL_H_
 
 #include <array>
+#include <type_traits>
+#include <tuple>
 
 #include "events.h"
 #include "jni_internal.h"
 #include "nativehelper/ScopedLocalRef.h"
+#include "scoped_thread_state_change-inl.h"
 #include "ti_breakpoint.h"
 
 #include "art_jvmti.h"
@@ -115,7 +118,6 @@ FORALL_EVENT_TYPES(GET_CALLBACK)
 
 // C++ does not allow partial template function specialization. The dispatch for our separated
 // ClassFileLoadHook event types is the same, so use this helper for code deduplication.
-// TODO Locking of some type!
 template <ArtJvmtiEvent kEvent>
 inline void EventHandler::DispatchClassFileLoadHookEvent(art::Thread* thread,
                                                          JNIEnv* jnienv,
@@ -137,37 +139,30 @@ inline void EventHandler::DispatchClassFileLoadHookEvent(art::Thread* thread,
     if (env == nullptr) {
       continue;
     }
-    if (ShouldDispatch<kEvent>(env, thread)) {
-      ScopedLocalRef<jthrowable> thr(jnienv, jnienv->ExceptionOccurred());
-      jnienv->ExceptionClear();
-      jint new_len = 0;
-      unsigned char* new_data = nullptr;
-      auto callback = impl::GetCallback<kEvent>(env);
-      callback(env,
-               jnienv,
-               class_being_redefined,
-               loader,
-               name,
-               protection_domain,
-               current_len,
-               current_class_data,
-               &new_len,
-               &new_data);
-      if (thr.get() != nullptr && !jnienv->ExceptionCheck()) {
-        jnienv->Throw(thr.get());
-      }
-      if (new_data != nullptr && new_data != current_class_data) {
-        // Destroy the data the last transformer made. We skip this if the previous state was the
-        // initial one since we don't know here which jvmtiEnv allocated it.
-        // NB Currently this doesn't matter since all allocations just go to malloc but in the
-        // future we might have jvmtiEnv's keep track of their allocations for leak-checking.
-        if (last_env != nullptr) {
-          last_env->Deallocate(current_class_data);
-        }
-        last_env = env;
-        current_class_data = new_data;
-        current_len = new_len;
+    jint new_len = 0;
+    unsigned char* new_data = nullptr;
+    DispatchEventOnEnv<kEvent>(env,
+                               thread,
+                               jnienv,
+                               class_being_redefined,
+                               loader,
+                               name,
+                               protection_domain,
+                               current_len,
+                               static_cast<const unsigned char*>(current_class_data),
+                               &new_len,
+                               &new_data);
+    if (new_data != nullptr && new_data != current_class_data) {
+      // Destroy the data the last transformer made. We skip this if the previous state was the
+      // initial one since we don't know here which jvmtiEnv allocated it.
+      // NB Currently this doesn't matter since all allocations just go to malloc but in the
+      // future we might have jvmtiEnv's keep track of their allocations for leak-checking.
+      if (last_env != nullptr) {
+        last_env->Deallocate(current_class_data);
       }
+      last_env = env;
+      current_class_data = new_data;
+      current_len = new_len;
     }
   }
   if (last_env != nullptr) {
@@ -180,97 +175,125 @@ inline void EventHandler::DispatchClassFileLoadHookEvent(art::Thread* thread,
 // exactly the argument types of the corresponding Jvmti kEvent function pointer.
 
 template <ArtJvmtiEvent kEvent, typename ...Args>
+inline void EventHandler::ExecuteCallback(ArtJvmTiEnv* env, Args... args) {
+  using FnType = typename impl::EventFnType<kEvent>::type;
+  FnType callback = impl::GetCallback<kEvent>(env);
+  if (callback != nullptr) {
+    (*callback)(env, args...);
+  }
+}
+
+template <ArtJvmtiEvent kEvent, typename ...Args>
 inline void EventHandler::DispatchEvent(art::Thread* thread, Args... args) const {
+  static_assert(!std::is_same<JNIEnv*,
+                              typename std::decay_t<
+                                  std::tuple_element_t<0, std::tuple<Args..., nullptr_t>>>>::value,
+                "Should be calling DispatchEvent with explicit JNIEnv* argument!");
+  DCHECK(thread == nullptr || !thread->IsExceptionPending());
   for (ArtJvmTiEnv* env : envs) {
     if (env != nullptr) {
-      DispatchEvent<kEvent, Args...>(env, thread, args...);
+      DispatchEventOnEnv<kEvent, Args...>(env, thread, args...);
     }
   }
 }
 
-// Events with JNIEnvs need to stash pending exceptions since they can cause new ones to be thrown.
-// In accordance with the JVMTI specification we allow exceptions originating from events to
-// overwrite the current exception, including exceptions originating from earlier events.
-// TODO It would be nice to add the overwritten exceptions to the suppressed exceptions list of the
-// newest exception.
+// Helper for ensuring that the dispatch environment is sane. Events with JNIEnvs need to stash
+// pending exceptions since they can cause new ones to be thrown. In accordance with the JVMTI
+// specification we allow exceptions originating from events to overwrite the current exception,
+// including exceptions originating from earlier events.
+class ScopedEventDispatchEnvironment FINAL : public art::ValueObject {
+ public:
+  explicit ScopedEventDispatchEnvironment(JNIEnv* env)
+      : env_(env),
+        thr_(env_, env_->ExceptionOccurred()),
+        suspend_(art::Thread::Current(), art::kNative) {
+    // The spec doesn't say how much local data should be there, so we just give 128 which seems
+    // likely to be enough for most cases.
+    env_->PushLocalFrame(128);
+    env_->ExceptionClear();
+    UNUSED(suspend_);
+  }
+
+  ~ScopedEventDispatchEnvironment() {
+    if (thr_.get() != nullptr && !env_->ExceptionCheck()) {
+      // TODO It would be nice to add the overwritten exceptions to the suppressed exceptions list
+      // of the newest exception.
+      env_->Throw(thr_.get());
+    }
+    env_->PopLocalFrame(nullptr);
+  }
+
+ private:
+  JNIEnv* env_;
+  ScopedLocalRef<jthrowable> thr_;
+  // Not actually unused. The destructor/constructor does important work.
+  art::ScopedThreadStateChange suspend_;
+
+  DISALLOW_COPY_AND_ASSIGN(ScopedEventDispatchEnvironment);
+};
+
 template <ArtJvmtiEvent kEvent, typename ...Args>
 inline void EventHandler::DispatchEvent(art::Thread* thread, JNIEnv* jnienv, Args... args) const {
   for (ArtJvmTiEnv* env : envs) {
     if (env != nullptr) {
-      ScopedLocalRef<jthrowable> thr(jnienv, jnienv->ExceptionOccurred());
-      jnienv->ExceptionClear();
-      DispatchEvent<kEvent, JNIEnv*, Args...>(env, thread, jnienv, args...);
-      if (thr.get() != nullptr && !jnienv->ExceptionCheck()) {
-        jnienv->Throw(thr.get());
-      }
+      DispatchEventOnEnv<kEvent, Args...>(env, thread, jnienv, args...);
     }
   }
 }
 
 template <ArtJvmtiEvent kEvent, typename ...Args>
-inline void EventHandler::DispatchEvent(ArtJvmTiEnv* env, art::Thread* thread, Args... args) const {
-  using FnType = void(jvmtiEnv*, Args...);
-  if (ShouldDispatch<kEvent>(env, thread)) {
-    FnType* callback = impl::GetCallback<kEvent>(env);
-    if (callback != nullptr) {
-      (*callback)(env, args...);
-    }
+inline void EventHandler::DispatchEventOnEnv(
+    ArtJvmTiEnv* env, art::Thread* thread, JNIEnv* jnienv, Args... args) const {
+  DCHECK(env != nullptr);
+  if (ShouldDispatch<kEvent, JNIEnv*, Args...>(env, thread, jnienv, args...)) {
+    ScopedEventDispatchEnvironment sede(jnienv);
+    ExecuteCallback<kEvent, JNIEnv*, Args...>(env, jnienv, args...);
   }
 }
 
+template <ArtJvmtiEvent kEvent, typename ...Args>
+inline void EventHandler::DispatchEventOnEnv(
+    ArtJvmTiEnv* env, art::Thread* thread, Args... args) const {
+  static_assert(!std::is_same<JNIEnv*,
+                              typename std::decay_t<
+                                  std::tuple_element_t<0, std::tuple<Args..., nullptr_t>>>>::value,
+                "Should be calling DispatchEventOnEnv with explicit JNIEnv* argument!");
+  if (ShouldDispatch<kEvent>(env, thread, args...)) {
+    ExecuteCallback<kEvent, Args...>(env, args...);
+  }
+}
+
+// Events that need custom logic for if we send the event but are otherwise normal. This includes
+// the kBreakpoint, kFramePop, kFieldAccess, and kFieldModification events.
+
 // Need to give custom specializations for Breakpoint since it needs to filter out which particular
 // methods/dex_pcs agents get notified on.
 template <>
-inline void EventHandler::DispatchEvent<ArtJvmtiEvent::kBreakpoint>(art::Thread* thread,
-                                                                    JNIEnv* jnienv,
-                                                                    jthread jni_thread,
-                                                                    jmethodID jmethod,
-                                                                    jlocation location) const {
+inline bool EventHandler::ShouldDispatch<ArtJvmtiEvent::kBreakpoint>(
+    ArtJvmTiEnv* env,
+    art::Thread* thread,
+    JNIEnv* jnienv ATTRIBUTE_UNUSED,
+    jthread jni_thread ATTRIBUTE_UNUSED,
+    jmethodID jmethod,
+    jlocation location) const {
   art::ArtMethod* method = art::jni::DecodeArtMethod(jmethod);
-  for (ArtJvmTiEnv* env : envs) {
-    // Search for a breakpoint on this particular method and location.
-    if (env != nullptr &&
-        ShouldDispatch<ArtJvmtiEvent::kBreakpoint>(env, thread) &&
-        env->breakpoints.find({method, location}) != env->breakpoints.end()) {
-      // We temporarily clear any pending exceptions so the event can call back into java code.
-      ScopedLocalRef<jthrowable> thr(jnienv, jnienv->ExceptionOccurred());
-      jnienv->ExceptionClear();
-      auto callback = impl::GetCallback<ArtJvmtiEvent::kBreakpoint>(env);
-      (*callback)(env, jnienv, jni_thread, jmethod, location);
-      if (thr.get() != nullptr && !jnienv->ExceptionCheck()) {
-        jnienv->Throw(thr.get());
-      }
-    }
-  }
+  return ShouldDispatchOnThread<ArtJvmtiEvent::kBreakpoint>(env, thread) &&
+      env->breakpoints.find({method, location}) != env->breakpoints.end();
 }
 
-// Need to give custom specializations for FramePop since it needs to filter out which particular
-// agents get the event. This specialization gets an extra argument so we can determine which (if
-// any) environments have the frame pop.
 template <>
-inline void EventHandler::DispatchEvent<ArtJvmtiEvent::kFramePop>(
+inline bool EventHandler::ShouldDispatch<ArtJvmtiEvent::kFramePop>(
+    ArtJvmTiEnv* env,
     art::Thread* thread,
-    JNIEnv* jnienv,
-    jthread jni_thread,
-    jmethodID jmethod,
-    jboolean is_exception,
+    JNIEnv* jnienv ATTRIBUTE_UNUSED,
+    jthread jni_thread ATTRIBUTE_UNUSED,
+    jmethodID jmethod ATTRIBUTE_UNUSED,
+    jboolean is_exception ATTRIBUTE_UNUSED,
     const art::ShadowFrame* frame) const {
-  for (ArtJvmTiEnv* env : envs) {
-    // Search for the frame. Do this before checking if we need to send the event so that we don't
-    // have to deal with use-after-free or the frames being reallocated later.
-    if (env != nullptr && env->notify_frames.erase(frame) != 0) {
-      if (ShouldDispatch<ArtJvmtiEvent::kFramePop>(env, thread)) {
-        // We temporarily clear any pending exceptions so the event can call back into java code.
-        ScopedLocalRef<jthrowable> thr(jnienv, jnienv->ExceptionOccurred());
-        jnienv->ExceptionClear();
-        auto callback = impl::GetCallback<ArtJvmtiEvent::kFramePop>(env);
-        (*callback)(env, jnienv, jni_thread, jmethod, is_exception);
-        if (thr.get() != nullptr && !jnienv->ExceptionCheck()) {
-          jnienv->Throw(thr.get());
-        }
-      }
-    }
-  }
+  // Search for the frame. Do this before checking if we need to send the event so that we don't
+  // have to deal with use-after-free or the frames being reallocated later.
+  return env->notify_frames.erase(frame) != 0 &&
+      ShouldDispatchOnThread<ArtJvmtiEvent::kFramePop>(env, thread);
 }
 
 // Need to give custom specializations for FieldAccess and FieldModification since they need to
@@ -278,64 +301,54 @@ inline void EventHandler::DispatchEvent<ArtJvmtiEvent::kFramePop>(
 // TODO The spec allows us to do shortcuts like only allow one agent to ever set these watches. This
 // could make the system more performant.
 template <>
-inline void EventHandler::DispatchEvent<ArtJvmtiEvent::kFieldModification>(art::Thread* thread,
-                                                                           JNIEnv* jnienv,
-                                                                           jthread jni_thread,
-                                                                           jmethodID method,
-                                                                           jlocation location,
-                                                                           jclass field_klass,
-                                                                           jobject object,
-                                                                           jfieldID field,
-                                                                           char type_char,
-                                                                           jvalue val) const {
-  for (ArtJvmTiEnv* env : envs) {
-    if (env != nullptr &&
-        ShouldDispatch<ArtJvmtiEvent::kFieldModification>(env, thread) &&
-        env->modify_watched_fields.find(
-            art::jni::DecodeArtField(field)) != env->modify_watched_fields.end()) {
-      ScopedLocalRef<jthrowable> thr(jnienv, jnienv->ExceptionOccurred());
-      jnienv->ExceptionClear();
-      auto callback = impl::GetCallback<ArtJvmtiEvent::kFieldModification>(env);
-      (*callback)(env,
-                  jnienv,
-                  jni_thread,
-                  method,
-                  location,
-                  field_klass,
-                  object,
-                  field,
-                  type_char,
-                  val);
-      if (thr.get() != nullptr && !jnienv->ExceptionCheck()) {
-        jnienv->Throw(thr.get());
-      }
-    }
-  }
+inline bool EventHandler::ShouldDispatch<ArtJvmtiEvent::kFieldModification>(
+    ArtJvmTiEnv* env,
+    art::Thread* thread,
+    JNIEnv* jnienv ATTRIBUTE_UNUSED,
+    jthread jni_thread ATTRIBUTE_UNUSED,
+    jmethodID method ATTRIBUTE_UNUSED,
+    jlocation location ATTRIBUTE_UNUSED,
+    jclass field_klass ATTRIBUTE_UNUSED,
+    jobject object ATTRIBUTE_UNUSED,
+    jfieldID field,
+    char type_char ATTRIBUTE_UNUSED,
+    jvalue val ATTRIBUTE_UNUSED) const {
+  return ShouldDispatchOnThread<ArtJvmtiEvent::kFieldModification>(env, thread) &&
+      env->modify_watched_fields.find(
+          art::jni::DecodeArtField(field)) != env->modify_watched_fields.end();
 }
 
 template <>
-inline void EventHandler::DispatchEvent<ArtJvmtiEvent::kFieldAccess>(art::Thread* thread,
-                                                                     JNIEnv* jnienv,
-                                                                     jthread jni_thread,
-                                                                     jmethodID method,
-                                                                     jlocation location,
-                                                                     jclass field_klass,
-                                                                     jobject object,
-                                                                     jfieldID field) const {
-  for (ArtJvmTiEnv* env : envs) {
-    if (env != nullptr &&
-        ShouldDispatch<ArtJvmtiEvent::kFieldAccess>(env, thread) &&
-        env->access_watched_fields.find(
-            art::jni::DecodeArtField(field)) != env->access_watched_fields.end()) {
-      ScopedLocalRef<jthrowable> thr(jnienv, jnienv->ExceptionOccurred());
-      jnienv->ExceptionClear();
-      auto callback = impl::GetCallback<ArtJvmtiEvent::kFieldAccess>(env);
-      (*callback)(env, jnienv, jni_thread, method, location, field_klass, object, field);
-      if (thr.get() != nullptr && !jnienv->ExceptionCheck()) {
-        jnienv->Throw(thr.get());
-      }
-    }
-  }
+inline bool EventHandler::ShouldDispatch<ArtJvmtiEvent::kFieldAccess>(
+    ArtJvmTiEnv* env,
+    art::Thread* thread,
+    JNIEnv* jnienv ATTRIBUTE_UNUSED,
+    jthread jni_thread ATTRIBUTE_UNUSED,
+    jmethodID method ATTRIBUTE_UNUSED,
+    jlocation location ATTRIBUTE_UNUSED,
+    jclass field_klass ATTRIBUTE_UNUSED,
+    jobject object ATTRIBUTE_UNUSED,
+    jfieldID field) const {
+  return ShouldDispatchOnThread<ArtJvmtiEvent::kFieldAccess>(env, thread) &&
+      env->access_watched_fields.find(
+          art::jni::DecodeArtField(field)) != env->access_watched_fields.end();
+}
+
+// Need to give custom specializations for FramePop since it needs to filter out which particular
+// agents get the event. This specialization gets an extra argument so we can determine which (if
+// any) environments have the frame pop.
+// TODO It might be useful to use more template magic to have this only define ShouldDispatch or
+// something.
+template <>
+inline void EventHandler::ExecuteCallback<ArtJvmtiEvent::kFramePop>(
+    ArtJvmTiEnv* env,
+    JNIEnv* jnienv,
+    jthread jni_thread,
+    jmethodID jmethod,
+    jboolean is_exception,
+    const art::ShadowFrame* frame ATTRIBUTE_UNUSED) {
+  ExecuteCallback<ArtJvmtiEvent::kFramePop>(
+      env, jnienv, jni_thread, jmethod, is_exception);
 }
 
 // Need to give a custom specialization for NativeMethodBind since it has to deal with an out
@@ -349,14 +362,21 @@ inline void EventHandler::DispatchEvent<ArtJvmtiEvent::kNativeMethodBind>(art::T
                                                                           void** new_method) const {
   *new_method = cur_method;
   for (ArtJvmTiEnv* env : envs) {
-    if (env != nullptr && ShouldDispatch<ArtJvmtiEvent::kNativeMethodBind>(env, thread)) {
-      auto callback = impl::GetCallback<ArtJvmtiEvent::kNativeMethodBind>(env);
-      (*callback)(env, jnienv, jni_thread, method, cur_method, new_method);
+    if (env != nullptr) {
+      *new_method = cur_method;
+      DispatchEventOnEnv<ArtJvmtiEvent::kNativeMethodBind>(env,
+                                                           thread,
+                                                           jnienv,
+                                                           jni_thread,
+                                                           method,
+                                                           cur_method,
+                                                           new_method);
       if (*new_method != nullptr) {
         cur_method = *new_method;
       }
     }
   }
+  *new_method = cur_method;
 }
 
 // C++ does not allow partial template function specialization. The dispatch for our separated
@@ -386,6 +406,7 @@ inline void EventHandler::DispatchEvent<ArtJvmtiEvent::kClassFileLoadHookRetrans
       new_class_data_len,
       new_class_data);
 }
+
 template <>
 inline void EventHandler::DispatchEvent<ArtJvmtiEvent::kClassFileLoadHookNonRetransformable>(
     art::Thread* thread,
@@ -412,8 +433,7 @@ inline void EventHandler::DispatchEvent<ArtJvmtiEvent::kClassFileLoadHookNonRetr
 }
 
 template <ArtJvmtiEvent kEvent>
-inline bool EventHandler::ShouldDispatch(ArtJvmTiEnv* env,
-                                         art::Thread* thread) {
+inline bool EventHandler::ShouldDispatchOnThread(ArtJvmTiEnv* env, art::Thread* thread) {
   bool dispatch = env->event_masks.global_event_mask.Test(kEvent);
 
   if (!dispatch && thread != nullptr && env->event_masks.unioned_thread_event_mask.Test(kEvent)) {
@@ -423,6 +443,17 @@ inline bool EventHandler::ShouldDispatch(ArtJvmTiEnv* env,
   return dispatch;
 }
 
+template <ArtJvmtiEvent kEvent, typename ...Args>
+inline bool EventHandler::ShouldDispatch(ArtJvmTiEnv* env,
+                                         art::Thread* thread,
+                                         Args... args ATTRIBUTE_UNUSED) const {
+  static_assert(std::is_same<typename impl::EventFnType<kEvent>::type,
+                             void(*)(jvmtiEnv*, Args...)>::value,
+                "Unexpected different type of shouldDispatch");
+
+  return ShouldDispatchOnThread<kEvent>(env, thread);
+}
+
 inline void EventHandler::RecalculateGlobalEventMask(ArtJvmtiEvent event) {
   bool union_value = false;
   for (const ArtJvmTiEnv* stored_env : envs) {
diff --git a/openjdkjvmti/events.cc b/openjdkjvmti/events.cc
index 5911a0373d..1cc82b7750 100644
--- a/openjdkjvmti/events.cc
+++ b/openjdkjvmti/events.cc
@@ -192,6 +192,25 @@ static bool IsThreadControllable(ArtJvmtiEvent event) {
   }
 }
 
+template<typename Type>
+static Type AddLocalRef(art::JNIEnvExt* e, art::mirror::Object* obj)
+    REQUIRES_SHARED(art::Locks::mutator_lock_) {
+  return (obj == nullptr) ? nullptr : e->AddLocalReference<Type>(obj);
+}
+
+template<ArtJvmtiEvent kEvent, typename ...Args>
+static void RunEventCallback(EventHandler* handler,
+                             art::Thread* self,
+                             art::JNIEnvExt* jnienv,
+                             Args... args)
+    REQUIRES_SHARED(art::Locks::mutator_lock_) {
+  ScopedLocalRef<jthread> thread_jni(jnienv, AddLocalRef<jthread>(jnienv, self->GetPeer()));
+  handler->DispatchEvent<kEvent>(self,
+                                 static_cast<JNIEnv*>(jnienv),
+                                 thread_jni.get(),
+                                 args...);
+}
+
 class JvmtiAllocationListener : public art::gc::AllocationListener {
  public:
   explicit JvmtiAllocationListener(EventHandler* handler) : handler_(handler) {}
@@ -211,26 +230,17 @@ class JvmtiAllocationListener : public art::gc::AllocationListener {
       //      jclass object_klass,
       //      jlong size
       art::JNIEnvExt* jni_env = self->GetJniEnv();
-
-      jthread thread_peer;
-      if (self->IsStillStarting()) {
-        thread_peer = nullptr;
-      } else {
-        thread_peer = jni_env->AddLocalReference<jthread>(self->GetPeer());
-      }
-
-      ScopedLocalRef<jthread> thread(jni_env, thread_peer);
       ScopedLocalRef<jobject> object(
           jni_env, jni_env->AddLocalReference<jobject>(*obj));
       ScopedLocalRef<jclass> klass(
           jni_env, jni_env->AddLocalReference<jclass>(obj->Ptr()->GetClass()));
 
-      handler_->DispatchEvent<ArtJvmtiEvent::kVmObjectAlloc>(self,
-                                                             reinterpret_cast<JNIEnv*>(jni_env),
-                                                             thread.get(),
-                                                             object.get(),
-                                                             klass.get(),
-                                                             static_cast<jlong>(byte_count));
+      RunEventCallback<ArtJvmtiEvent::kVmObjectAlloc>(handler_,
+                                                      self,
+                                                      jni_env,
+                                                      object.get(),
+                                                      klass.get(),
+                                                      static_cast<jlong>(byte_count));
     }
   }
 
@@ -250,38 +260,6 @@ static void SetupObjectAllocationTracking(art::gc::AllocationListener* listener,
   }
 }
 
-template<typename Type>
-static Type AddLocalRef(art::JNIEnvExt* e, art::mirror::Object* obj)
-    REQUIRES_SHARED(art::Locks::mutator_lock_) {
-  return (obj == nullptr) ? nullptr : e->AddLocalReference<Type>(obj);
-}
-
-template<ArtJvmtiEvent kEvent, typename ...Args>
-static void RunEventCallback(EventHandler* handler,
-                             art::Thread* self,
-                             art::JNIEnvExt* jnienv,
-                             Args... args)
-    REQUIRES_SHARED(art::Locks::mutator_lock_) {
-  ScopedLocalRef<jthread> thread_jni(jnienv, AddLocalRef<jthread>(jnienv, self->GetPeer()));
-  art::StackHandleScope<1> hs(self);
-  art::Handle<art::mirror::Throwable> old_exception(hs.NewHandle(self->GetException()));
-  self->ClearException();
-  // Just give the event a good sized JNI frame. 100 should be fine.
-  jnienv->PushFrame(100);
-  {
-    // Need to do trampoline! :(
-    art::ScopedThreadSuspension sts(self, art::ThreadState::kNative);
-    handler->DispatchEvent<kEvent>(self,
-                                   static_cast<JNIEnv*>(jnienv),
-                                   thread_jni.get(),
-                                   args...);
-  }
-  jnienv->PopFrame();
-  if (!self->IsExceptionPending() && !old_exception.IsNull()) {
-    self->SetException(old_exception.Get());
-  }
-}
-
 class JvmtiMonitorListener : public art::MonitorCallback {
  public:
   explicit JvmtiMonitorListener(EventHandler* handler) : handler_(handler) {}
@@ -649,17 +627,15 @@ class JvmtiMethodTraceListener FINAL : public art::instrumentation::Instrumentat
 
   void WatchedFramePop(art::Thread* self, const art::ShadowFrame& frame)
       REQUIRES_SHARED(art::Locks::mutator_lock_) OVERRIDE {
-    if (event_handler_->IsEventEnabledAnywhere(ArtJvmtiEvent::kFramePop)) {
       art::JNIEnvExt* jnienv = self->GetJniEnv();
-      jboolean is_exception_pending = self->IsExceptionPending();
-      RunEventCallback<ArtJvmtiEvent::kFramePop>(
-          event_handler_,
-          self,
-          jnienv,
-          art::jni::EncodeArtMethod(frame.GetMethod()),
-          is_exception_pending,
-          &frame);
-    }
+    jboolean is_exception_pending = self->IsExceptionPending();
+    RunEventCallback<ArtJvmtiEvent::kFramePop>(
+        event_handler_,
+        self,
+        jnienv,
+        art::jni::EncodeArtMethod(frame.GetMethod()),
+        is_exception_pending,
+        &frame);
   }
 
   static void FindCatchMethodsFromThrow(art::Thread* self,
diff --git a/openjdkjvmti/events.h b/openjdkjvmti/events.h
index 0d36c46fa1..a062e1589e 100644
--- a/openjdkjvmti/events.h
+++ b/openjdkjvmti/events.h
@@ -161,20 +161,21 @@ class EventHandler {
                       ArtJvmtiEvent event,
                       jvmtiEventMode mode);
 
-  // Dispatch event to all registered environments.
+  // Dispatch event to all registered environments. Since this one doesn't have a JNIEnv* it doesn't
+  // matter if it has the mutator_lock.
   template <ArtJvmtiEvent kEvent, typename ...Args>
   ALWAYS_INLINE
   inline void DispatchEvent(art::Thread* thread, Args... args) const;
+
   // Dispatch event to all registered environments stashing exceptions as needed. This works since
   // JNIEnv* is always the second argument if it is passed to an event. Needed since C++ does not
   // allow partial template function specialization.
+  //
+  // We need both of these since we want to make sure to push a stack frame when it is possible for
+  // the event to allocate local references.
   template <ArtJvmtiEvent kEvent, typename ...Args>
   ALWAYS_INLINE
-  void DispatchEvent(art::Thread* thread, JNIEnv* jnienv, Args... args) const;
-  // Dispatch event to the given environment, only.
-  template <ArtJvmtiEvent kEvent, typename ...Args>
-  ALWAYS_INLINE
-  inline void DispatchEvent(ArtJvmTiEnv* env, art::Thread* thread, Args... args) const;
+  inline void DispatchEvent(art::Thread* thread, JNIEnv* jnienv, Args... args) const;
 
   // Tell the event handler capabilities were added/lost so it can adjust the sent events.If
   // caps_added is true then caps is all the newly set capabilities of the jvmtiEnv. If it is false
@@ -184,10 +185,29 @@ class EventHandler {
                                         const jvmtiCapabilities& caps,
                                         bool added);
 
+  // Dispatch event to the given environment, only.
+  template <ArtJvmtiEvent kEvent, typename ...Args>
+  ALWAYS_INLINE
+  inline void DispatchEventOnEnv(
+      ArtJvmTiEnv* env, art::Thread* thread, JNIEnv* jnienv, Args... args) const;
+
+  // Dispatch event to the given environment, only.
+  template <ArtJvmtiEvent kEvent, typename ...Args>
+  ALWAYS_INLINE
+  inline void DispatchEventOnEnv(ArtJvmTiEnv* env, art::Thread* thread, Args... args) const;
+
  private:
   template <ArtJvmtiEvent kEvent>
   ALWAYS_INLINE
-  static inline bool ShouldDispatch(ArtJvmTiEnv* env, art::Thread* thread);
+  static inline bool ShouldDispatchOnThread(ArtJvmTiEnv* env, art::Thread* thread);
+
+  template <ArtJvmtiEvent kEvent, typename ...Args>
+  ALWAYS_INLINE
+  static inline void ExecuteCallback(ArtJvmTiEnv* env, Args... args);
+
+  template <ArtJvmtiEvent kEvent, typename ...Args>
+  ALWAYS_INLINE
+  inline bool ShouldDispatch(ArtJvmTiEnv* env, art::Thread* thread, Args... args) const;
 
   ALWAYS_INLINE
   inline bool NeedsEventUpdate(ArtJvmTiEnv* env,
diff --git a/openjdkjvmti/object_tagging.cc b/openjdkjvmti/object_tagging.cc
index dcdd3ede13..6ba7165577 100644
--- a/openjdkjvmti/object_tagging.cc
+++ b/openjdkjvmti/object_tagging.cc
@@ -61,7 +61,7 @@ bool ObjectTagTable::DoesHandleNullOnSweep() {
   return event_handler_->IsEventEnabledAnywhere(ArtJvmtiEvent::kObjectFree);
 }
 void ObjectTagTable::HandleNullSweep(jlong tag) {
-  event_handler_->DispatchEvent<ArtJvmtiEvent::kObjectFree>(jvmti_env_, nullptr, tag);
+  event_handler_->DispatchEventOnEnv<ArtJvmtiEvent::kObjectFree>(jvmti_env_, nullptr, tag);
 }
 
 }  // namespace openjdkjvmti
diff --git a/openjdkjvmti/ti_class.cc b/openjdkjvmti/ti_class.cc
index 954b5d1d03..daf4a8b7f2 100644
--- a/openjdkjvmti/ti_class.cc
+++ b/openjdkjvmti/ti_class.cc
@@ -318,14 +318,11 @@ struct ClassCallback : public art::ClassLoadCallback {
       ScopedLocalRef<jthread> thread_jni(
           thread->GetJniEnv(),
           peer.IsNull() ? nullptr : thread->GetJniEnv()->AddLocalReference<jthread>(peer));
-      {
-        art::ScopedThreadSuspension sts(thread, art::ThreadState::kNative);
-        event_handler->DispatchEvent<ArtJvmtiEvent::kClassLoad>(
-            thread,
-            static_cast<JNIEnv*>(thread->GetJniEnv()),
-            thread_jni.get(),
-            jklass.get());
-      }
+      event_handler->DispatchEvent<ArtJvmtiEvent::kClassLoad>(
+          thread,
+          static_cast<JNIEnv*>(thread->GetJniEnv()),
+          thread_jni.get(),
+          jklass.get());
       if (klass->IsTemp()) {
         AddTempClass(thread, jklass.get());
       }
@@ -348,7 +345,6 @@ struct ClassCallback : public art::ClassLoadCallback {
       ScopedLocalRef<jthread> thread_jni(
           thread->GetJniEnv(),
           peer.IsNull() ? nullptr : thread->GetJniEnv()->AddLocalReference<jthread>(peer));
-      art::ScopedThreadSuspension sts(thread, art::ThreadState::kNative);
       event_handler->DispatchEvent<ArtJvmtiEvent::kClassPrepare>(
           thread,
           static_cast<JNIEnv*>(thread->GetJniEnv()),
diff --git a/test/912-classes/classes.cc b/test/912-classes/classes.cc
index 869eacd82c..9727057668 100644
--- a/test/912-classes/classes.cc
+++ b/test/912-classes/classes.cc
@@ -16,6 +16,7 @@
 
 #include <stdio.h>
 
+#include <condition_variable>
 #include <mutex>
 #include <vector>
 
@@ -378,6 +379,48 @@ extern "C" JNIEXPORT void JNICALL Java_art_Test912_enableClassLoadPreparePrintEv
                ClassLoadPreparePrinter::ClassPrepareCallback);
 }
 
+template<typename T>
+static jthread RunEventThread(const std::string& name,
+                              jvmtiEnv* jvmti,
+                              JNIEnv* env,
+                              void (*func)(jvmtiEnv*, JNIEnv*, T*),
+                              T* data) {
+  // Create a Thread object.
+  std::string name_str = name;
+  name_str += ": JVMTI_THREAD-Test912";
+  ScopedLocalRef<jobject> thread_name(env, env->NewStringUTF(name_str.c_str()));
+  CHECK(thread_name.get() != nullptr);
+
+  ScopedLocalRef<jclass> thread_klass(env, env->FindClass("java/lang/Thread"));
+  CHECK(thread_klass.get() != nullptr);
+
+  ScopedLocalRef<jobject> thread(env, env->AllocObject(thread_klass.get()));
+  CHECK(thread.get() != nullptr);
+
+  jmethodID initID = env->GetMethodID(thread_klass.get(), "<init>", "(Ljava/lang/String;)V");
+  CHECK(initID != nullptr);
+
+  env->CallNonvirtualVoidMethod(thread.get(), thread_klass.get(), initID, thread_name.get());
+  CHECK(!env->ExceptionCheck());
+
+  // Run agent thread.
+  CheckJvmtiError(jvmti, jvmti->RunAgentThread(thread.get(),
+                                               reinterpret_cast<jvmtiStartFunction>(func),
+                                               reinterpret_cast<void*>(data),
+                                               JVMTI_THREAD_NORM_PRIORITY));
+  return thread.release();
+}
+
+static void JoinTread(JNIEnv* env, jthread thr) {
+  ScopedLocalRef<jclass> thread_klass(env, env->FindClass("java/lang/Thread"));
+  CHECK(thread_klass.get() != nullptr);
+
+  jmethodID joinID = env->GetMethodID(thread_klass.get(), "join", "()V");
+  CHECK(joinID != nullptr);
+
+  env->CallVoidMethod(thr, joinID);
+}
+
 class ClassLoadPrepareEquality {
  public:
   static constexpr const char* kClassName = "Lart/Test912$ClassE;";
@@ -389,6 +432,21 @@ class ClassLoadPrepareEquality {
   static constexpr const char* kWeakInitSig = "(Ljava/lang/Object;)V";
   static constexpr const char* kWeakGetSig = "()Ljava/lang/Object;";
 
+  static void AgentThreadTest(jvmtiEnv* jvmti ATTRIBUTE_UNUSED,
+                              JNIEnv* env,
+                              jobject* obj_global) {
+    jobject target = *obj_global;
+    jobject target_local = env->NewLocalRef(target);
+    {
+      std::unique_lock<std::mutex> lk(mutex_);
+      started_ = true;
+      cond_started_.notify_all();
+      cond_finished_.wait(lk, [] { return finished_; });
+      CHECK(finished_);
+    }
+    CHECK(env->IsSameObject(target, target_local));
+  }
+
   static void JNICALL ClassLoadCallback(jvmtiEnv* jenv,
                                         JNIEnv* jni_env,
                                         jthread thread ATTRIBUTE_UNUSED,
@@ -398,9 +456,13 @@ class ClassLoadPrepareEquality {
       found_ = true;
       stored_class_ = jni_env->NewGlobalRef(klass);
       weakly_stored_class_ = jni_env->NewWeakGlobalRef(klass);
-      // The following is bad and relies on implementation details. But otherwise a test would be
-      // a lot more complicated.
-      local_stored_class_ = jni_env->NewLocalRef(klass);
+      // Check that we update the local refs.
+      agent_thread_ = static_cast<jthread>(jni_env->NewGlobalRef(RunEventThread<jobject>(
+          "local-ref", jenv, jni_env, &AgentThreadTest, static_cast<jobject*>(&stored_class_))));
+      {
+        std::unique_lock<std::mutex> lk(mutex_);
+        cond_started_.wait(lk, [] { return started_; });
+      }
       // Store the value into a field in the heap.
       SetOrCompare(jni_env, klass, true);
     }
@@ -415,9 +477,14 @@ class ClassLoadPrepareEquality {
       CHECK(stored_class_ != nullptr);
       CHECK(jni_env->IsSameObject(stored_class_, klass));
       CHECK(jni_env->IsSameObject(weakly_stored_class_, klass));
-      CHECK(jni_env->IsSameObject(local_stored_class_, klass));
+      {
+        std::unique_lock<std::mutex> lk(mutex_);
+        finished_ = true;
+        cond_finished_.notify_all();
+      }
       // Look up the value in a field in the heap.
       SetOrCompare(jni_env, klass, false);
+      JoinTread(jni_env, agent_thread_);
       compared_ = true;
     }
   }
@@ -487,14 +554,25 @@ class ClassLoadPrepareEquality {
  private:
   static jobject stored_class_;
   static jweak weakly_stored_class_;
-  static jobject local_stored_class_;
+  static jthread agent_thread_;
+  static std::mutex mutex_;
+  static bool started_;
+  static std::condition_variable cond_finished_;
+  static bool finished_;
+  static std::condition_variable cond_started_;
   static bool found_;
   static bool compared_;
 };
+
 jclass ClassLoadPrepareEquality::storage_class_ = nullptr;
 jobject ClassLoadPrepareEquality::stored_class_ = nullptr;
 jweak ClassLoadPrepareEquality::weakly_stored_class_ = nullptr;
-jobject ClassLoadPrepareEquality::local_stored_class_ = nullptr;
+jthread ClassLoadPrepareEquality::agent_thread_ = nullptr;
+std::mutex ClassLoadPrepareEquality::mutex_;
+bool ClassLoadPrepareEquality::started_ = false;
+std::condition_variable ClassLoadPrepareEquality::cond_started_;
+bool ClassLoadPrepareEquality::finished_ = false;
+std::condition_variable ClassLoadPrepareEquality::cond_finished_;
 bool ClassLoadPrepareEquality::found_ = false;
 bool ClassLoadPrepareEquality::compared_ = false;