PartitionAlloc: PCScan: Use slot size to retrieve allocation start

This reduces the average PCScan time on the renderer process (with PA as
a malloc) from 269ms to 29ms (9x).

Bug: 11297512
Change-Id: I445f609751e25d943ed779453aa07a0142cbf5b0
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2498464
Commit-Queue: Anton Bikineev <bikineev@chromium.org>
Reviewed-by: Bartek Nowierski <bartekn@chromium.org>
Reviewed-by: Michael Lippautz <mlippautz@chromium.org>
Reviewed-by: Kentaro Hara <haraken@chromium.org>
Cr-Commit-Position: refs/heads/master@{#821696}

base/allocator/partition_allocator/object_bitmap.h

@@ -45,12 +45,6 @@ class ObjectBitmap final {
 
   inline ObjectBitmap();
 
-  // Finds the beginning of the closest object that starts at or before
-  // |address|. It may return an object from another slot if the slot where
-  // |address| lies in is unallocated. The caller is responsible for range
-  // checking. Returns |kSentinel| if no object was found.
-  inline uintptr_t FindPotentialObjectBeginning(uintptr_t address) const;
-
   inline void SetBit(uintptr_t address);
   inline void ClearBit(uintptr_t address);
   inline bool CheckBit(uintptr_t address) const;
@@ -89,35 +83,6 @@ template <size_t PageSize, size_t PageAlignment, size_t ObjectAlignment>
 inline ObjectBitmap<PageSize, PageAlignment, ObjectAlignment>::ObjectBitmap() =
     default;
 
-template <size_t PageSize, size_t PageAlignment, size_t ObjectAlignment>
-uintptr_t ObjectBitmap<PageSize, PageAlignment, ObjectAlignment>::
-    FindPotentialObjectBeginning(uintptr_t address) const {
-  const uintptr_t page_base = reinterpret_cast<uintptr_t>(this) & kPageBaseMask;
-  PA_DCHECK(page_base <= address && address < page_base + kPageSize);
-
-  size_t cell_index, bit;
-  std::tie(cell_index, bit) = ObjectIndexAndBit(address);
-
-  // Find the first set bit at or before |bit|.
-  uint8_t byte = LoadCell(cell_index) & ((1 << (bit + 1)) - 1);
-  while (!byte && cell_index) {
-    PA_DCHECK(0u < cell_index);
-    byte = LoadCell(--cell_index);
-  }
-
-  if (!byte) {
-    // No object was found.
-    return kSentinel;
-  }
-
-  const int leading_zeroes = base::bits::CountLeadingZeroBits(byte);
-  const size_t object_number =
-      (cell_index * kBitsPerCell) + (kBitsPerCell - 1) - leading_zeroes;
-  const size_t offset_in_page = object_number * kObjectAlignment;
-
-  return offset_in_page + page_base;
-}
-
 template <size_t PageSize, size_t PageAlignment, size_t ObjectAlignment>
 void ObjectBitmap<PageSize, PageAlignment, ObjectAlignment>::SetBit(
     uintptr_t address) {

base/allocator/partition_allocator/object_bitmap_unittest.cc

@@ -141,30 +141,5 @@ TEST_F(ObjectBitmapTest, AdjacentObjectsAtEnd) {
   EXPECT_EQ(2u, count);
 }
 
-TEST_F(ObjectBitmapTest, FindElementSentinel) {
-  EXPECT_EQ(TestBitmap::kSentinel,
-            bitmap().FindPotentialObjectBeginning(ObjectAddress(654)));
-}
-
-TEST_F(ObjectBitmapTest, FindElementExact) {
-  SetBitForObject(654);
-  EXPECT_EQ(ObjectAddress(654),
-            bitmap().FindPotentialObjectBeginning(ObjectAddress(654)));
-}
-
-TEST_F(ObjectBitmapTest, FindElementApproximate) {
-  static const size_t kInternalDelta = 37;
-  SetBitForObject(654);
-  EXPECT_EQ(ObjectAddress(654), bitmap().FindPotentialObjectBeginning(
-                                    ObjectAddress(654 + kInternalDelta)));
-}
-
-TEST_F(ObjectBitmapTest, FindElementIteratingWholeBitmap) {
-  SetBitForObject(0);
-  const uintptr_t hint_index = LastIndex();
-  EXPECT_EQ(ObjectAddress(0),
-            bitmap().FindPotentialObjectBeginning(ObjectAddress(hint_index)));
-}
-
 }  // namespace internal
 }  // namespace base

base/allocator/partition_allocator/partition_alloc_unittest.cc

@@ -2555,6 +2555,18 @@ TEST_F(PartitionAllocTest, OptimizedGetSlotOffset) {
   }
 }
 
