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
diff --git a/test/904-object-allocation/tracking.cc b/test/904-object-allocation/tracking.cc
index b22fc6c..57bfed5 100644
--- a/test/904-object-allocation/tracking.cc
+++ b/test/904-object-allocation/tracking.cc
@@ -58,10 +58,10 @@
 }
 
 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 @@
     printf("Unable to get jvmti env!\n");
     return 1;
   }
+  jvmti_env->SetEventNotificationMode(JVMTI_ENABLE, JVMTI_EVENT_VM_OBJECT_ALLOC, nullptr);
   return 0;
 }