ART: Fix and extend test 904

Fix using the wrong thread when unregistering. The inner thread
should not use Thread.currentThread (copy-n-paste mistake).

Add test for registering in OnLoad.

Add test for keeping registered on exit.

Bug: 31684277
Bug: 31684920
Test: m test-art-host-run-test-904-object-allocation
Change-Id: If0d5eeeb864684ab42eeaf62698d1e639f2071d6

2 files changed, 17 insertions, 5 deletions

diff --git a/test/904-object-allocation/src/Main.java b/test/904-object-allocation/src/Main.java
index 9a089bdd78..fc8a112e49 100644
--- a/test/904-object-allocation/src/Main.java
+++ b/test/904-object-allocation/src/Main.java
@@ -40,7 +40,11 @@ public class Main {
   }
 
   public static void doTest(ArrayList<Object> l) throws Exception {
-    setupObjectAllocCallback();
+    // Disable the global registration from OnLoad, to get into a known state.
+    enableAllocationTracking(null, false);
+
+    // Enable actual logging callback.
+    setupObjectAllocCallback(true);
 
     enableAllocationTracking(null, true);
 
@@ -74,6 +78,11 @@ public class Main {
     testThread(l, false, true);
 
     l.add(new Byte((byte)0));
+
+    // Disable actual logging callback and re-enable tracking, so we can keep the event enabled and
+    // check that shutdown works correctly.
+    setupObjectAllocCallback(false);
+    enableAllocationTracking(null, true);
   }
 
   private static void testThread(final ArrayList<Object> l, final boolean sameThread,
@@ -82,6 +91,8 @@ public class Main {
     final SimpleBarrier trackBarrier = new SimpleBarrier(1);
     final SimpleBarrier disableBarrier = new SimpleBarrier(1);
 
+    final Thread thisThread = Thread.currentThread();
+
     Thread t = new Thread() {
       public void run() {
         try {
@@ -95,7 +106,7 @@ public class Main {
         l.add(new Double(0.0));
 
         if (disableTracking) {
-          enableAllocationTracking(sameThread ? this : Thread.currentThread(), false);
+          enableAllocationTracking(sameThread ? this : thisThread, false);
         }
       }
     };
@@ -127,6 +138,6 @@ public class Main {
     }
   }
 
-  private static native void setupObjectAllocCallback();
+  private static native void setupObjectAllocCallback(boolean enable);
   private static native void enableAllocationTracking(Thread thread, boolean enable);
 }
diff --git a/test/904-object-allocation/tracking.cc b/test/904-object-allocation/tracking.cc
index b22fc6c8c0..57bfed544e 100644
--- a/test/904-object-allocation/tracking.cc
+++ b/test/904-object-allocation/tracking.cc
@@ -58,10 +58,10 @@ static void JNICALL ObjectAllocated(jvmtiEnv* ti_env ATTRIBUTE_UNUSED,
 }
 
 extern "C" JNIEXPORT void JNICALL Java_Main_setupObjectAllocCallback(
-    JNIEnv* env ATTRIBUTE_UNUSED, jclass klass ATTRIBUTE_UNUSED) {
+    JNIEnv* env ATTRIBUTE_UNUSED, jclass klass ATTRIBUTE_UNUSED, jboolean enable) {
   jvmtiEventCallbacks callbacks;
   memset(&callbacks, 0, sizeof(jvmtiEventCallbacks));
-  callbacks.VMObjectAlloc = ObjectAllocated;
+  callbacks.VMObjectAlloc = enable ? ObjectAllocated : nullptr;
 
   jvmtiError ret = jvmti_env->SetEventCallbacks(&callbacks, sizeof(callbacks));
   if (ret != JVMTI_ERROR_NONE) {
@@ -94,6 +94,7 @@ jint OnLoad(JavaVM* vm,
     printf("Unable to get jvmti env!\n");
     return 1;
   }
+  jvmti_env->SetEventNotificationMode(JVMTI_ENABLE, JVMTI_EVENT_VM_OBJECT_ALLOC, nullptr);
   return 0;
 }