Enable sharding of AllocationRegister on desktop.

third_party/WebKit/Source/platform/PartitionAllocMemoryDumpProvider.cpp

 // Copyright 2015 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "platform/PartitionAllocMemoryDumpProvider.h"
 
 #include <unordered_map>
 
 #include "base/strings/stringprintf.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/process_memory_dump.h"
+#include "base/trace_event/sharded_allocation_register.h"
 #include "base/trace_event/trace_event_memory_overhead.h"
 #include "platform/wtf/allocator/Partitions.h"
 #include "platform/wtf/text/WTFString.h"
 
 namespace blink {
 
 namespace {
 
 using namespace WTF;
 
@@ skipping 104 matching lines @@
   DEFINE_STATIC_LOCAL(PartitionAllocMemoryDumpProvider, instance, ());
   return &instance;
 }
 
 bool PartitionAllocMemoryDumpProvider::OnMemoryDump(
     const base::trace_event::MemoryDumpArgs& args,
     base::trace_event::ProcessMemoryDump* memory_dump) {
   using base::trace_event::MemoryDumpLevelOfDetail;
 
   MemoryDumpLevelOfDetail level_of_detail = args.level_of_detail;
-  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_);
-      // Dump only the overhead estimation in non-detailed dumps.
-      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);
+    // Dump only the overhead estimation in non-detailed dumps.
+    if (level_of_detail == MemoryDumpLevelOfDetail::DETAILED) {
+      allocation_register_->UpdateAndReturnsMetrics(metrics_by_context);
     }
+    allocation_register_->EstimateTraceMemoryOverhead(&overhead);
     memory_dump->DumpHeapUsage(metrics_by_context, overhead, "partition_alloc");
   }
 
   PartitionStatsDumperImpl partition_stats_dumper(memory_dump, level_of_detail);
 
   base::trace_event::MemoryAllocatorDump* partitions_dump =
       memory_dump->CreateAllocatorDump(base::StringPrintf(
           "%s/%s", kPartitionAllocDumpName, kPartitionsDumpName));
 
   // This method calls memoryStats.partitionsDumpBucketStats with memory
   // statistics.
   WTF::Partitions::DumpMemoryStats(
       level_of_detail != MemoryDumpLevelOfDetail::DETAILED,
       &partition_stats_dumper);
 
   base::trace_event::MemoryAllocatorDump* allocated_objects_dump =
       memory_dump->CreateAllocatorDump(Partitions::kAllocatedObjectPoolName);
   allocated_objects_dump->AddScalar("size", "bytes",
                                     partition_stats_dumper.TotalActiveBytes());
   memory_dump->AddOwnershipEdge(allocated_objects_dump->guid(),
                                 partitions_dump->guid());
 
   return true;
 }
 
 // |m_allocationRegister| should be initialized only when necessary to avoid
 // waste of memory.
 PartitionAllocMemoryDumpProvider::PartitionAllocMemoryDumpProvider()
-    : allocation_register_(nullptr), is_heap_profiling_enabled_(false) {}
+    : allocation_register_(new base::trace_event::ShardedAllocationRegister) {}
 
 PartitionAllocMemoryDumpProvider::~PartitionAllocMemoryDumpProvider() {}
 
 void PartitionAllocMemoryDumpProvider::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();
     WTF::PartitionAllocHooks::SetAllocationHook(ReportAllocation);
     WTF::PartitionAllocHooks::SetFreeHook(ReportFree);
   } else {
     WTF::PartitionAllocHooks::SetAllocationHook(nullptr);
     WTF::PartitionAllocHooks::SetFreeHook(nullptr);
   }
-  is_heap_profiling_enabled_ = enabled;
+  allocation_register_->SetEnabled(enabled);
 }
 
 void PartitionAllocMemoryDumpProvider::insert(void* address,
                                               size_t size,
                                               const char* type_name) {
   base::trace_event::AllocationContext context;
   if (!base::trace_event::AllocationContextTracker::
            GetInstanceForCurrentThread()
                ->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 PartitionAllocMemoryDumpProvider::Remove(void* address) {
-  MutexLocker locker(allocation_register_mutex_);
-  if (allocation_register_)
-    allocation_register_->Remove(address);
+  if (!allocation_register_->IsEnabled())
+    return;
+  allocation_register_->Remove(address);
 }
 
 }  // namespace blink