Track jvmti allocations related to object tagging

Object tagging overhead can be significant. We now surface that
overhead via the jvmti-extension for memory use.

Test: ./test.py --host -j40
Bug: 62065509

Change-Id: Id0b98e74d66a1a99ac89186176ade39c922569cd
diff --git a/runtime/openjdkjvmti/jvmti_allocator.h b/runtime/openjdkjvmti/jvmti_allocator.h
index 1225c14..44b1cb1 100644
--- a/runtime/openjdkjvmti/jvmti_allocator.h
+++ b/runtime/openjdkjvmti/jvmti_allocator.h
@@ -36,6 +36,8 @@
 #include "base/macros.h"
 #include "jvmti.h"
 
+#include "ti_allocator.h"
+
 namespace openjdkjvmti {
 
 template <typename T> class JvmtiAllocator;
@@ -53,6 +55,7 @@
   };
 
   explicit JvmtiAllocator(jvmtiEnv* env) : env_(env) {}
+  explicit JvmtiAllocator() : env_(nullptr) {}
 
   template <typename U>
   JvmtiAllocator(const JvmtiAllocator<U>& other)  // NOLINT, implicit
@@ -89,6 +92,7 @@
   };
 
   explicit JvmtiAllocator(jvmtiEnv* env) : env_(env) {}
+  explicit JvmtiAllocator() : env_(nullptr) {}
 
   template <typename U>
   JvmtiAllocator(const JvmtiAllocator<U>& other)  // NOLINT, implicit
@@ -108,8 +112,8 @@
   pointer allocate(size_type n, JvmtiAllocator<void>::pointer hint ATTRIBUTE_UNUSED = nullptr) {
     DCHECK_LE(n, max_size());
     if (env_ == nullptr) {
-      T* result = reinterpret_cast<T*>(malloc(n * sizeof(T)));
-      CHECK(result != nullptr || n == 0u);  // Abort if malloc() fails.
+      T* result = reinterpret_cast<T*>(AllocUtil::AllocateImpl(n * sizeof(T)));
+      CHECK(result != nullptr || n == 0u);  // Abort if AllocateImpl() fails.
       return result;
     } else {
       unsigned char* result;
@@ -120,7 +124,7 @@
   }
   void deallocate(pointer p, size_type n ATTRIBUTE_UNUSED) {
     if (env_ == nullptr) {
-      free(p);
+      AllocUtil::DeallocateImpl(reinterpret_cast<unsigned char*>(p));
     } else {
       jvmtiError dealloc_error = env_->Deallocate(reinterpret_cast<unsigned char*>(p));
       CHECK(dealloc_error == JVMTI_ERROR_NONE);
diff --git a/runtime/openjdkjvmti/jvmti_weak_table.h b/runtime/openjdkjvmti/jvmti_weak_table.h
index 01c24b1..a5175a4 100644
--- a/runtime/openjdkjvmti/jvmti_weak_table.h
+++ b/runtime/openjdkjvmti/jvmti_weak_table.h
@@ -40,6 +40,7 @@
 #include "gc_root-inl.h"
 #include "globals.h"
 #include "jvmti.h"
+#include "jvmti_allocator.h"
 #include "mirror/object.h"
 #include "thread-current-inl.h"
 
@@ -191,7 +192,7 @@
       REQUIRES_SHARED(art::Locks::mutator_lock_)
       REQUIRES(allow_disallow_lock_);
 
-  template <typename Storage, class Allocator = std::allocator<T>>
+  template <typename Storage, class Allocator = JvmtiAllocator<T>>
   struct ReleasableContainer;
 
   struct HashGcRoot {
@@ -209,10 +210,12 @@
     }
   };
 
+  using TagAllocator = JvmtiAllocator<std::pair<const art::GcRoot<art::mirror::Object>, T>>;
   std::unordered_map<art::GcRoot<art::mirror::Object>,
                      T,
                      HashGcRoot,
-                     EqGcRoot> tagged_objects_
+                     EqGcRoot,
+                     TagAllocator> tagged_objects_
       GUARDED_BY(allow_disallow_lock_)
       GUARDED_BY(art::Locks::mutator_lock_);
   // To avoid repeatedly scanning the whole table, remember if we did that since the last sweep.
diff --git a/runtime/openjdkjvmti/ti_allocator.cc b/runtime/openjdkjvmti/ti_allocator.cc
index b82c4f4..8a0237d 100644
--- a/runtime/openjdkjvmti/ti_allocator.cc
+++ b/runtime/openjdkjvmti/ti_allocator.cc
@@ -59,20 +59,31 @@
     *mem_ptr = nullptr;
     return OK;
   }
-  *mem_ptr = reinterpret_cast<unsigned char*>(malloc(size));
+  *mem_ptr = AllocateImpl(size);
   if (UNLIKELY(*mem_ptr == nullptr)) {
     return ERR(OUT_OF_MEMORY);
   }
-  allocated += malloc_usable_size(*mem_ptr);
   return OK;
 }
 
