Enable sharding of AllocationRegister on desktop.

base/trace_event/malloc_dump_provider.cc

Index: base/trace_event/malloc_dump_provider.cc
diff --git a/base/trace_event/malloc_dump_provider.cc b/base/trace_event/malloc_dump_provider.cc
index 7b3831c08d64a723d3773ba754708a3121c69e49..21357c7ab0d5baf9e85896985ce710ac8ddb561b 100644
--- a/base/trace_event/malloc_dump_provider.cc
+++ b/base/trace_event/malloc_dump_provider.cc
@@ -14,7 +14,6 @@
 #include "base/debug/profiler.h"
 #include "base/trace_event/heap_profiler_allocation_context.h"
 #include "base/trace_event/heap_profiler_allocation_context_tracker.h"
-#include "base/trace_event/heap_profiler_allocation_register.h"
 #include "base/trace_event/heap_profiler_heap_dump_writer.h"
 #include "base/trace_event/process_memory_dump.h"
 #include "base/trace_event/trace_event_argument.h"
@@ -190,7 +189,7 @@ MallocDumpProvider* MallocDumpProvider::GetInstance() {
 }
 
 MallocDumpProvider::MallocDumpProvider()
-    : heap_profiler_enabled_(false), tid_dumping_heap_(kInvalidThreadId) {}
+    : tid_dumping_heap_(kInvalidThreadId) {}
 
 MallocDumpProvider::~MallocDumpProvider() {}
 
@@ -286,7 +285,7 @@ bool MallocDumpProvider::OnMemoryDump(const MemoryDumpArgs& args,
   }
 
   // Heap profiler dumps.
-  if (!heap_profiler_enabled_)
+  if (!allocation_register_.is_enabled())
     return true;
 
   // The dumps of the heap profiler should be created only when heap profiling
@@ -303,30 +302,22 @@ bool MallocDumpProvider::OnMemoryDump(const MemoryDumpArgs& args,
     size_t shim_allocated_objects_count = 0;
     TraceEventMemoryOverhead overhead;
     std::unordered_map<AllocationContext, AllocationMetrics> metrics_by_context;
-    {
-      AutoLock lock(allocation_register_lock_);
-      if (allocation_register_) {
-        if (args.level_of_detail == MemoryDumpLevelOfDetail::DETAILED) {
-          for (const auto& alloc_size : *allocation_register_) {
-            AllocationMetrics& metrics = metrics_by_context[alloc_size.context];
-            metrics.size += alloc_size.size;
-            metrics.count++;
-
-            // Aggregate data for objects allocated through the shim.
-            shim_allocated_objects_size += alloc_size.size;
-            shim_allocated_objects_count++;
-          }
-        }
-        allocation_register_->EstimateTraceMemoryOverhead(&overhead);
-      }
-
-      inner_dump->AddScalar("shim_allocated_objects_size",
-                             MemoryAllocatorDump::kUnitsBytes,
-                             shim_allocated_objects_size);
-      inner_dump->AddScalar("shim_allocator_object_count",
-                             MemoryAllocatorDump::kUnitsObjects,
-                             shim_allocated_objects_count);
-    }  // lock(allocation_register_lock_)
+    if (args.level_of_detail == MemoryDumpLevelOfDetail::DETAILED) {
+      ShardedAllocationRegister::OutputMetrics metrics =
+          allocation_register_.UpdateAndReturnsMetrics(metrics_by_context);
+
+      // Aggregate data for objects allocated through the shim.
+      shim_allocated_objects_size += metrics.size;
+      shim_allocated_objects_count += metrics.count;
+    }
+    allocation_register_.EstimateTraceMemoryOverhead(&overhead);
+
+    inner_dump->AddScalar("shim_allocated_objects_size",
+                          MemoryAllocatorDump::kUnitsBytes,
+                          shim_allocated_objects_size);
+    inner_dump->AddScalar("shim_allocator_object_count",
+                          MemoryAllocatorDump::kUnitsObjects,
+                          shim_allocated_objects_count);
     pmd->DumpHeapUsage(metrics_by_context, overhead, "malloc");
   }
   tid_dumping_heap_ = kInvalidThreadId;
@@ -337,20 +328,12 @@ bool MallocDumpProvider::OnMemoryDump(const MemoryDumpArgs& args,
 void MallocDumpProvider::OnHeapProfilingEnabled(bool enabled) {
 #if BUILDFLAG(USE_EXPERIMENTAL_ALLOCATOR_SHIM)
   if (enabled) {
-    {
-      AutoLock lock(allocation_register_lock_);
-      allocation_register_.reset(new AllocationRegister());
-    }
+    allocation_register_.SetEnabled();
     allocator::InsertAllocatorDispatch(&g_allocator_hooks);
   } else {
-    AutoLock lock(allocation_register_lock_);
-    allocation_register_.reset();
-    // Insert/RemoveAllocation below will no-op if the register is torn down.
-    // Once disabled, heap profiling will not re-enabled anymore for the
-    // lifetime of the process.
+    allocation_register_.SetDisabled();
   }
 #endif
-  heap_profiler_enabled_ = enabled;
 }
 
 void MallocDumpProvider::InsertAllocation(void* address, size_t size) {
@@ -372,11 +355,10 @@ void MallocDumpProvider::InsertAllocation(void* address, size_t size) {
   if (!tracker->GetContextSnapshot(&context))
     return;
 
-  AutoLock lock(allocation_register_lock_);
-  if (!allocation_register_)
+  if (!allocation_register_.is_enabled())
     return;
 
-  allocation_register_->Insert(address, size, context);
+  allocation_register_.Insert(address, size, context);
 }
 
 void MallocDumpProvider::RemoveAllocation(void* address) {
@@ -385,10 +367,9 @@ void MallocDumpProvider::RemoveAllocation(void* address) {
   if (tid_dumping_heap_ != kInvalidThreadId &&
       tid_dumping_heap_ == PlatformThread::CurrentId())
     return;
-  AutoLock lock(allocation_register_lock_);
-  if (!allocation_register_)
+  if (!allocation_register_.is_enabled())
     return;
-  allocation_register_->Remove(address);
+  allocation_register_.Remove(address);
 }
 
 }  // namespace trace_event