ART: Add allocation callback

Bug: 31684277
Test: m test-art-host
Change-Id: I959f44e23ca5fe55ed678315708895faf0aadb04
diff --git a/runtime/gc/allocation_listener.h b/runtime/gc/allocation_listener.h
new file mode 100644
index 0000000..6fb74d3
--- /dev/null
+++ b/runtime/gc/allocation_listener.h
@@ -0,0 +1,49 @@
+/*
+ * 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.
+ */
+
+#ifndef ART_RUNTIME_GC_ALLOCATION_LISTENER_H_
+#define ART_RUNTIME_GC_ALLOCATION_LISTENER_H_
+
+#include <list>
+#include <memory>
+
+#include "base/macros.h"
+#include "base/mutex.h"
+#include "object_callbacks.h"
+#include "gc_root.h"
+
+namespace art {
+
+namespace mirror {
+  class Object;
+}
+
+class Thread;
+
+namespace gc {
+
+class AllocationListener {
+ public:
+  virtual ~AllocationListener() {}
+
+  virtual void ObjectAllocated(Thread* self, mirror::Object** obj, size_t byte_count)
+      REQUIRES_SHARED(Locks::mutator_lock_) = 0;
+};
+
+}  // namespace gc
+}  // namespace art
+
+#endif  // ART_RUNTIME_GC_ALLOCATION_LISTENER_H_
diff --git a/runtime/gc/heap-inl.h b/runtime/gc/heap-inl.h
index 83789cc..00adefb 100644
--- a/runtime/gc/heap-inl.h
+++ b/runtime/gc/heap-inl.h
@@ -19,6 +19,7 @@
 
 #include "heap.h"
 
+#include "allocation_listener.h"
 #include "base/time_utils.h"
 #include "gc/accounting/card_table-inl.h"
 #include "gc/allocation_record.h"
@@ -184,6 +185,12 @@
       DCHECK(allocation_records_ != nullptr);
       allocation_records_->RecordAllocation(self, &obj, bytes_allocated);
     }
+    AllocationListener* l = alloc_listener_.LoadSequentiallyConsistent();
+    if (l != nullptr) {
+      // Same as above. We assume that a listener that was once stored will never be deleted.
+      // Otherwise we'd have to perform this under a lock.
+      l->ObjectAllocated(self, &obj, bytes_allocated);
+    }
   } else {
     DCHECK(!IsAllocTrackingEnabled());
   }
diff --git a/runtime/gc/heap.cc b/runtime/gc/heap.cc
index 01ad8d0..33f849a 100644
--- a/runtime/gc/heap.cc
+++ b/runtime/gc/heap.cc
@@ -21,6 +21,7 @@
 #include <unwind.h>  // For GC verification.
 #include <vector>
 
+#include "allocation_listener.h"
 #include "art_field-inl.h"
 #include "base/allocator.h"
 #include "base/arena_allocator.h"
@@ -1287,6 +1288,16 @@
   }
 }
 
+ALWAYS_INLINE
+static inline AllocationListener* GetAndOverwriteAllocationListener(
+    Atomic<AllocationListener*>* storage, AllocationListener* new_value) {
+  AllocationListener* old;
+  do {
+    old = storage->LoadSequentiallyConsistent();
+  } while (!storage->CompareExchangeStrongSequentiallyConsistent(old, new_value));
+  return old;
+}
+
 Heap::~Heap() {
   VLOG(heap) << "Starting ~Heap()";
   STLDeleteElements(&garbage_collectors_);
@@ -1307,6 +1318,10 @@
         << " total=" << seen_backtrace_count_.LoadRelaxed() +
             unique_backtrace_count_.LoadRelaxed();
   }
+  // Delete any still registered allocation listener.
+  AllocationListener* l = GetAndOverwriteAllocationListener(&alloc_listener_, nullptr);
+  delete l;
+
   VLOG(heap) << "Finished ~Heap()";
 }
 
@@ -4223,5 +4238,22 @@
   }
 }
 
