[Tracing] Add lookup support to AllocationRegister

This adds a |Get| method to |AllocationRegister| which can be used to
change the context and size of an allocation after it has been inserted.
One possible use case of this is to post-annotate allocations with type
information, when this information is not available to the allocator,
but it is later in the program.

Because |AllocationRegister::Allocation| is now exposed mutably, its
|address| member has been made const. The hash of the address has to
match the bucket in which the |Allocation| is stored, so after insertion
the address must not be changed. This is enforced by making the member
const, at the cost of a few const violations in the implementation of
|AllocationRegister|.

BUG=574756

Review URL: https://codereview.chromium.org/1574493002

Cr-Commit-Position: refs/heads/master@{#368583}

3 files changed, 77 insertions, 11 deletions

diff --git a/base/trace_event/heap_profiler_allocation_register.cc b/base/trace_event/heap_profiler_allocation_register.cc
index 825a0d0..9a846c5 100644
--- a/base/trace_event/heap_profiler_allocation_register.cc
+++ b/base/trace_event/heap_profiler_allocation_register.cc
@@ -43,7 +43,11 @@ void AllocationRegister::Insert(void* address,
   if (*idx_ptr == 0) {
     *idx_ptr = GetFreeCell();
 
-    cells_[*idx_ptr].allocation.address = address;
+    // The address stored in a cell is const as long as it is exposed (via the
+    // iterators or |Get|), but because cells are re-used, a const cast is
+    // required to set it on insert and remove.
+    void* const& allocation_address = cells_[*idx_ptr].allocation.address;
+    const_cast<void*&>(allocation_address) = address;
     cells_[*idx_ptr].next = 0;
   }
 
@@ -71,7 +75,14 @@ void AllocationRegister::Remove(void* address) {
   free_list_ = freed_idx;
 
   // Reset the address, so that on iteration the free cell is ignored.
-  freed_cell->allocation.address = nullptr;
+  const_cast<void*&>(freed_cell->allocation.address) = nullptr;
+}
+
+AllocationRegister::Allocation* AllocationRegister::Get(void* address) {
+  CellIndex* idx_ptr = Lookup(address);
+
+  // If the index is 0, the address is not present in the table.
+  return *idx_ptr == 0 ? nullptr : &cells_[*idx_ptr].allocation;
 }
 
 AllocationRegister::ConstIterator AllocationRegister::begin() const {
diff --git a/base/trace_event/heap_profiler_allocation_register.h b/base/trace_event/heap_profiler_allocation_register.h
index b8afa6b..804be89 100644
--- a/base/trace_event/heap_profiler_allocation_register.h
+++ b/base/trace_event/heap_profiler_allocation_register.h
@@ -34,7 +34,7 @@ class BASE_EXPORT AllocationRegister {
   // The data stored in the hash table;
   // contains the details about an allocation.
   struct Allocation {
-    void* address;
+    void* const address;
     size_t size;
     AllocationContext context;
   };
@@ -75,6 +75,11 @@ class BASE_EXPORT AllocationRegister {
   // with a null pointer.
   void Remove(void* address);
 
+  // Returns a pointer to the allocation at the address, or null if there is no
+  // allocation at that address. This can be used to change the allocation
+  // context after insertion, for example to change the type name.
+  Allocation* Get(void* address);
+
   ConstIterator begin() const;
   ConstIterator end() const;
 
diff --git a/base/trace_event/heap_profiler_allocation_register_unittest.cc b/base/trace_event/heap_profiler_allocation_register_unittest.cc
index b74f052..3ec4580 100644
--- a/base/trace_event/heap_profiler_allocation_register_unittest.cc
+++ b/base/trace_event/heap_profiler_allocation_register_unittest.cc
@@ -110,20 +110,24 @@ TEST_F(AllocationRegisterTest, DoubleInsertOverwrites) {
   ctx.backtrace.frames[0] = frame1;
   reg.Insert(reinterpret_cast<void*>(1), 11, ctx);
 
-  auto elem = *reg.begin();
+  {
+    AllocationRegister::Allocation elem = *reg.begin();
 
-  EXPECT_EQ(frame1, elem.context.backtrace.frames[0]);
-  EXPECT_EQ(11u, elem.size);
-  EXPECT_EQ(reinterpret_cast<void*>(1), elem.address);
+    EXPECT_EQ(frame1, elem.context.backtrace.frames[0]);
+    EXPECT_EQ(11u, elem.size);
+    EXPECT_EQ(reinterpret_cast<void*>(1), elem.address);
+  }
 
   ctx.backtrace.frames[0] = frame2;
   reg.Insert(reinterpret_cast<void*>(1), 13, ctx);
 
-  elem = *reg.begin();
+  {
+    AllocationRegister::Allocation elem = *reg.begin();
 
-  EXPECT_EQ(frame2, elem.context.backtrace.frames[0]);
-  EXPECT_EQ(13u, elem.size);
-  EXPECT_EQ(reinterpret_cast<void*>(1), elem.address);
+    EXPECT_EQ(frame2, elem.context.backtrace.frames[0]);
+    EXPECT_EQ(13u, elem.size);
+    EXPECT_EQ(reinterpret_cast<void*>(1), elem.address);
+  }
 }
 
 // Check that even if more entries than the number of buckets are inserted, the
@@ -205,6 +209,52 @@ TEST_F(AllocationRegisterTest, InsertRemoveRandomOrder) {
   ASSERT_EQ(prime - 1, GetHighWaterMark(reg) - initial_water_mark);
 }
 
+TEST_F(AllocationRegisterTest, ChangeContextAfterInsertion) {
+  using Allocation = AllocationRegister::Allocation;
+  const char kStdString[] = "std::string";
+  AllocationRegister reg;
+  AllocationContext ctx = AllocationContext::Empty();
+
+  reg.Insert(reinterpret_cast<void*>(17), 1, ctx);
+  reg.Insert(reinterpret_cast<void*>(19), 2, ctx);
+  reg.Insert(reinterpret_cast<void*>(23), 3, ctx);
+
+  // Looking up addresses that were not inserted should return null.
+  // A null pointer lookup is a valid thing to do.
+  EXPECT_EQ(nullptr, reg.Get(nullptr));
+  EXPECT_EQ(nullptr, reg.Get(reinterpret_cast<void*>(13)));
+
+  Allocation* a17 = reg.Get(reinterpret_cast<void*>(17));
+  Allocation* a19 = reg.Get(reinterpret_cast<void*>(19));
+  Allocation* a23 = reg.Get(reinterpret_cast<void*>(23));
+
+  EXPECT_NE(nullptr, a17);
+  EXPECT_NE(nullptr, a19);
+  EXPECT_NE(nullptr, a23);
+
+  a17->size = 100;
+  a19->context.type_name = kStdString;
+
+  reg.Remove(reinterpret_cast<void*>(23));
+
+  // Lookup should not find any garbage after removal.
+  EXPECT_EQ(nullptr, reg.Get(reinterpret_cast<void*>(23)));
+
+  // Mutating allocations should have modified the allocations in the register.
+  for (const Allocation& allocation : reg) {
+     if (allocation.address == reinterpret_cast<void*>(17))
+       EXPECT_EQ(100u, allocation.size);
+     if (allocation.address == reinterpret_cast<void*>(19))
+       EXPECT_EQ(kStdString, allocation.context.type_name);
+  }
+
+  reg.Remove(reinterpret_cast<void*>(17));
+  reg.Remove(reinterpret_cast<void*>(19));
+
+  EXPECT_EQ(nullptr, reg.Get(reinterpret_cast<void*>(17)));
+  EXPECT_EQ(nullptr, reg.Get(reinterpret_cast<void*>(19)));
+}
+
 // Check that the process aborts due to hitting the guard page when inserting
 // too many elements.
 #if GTEST_HAS_DEATH_TEST