+// Test that the optimized `GetSlotNumber` implementation produces valid
+// results.
+TEST_F(PartitionAllocTest, OptimizedGetSlotNumber) {
+  for (auto& bucket : allocator.root()->buckets) {
+    for (size_t slot = 0, offset = bucket.slot_size / 2;
+         slot < bucket.get_slots_per_span();
+         ++slot, offset += bucket.slot_size) {
+      EXPECT_EQ(slot, bucket.GetSlotNumber(offset));
+    }
+  }
+}
+
 TEST_F(PartitionAllocTest, GetUsableSize) {
   size_t delta = SystemPageSize() + 1;
   for (size_t size = 1; size <= kMinDirectMappedDownsize; size += delta) {

base/allocator/partition_allocator/partition_bucket.h

@@ -128,6 +128,19 @@ struct PartitionBucket {
     return static_cast<size_t>(offset_in_slot);
   }
 
+  // Returns a slot number starting from the beginning of the slot span.
+  ALWAYS_INLINE size_t GetSlotNumber(size_t offset_in_slot_span) {
+    // See the static assertion for `kReciprocalShift` above.
+    PA_DCHECK(offset_in_slot_span <= kMaxBucketed);
+    PA_DCHECK(slot_size <= kMaxBucketed);
+
+    const size_t offset_in_slot =
+        ((offset_in_slot_span * slot_size_reciprocal) >> kReciprocalShift);
+    PA_DCHECK(offset_in_slot_span / slot_size == offset_in_slot);
+
+    return offset_in_slot;
+  }
+
  private:
   static NOINLINE void OnFull();
 

base/allocator/partition_allocator/partition_page.h

@@ -254,6 +254,40 @@ ALWAYS_INLINE bool IsWithinSuperPagePayload(char* ptr, bool with_pcscan) {
   return ptr >= payload_start && ptr < payload_end;
 }
 
+// Returns the start of a slot or nullptr if |maybe_inner_ptr| is not inside of
+// an existing slot span. The function may return a pointer even inside a
+// decommitted or free slot span, it's the caller responsibility to check if
+// memory is actually allocated.
+// The precondition is that |maybe_inner_ptr| must point to payload of a valid
+// super page.
+template <bool thread_safe>
+ALWAYS_INLINE char* GetSlotStartInSuperPage(char* maybe_inner_ptr) {
+  PA_DCHECK(!IsManagedByPartitionAllocDirectMap(maybe_inner_ptr));
+  char* super_page_ptr = reinterpret_cast<char*>(
+      reinterpret_cast<uintptr_t>(maybe_inner_ptr) & kSuperPageBaseMask);
+  auto* extent = reinterpret_cast<PartitionSuperPageExtentEntry<thread_safe>*>(
+      PartitionSuperPageToMetadataArea(super_page_ptr));
+  PA_DCHECK(IsWithinSuperPagePayload(maybe_inner_ptr,
+                                     extent->root->pcscan.has_value()));
+  auto* slot_span = SlotSpanMetadata<thread_safe>::FromPointerNoAlignmentCheck(
+      maybe_inner_ptr);
+  // Check if the slot span is actually used and valid.
+  if (!slot_span->bucket)
+    return nullptr;
+  PA_DCHECK(PartitionRoot<thread_safe>::IsValidSlotSpan(slot_span));
+  char* const slot_span_begin =
+      static_cast<char*>(SlotSpanMetadata<thread_safe>::ToPointer(slot_span));
+  const ptrdiff_t ptr_offset = maybe_inner_ptr - slot_span_begin;
+  PA_DCHECK(0 <= ptr_offset &&
+            ptr_offset < static_cast<ptrdiff_t>(
+                             slot_span->bucket->get_bytes_per_span()));
+  const size_t slot_size = slot_span->bucket->slot_size;
+  const size_t slot_number = slot_span->bucket->GetSlotNumber(ptr_offset);
+  char* const result = slot_span_begin + (slot_number * slot_size);
+  PA_DCHECK(result <= maybe_inner_ptr && maybe_inner_ptr < result + slot_size);
+  return result;
+}
+
 // See the comment for |FromPointer|.
 template <bool thread_safe>
 ALWAYS_INLINE SlotSpanMetadata<thread_safe>*

base/allocator/partition_allocator/pcscan.cc

@@ -75,6 +75,18 @@ void ReportStats(size_t swept_bytes, size_t last_size, size_t new_size) {
           << static_cast<double>(new_size) / last_size;
 }
 
+template <bool thread_safe>
+uintptr_t GetObjectStartInSuperPage(uintptr_t maybe_ptr, bool allow_extras) {
+  char* allocation_start =
+      GetSlotStartInSuperPage<thread_safe>(reinterpret_cast<char*>(maybe_ptr));
+  if (!allocation_start) {
+    // |maybe_ptr| refers to a garbage or is outside of the payload region.
+    return 0;
+  }
+  return reinterpret_cast<uintptr_t>(
+      PartitionPointerAdjustAdd(allow_extras, allocation_start));
+}
+
 }  // namespace
 
 // This class is responsible for performing the entire PCScan task.
@@ -171,8 +183,9 @@ size_t PCScan<thread_safe>::PCScanTask::TryMarkObjectInNormalBucketPool(
     return 0;
 
   // Check if pointer was in the quarantine bitmap.
-  const uintptr_t base = bitmap->FindPotentialObjectBeginning(maybe_ptr);
-  if (!base)
+  const uintptr_t base =
+      GetObjectStartInSuperPage<thread_safe>(maybe_ptr, root_.allow_extras);
+  if (!bitmap->CheckBit(base))
     return 0;
 
   PA_DCHECK((maybe_ptr & kSuperPageBaseMask) == (base & kSuperPageBaseMask));

base/allocator/partition_allocator/pcscan_unittest.cc

@@ -163,7 +163,6 @@ template <typename SourceList, typename ValueList>
 void TestDanglingReference(PCScanTest& test,
                            SourceList* source,
                            ValueList* value) {
-  CHECK_EQ(value, source->next);
   auto& root = test.root();
   {
     // Free |value| and leave the dangling reference in |source|.
@@ -270,6 +269,39 @@ TEST_F(PCScanTest, DanglingReferenceFromFullPage) {
   TestDanglingReference(*this, source, value);
 }
 
+namespace {
+
+template <size_t Size>
+struct ListWithInnerReference {
+  char buffer1[Size];
+  char* next = nullptr;
+  char buffer2[Size];
+
+  static ListWithInnerReference* Create(ThreadSafePartitionRoot& root) {
+    auto* list = static_cast<ListWithInnerReference*>(
+        root.Alloc(sizeof(ListWithInnerReference), nullptr));
+    return list;
+  }
+
+  static void Destroy(ThreadSafePartitionRoot& root,
+                      ListWithInnerReference* list) {
+    root.Free(list);
+  }
+};
+
+}  // namespace
+
+TEST_F(PCScanTest, DanglingInnerReference) {
+  using SourceList = ListWithInnerReference<64>;
+  using ValueList = SourceList;
+
+  auto* source = SourceList::Create(root());
+  auto* value = ValueList::Create(root());
+  source->next = value->buffer2;
+
+  TestDanglingReference(*this, source, value);
+}
+
 }  // namespace internal
 }  // namespace base