+void Heap::SetAllocationListener(AllocationListener* l) {
+  AllocationListener* old = GetAndOverwriteAllocationListener(&alloc_listener_, l);
+
+  if (old == nullptr) {
+    Runtime::Current()->GetInstrumentation()->InstrumentQuickAllocEntryPoints();
+  }
+}
+
+void Heap::RemoveAllocationListener() {
+  AllocationListener* old = GetAndOverwriteAllocationListener(&alloc_listener_, nullptr);
+
+  if (old != nullptr) {
+    Runtime::Current()->GetInstrumentation()->InstrumentQuickAllocEntryPoints();
+  }
+}
+
+
 }  // namespace gc
 }  // namespace art
diff --git a/runtime/gc/heap.h b/runtime/gc/heap.h
index 678edff..5e17a52 100644
--- a/runtime/gc/heap.h
+++ b/runtime/gc/heap.h
@@ -57,6 +57,7 @@
 
 namespace gc {
 
+class AllocationListener;
 class AllocRecordObjectMap;
 class ReferenceProcessor;
 class TaskProcessor;
@@ -784,6 +785,12 @@
   HomogeneousSpaceCompactResult PerformHomogeneousSpaceCompact() REQUIRES(!*gc_complete_lock_);
   bool SupportHomogeneousSpaceCompactAndCollectorTransitions() const;
 
+  // Install an allocation listener.
+  void SetAllocationListener(AllocationListener* l);
+  // Remove an allocation listener. Note: the listener must not be deleted, as for performance
+  // reasons, we assume it stays valid when we read it (so that we don't require a lock).
+  void RemoveAllocationListener();
+
  private:
   class ConcurrentGCTask;
   class CollectorTransitionTask;
@@ -1352,6 +1359,9 @@
   // Boot image spaces.
   std::vector<space::ImageSpace*> boot_image_spaces_;
 
+  // An installed allocation listener.
+  Atomic<AllocationListener*> alloc_listener_;
+
   friend class CollectorTransitionTask;
   friend class collector::GarbageCollector;
   friend class collector::MarkCompact;
diff --git a/runtime/openjdkjvmti/events-inl.h b/runtime/openjdkjvmti/events-inl.h
index b298301..d027201 100644
--- a/runtime/openjdkjvmti/events-inl.h
+++ b/runtime/openjdkjvmti/events-inl.h
@@ -29,6 +29,9 @@
     return nullptr;
   }
 
+  // TODO: Add a type check. Can be done, for example, by an explicitly instantiated template
+  //       function.
+
   switch (event) {
     case JVMTI_EVENT_VM_INIT:
       return reinterpret_cast<FnType*>(env->event_callbacks->VMInit);
diff --git a/runtime/openjdkjvmti/events.cc b/runtime/openjdkjvmti/events.cc
index 48f3de4..4d5b7e0 100644
--- a/runtime/openjdkjvmti/events.cc
+++ b/runtime/openjdkjvmti/events.cc
@@ -29,9 +29,17 @@
  * questions.
  */
 
-#include "events.h"
+#include "events-inl.h"
 
 #include "art_jvmti.h"
+#include "base/logging.h"
+#include "gc/allocation_listener.h"
+#include "instrumentation.h"
+#include "jni_env_ext-inl.h"
+#include "mirror/class.h"
+#include "mirror/object.h"
+#include "runtime.h"
+#include "ScopedLocalRef.h"
 
 namespace openjdkjvmti {
 
@@ -116,8 +124,65 @@
   }
 }
 
