[Tracing] Add lookup support to AllocationRegister

base/trace_event/heap_profiler_allocation_register_unittest.cc

 // 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 "base/trace_event/heap_profiler_allocation_register.h"
 
 #include <stddef.h>
 #include <stdint.h>
 
 #include "base/process/process_metrics.h"
@@ skipping 92 matching lines @@
   // TODO(ruuda): Although double insert happens in practice, it should not.
   // Find out the cause and ban double insert if possible.
   AllocationRegister reg;
   AllocationContext ctx = AllocationContext::Empty();
   StackFrame frame1 = "Foo";
   StackFrame frame2 = "Bar";
 
   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
 // register still behaves correctly.
 TEST_F(AllocationRegisterTest, InsertRemoveCollisions) {
   size_t expected_sum = 0;
   AllocationRegister reg;
   AllocationContext ctx = AllocationContext::Empty();
 
   // By inserting 100 more entries than the number of buckets, there will be at
@@ skipping 61 matching lines @@
   for (uintptr_t i = 2; i < prime; i++)
     reg.Insert(reinterpret_cast<void*>(i), 0, ctx);
 
   ASSERT_EQ(prime - 2, GetHighWaterMark(reg) - initial_water_mark);
 
   // Inserting one more entry should use a fresh cell again.
   reg.Insert(reinterpret_cast<void*>(prime), 0, ctx);
   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
 TEST_F(AllocationRegisterTest, OverflowDeathTest) {
   // Use a smaller register to prevent OOM errors on low-end devices.
   AllocationRegister reg(static_cast<uint32_t>(GetNumCellsPerPage()));
   AllocationContext ctx = AllocationContext::Empty();
   uintptr_t i;
 
   // Fill up all of the memory allocated for the register. |GetNumCells(reg)|
   // minus 1 elements are inserted, because cell 0 is unused, so this should
   // fill up the available cells exactly.
   for (i = 1; i < GetNumCells(reg); i++) {
     reg.Insert(reinterpret_cast<void*>(i), 0, ctx);
   }
 
   // Adding just one extra element might still work because the allocated memory
   // is rounded up to the page size. Adding a page full of elements should cause
   // overflow.
   const size_t cells_per_page = GetNumCellsPerPage();
 
   ASSERT_DEATH(for (size_t j = 0; j < cells_per_page; j++) {
     reg.Insert(reinterpret_cast<void*>(i + j), 0, ctx);
   }, "");
 }
 #endif
 
 }  // namespace trace_event
 }  // namespace base