[Tracing] Add hook to PartitionAlloc for heap profiling

content/child/web_memory_dump_provider_adapter.cc

Index: content/child/web_memory_dump_provider_adapter.cc
diff --git a/content/child/web_memory_dump_provider_adapter.cc b/content/child/web_memory_dump_provider_adapter.cc
index ef2cab810f0ab628f41333285b72440cd991f46b..3ec5fe86074205b3ac33a48e012705251435bee5 100644
--- a/content/child/web_memory_dump_provider_adapter.cc
+++ b/content/child/web_memory_dump_provider_adapter.cc
@@ -4,9 +4,39 @@
 
 #include "content/child/web_memory_dump_provider_adapter.h"
 
+#include "base/lazy_instance.h"
+#include "base/synchronization/lock.h"
+#include "base/trace_event/memory_profiler_allocation_context.h"
+#include "base/trace_event/memory_profiler_allocation_register.h"
 #include "content/child/web_process_memory_dump_impl.h"
 #include "third_party/WebKit/public/platform/WebMemoryDumpProvider.h"
 
+using namespace base;
+using namespace base::trace_event;
+
+namespace {
+
+AllocationRegister* g_allocation_register = nullptr;
+LazyInstance<Lock>::Leaky g_allocation_register_lock =
+    LAZY_INSTANCE_INITIALIZER;
+
+void ReportAllocation(void* address, size_t size) {
+  AllocationContext context = AllocationContextTracker::GetContextSnapshot();
+  AutoLock lock(g_allocation_register_lock.Get());
+
+  if (g_allocation_register)
+    g_allocation_register->Insert(address, size, context);
+}
+
+void ReportFree(void* address) {
+  AutoLock lock(g_allocation_register_lock.Get());
+
+  if (g_allocation_register)
+    g_allocation_register->Remove(address);
+}
+
+}  // namespace
+
 namespace content {
 
 WebMemoryDumpProviderAdapter::WebMemoryDumpProviderAdapter(
@@ -35,7 +65,37 @@ bool WebMemoryDumpProviderAdapter::OnMemoryDump(
   }
   WebProcessMemoryDumpImpl web_pmd_impl(args.level_of_detail, pmd);
 
+  if (web_memory_dump_provider_->supportsHeapProfiling()) {
+    // TODO(ruuda): Dump |g_allocation_register| into the |ProcessMemoryDump|.
+  }
+
   return web_memory_dump_provider_->onMemoryDump(level, &web_pmd_impl);
 }
 
+void WebMemoryDumpProviderAdapter::OnHeapProfilingEnabled(bool enabled) {
+  if (!web_memory_dump_provider_->supportsHeapProfiling())
+    return;
+
+  if (enabled) {
+    {
+      AutoLock lock(g_allocation_register_lock.Get());
+      if (!g_allocation_register)
+        g_allocation_register = new AllocationRegister();
+    }
+
+    // Make this dump provider call the global hooks on every allocation / free.
+    // TODO(ruuda): Because bookkeeping is done here in the adapter, and not in
+    // the dump providers themselves, all dump providers in Blink share the
+    // same global allocation register. At the moment this is not a problem,
+    // because the only dump provider that supports heap profiling is the
+    // PartitionAlloc dump provider. When Blink can depend on base and this
+    // glue layer is removed, dump providers can have their own instance of the
+    // allocation register.
+    web_memory_dump_provider_->onHeapProfilingEnabled(ReportAllocation,
+                                                      ReportFree);
+  } else {
+    web_memory_dump_provider_->onHeapProfilingEnabled(nullptr, nullptr);
+  }
+}
+
 }  // namespace content