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

11 files changed, 388 insertions, 15 deletions

diff --git a/runtime/openjdkjvmti/OpenjdkJvmTi.cc b/runtime/openjdkjvmti/OpenjdkJvmTi.cc
index 505e844c92..d3e8798bd6 100644
--- a/runtime/openjdkjvmti/OpenjdkJvmTi.cc
+++ b/runtime/openjdkjvmti/OpenjdkJvmTi.cc
@@ -1208,6 +1208,23 @@ class JvmtiFunctions {
       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 603a43ffd2..b82c4f4122 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 @@ jvmtiError AllocUtil::Allocate(jvmtiEnv* env ATTRIBUTE_UNUSED,
     *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 7f1aa6d9c0..aba77ae2fb 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
diff --git a/test/1900-track-alloc/alloc.cc b/test/1900-track-alloc/alloc.cc
new file mode 100644
index 0000000000..db5617c54c
--- /dev/null
+++ b/test/1900-track-alloc/alloc.cc
@@ -0,0 +1,159 @@
+/*
+ * 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 "jvmti.h"
+
+// Test infrastructure
+#include "jvmti_helper.h"
+#include "scoped_local_ref.h"
+#include "test_env.h"
+
+namespace art {
+namespace Test1900TrackAlloc {
+
+typedef jvmtiError (*GetGlobalState)(jvmtiEnv* env, jlong* allocated);
+
+struct AllocTrackingData {
+  GetGlobalState get_global_state;
+};
+
+template <typename T>
+static void Dealloc(T* t) {
+  jvmti_env->Deallocate(reinterpret_cast<unsigned char*>(t));
+}
+
+template <typename T, typename ...Rest>
+static void Dealloc(T* t, Rest... rs) {
+  Dealloc(t);
+  Dealloc(rs...);
+}
+
+extern "C" JNIEXPORT void JNICALL Java_art_Test1900_doDeallocate(JNIEnv* env,
+                                                                 jclass,
+                                                                 jlong jvmti_env_ptr,
+                                                                 jlong ptr) {
+  JvmtiErrorToException(env,
+                        reinterpret_cast<jvmtiEnv*>(jvmti_env_ptr),
+                        reinterpret_cast<jvmtiEnv*>(jvmti_env_ptr)->Deallocate(
+                            reinterpret_cast<unsigned char*>(static_cast<intptr_t>(ptr))));
+}
+
+extern "C" JNIEXPORT jlong JNICALL Java_art_Test1900_doAllocate(JNIEnv* env,
+                                                                jclass,
+                                                                jlong jvmti_env_ptr,
+                                                                jlong size) {
+  unsigned char* res = nullptr;
+  JvmtiErrorToException(env,
+                        reinterpret_cast<jvmtiEnv*>(jvmti_env_ptr),
+                        reinterpret_cast<jvmtiEnv*>(jvmti_env_ptr)->Allocate(size, &res));
+  return static_cast<jlong>(reinterpret_cast<intptr_t>(res));
+}
+
+extern "C" JNIEXPORT jlong JNICALL Java_art_Test1900_getAmountAllocated(JNIEnv* env, jclass) {
+  AllocTrackingData* data = nullptr;
+  if (JvmtiErrorToException(
+      env, jvmti_env, jvmti_env->GetEnvironmentLocalStorage(reinterpret_cast<void**>(&data)))) {
+    return -1;
+  }
+  if (data == nullptr || data->get_global_state == nullptr) {
+    ScopedLocalRef<jclass> rt_exception(env, env->FindClass("java/lang/RuntimeException"));
+    env->ThrowNew(rt_exception.get(), "Alloc tracking data not initialized.");
+    return -1;
+  }
+  jlong allocated = -1;
+  JvmtiErrorToException(env, jvmti_env, data->get_global_state(jvmti_env, &allocated));
+  return allocated;
+}
+
+static void DeallocParams(jvmtiParamInfo* params, jint n_params) {
+  for (jint i = 0; i < n_params; i++) {
+    Dealloc(params[i].name);
+  }
+}
+
+extern "C" JNIEXPORT jlong JNICALL Java_art_Test1900_getDefaultJvmtiEnv(JNIEnv*, jclass) {
+  return static_cast<jlong>(reinterpret_cast<intptr_t>(jvmti_env));
+}
+
+extern "C" JNIEXPORT void Java_art_Test1900_destroyJvmtiEnv(JNIEnv* env,
+                                                            jclass,
+                                                            jlong jvmti_env_ptr) {
+  JvmtiErrorToException(env,
+                        jvmti_env,
+                        reinterpret_cast<jvmtiEnv*>(jvmti_env_ptr)->DisposeEnvironment());
+}
+
+extern "C" JNIEXPORT jlong Java_art_Test1900_newJvmtiEnv(JNIEnv* env, jclass) {
+  JavaVM* vm = nullptr;
+  if (env->GetJavaVM(&vm) != 0) {
+    ScopedLocalRef<jclass> rt_exception(env, env->FindClass("java/lang/RuntimeException"));
+    env->ThrowNew(rt_exception.get(), "Unable to get JavaVM");
+    return -1;
+  }
+  jvmtiEnv* new_env = nullptr;
+  if (vm->GetEnv(reinterpret_cast<void**>(&new_env), JVMTI_VERSION_1_0) != 0) {
+    ScopedLocalRef<jclass> rt_exception(env, env->FindClass("java/lang/RuntimeException"));
+    env->ThrowNew(rt_exception.get(), "Unable to create new jvmtiEnv");
+    return -1;
+  }
+  return static_cast<jlong>(reinterpret_cast<intptr_t>(new_env));
+}
+
+extern "C" JNIEXPORT void JNICALL Java_art_Test1900_initializeTest(JNIEnv* env, jclass) {
+  void* old_data = nullptr;
+  if (JvmtiErrorToException(env, jvmti_env, jvmti_env->GetEnvironmentLocalStorage(&old_data))) {
+    return;
+  } else if (old_data != nullptr) {
+    ScopedLocalRef<jclass> rt_exception(env, env->FindClass("java/lang/RuntimeException"));
+    env->ThrowNew(rt_exception.get(), "Environment already has local storage set!");
+    return;
+  }
+  AllocTrackingData* data = nullptr;
+  if (JvmtiErrorToException(env,
+                            jvmti_env,
+                            jvmti_env->Allocate(sizeof(AllocTrackingData),
+                                                reinterpret_cast<unsigned char**>(&data)))) {
+    return;
+  }
+  memset(data, 0, sizeof(AllocTrackingData));
+  // Get the extensions.
+  jint n_ext = 0;
+  jvmtiExtensionFunctionInfo* infos = nullptr;
+  if (JvmtiErrorToException(env, jvmti_env, jvmti_env->GetExtensionFunctions(&n_ext, &infos))) {
+    return;
+  }
+  for (jint i = 0; i < n_ext; i++) {
+    jvmtiExtensionFunctionInfo* cur_info = &infos[i];
+    if (strcmp("com.android.art.alloc.get_global_jvmti_allocation_state", cur_info->id) == 0) {
+      data->get_global_state = reinterpret_cast<GetGlobalState>(cur_info->func);
+    }
+    // Cleanup the cur_info
+    DeallocParams(cur_info->params, cur_info->param_count);
+    Dealloc(cur_info->id, cur_info->short_description, cur_info->params, cur_info->errors);
+  }
+  // Cleanup the array.
+  Dealloc(infos);
+  if (data->get_global_state == nullptr) {
+    ScopedLocalRef<jclass> rt_exception(env, env->FindClass("java/lang/RuntimeException"));
+    env->ThrowNew(rt_exception.get(), "Unable to find memory tracking extensions.");
+    return;
+  }
+  JvmtiErrorToException(env, jvmti_env, jvmti_env->SetEnvironmentLocalStorage(data));
+  return;
+}
+
+}  // namespace Test1900TrackAlloc
+}  // namespace art
diff --git a/test/1900-track-alloc/expected.txt b/test/1900-track-alloc/expected.txt
new file mode 100644
index 0000000000..e69de29bb2
--- /dev/null
+++ b/test/1900-track-alloc/expected.txt
diff --git a/test/1900-track-alloc/info.txt b/test/1900-track-alloc/info.txt
new file mode 100644
index 0000000000..e1d35ae026
--- /dev/null
+++ b/test/1900-track-alloc/info.txt
@@ -0,0 +1 @@
+Tests the jvmti-extension to get allocated memory snapshot.
diff --git a/test/1900-track-alloc/run b/test/1900-track-alloc/run
new file mode 100755
index 0000000000..c6e62ae6cd
--- /dev/null
+++ b/test/1900-track-alloc/run
@@ -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-run "$@" --jvmti
diff --git a/test/1900-track-alloc/src/Main.java b/test/1900-track-alloc/src/Main.java
new file mode 100644
index 0000000000..0dab4ef726
--- /dev/null
+++ b/test/1900-track-alloc/src/Main.java
@@ -0,0 +1,21 @@
+/*
+ * 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.
+ */
+
+public class Main {
+  public static void main(String[] args) throws Exception {
+    art.Test1900.run();
+  }
+}
diff --git a/test/1900-track-alloc/src/art/Test1900.java b/test/1900-track-alloc/src/art/Test1900.java
new file mode 100644
index 0000000000..717999bed5
--- /dev/null
+++ b/test/1900-track-alloc/src/art/Test1900.java
@@ -0,0 +1,144 @@
+/*
+ * 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.
+ */
+
+package art;
+
+import java.lang.ref.WeakReference;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.Semaphore;
+
+public class Test1900 {
+  public static void checkLE(long exp, long o) {
+    if (exp > o) {
+      throw new Error("Expected: " + exp + " Got: " + o);
+    }
+  }
+  public static void checkEq(long exp, long o) {
+    if (exp != o) {
+      throw new Error("Expected: " + exp + " Got: " + o);
+    }
+  }
+
+  public static void runConcurrent(Runnable... rs) throws Exception {
+    final CountDownLatch latch = new CountDownLatch(rs.length);
+    Thread[] thrs = new Thread[rs.length];
+    for (int i = 0; i < rs.length; i++) {
+      final Runnable r = rs[i];
+      thrs[i] = new Thread(() -> {
+        latch.countDown();
+        r.run();
+      });
+      thrs[i].start();
+    }
+    for (Thread thr : thrs) {
+      thr.join();
+    }
+  }
+  static class Holder {
+    public long val;
+  }
+
+  public static void run() throws Exception {
+    initializeTest();
+    // Get the overhead for the native part of this test.
+    final long base_state = getAmountAllocated();
+
+    // Basic alloc-dealloc
+    checkEq(base_state + 0, getAmountAllocated());
+    long abc = doAllocate(10);
+    checkLE(base_state + 10, getAmountAllocated());
+    long def = doAllocate(10);
+    checkLE(base_state + 20, getAmountAllocated());
+    doDeallocate(abc);
+    checkLE(base_state + 10, getAmountAllocated());
+
+    doDeallocate(def);
+
+    checkEq(base_state + 0, getAmountAllocated());
+
+    // Try doing it concurrently.
+    Runnable add10 = () -> { long x = doAllocate(10); doDeallocate(x); };
+    Runnable[] rs = new Runnable[100];
+    Arrays.fill(rs, add10);
+    runConcurrent(rs);
+    checkEq(base_state + 0, getAmountAllocated());
+
+    // Try doing it concurrently with different threads to allocate and deallocate.
+    final Semaphore sem = new Semaphore(0);
+    final Holder h = new Holder();
+    runConcurrent(
+        () -> {
+          try {
+            h.val = doAllocate(100);
+            checkLE(base_state + 100, getAmountAllocated());
+            sem.release();
+          } catch (Exception e) { throw new Error("exception!", e); }
+        },
+        () -> {
+          try {
+            sem.acquire();
+            long after_acq = getAmountAllocated();
+            doDeallocate(h.val);
+            checkLE(base_state + 100, after_acq);
+          } catch (Exception e) { throw new Error("exception!", e); }
+        }
+    );
+    checkEq(base_state + 0, getAmountAllocated());
+
+    // Try doing it with multiple jvmtienvs.
+    long env1 = newJvmtiEnv();
+    long env2 = newJvmtiEnv();
+
+    final long new_base_state = getAmountAllocated();
+    // new jvmtienvs shouldn't save us memory.
+    checkLE(base_state, new_base_state);
+    // Make sure we track both.
+    abc = doAllocate(env1, 10);
+    checkLE(new_base_state + 10, getAmountAllocated());
+    def = doAllocate(env2, 10);
+    checkLE(new_base_state + 20, getAmountAllocated());
+    doDeallocate(env1, abc);
+    checkLE(new_base_state + 10, getAmountAllocated());
+
+    doDeallocate(env2, def);
+
+    checkEq(new_base_state + 0, getAmountAllocated());
+
+    destroyJvmtiEnv(env1);
+    destroyJvmtiEnv(env2);
+
+    // Back to normal after getting rid of the envs.
+    checkEq(base_state + 0, getAmountAllocated());
+  }
+
+  private static native long doAllocate(long jvmtienv, long size);
+  private static long doAllocate(long size) {
+    return doAllocate(getDefaultJvmtiEnv(), size);
+  }
+
+  private static native void doDeallocate(long jvmtienv, long ptr);
+  private static void doDeallocate(long size) {
+    doDeallocate(getDefaultJvmtiEnv(), size);
+  }
+
+  private static native long getDefaultJvmtiEnv();
+  private static native long newJvmtiEnv();
+  private static native void destroyJvmtiEnv(long jvmtienv);
+  private static native long getAmountAllocated();
+  private static native void initializeTest();
+}
diff --git a/test/Android.bp b/test/Android.bp
index e5d2ca0863..bf3df59da4 100644
--- a/test/Android.bp
+++ b/test/Android.bp
@@ -286,6 +286,7 @@ art_cc_defaults {
         "992-source-data/source_file.cc",
         "993-breakpoints/breakpoints.cc",
         "996-breakpoint-obsolete/obsolete_breakpoints.cc",
+        "1900-track-alloc/alloc.cc",
         "1901-get-bytecodes/bytecodes.cc",
     ],
     shared_libs: [
diff --git a/test/knownfailures.json b/test/knownfailures.json
index d3d92c651e..3ee4299472 100644
--- a/test/knownfailures.json
+++ b/test/knownfailures.json
@@ -538,7 +538,8 @@
             "595-profile-saving",
             "900-hello-plugin",
             "909-attach-agent",
-            "981-dedup-original-dex"
+            "981-dedup-original-dex",
+            "1900-track-alloc"
         ],
         "description": ["Tests that require exact knowledge of the number of plugins and agents."],
         "variant": "jvmti-stress | redefine-stress | trace-stress | field-stress | step-stress"