Enable sharding of AllocationRegister on desktop.

third_party/WebKit/Source/platform/heap/BlinkGCMemoryDumpProvider.cpp

Index: third_party/WebKit/Source/platform/heap/BlinkGCMemoryDumpProvider.cpp
diff --git a/third_party/WebKit/Source/platform/heap/BlinkGCMemoryDumpProvider.cpp b/third_party/WebKit/Source/platform/heap/BlinkGCMemoryDumpProvider.cpp
index baefa38a3657d96e2ca457172b38070de389e552..12385be3f9e4d662c1827b790aeb7b59c70089f4 100644
--- a/third_party/WebKit/Source/platform/heap/BlinkGCMemoryDumpProvider.cpp
+++ b/third_party/WebKit/Source/platform/heap/BlinkGCMemoryDumpProvider.cpp
@@ -7,9 +7,9 @@
 #include <unordered_map>
 
 #include "base/trace_event/heap_profiler_allocation_context_tracker.h"
-#include "base/trace_event/heap_profiler_allocation_register.h"
 #include "base/trace_event/memory_allocator_dump.h"
 #include "base/trace_event/process_memory_dump.h"
+#include "base/trace_event/sharded_allocation_register.h"
 #include "base/trace_event/trace_event_memory_overhead.h"
 #include "platform/heap/Handle.h"
 #include "platform/instrumentation/tracing/web_memory_allocator_dump.h"
@@ -66,24 +66,16 @@ bool BlinkGCMemoryDumpProvider::OnMemoryDump(
                                            BlinkGC::kForcedGC);
   DumpMemoryTotals(memory_dump);
 
-  if (is_heap_profiling_enabled_) {
+  if (allocation_register_->IsEnabled()) {
     // Overhead should always be reported, regardless of light vs. heavy.
     base::trace_event::TraceEventMemoryOverhead overhead;
     std::unordered_map<base::trace_event::AllocationContext,
                        base::trace_event::AllocationMetrics>
         metrics_by_context;
-    {
-      MutexLocker locker(allocation_register_mutex_);
-      if (level_of_detail == MemoryDumpLevelOfDetail::DETAILED) {
-        for (const auto& alloc_size : *allocation_register_) {
-          base::trace_event::AllocationMetrics& metrics =
-              metrics_by_context[alloc_size.context];
-          metrics.size += alloc_size.size;
-          metrics.count++;
-        }
-      }
-      allocation_register_->EstimateTraceMemoryOverhead(&overhead);
+    if (level_of_detail == MemoryDumpLevelOfDetail::DETAILED) {
+      allocation_register_->UpdateAndReturnsMetrics(metrics_by_context);
     }
+    allocation_register_->EstimateTraceMemoryOverhead(&overhead);
     memory_dump->DumpHeapUsage(metrics_by_context, overhead, "blink_gc");
   }
 
@@ -95,18 +87,15 @@ bool BlinkGCMemoryDumpProvider::OnMemoryDump(
 
 void BlinkGCMemoryDumpProvider::OnHeapProfilingEnabled(bool enabled) {
   if (enabled) {
-    {
-      MutexLocker locker(allocation_register_mutex_);
-      if (!allocation_register_)
-        allocation_register_.reset(new base::trace_event::AllocationRegister());
-    }
+    if (!allocation_register_->IsInitialized())
+      allocation_register_->Initialize();
     HeapAllocHooks::SetAllocationHook(ReportAllocation);
     HeapAllocHooks::SetFreeHook(ReportFree);
   } else {
     HeapAllocHooks::SetAllocationHook(nullptr);
     HeapAllocHooks::SetFreeHook(nullptr);
   }
-  is_heap_profiling_enabled_ = enabled;
+  allocation_register_->SetEnabled(enabled);
 }
 
 base::trace_event::MemoryAllocatorDump*
@@ -122,10 +111,10 @@ void BlinkGCMemoryDumpProvider::ClearProcessDumpForCurrentGC() {
 }
 
 BlinkGCMemoryDumpProvider::BlinkGCMemoryDumpProvider()
-    : current_process_memory_dump_(new base::trace_event::ProcessMemoryDump(
+    : allocation_register_(new base::trace_event::ShardedAllocationRegister),
+      current_process_memory_dump_(new base::trace_event::ProcessMemoryDump(
           nullptr,
-          {base::trace_event::MemoryDumpLevelOfDetail::DETAILED})),
-      is_heap_profiling_enabled_(false) {}
+          {base::trace_event::MemoryDumpLevelOfDetail::DETAILED})) {}
 
 void BlinkGCMemoryDumpProvider::insert(Address address,
                                        size_t size,
@@ -136,15 +125,15 @@ void BlinkGCMemoryDumpProvider::insert(Address address,
                ->GetContextSnapshot(&context))
     return;
   context.type_name = type_name;
-  MutexLocker locker(allocation_register_mutex_);
-  if (allocation_register_)
-    allocation_register_->Insert(address, size, context);
+  if (!allocation_register_->IsEnabled())
+    return;
+  allocation_register_->Insert(address, size, context);
 }
 
 void BlinkGCMemoryDumpProvider::Remove(Address address) {
-  MutexLocker locker(allocation_register_mutex_);
-  if (allocation_register_)
-    allocation_register_->Remove(address);
+  if (!allocation_register_->IsEnabled())
+    return;
+  allocation_register_->Remove(address);
 }
 
 }  // namespace blink