From aaed4feb28d99bc05a9acc6ca7e66295f62a6b9a Mon Sep 17 00:00:00 2001
From: Aliaksei Kandratsenka <alkondratenko@gmail.com>
Date: Fri, 27 Sep 2024 14:42:22 -0400
Subject: [PATCH] [heap-profiler] use regular stacktracing API

Instead of MallocHook::GetCallerStackTrace.

Thing is, GetCallerStackTrace isn't reliable beyond ELF systems, like
OSX. And, yet, things just work without it for e.g. heap
sampling. Why? Because pprof already knows how to exclude
tcmalloc-internal stack frames (by looking at e.g. tcmalloc::
namespace). So we do the same for heap profiler.

This fixes heap profiling unit tests on OSX.
---
 src/heap-profile-table.cc |  7 -------
 src/heap-profile-table.h  | 12 ------------
 src/heap-profiler.cc      | 32 ++++++++++++++------------------
 3 files changed, 14 insertions(+), 37 deletions(-)

diff --git a/src/heap-profile-table.cc b/src/heap-profile-table.cc
index 557d214..0f4ba20 100644
--- a/src/heap-profile-table.cc
+++ b/src/heap-profile-table.cc
@@ -96,7 +96,6 @@ const char HeapProfileTable::kFileExt[] = ".heap";
 //----------------------------------------------------------------------
 
 static const int kHashTableSize = 179999;   // Size for bucket_table_.
-/*static*/ const int HeapProfileTable::kMaxStackDepth;
 
 //----------------------------------------------------------------------
 
@@ -187,12 +186,6 @@ HeapProfileTable::Bucket* HeapProfileTable::GetBucket(int depth,
   return b;
 }
 
-int HeapProfileTable::GetCallerStackTrace(
-    int skip_count, void* stack[kMaxStackDepth]) {
-  return MallocHook::GetCallerStackTrace(
-      stack, kMaxStackDepth, kStripFrames + skip_count + 1);
-}
-
 void HeapProfileTable::RecordAlloc(
     const void* ptr, size_t bytes, int stack_depth,
     const void* const call_stack[]) {
diff --git a/src/heap-profile-table.h b/src/heap-profile-table.h
index c074f97..7c266c3 100644
--- a/src/heap-profile-table.h
+++ b/src/heap-profile-table.h
@@ -55,9 +55,6 @@ class HeapProfileTable {
   // Extension to be used for heap pforile files.
   static const char kFileExt[];
 
-  // Longest stack trace we record.
-  static const int kMaxStackDepth = 32;
-
   // data types ----------------------------
 
   // Profile stats.
@@ -81,15 +78,6 @@ class HeapProfileTable {
   HeapProfileTable(Allocator alloc, DeAllocator dealloc);
   ~HeapProfileTable();
 
-  // Collect the stack trace for the function that asked to do the
-  // allocation for passing to RecordAlloc() below.
-  //
-  // The stack trace is stored in 'stack'. The stack depth is returned.
-  //
-  // 'skip_count' gives the number of stack frames between this call
-  // and the memory allocation function.
-  static int GetCallerStackTrace(int skip_count, void* stack[kMaxStackDepth]);
-
   // Record an allocation at 'ptr' of 'bytes' bytes.  'stack_depth'
   // and 'call_stack' identifying the function that requested the
   // allocation. They can be generated using GetCallerStackTrace() above.
diff --git a/src/heap-profiler.cc b/src/heap-profiler.cc
index 7151625..9c7045d 100644
--- a/src/heap-profiler.cc
+++ b/src/heap-profiler.cc
@@ -67,6 +67,7 @@
 #include "base/low_level_alloc.h"
 #include "base/sysinfo.h"      // for GetUniquePathFromEnv()
 #include "heap-profile-table.h"
+#include "malloc_backtrace.h"
 
 #ifndef	PATH_MAX
 #ifdef MAXPATHLEN
@@ -273,11 +274,18 @@ static void MaybeDumpProfileLocked() {
   }
 }
 
+//----------------------------------------------------------------------
+// Allocation/deallocation hooks for MallocHook
+//----------------------------------------------------------------------
+
 // Record an allocation in the profile.
-static void RecordAlloc(const void* ptr, size_t bytes, int skip_count) {
+static void NewHook(const void* ptr, size_t bytes) {
+  if (!ptr) return;
+
   // Take the stack trace outside the critical section.
-  void* stack[HeapProfileTable::kMaxStackDepth];
-  int depth = HeapProfileTable::GetCallerStackTrace(skip_count + 1, stack);
+  static constexpr int kDepth = 32;
+  void* stack[kDepth];
+  int depth = tcmalloc::GrabBacktrace(stack, kDepth, 1);
   SpinLockHolder l(&heap_lock);
   if (is_on) {
     heap_profile->RecordAlloc(ptr, bytes, depth, stack);
@@ -286,7 +294,9 @@ static void RecordAlloc(const void* ptr, size_t bytes, int skip_count) {
 }
 
 // Record a deallocation in the profile.
-static void RecordFree(const void* ptr) {
+static void DeleteHook(const void* ptr) {
+  if (!ptr) return;
+
   SpinLockHolder l(&heap_lock);
   if (is_on) {
     heap_profile->RecordFree(ptr);
@@ -294,20 +304,6 @@ static void RecordFree(const void* ptr) {
   }
 }
 
-//----------------------------------------------------------------------
-// Allocation/deallocation hooks for MallocHook
-//----------------------------------------------------------------------
-
-// static
-void NewHook(const void* ptr, size_t size) {
-  if (ptr != nullptr) RecordAlloc(ptr, size, 0);
-}
-
-// static
-void DeleteHook(const void* ptr) {
-  if (ptr != nullptr) RecordFree(ptr);
-}
-
 //----------------------------------------------------------------------
 // Starting/stopping/dumping
 //----------------------------------------------------------------------