Add API to ART/Perfetto Java Heap Profiler

Add Perfetto API calls to ART Sampling Java Heap Profiler.

Test: Passing Tests
Local Testing

Bug: 160214819
Change-Id: I4f4740eed89bd7c762888d9f032e6bed970db634
diff --git a/runtime/Android.bp b/runtime/Android.bp
index 913a32f..403a4c3 100644
--- a/runtime/Android.bp
+++ b/runtime/Android.bp
@@ -380,6 +380,7 @@
                 "libicu",
                 "libstatssocket",
                 "libz", // For adler32.
+                "heapprofd_client_api",
             ],
             static_libs: [
                 "libstatslog_art",
diff --git a/runtime/gc/heap.cc b/runtime/gc/heap.cc
index 4ca6bf7..7cfa0fa 100644
--- a/runtime/gc/heap.cc
+++ b/runtime/gc/heap.cc
@@ -92,6 +92,9 @@
 #include "mirror/var_handle.h"
 #include "nativehelper/scoped_local_ref.h"
 #include "obj_ptr-inl.h"
+#ifdef ART_TARGET_ANDROID
+#include "perfetto/heap_profile.h"
+#endif
 #include "reflection.h"
 #include "runtime.h"
 #include "javaheapprof/javaheapsampler.h"
@@ -102,6 +105,36 @@
 
 namespace art {
 
+#ifdef ART_TARGET_ANDROID
+namespace {
+
+// Enable the heap sampler Callback function used by Perfetto.
+void EnableHeapSamplerCallback(void* enable_ptr,
+                               const AHeapProfileEnableCallbackInfo* enable_info_ptr) {
+  HeapSampler* sampler_self = reinterpret_cast<HeapSampler*>(enable_ptr);
+  // Set the ART profiler sampling interval to the value from Perfetto.
+  uint64_t interval = AHeapProfileEnableCallbackInfo_getSamplingInterval(enable_info_ptr);
+  if (interval > 0) {
+    sampler_self->SetSamplingInterval(interval);
+  }
+  // Else default is 4K sampling interval. However, default case shouldn't happen for Perfetto API.
+  // AHeapProfileEnableCallbackInfo_getSamplingInterval should always give the requested
+  // (non-negative) sampling interval. It is a uint64_t and gets checked for != 0
+  // Do not call heap as a temp here, it will build but test run will silently fail.
+  // Heap is not fully constructed yet in some cases.
+  sampler_self->EnableHeapSampler();
+}
+
+// Disable the heap sampler Callback function used by Perfetto.
+void DisableHeapSamplerCallback(void* disable_ptr,
+                                const AHeapProfileDisableCallbackInfo* info_ptr ATTRIBUTE_UNUSED) {
+  HeapSampler* sampler_self = reinterpret_cast<HeapSampler*>(disable_ptr);
+  sampler_self->DisableHeapSampler();
+}
+
+}  // namespace
+#endif
+
 namespace gc {
 
 DEFINE_RUNTIME_DEBUG_FLAG(Heap, kStressCollectorTransition);
@@ -349,7 +382,6 @@
                                         kGcCountRateMaxBucketCount),
       alloc_tracking_enabled_(false),
       alloc_record_depth_(AllocRecordObjectMap::kDefaultAllocStackDepth),
-      perfetto_javaheapprof_heapid_(0),
       backtrace_lock_(nullptr),
       seen_backtrace_count_(0u),
       unique_backtrace_count_(0u),
@@ -750,7 +782,7 @@
     InitPerfettoJavaHeapProf();
   } else {
     // Disable the Java Heap Profiler.
-    GetHeapSampler().DisableHeapSampler(/*disable_ptr=*/nullptr, /*disable_info_ptr=*/nullptr);
+    GetHeapSampler().DisableHeapSampler();
   }
 
   instrumentation::Instrumentation* const instrumentation = runtime->GetInstrumentation();
@@ -4046,15 +4078,25 @@
 
 // Perfetto initialization.
 void Heap::InitPerfettoJavaHeapProf() {
-  // Register the heap and create the heapid.
-  // Use a heap name = "HeapSampler".
   // Initialize Perfetto Heap info and Heap id.
-  static uint32_t heap_id = 1;  // Initialize to 1, to be overwritten by Perfetto heap id.
-  SetPerfettoJavaHeapProfHeapID(heap_id);
-  // Enable the Java Heap Profiler.
-  GetHeapSampler().EnableHeapSampler(/*enable_ptr=*/nullptr, /*enable_info_ptr=*/nullptr);
+  uint32_t heap_id = 1;  // Initialize to 1, to be overwritten by Perfetto heap id.
+#ifdef ART_TARGET_ANDROID
+  // Register the heap and create the heapid.
+  // Use a Perfetto heap name = "com.android.art" for the Java Heap Profiler.
+  AHeapInfo* info = AHeapInfo_create("com.android.art");
   // Set the Enable Callback, there is no callback data ("nullptr").
+  AHeapInfo_setEnabledCallback(info, &EnableHeapSamplerCallback, &heap_sampler_);
   // Set the Disable Callback.
+  AHeapInfo_setDisabledCallback(info, &DisableHeapSamplerCallback, &heap_sampler_);
+  heap_id = AHeapProfile_registerHeap(info);
+  // Do not enable the Java Heap Profiler in this case, wait for Perfetto to enable it through
+  // the callback function.
+#else
+  // This is the host case, enable the Java Heap Profiler for host testing.
+  // Perfetto API is currently not available on host.
+  heap_sampler_.EnableHeapSampler();
+#endif
+  heap_sampler_.SetHeapID(heap_id);
   VLOG(heap) << "Java Heap Profiler Initialized";
 }
 
diff --git a/runtime/gc/heap.h b/runtime/gc/heap.h
index d45f276..35dce5e 100644
--- a/runtime/gc/heap.h
+++ b/runtime/gc/heap.h
@@ -836,16 +836,6 @@
   void DumpBlockingGcCountRateHistogram(std::ostream& os) const REQUIRES(!*gc_complete_lock_);
 
   // Perfetto Art Heap Profiler Support.
-  // HeapID is a heap identifier used by the Perfetto API and is used in allocation reporting
-  // to Perfetto API.
-  void SetPerfettoJavaHeapProfHeapID(uint32_t heapid) {
-    perfetto_javaheapprof_heapid_ = heapid;
-  }
-
-  uint32_t GetPerfettoJavaHeapProfHeapID() const {
-    return perfetto_javaheapprof_heapid_;
-  }
-
   HeapSampler& GetHeapSampler() {
     return heap_sampler_;
   }
@@ -1605,7 +1595,6 @@
   size_t alloc_record_depth_;
 
   // Perfetto Java Heap Profiler support.
-  uint32_t perfetto_javaheapprof_heapid_;
   HeapSampler heap_sampler_;
 
   // GC stress related data structures.
diff --git a/runtime/javaheapprof/javaheapsampler.cc b/runtime/javaheapprof/javaheapsampler.cc
index a1c58d8..a73ed0b 100644
--- a/runtime/javaheapprof/javaheapsampler.cc
+++ b/runtime/javaheapprof/javaheapsampler.cc
@@ -18,6 +18,9 @@
 #include "base/locks.h"
 #include "gc/heap.h"
 #include "javaheapprof/javaheapsampler.h"
+#ifdef ART_TARGET_ANDROID
+#include "perfetto/heap_profile.h"
+#endif
 #include "runtime.h"
 
 namespace art {
@@ -58,8 +61,13 @@
 // Also bytes_until_sample can only be updated after the allocation and reporting is done.
 // Thus next bytes_until_sample is previously calculated (before allocation) to be able to
 // get the next tlab_size, but only saved/updated here.
-void HeapSampler::ReportSample(art::mirror::Object* obj ATTRIBUTE_UNUSED, size_t allocation_size) {
+void HeapSampler::ReportSample(art::mirror::Object* obj, size_t allocation_size) {
   VLOG(heap) << "JHP:***Report Perfetto Allocation: alloc_size: " << allocation_size;
+  uint64_t perf_alloc_id = reinterpret_cast<uint64_t>(obj);
+  VLOG(heap) << "JHP:***Report Perfetto Allocation: obj: " << perf_alloc_id;
+#ifdef ART_TARGET_ANDROID
+  AHeapProfile_reportSample(perfetto_heap_id_, perf_alloc_id, allocation_size);
+#endif
 }
 
 // Check whether we should take a sample or not at this allocation and calculate the sample
@@ -123,51 +131,19 @@
              << " next_bytes_until_sample = " << next_bytes_until_sample;
 }
 
-// Enable the heap sampler and initialize/set the sampling interval.
-void HeapSampler::EnableHeapSampler(void* enable_ptr ATTRIBUTE_UNUSED,
-                                    const void* enable_info_ptr ATTRIBUTE_UNUSED) {
-  uint64_t interval = 4 * 1024;
-  // Set the ART profiler sampling interval to the value from AHeapProfileSessionInfo
-  // Set interval to sampling interval from AHeapProfileSessionInfo
-  if (interval > 0) {
-    // Make sure that rng_ and geo_dist are thread safe by acquiring a lock to access.
-    art::MutexLock mu(art::Thread::Current(), geo_dist_rng_lock_);
-    SetSamplingInterval(interval);
-  }
-  // Else default is 4K sampling interval. However, default case shouldn't happen for Perfetto API.
-  // AHeapProfileEnableCallbackInfo_getSamplingInterval should always give the requested
-  // (non-negative) sampling interval. It is a uint64_t and gets checked for != 0
-  // Do not call heap->GetPerfettoJavaHeapProfHeapID() as a temp here, it will build but test run
-  // will silently fail. Heap is not fully constructed yet.
-  // heap_id will be set through the Perfetto API.
-  perfetto_heap_id_ = 1;  // To be set by Perfetto API
-  enabled_.store(true, std::memory_order_release);
-}
-
 bool HeapSampler::IsEnabled() {
   return enabled_.load(std::memory_order_acquire);
 }
 
-void HeapSampler::DisableHeapSampler(void* disable_ptr ATTRIBUTE_UNUSED,
-                                     const void* disable_info_ptr ATTRIBUTE_UNUSED) {
-  enabled_.store(false, std::memory_order_release);
-}
-
 int HeapSampler::GetSamplingInterval() {
   return p_sampling_interval_.load(std::memory_order_acquire);
 }
 
 void HeapSampler::SetSamplingInterval(int sampling_interval) {
+  // Make sure that rng_ and geo_dist are thread safe by acquiring a lock to access.
+  art::MutexLock mu(art::Thread::Current(), geo_dist_rng_lock_);
   p_sampling_interval_.store(sampling_interval, std::memory_order_release);
   geo_dist_.param(std::geometric_distribution<size_t>::param_type(1.0/p_sampling_interval_));
 }
 
-void HeapSampler::SetSessionInfo(void* info) {
-  perfetto_session_info_ = info;
-}
-
-void* HeapSampler::GetSessionInfo() {
-  return perfetto_session_info_;
-}
-
 }  // namespace art
diff --git a/runtime/javaheapprof/javaheapsampler.h b/runtime/javaheapprof/javaheapsampler.h
index 02cb7b7..618893c 100644
--- a/runtime/javaheapprof/javaheapsampler.h
+++ b/runtime/javaheapprof/javaheapsampler.h
@@ -42,6 +42,15 @@
     thread_local size_t bytes_until_sample = 0;
     return &bytes_until_sample;
   }