+class JvmtiAllocationListener : public art::gc::AllocationListener {
+ public:
+  explicit JvmtiAllocationListener(EventHandler* handler) : handler_(handler) {}
+
+  void ObjectAllocated(art::Thread* self, art::mirror::Object** obj, size_t byte_count)
+      REQUIRES_SHARED(art::Locks::mutator_lock_) {
+    DCHECK_EQ(self, art::Thread::Current());
+
+    if (handler_->IsEventEnabledAnywhere(JVMTI_EVENT_VM_OBJECT_ALLOC)) {
+      // jvmtiEventVMObjectAlloc parameters:
+      //      jvmtiEnv *jvmti_env,
+      //      JNIEnv* jni_env,
+      //      jthread thread,
+      //      jobject object,
+      //      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)->GetClass()));
+
+      handler_->DispatchEvent(self,
+                              JVMTI_EVENT_VM_OBJECT_ALLOC,
+                              jni_env,
+                              thread.get(),
+                              object.get(),
+                              klass.get(),
+                              byte_count);
+    }
+  }
+
+ private:
+  EventHandler* handler_;
+};
+
+static void SetupObjectAllocationTracking(art::gc::AllocationListener* listener, bool enable) {
+  if (enable) {
+    art::Runtime::Current()->GetHeap()->SetAllocationListener(listener);
+  } else {
+    art::Runtime::Current()->GetHeap()->RemoveAllocationListener();
+  }
+}
+
 // Handle special work for the given event type, if necessary.
-static void HandleEventType(jvmtiEvent event ATTRIBUTE_UNUSED, bool enable ATTRIBUTE_UNUSED) {
+void EventHandler::HandleEventType(jvmtiEvent event, bool enable) {
+  if (event == JVMTI_EVENT_VM_OBJECT_ALLOC) {
+    SetupObjectAllocationTracking(alloc_listener_.get(), enable);
+    return;
+  }
 }
 
 jvmtiError EventHandler::SetEvent(ArtJvmTiEnv* env,
@@ -172,4 +237,11 @@
   return ERR(NONE);
 }
 
+EventHandler::EventHandler() {
+  alloc_listener_.reset(new JvmtiAllocationListener(this));
+}
+
+EventHandler::~EventHandler() {
+}
+
 }  // namespace openjdkjvmti
