ART: Add allocation callback

Bug: 31684277
Test: m test-art-host
Change-Id: I959f44e23ca5fe55ed678315708895faf0aadb04
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