Add support for tracking jvmti allocations

Adds a jvmti extension for getting the amount of memory allocated and
by jvmtiEnvs and currently live.

Bug: 62065509
Test: ./test.py --host -j40
Change-Id: I5d05b171260d91f9c415f7a5cb40cc01b48d7d07
diff --git a/runtime/openjdkjvmti/OpenjdkJvmTi.cc b/runtime/openjdkjvmti/OpenjdkJvmTi.cc
index 505e844..d3e8798 100644
--- a/runtime/openjdkjvmti/OpenjdkJvmTi.cc
+++ b/runtime/openjdkjvmti/OpenjdkJvmTi.cc
@@ -1208,6 +1208,23 @@
       return error;
     }
 
+    error = add_extension(
+        reinterpret_cast<jvmtiExtensionFunction>(AllocUtil::GetGlobalJvmtiAllocationState),
+        "com.android.art.alloc.get_global_jvmti_allocation_state",
+        "Returns the total amount of memory currently allocated by all jvmtiEnvs through the"
+        " 'Allocate' jvmti function. This does not include any memory that has been deallocated"
+        " through the 'Deallocate' function. This number is approximate and might not correspond"
+        " exactly to the sum of the sizes of all not freed allocations.",
+        1,
+        {                                                          // NOLINT [whitespace/braces] [4]
+            { "currently_allocated", JVMTI_KIND_OUT, JVMTI_TYPE_JLONG, false},
+        },
+        1,
+        { ERR(NULL_POINTER) });
+    if (error != ERR(NONE)) {
+      return error;
+    }
+
     // Copy into output buffer.
 
     *extension_count_ptr = ext_vector.size();
diff --git a/runtime/openjdkjvmti/ti_allocator.cc b/runtime/openjdkjvmti/ti_allocator.cc
index 603a43f..b82c4f4 100644
--- a/runtime/openjdkjvmti/ti_allocator.cc
+++ b/runtime/openjdkjvmti/ti_allocator.cc
@@ -31,23 +31,25 @@
 
 #include "ti_allocator.h"
 
+#include <malloc.h>
+#include <atomic>
+
 #include "art_jvmti.h"
-#include "art_method-inl.h"
 #include "base/enums.h"
-#include "dex_file_annotations.h"
-#include "events-inl.h"
-#include "jni_internal.h"
-#include "mirror/object_array-inl.h"
-#include "modifiers.h"
-#include "runtime_callbacks.h"
-#include "scoped_thread_state_change-inl.h"
-#include "ScopedLocalRef.h"
-#include "thread-current-inl.h"
-#include "thread_list.h"
-#include "ti_phase.h"
 
 namespace openjdkjvmti {
 
+std::atomic<jlong> AllocUtil::allocated;
+
+jvmtiError AllocUtil::GetGlobalJvmtiAllocationState(jvmtiEnv* env ATTRIBUTE_UNUSED,
+                                                    jlong* allocated_ptr) {
+  if (allocated_ptr == nullptr) {
+    return ERR(NULL_POINTER);
+  }
+  *allocated_ptr = allocated.load();
+  return OK;
+}
+
 jvmtiError AllocUtil::Allocate(jvmtiEnv* env ATTRIBUTE_UNUSED,
                                jlong size,
                                unsigned char** mem_ptr) {
@@ -57,12 +59,17 @@
     *mem_ptr = nullptr;
     return OK;
   }
-  *mem_ptr = static_cast<unsigned char*>(malloc(size));
-  return (*mem_ptr != nullptr) ? OK : ERR(OUT_OF_MEMORY);
+  *mem_ptr = reinterpret_cast<unsigned char*>(malloc(size));
+  if (UNLIKELY(*mem_ptr == nullptr)) {
+    return ERR(OUT_OF_MEMORY);
+  }
+  allocated += malloc_usable_size(*mem_ptr);
+  return OK;
 }
 
 jvmtiError AllocUtil::Deallocate(jvmtiEnv* env ATTRIBUTE_UNUSED, unsigned char* mem) {
   if (mem != nullptr) {
+    allocated -= malloc_usable_size(mem);
     free(mem);
   }
   return OK;
diff --git a/runtime/openjdkjvmti/ti_allocator.h b/runtime/openjdkjvmti/ti_allocator.h
index 7f1aa6d..aba77ae 100644
--- a/runtime/openjdkjvmti/ti_allocator.h
+++ b/runtime/openjdkjvmti/ti_allocator.h
@@ -35,12 +35,17 @@
 #include "jni.h"
 #include "jvmti.h"
 
+#include <atomic>
+
 namespace openjdkjvmti {
 
 class AllocUtil {
  public:
   static jvmtiError Allocate(jvmtiEnv* env, jlong size, unsigned char** mem_ptr);
   static jvmtiError Deallocate(jvmtiEnv* env, unsigned char* mem);
+  static jvmtiError GetGlobalJvmtiAllocationState(jvmtiEnv* env, jlong* total_allocated);
+ private:
+  static std::atomic<jlong> allocated;
 };
 
 }  // namespace openjdkjvmti