diff --git a/runtime/openjdkjvmti/events.h b/runtime/openjdkjvmti/events.h
index 5716d03..3212b12 100644
--- a/runtime/openjdkjvmti/events.h
+++ b/runtime/openjdkjvmti/events.h
@@ -27,6 +27,7 @@
 namespace openjdkjvmti {
 
 struct ArtJvmTiEnv;
+class JvmtiAllocationListener;
 
 struct EventMask {
   static constexpr size_t kEventsSize = JVMTI_MAX_EVENT_TYPE_VAL - JVMTI_MIN_EVENT_TYPE_VAL + 1;
@@ -71,12 +72,10 @@
 };
 
 // Helper class for event handling.
-struct EventHandler {
-  // List of all JvmTiEnv objects that have been created, in their creation order.
-  std::vector<ArtJvmTiEnv*> envs;
-
-  // A union of all enabled events, anywhere.
-  EventMask global_mask;
+class EventHandler {
+ public:
+  EventHandler();
+  ~EventHandler();
 
   // Register an env. It is assumed that this happens on env creation, that is, no events are
   // enabled, yet.
@@ -93,6 +92,17 @@
 
   template <typename ...Args>
   ALWAYS_INLINE inline void DispatchEvent(art::Thread* thread, jvmtiEvent event, Args... args);
+
+ private:
+  void HandleEventType(jvmtiEvent event, bool enable);
+
+  // List of all JvmTiEnv objects that have been created, in their creation order.
+  std::vector<ArtJvmTiEnv*> envs;
+
+  // A union of all enabled events, anywhere.
+  EventMask global_mask;
+
+  std::unique_ptr<JvmtiAllocationListener> alloc_listener_;
 };
 
 }  // namespace openjdkjvmti
diff --git a/test/904-object-allocation/build b/test/904-object-allocation/build
new file mode 100755
index 0000000..898e2e5
--- /dev/null
+++ b/test/904-object-allocation/build
@@ -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-build "$@" --experimental agents
diff --git a/test/904-object-allocation/expected.txt b/test/904-object-allocation/expected.txt
new file mode 100644
index 0000000..371d2b7
--- /dev/null
+++ b/test/904-object-allocation/expected.txt
@@ -0,0 +1,8 @@
+ObjectAllocated type java.lang.Object/java.lang.Object size 8
+ObjectAllocated type java.lang.Integer/java.lang.Integer size 16
+ObjectAllocated type java.lang.Short/java.lang.Short size 16
+Tracking on same thread
+ObjectAllocated type java.lang.Double/java.lang.Double size 16
+Tracking on same thread, not disabling tracking
+ObjectAllocated type java.lang.Double/java.lang.Double size 16
+Tracking on different thread
diff --git a/test/904-object-allocation/info.txt b/test/904-object-allocation/info.txt
new file mode 100644
index 0000000..875a5f6
--- /dev/null
+++ b/test/904-object-allocation/info.txt
@@ -0,0 +1 @@
+Tests basic functions in the jvmti plugin.
diff --git a/test/904-object-allocation/run b/test/904-object-allocation/run
new file mode 100755
index 0000000..2f7ad21
--- /dev/null
+++ b/test/904-object-allocation/run
@@ -0,0 +1,43 @@
+#!/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.
+
+plugin=libopenjdkjvmtid.so
+agent=libtiagentd.so
+lib=tiagentd
+if  [[ "$@" == *"-O"* ]]; then
+  agent=libtiagent.so
+  plugin=libopenjdkjvmti.so
+  lib=tiagent
+fi
+
+if [[ "$@" == *"--jvm"* ]]; then
+  arg="jvm"
+else
+  arg="art"
+fi
+
+if [[ "$@" != *"--debuggable"* ]]; then
+  other_args=" -Xcompiler-option --debuggable "
+else
+  other_args=""
+fi
+
+./default-run "$@" --experimental agents \
+                   --experimental runtime-plugins \
+                   --runtime-option -agentpath:${agent}=904-object-allocation,${arg} \
+                   --android-runtime-option -Xplugin:${plugin} \
+                   ${other_args} \
+                   --args ${lib}
diff --git a/test/904-object-allocation/src/Main.java b/test/904-object-allocation/src/Main.java
new file mode 100644
index 0000000..63f7f89
--- /dev/null
+++ b/test/904-object-allocation/src/Main.java
@@ -0,0 +1,119 @@
+/*
+ * 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.
+ */
+
+import java.util.ArrayList;
+
+public class Main {
+  public static void main(String[] args) throws Exception {
+    System.loadLibrary(args[1]);
+
+    // Use a list to ensure objects must be allocated.
+    ArrayList<Object> l = new ArrayList<>(100);
+
+    doTest(l);
+  }
+
+  public static void doTest(ArrayList<Object> l) throws Exception {
+    setupCallback();
+
+    enableAllocationTracking(null, true);
+
+    l.add(new Object());
+    l.add(new Integer(1));
+
+    enableAllocationTracking(null, false);
+
+    l.add(new Float(1.0f));
+
+    enableAllocationTracking(Thread.currentThread(), true);
+
+    l.add(new Short((short)0));
+
+    enableAllocationTracking(Thread.currentThread(), false);
+
+    l.add(new Byte((byte)0));
+
+    System.out.println("Tracking on same thread");
+
+    testThread(l, true, true);
+
+    l.add(new Byte((byte)0));
+
+    System.out.println("Tracking on same thread, not disabling tracking");
+
+    testThread(l, true, false);
+
+    System.out.println("Tracking on different thread");
+
+    testThread(l, false, true);
+
+    l.add(new Byte((byte)0));
+  }
+
+  private static void testThread(final ArrayList<Object> l, final boolean sameThread,
+      final boolean disableTracking) throws Exception {
+    final SimpleBarrier startBarrier = new SimpleBarrier(1);
+    final SimpleBarrier trackBarrier = new SimpleBarrier(1);
+    final SimpleBarrier disableBarrier = new SimpleBarrier(1);
+
+    Thread t = new Thread() {
+      public void run() {
+        try {
+          startBarrier.dec();
+          trackBarrier.waitFor();
+        } catch (Exception e) {
+          e.printStackTrace(System.out);
+          System.exit(1);
+        }
+
+        l.add(new Double(0.0));
+
+        if (disableTracking) {
+          enableAllocationTracking(sameThread ? this : Thread.currentThread(), false);
+        }
+      }
+    };
+
+    t.start();
+    startBarrier.waitFor();
+    enableAllocationTracking(sameThread ? t : Thread.currentThread(), true);
+    trackBarrier.dec();
+
+    t.join();
+  }
+
+  private static class SimpleBarrier {
+    int count;
+
+    public SimpleBarrier(int i) {
+      count = i;
+    }
+
+    public synchronized void dec() throws Exception {
+      count--;
+      notifyAll();
+    }
+
+    public synchronized void waitFor() throws Exception  {
+      while (count != 0) {
+        wait();
+      }
+    }
+  }
+
+  private static native void setupCallback();
+  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
new file mode 100644
index 0000000..c392cd4
--- /dev/null
+++ b/test/904-object-allocation/tracking.cc
@@ -0,0 +1,102 @@
+/*
+ * 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 "tracking.h"
+
+#include <iostream>
+#include <pthread.h>
+#include <stdio.h>
+#include <vector>
+
+#include "base/logging.h"
+#include "jni.h"
+#include "openjdkjvmti/jvmti.h"
+#include "ScopedLocalRef.h"
+#include "ScopedUtfChars.h"
+#include "utils.h"
+
+namespace art {
+namespace Test904ObjectAllocation {
+
+static jvmtiEnv* jvmti_env;
+
+static std::string GetClassName(JNIEnv* jni_env, jclass cls) {
+  ScopedLocalRef<jclass> class_class(jni_env, jni_env->GetObjectClass(cls));
+  jmethodID mid = jni_env->GetMethodID(class_class.get(), "getName", "()Ljava/lang/String;");
+  ScopedLocalRef<jstring> str(
+      jni_env, reinterpret_cast<jstring>(jni_env->CallObjectMethod(cls, mid)));
+  ScopedUtfChars utf_chars(jni_env, str.get());
+  return utf_chars.c_str();
+}
+
+static void JNICALL ObjectAllocated(jvmtiEnv* ti_env ATTRIBUTE_UNUSED,
+                                    JNIEnv* jni_env,
+                                    jthread thread ATTRIBUTE_UNUSED,
+                                    jobject object,
+                                    jclass object_klass,
+                                    jlong size) {
+  std::string object_klass_descriptor = GetClassName(jni_env, object_klass);
+  ScopedLocalRef<jclass> object_klass2(jni_env, jni_env->GetObjectClass(object));
+  std::string object_klass_descriptor2 = GetClassName(jni_env, object_klass2.get());
+
+  printf("ObjectAllocated type %s/%s size %zu\n",
+         object_klass_descriptor.c_str(),
+         object_klass_descriptor2.c_str(),
+         static_cast<size_t>(size));
+}
+
+extern "C" JNIEXPORT void JNICALL Java_Main_setupCallback(JNIEnv* env ATTRIBUTE_UNUSED, jclass) {
+  jvmtiEventCallbacks callbacks;
+  memset(&callbacks, 0, sizeof(jvmtiEventCallbacks));
+  callbacks.VMObjectAlloc = ObjectAllocated;
+
+  jvmtiError ret = jvmti_env->SetEventCallbacks(&callbacks, sizeof(callbacks));
+  if (ret != JVMTI_ERROR_NONE) {
+    char* err;
+    jvmti_env->GetErrorName(ret, &err);
+    printf("Error setting callbacks: %s\n", err);
+  }
+}
+
+extern "C" JNIEXPORT void JNICALL Java_Main_enableAllocationTracking(JNIEnv* env ATTRIBUTE_UNUSED,
+                                                                     jclass,
+                                                                     jthread thread,
+                                                                     jboolean enable) {
+  jvmtiError ret = jvmti_env->SetEventNotificationMode(
+      enable ? JVMTI_ENABLE : JVMTI_DISABLE,
+      JVMTI_EVENT_VM_OBJECT_ALLOC,
+      thread);
+  if (ret != JVMTI_ERROR_NONE) {
+    char* err;
+    jvmti_env->GetErrorName(ret, &err);
+    printf("Error getting tag: %s\n", err);
+  }
+}
+
+// Don't do anything
+jint OnLoad(JavaVM* vm,
+            char* options ATTRIBUTE_UNUSED,
+            void* reserved ATTRIBUTE_UNUSED) {
+  if (vm->GetEnv(reinterpret_cast<void**>(&jvmti_env), JVMTI_VERSION_1_0)) {
+    printf("Unable to get jvmti env!\n");
+    return 1;
+  }
+  return 0;
+}
+
+}  // namespace Test904ObjectAllocation
+}  // namespace art
+
diff --git a/test/904-object-allocation/tracking.h b/test/904-object-allocation/tracking.h
new file mode 100644
index 0000000..21c1837
--- /dev/null
+++ b/test/904-object-allocation/tracking.h
@@ -0,0 +1,30 @@
+/*
+ * 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.
+ */
+
+#ifndef ART_TEST_904_OBJECT_ALLOCATION_TRACKING_H_
+#define ART_TEST_904_OBJECT_ALLOCATION_TRACKING_H_
+
+#include <jni.h>
+
+namespace art {
+namespace Test904ObjectAllocation {
+
+jint OnLoad(JavaVM* vm, char* options, void* reserved);
+
+}  // namespace Test904ObjectAllocation
+}  // namespace art
+
+#endif  // ART_TEST_904_OBJECT_ALLOCATION_TRACKING_H_
diff --git a/test/Android.bp b/test/Android.bp
index d17261c..b2152e8 100644
--- a/test/Android.bp
+++ b/test/Android.bp
@@ -246,6 +246,7 @@
         "901-hello-ti-agent/basics.cc",
         "902-hello-transformation/transform.cc",
         "903-hello-tagging/tagging.cc",
+        "904-object-allocation/tracking.cc",
     ],
     shared_libs: [
         "libart",
@@ -265,6 +266,7 @@
         "901-hello-ti-agent/basics.cc",
         "902-hello-transformation/transform.cc",
         "903-hello-tagging/tagging.cc",
+        "904-object-allocation/tracking.cc",
     ],
     shared_libs: [
         "libartd",
diff --git a/test/Android.run-test.mk b/test/Android.run-test.mk
index a858c75..935d6ad 100644
--- a/test/Android.run-test.mk
+++ b/test/Android.run-test.mk
@@ -263,14 +263,14 @@
 # 147-stripped-dex-fallback isn't supported on device because --strip-dex
 # requires the zip command.
 # 569-checker-pattern-replacement tests behaviour present only on host.
-# 902-hello-transformation and 903-hello-tagging
-# isn't supported in current form due to linker
+# 90{2,3,4} are not supported in current form due to linker
 # restrictions. See b/31681198
 TEST_ART_BROKEN_TARGET_TESTS := \
   147-stripped-dex-fallback \
   569-checker-pattern-replacement \
   902-hello-transformation \
-  903-hello-tagging
+  903-hello-tagging \
+  904-object-allocation \
 
 ifneq (,$(filter target,$(TARGET_TYPES)))
   ART_TEST_KNOWN_BROKEN += $(call all-run-test-names,target,$(RUN_TYPES),$(PREBUILD_TYPES), \
diff --git a/test/ti-agent/common_load.cc b/test/ti-agent/common_load.cc
index 4c7df97..fde3c6b 100644
--- a/test/ti-agent/common_load.cc
+++ b/test/ti-agent/common_load.cc
@@ -26,6 +26,7 @@
 #include "901-hello-ti-agent/basics.h"
 #include "902-hello-transformation/transform.h"
 #include "903-hello-tagging/tagging.h"
+#include "904-object-allocation/tracking.h"
 
 namespace art {
 
@@ -43,6 +44,7 @@
   { "901-hello-ti-agent", Test901HelloTi::OnLoad, nullptr },
   { "902-hello-transformation", Test902HelloTransformation::OnLoad, nullptr },
   { "903-hello-tagging", Test903HelloTagging::OnLoad, nullptr },
+  { "904-object-allocation", Test904ObjectAllocation::OnLoad, nullptr },
 };
 
 static AgentLib* FindAgent(char* name) {