Fix ddm.publish_chunk extension to prevent deadlock

The ddm.publish_chunk jvmti extension event could could cause
deadlocks with the GC. This is not really fixable so instead remove
the JNIenv from the event function to prevent the dangerous actions
from being performed. Since this is an internal API that isn't really
used this should be safe enough. Rename the extension to
publish_chunk_safe so users can be notified of new behavior. Also fix
the test to stop relying on the buggy behavior.

Test: ./test.py
Bug: 161207646
Change-Id: Ib67d6607a8f359167069bc65cf5e3adaa70f0fc5

5 files changed, 140 insertions, 51 deletions

diff --git a/openjdkjvmti/events.cc b/openjdkjvmti/events.cc
index 64a02e874c..cd7155fa2f 100644
--- a/openjdkjvmti/events.cc
+++ b/openjdkjvmti/events.cc
@@ -300,7 +300,6 @@ class JvmtiDdmChunkListener : public art::DdmCallback {
       art::Thread* self = art::Thread::Current();
       handler_->DispatchEvent<ArtJvmtiEvent::kDdmPublishChunk>(
           self,
-          static_cast<JNIEnv*>(self->GetJniEnv()),
           static_cast<jint>(type),
           static_cast<jint>(data.size()),
           reinterpret_cast<const jbyte*>(data.data()));
diff --git a/openjdkjvmti/events.h b/openjdkjvmti/events.h
index d4eb17137e..be0839bdc5 100644
--- a/openjdkjvmti/events.h
+++ b/openjdkjvmti/events.h
@@ -99,7 +99,6 @@ constexpr jint kInternalEventCount = static_cast<jint>(ArtJvmtiEvent::kMaxIntern
                                      static_cast<jint>(ArtJvmtiEvent::kMinInternalEventTypeVal) + 1;
 
 using ArtJvmtiEventDdmPublishChunk = void (*)(jvmtiEnv *jvmti_env,
-                                              JNIEnv* jni_env,
                                               jint data_type,
                                               jint data_len,
                                               const jbyte* data);
diff --git a/openjdkjvmti/ti_extension.cc b/openjdkjvmti/ti_extension.cc
index f3f6e18018..dd68533523 100644
--- a/openjdkjvmti/ti_extension.cc
+++ b/openjdkjvmti/ti_extension.cc
@@ -599,15 +599,16 @@ jvmtiError ExtensionUtil::GetExtensionEvents(jvmtiEnv* env,
   jvmtiError error;
   error = add_extension(
       ArtJvmtiEvent::kDdmPublishChunk,
-      "com.android.art.internal.ddm.publish_chunk",
+      "com.android.art.internal.ddm.publish_chunk_safe",
       "Called when there is new ddms information that the agent or other clients can use. The"
       " agent is given the 'type' of the ddms chunk and a 'data_size' byte-buffer in 'data'."
       " The 'data' pointer is only valid for the duration of the publish_chunk event. The agent"
       " is responsible for interpreting the information present in the 'data' buffer. This is"
       " provided for backwards-compatibility support only. Agents should prefer to use relevant"
-      " JVMTI events and functions above listening for this event.",
+      " JVMTI events and functions above listening for this event. Previous publish_chunk"
+      " event was inherently unsafe since using the JNIEnv could cause deadlocks in some scenarios."
+      " The current version does not have these issues.",
       {
-        { "jni_env", JVMTI_KIND_IN_PTR, JVMTI_TYPE_JNIENV, false },
         { "type", JVMTI_KIND_IN, JVMTI_TYPE_JINT, false },
         { "data_size", JVMTI_KIND_IN, JVMTI_TYPE_JINT, false },
         { "data",  JVMTI_KIND_IN_BUF, JVMTI_TYPE_JBYTE, false },
diff --git a/test/1940-ddms-ext/ddm_ext.cc b/test/1940-ddms-ext/ddm_ext.cc
index 452187bdcb..110ad64ac6 100644
--- a/test/1940-ddms-ext/ddm_ext.cc
+++ b/test/1940-ddms-ext/ddm_ext.cc
@@ -14,6 +14,9 @@
  * limitations under the License.
  */
 
+#include <queue>
+#include <vector>
+
 #include "jvmti.h"
 
 // Test infrastructure
@@ -33,10 +36,16 @@ using DdmHandleChunk = jvmtiError(*)(jvmtiEnv* env,
                                      jint* len_data_out,
                                      jbyte** data_out);
 
+struct DdmCallbackData {
+ public:
+  DdmCallbackData(jint type, jint size, jbyte* data) : type_(type), data_(data, data + size) {}
+  jint type_;
+  std::vector<jbyte> data_;
+};
 struct DdmsTrackingData {
   DdmHandleChunk send_ddm_chunk;
-  jclass test_klass;
-  jmethodID publish_method;
+  jrawMonitorID callback_mon;
+  std::queue<DdmCallbackData> callbacks_received;
 };
 
 template <typename T>
@@ -111,21 +120,52 @@ static void DeallocParams(jvmtiParamInfo* params, jint n_params) {
   }
 }
 
-static void JNICALL PublishCB(jvmtiEnv* jvmti, JNIEnv* jnienv, jint type, jint size, jbyte* bytes) {
+extern "C" JNIEXPORT void JNICALL Java_art_Test1940_publishListen(JNIEnv* env,
+                                                                  jclass test_klass,
+                                                                  jobject publish) {
+  jmethodID publish_method = env->FromReflectedMethod(publish);
   DdmsTrackingData* data = nullptr;
-  if (JvmtiErrorToException(jnienv, jvmti,
-                            jvmti->GetEnvironmentLocalStorage(reinterpret_cast<void**>(&data)))) {
+  if (JvmtiErrorToException(
+          env, jvmti_env, jvmti_env->GetEnvironmentLocalStorage(reinterpret_cast<void**>(&data)))) {
     return;
   }
-  ScopedLocalRef<jbyteArray> res(jnienv, jnienv->NewByteArray(size));
-  jnienv->SetByteArrayRegion(res.get(), 0, size, bytes);
-  jnienv->CallStaticVoidMethod(data->test_klass, data->publish_method, type, res.get());
+  std::vector<DdmCallbackData> callbacks;
+  while (true) {
+    if (JvmtiErrorToException(env, jvmti_env, jvmti_env->RawMonitorEnter(data->callback_mon))) {
+      return;
+    }
+    while (data->callbacks_received.empty()) {
+      if (JvmtiErrorToException(env, jvmti_env, jvmti_env->RawMonitorWait(data->callback_mon, 0))) {
+        CHECK_EQ(JVMTI_ERROR_NONE, jvmti_env->RawMonitorExit(data->callback_mon));
+        return;
+      }
+    }
+    while (!data->callbacks_received.empty()) {
+      callbacks.emplace_back(std::move(data->callbacks_received.front()));
+      data->callbacks_received.pop();
+    }
+    if (JvmtiErrorToException(env, jvmti_env, jvmti_env->RawMonitorExit(data->callback_mon))) {
+      return;
+    }
+    for (auto cb : callbacks) {
+      ScopedLocalRef<jbyteArray> res(env, env->NewByteArray(cb.data_.size()));
+      env->SetByteArrayRegion(res.get(), 0, cb.data_.size(), cb.data_.data());
+      env->CallStaticVoidMethod(test_klass, publish_method, cb.type_, res.get());
+    }
+    callbacks.clear();
+  }
+}
+
+static void JNICALL PublishCB(jvmtiEnv* jvmti, jint type, jint size, jbyte* bytes) {
+  DdmsTrackingData* data = nullptr;
+  CHECK_EQ(JVMTI_ERROR_NONE, jvmti->GetEnvironmentLocalStorage(reinterpret_cast<void**>(&data)));
+  CHECK_EQ(JVMTI_ERROR_NONE, jvmti->RawMonitorEnter(data->callback_mon));
+  data->callbacks_received.emplace(type, size, bytes);
+  CHECK_EQ(JVMTI_ERROR_NONE, jvmti->RawMonitorNotifyAll(data->callback_mon));
+  CHECK_EQ(JVMTI_ERROR_NONE, jvmti->RawMonitorExit(data->callback_mon));
 }
 
-extern "C" JNIEXPORT void JNICALL Java_art_Test1940_initializeTest(JNIEnv* env,
-                                                                   jclass,
-                                                                   jclass method_klass,
-                                                                   jobject publish_method) {
+extern "C" JNIEXPORT void JNICALL Java_art_Test1940_initializeTest(JNIEnv* env, jclass) {
   void* old_data = nullptr;
   if (JvmtiErrorToException(env, jvmti_env, jvmti_env->GetEnvironmentLocalStorage(&old_data))) {
     return;
@@ -134,17 +174,16 @@ extern "C" JNIEXPORT void JNICALL Java_art_Test1940_initializeTest(JNIEnv* env,
     env->ThrowNew(rt_exception.get(), "Environment already has local storage set!");
     return;
   }
-  DdmsTrackingData* data = nullptr;
+  void* mem = nullptr;
   if (JvmtiErrorToException(env,
                             jvmti_env,
                             jvmti_env->Allocate(sizeof(DdmsTrackingData),
-                                                reinterpret_cast<unsigned char**>(&data)))) {
+                                                reinterpret_cast<unsigned char**>(&mem)))) {
     return;
   }
-  memset(data, 0, sizeof(DdmsTrackingData));
-  data->test_klass = reinterpret_cast<jclass>(env->NewGlobalRef(method_klass));
-  data->publish_method = env->FromReflectedMethod(publish_method);
-  if (env->ExceptionCheck()) {
+  DdmsTrackingData* data = new (mem) DdmsTrackingData{};
+  if (JvmtiErrorToException(
+          env, jvmti_env, jvmti_env->CreateRawMonitor("callback-mon", &data->callback_mon))) {
     return;
   }
   // Get the extensions.
@@ -181,7 +220,7 @@ extern "C" JNIEXPORT void JNICALL Java_art_Test1940_initializeTest(JNIEnv* env,
   }
   for (jint i = 0; i < n_ext; i++) {
     jvmtiExtensionEventInfo* cur_info = &events[i];
-    if (strcmp("com.android.art.internal.ddm.publish_chunk", cur_info->id) == 0) {
+    if (strcmp("com.android.art.internal.ddm.publish_chunk_safe", cur_info->id) == 0) {
       found_event = true;
       event_index = cur_info->extension_event_index;
     }
diff --git a/test/1940-ddms-ext/src-art/art/Test1940.java b/test/1940-ddms-ext/src-art/art/Test1940.java
index 2957f632ba..6fa31cab15 100644
--- a/test/1940-ddms-ext/src-art/art/Test1940.java
+++ b/test/1940-ddms-ext/src-art/art/Test1940.java
@@ -21,6 +21,7 @@ import dalvik.system.VMDebug;
 
 import java.lang.reflect.Method;
 import java.util.Arrays;
+import java.util.function.*;
 import java.util.zip.Adler32;
 import java.nio.*;
 
@@ -36,7 +37,7 @@ public class Test1940 {
   public static final boolean PRINT_ALL_CHUNKS = false;
 
   public static interface DdmHandler {
-    public void HandleChunk(int type, byte[] data);
+    public void HandleChunk(int type, byte[] data) throws Exception;
   }
 
   public static final class TestError extends Error {
@@ -49,6 +50,13 @@ public class Test1940 {
     }
   }
 
+  private static boolean chunkEq(Chunk a, Chunk b) {
+    return a.type == b.type &&
+           a.offset == b.offset &&
+           a.length == b.length &&
+           Arrays.equals(a.data, b.data);
+  }
+
   private static String printChunk(Chunk k) {
     byte[] out = new byte[k.length];
     System.arraycopy(k.data, k.offset, out, 0, k.length);
@@ -88,7 +96,7 @@ public class Test1940 {
 
   public static DdmHandler CURRENT_HANDLER;
 
-  public static void HandlePublish(int type, byte[] data) {
+  public static void HandlePublish(int type, byte[] data) throws Exception {
     if (PRINT_ALL_CHUNKS) {
       System.out.println(
           "Unknown Chunk published: " + printChunk(new Chunk(type, data, 0, data.length)));
@@ -114,13 +122,43 @@ public class Test1940 {
     return b.getInt() == (int) t.getId();
   }
 
+  public static final class AwaitChunkHandler implements DdmHandler {
+    public final Predicate<Chunk> needle;
+    public final DdmHandler chain;
+    private boolean found = false;
+    public AwaitChunkHandler(Predicate<Chunk> needle, DdmHandler chain) {
+      this.needle = needle;
+      this.chain = chain;
+    }
+    public void HandleChunk(int type, byte[] data) throws Exception {
+      chain.HandleChunk(type, data);
+      Chunk c = new Chunk(type, data, 0, data.length);
+      if (needle.test(c)) {
+        synchronized (this) {
+          found = true;
+          notifyAll();
+        }
+      }
+    }
+    public synchronized void Await() throws Exception {
+      while (!found) {
+        wait();
+      }
+    }
+  }
+
   public static void run() throws Exception {
-    CURRENT_HANDLER = (type, data) -> {
+    DdmHandler BaseHandler = (type, data) -> {
       System.out.println("Chunk published: " + printChunk(new Chunk(type, data, 0, data.length)));
     };
-    initializeTest(
-        Test1940.class,
-        Test1940.class.getDeclaredMethod("HandlePublish", Integer.TYPE, new byte[0].getClass()));
+    CURRENT_HANDLER = BaseHandler;
+    initializeTest();
+    Method publish = Test1940.class.getDeclaredMethod("HandlePublish",
+                                                      Integer.TYPE,
+                                                      new byte[0].getClass());
+    Thread listener = new Thread(() -> { Test1940.publishListen(publish); });
+    listener.setDaemon(true);
+    listener.start();
     // Test sending chunk directly.
     DdmServer.registerHandler(MY_DDMS_TYPE, SINGLE_HANDLER);
     DdmServer.registerHandler(MY_EMPTY_DDMS_TYPE, SINGLE_HANDLER);
@@ -139,8 +177,12 @@ public class Test1940 {
     // Test sending chunk through DdmServer#sendChunk
     Chunk c = new Chunk(
         MY_DDMS_TYPE, new byte[] { 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f, 0x10 }, 0, 8);
+    AwaitChunkHandler h = new AwaitChunkHandler((x) -> chunkEq(c, x), CURRENT_HANDLER);
+    CURRENT_HANDLER = h;
     System.out.println("Sending chunk: " + printChunk(c));
     DdmServer.sendChunk(c);
+    h.Await();
+    CURRENT_HANDLER = BaseHandler;
 
     // Test getting back an empty chunk.
     data = new byte[] { 0x1 };
@@ -161,22 +203,25 @@ public class Test1940 {
 
     // Test thread chunks are sent.
     final boolean[] types_seen = new boolean[] { false, false, false };
-    CURRENT_HANDLER = (type, cdata) -> {
-      switch (type) {
-        case TYPE_THCR:
-          types_seen[0] = true;
-          break;
-        case TYPE_THNM:
-          types_seen[1] = true;
-          break;
-        case TYPE_THDE:
-          types_seen[2] = true;
-          break;
-        default:
-          // We don't want to print other types.
-          break;
-      }
-    };
+    AwaitChunkHandler wait_thd= new AwaitChunkHandler(
+      (x) -> types_seen[0] && types_seen[1] && types_seen[2],
+      (type, cdata) -> {
+        switch (type) {
+          case TYPE_THCR:
+            types_seen[0] = true;
+            break;
+          case TYPE_THNM:
+            types_seen[1] = true;
+            break;
+          case TYPE_THDE:
+            types_seen[2] = true;
+            break;
+          default:
+            // We don't want to print other types.
+            break;
+        }
+      });
+    CURRENT_HANDLER = wait_thd;
     DdmVmInternal.threadNotify(true);
     System.out.println("threadNotify started!");
     final Thread thr = new Thread(() -> { return; }, "THREAD");
@@ -186,6 +231,7 @@ public class Test1940 {
     System.out.println("Target thread finished!");
     DdmVmInternal.threadNotify(false);
     System.out.println("threadNotify Disabled!");
+    wait_thd.Await();
     // Make sure we saw at least one of Thread-create, Thread name, & thread death.
     if (!types_seen[0] || !types_seen[1] || !types_seen[2]) {
       System.out.println("Didn't see expected chunks for thread creation! got: " +
@@ -195,25 +241,28 @@ public class Test1940 {
     }
 
     // Test heap chunks are sent.
-    CURRENT_HANDLER = (type, cdata) -> {
+    AwaitChunkHandler hpif = new AwaitChunkHandler((x) -> x.type == TYPE_HPIF, (type, cdata) -> {
       // The actual data is far to noisy for this test as it includes information about global heap
       // state.
       if (type == TYPE_HPIF) {
         System.out.println("Expected chunk type published: " + type);
       }
-    };
+    });
+    CURRENT_HANDLER = hpif;
     final int HPIF_WHEN_NOW = 1;
     if (!DdmVmInternal.heapInfoNotify(HPIF_WHEN_NOW)) {
       System.out.println("Unexpected failure for heapInfoNotify!");
     }
+    hpif.Await();
 
     // method Tracing
-    CURRENT_HANDLER = (type, cdata) -> {
+    AwaitChunkHandler mpse = new AwaitChunkHandler((x) -> x.type == TYPE_MPSE, (type, cdata) -> {
       // This chunk includes timing and thread information so we just check the type.
       if (type == TYPE_MPSE) {
         System.out.println("Expected chunk type published: " + type);
       }
-    };
+    });
+    CURRENT_HANDLER = mpse;
     VMDebug.startMethodTracingDdms(/*size: default*/0,
                                    /*flags: none*/ 0,
                                    /*sampling*/ false,
@@ -223,6 +272,7 @@ public class Test1940 {
     doNothing();
     doNothing();
     VMDebug.stopMethodTracing();
+    mpse.Await();
   }
 
   private static void doNothing() {}
@@ -230,6 +280,7 @@ public class Test1940 {
     return processChunk(new Chunk(MY_DDMS_TYPE, val, 0, val.length));
   }
 
-  private static native void initializeTest(Class<?> k, Method m);
+  private static native void initializeTest();
   private static native Chunk processChunk(Chunk val);
+  private static native void publishListen(Method publish);
 }