+unsigned char* AllocUtil::AllocateImpl(jlong size) {
+  unsigned char* ret = size != 0 ? reinterpret_cast<unsigned char*>(malloc(size)) : nullptr;
+  if (LIKELY(ret != nullptr)) {
+    allocated += malloc_usable_size(ret);
+  }
+  return ret;
+}
+
 jvmtiError AllocUtil::Deallocate(jvmtiEnv* env ATTRIBUTE_UNUSED, unsigned char* mem) {
+  DeallocateImpl(mem);
+  return OK;
+}
+
+void AllocUtil::DeallocateImpl(unsigned char* mem) {
   if (mem != nullptr) {
     allocated -= malloc_usable_size(mem);
     free(mem);
   }
-  return OK;
 }
 
 }  // namespace openjdkjvmti
diff --git a/runtime/openjdkjvmti/ti_allocator.h b/runtime/openjdkjvmti/ti_allocator.h
index aba77ae..35575c3 100644
--- a/runtime/openjdkjvmti/ti_allocator.h
+++ b/runtime/openjdkjvmti/ti_allocator.h
@@ -36,16 +36,27 @@
 #include "jvmti.h"
 
 #include <atomic>
+#include <memory>
 
 namespace openjdkjvmti {
 
+template<typename T>
+class JvmtiAllocator;
+
 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 void DeallocateImpl(unsigned char* mem);
+  static unsigned char* AllocateImpl(jlong size);
+
   static std::atomic<jlong> allocated;
+
+  template <typename T>
+  friend class JvmtiAllocator;
 };
 
 }  // namespace openjdkjvmti
diff --git a/test/1900-track-alloc/src/art/Main.java b/test/1900-track-alloc/src/art/Main.java
new file mode 100644
index 0000000..aa5498b
--- /dev/null
+++ b/test/1900-track-alloc/src/art/Main.java
@@ -0,0 +1,32 @@
+/*
+ * Copyright (C) 2017 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;
+
+// Binder class so the agent's C code has something that can be bound and exposed to tests.
+// In a package to separate cleanly and work around CTS reference issues (though this class
+// should be replaced in the CTS version).
+public class Main {
+  // Load the given class with the given classloader, and bind all native methods to corresponding
+  // C methods in the agent. Will abort if any of the steps fail.
+  public static native void bindAgentJNI(String className, ClassLoader classLoader);
+  // Same as above, giving the class directly.
+  public static native void bindAgentJNIForClass(Class<?> klass);
+
+  // Common infrastructure.
+  public static native void setTag(Object o, long tag);
+  public static native long getTag(Object o);
+}
diff --git a/test/1900-track-alloc/src/art/Test1900.java b/test/1900-track-alloc/src/art/Test1900.java
index 717999b..becee1b 100644
--- a/test/1900-track-alloc/src/art/Test1900.java
+++ b/test/1900-track-alloc/src/art/Test1900.java
@@ -124,6 +124,15 @@
 
     // Back to normal after getting rid of the envs.
     checkEq(base_state + 0, getAmountAllocated());
+
+    // Try adding some tags
+    Object a = new Object();
+    Object b = new Object();
+    Main.setTag(a, 100);
+    Main.setTag(b, 200);
+
+    // tags should be counted and should have some data associated with them.
+    checkLE(base_state + 1, getAmountAllocated());
   }
 
   private static native long doAllocate(long jvmtienv, long size);