+  void SetHeapID(uint32_t heap_id) {
+    perfetto_heap_id_ = heap_id;
+  }
+  void EnableHeapSampler() {
+    enabled_.store(true, std::memory_order_release);
+  }
+  void DisableHeapSampler() {
+    enabled_.store(false, std::memory_order_release);
+  }
   // Report a sample to Perfetto.
   void ReportSample(art::mirror::Object* obj, size_t allocation_size);
   // Check whether we should take a sample or not at this allocation, and return the
@@ -60,16 +69,10 @@
   void AdjustSampleOffset(size_t adjustment);
   // Is heap sampler enabled?
   bool IsEnabled();
-  void EnableHeapSampler(void* enable_ptr, const void* enable_info_ptr);
-  void DisableHeapSampler(void* disable_ptr, const void* disable_info_ptr);
   // Set the sampling interval.
-  void SetSamplingInterval(int sampling_interval) REQUIRES(geo_dist_rng_lock_);
+  void SetSamplingInterval(int sampling_interval) REQUIRES(!geo_dist_rng_lock_);
   // Return the sampling interval.
   int GetSamplingInterval();
-  // Set the Perfetto Session Info.
-  void SetSessionInfo(void* info);
-  // Get the Perfetto Session Info.
-  void* GetSessionInfo();
 
  private:
   size_t NextGeoDistRandSample() REQUIRES(!geo_dist_rng_lock_);
@@ -81,7 +84,6 @@
   // Default sampling interval is 4kb.
   // Writes guarded by geo_dist_rng_lock_.
   std::atomic<int> p_sampling_interval_{4 * 1024};
-  void* perfetto_session_info_;
   uint32_t perfetto_heap_id_ = 0;
   // std random number generator.
   std::minstd_rand rng_ GUARDED_BY(geo_dist_rng_lock_);  // Holds the state