[PA] Add PartitionAllocatorForTesting for PartitionAllocTest.

Since PartitionAllocGlobalUninitForTesting() releases the reserved
address space, all allocators will not work correctly after
invoking PartitionAllocGlobalUninitForTesting().
Instead make each allocator for testing release its allocated pages
(and also address spaces) and make
PartitionAllocGlobalUninitForTesting() do nothing for
AddressPoolManager or PartitionAddressSpace.

Change-Id: I242ecb98a4ed9647830bb1b83644619a3affd5b2
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/4075584
Reviewed-by: Daniel Cheng <dcheng@chromium.org>
Reviewed-by: Etienne Bergeron <etienneb@chromium.org>
Reviewed-by: Bartek Nowierski <bartekn@chromium.org>
Commit-Queue: Takashi Sakamoto <tasak@google.com>
Cr-Commit-Position: refs/heads/main@{#1103652}

base/BUILD.gn

@@ -4748,6 +4748,7 @@ source_set("partition_alloc_test_support") {
     "allocator/partition_allocator/extended_api.cc",
     "allocator/partition_allocator/extended_api.h",
     "allocator/partition_allocator/partition_alloc_base/threading/platform_thread_for_testing.h",
+    "allocator/partition_allocator/partition_alloc_for_testing.h",
   ]
   if (is_posix) {
     sources += [ "allocator/partition_allocator/partition_alloc_base/threading/platform_thread_posix_for_testing.cc" ]

base/allocator/dispatcher/dispatcher_unittest.cc

@@ -13,7 +13,7 @@
 #include "build/build_config.h"
 
 #if BUILDFLAG(USE_PARTITION_ALLOC)
-#include "base/allocator/partition_allocator/partition_alloc.h"
+#include "base/allocator/partition_allocator/partition_alloc_for_testing.h"  // nogncheck
 #endif
 
 #if BUILDFLAG(USE_ALLOCATOR_SHIM)
@@ -101,6 +101,11 @@ TEST_F(BaseAllocatorDispatcherTest, VerifyInitialization) {
 struct PartitionAllocator {
   void* Alloc(size_t size) { return alloc_.AllocWithFlags(0, size, nullptr); }
   void Free(void* data) { alloc_.Free(data); }
+  ~PartitionAllocator() {
+    // Use |DisallowLeaks| to confirm that there is no memory allocated and
+    // not yet freed.
+    alloc_.ResetForTesting(::partition_alloc::internal::DisallowLeaks);
+  }
 
  private:
   ThreadSafePartitionRoot alloc_{{

base/allocator/partition_allocator/memory_reclaimer_unittest.cc

@@ -7,11 +7,11 @@
 #include <memory>
 #include <utility>
 
-#include "base/allocator/partition_allocator/partition_alloc.h"
 #include "base/allocator/partition_allocator/partition_alloc_base/compiler_specific.h"
 #include "base/allocator/partition_allocator/partition_alloc_base/logging.h"
 #include "base/allocator/partition_allocator/partition_alloc_buildflags.h"
 #include "base/allocator/partition_allocator/partition_alloc_config.h"
+#include "base/allocator/partition_allocator/partition_alloc_for_testing.h"
 #include "base/allocator/partition_allocator/shim/allocator_shim_default_dispatch_to_partition_alloc.h"
 #include "build/build_config.h"
 #include "testing/gtest/include/gtest/gtest.h"
@@ -37,28 +37,31 @@ void HandleOOM(size_t unused_size) {
 
 class MemoryReclaimerTest : public ::testing::Test {
  public:
-  MemoryReclaimerTest() = default;
-
- protected:
-  void SetUp() override {
-    PartitionAllocGlobalInit(HandleOOM);
+  MemoryReclaimerTest() {
+    // Since MemoryReclaimer::ResetForTesting() clears partitions_,
+    // we need to make PartitionAllocator after this ResetForTesting().
+    // Otherwise, we will see no PartitionAllocator is registered.
     MemoryReclaimer::Instance()->ResetForTesting();
-    allocator_ = std::make_unique<PartitionAllocator>();
-    allocator_->init({
-        PartitionOptions::AlignedAlloc::kDisallowed,
-        PartitionOptions::ThreadCache::kDisabled,
-        PartitionOptions::Quarantine::kAllowed,
-        PartitionOptions::Cookie::kAllowed,
-        PartitionOptions::BackupRefPtr::kDisabled,
-        PartitionOptions::BackupRefPtrZapping::kDisabled,
-        PartitionOptions::UseConfigurablePool::kNo,
-    });
+
+    allocator_ =
+        std::make_unique<PartitionAllocatorForTesting>(PartitionOptions{
+            PartitionOptions::AlignedAlloc::kDisallowed,
+            PartitionOptions::ThreadCache::kDisabled,
+            PartitionOptions::Quarantine::kAllowed,
+            PartitionOptions::Cookie::kAllowed,
+            PartitionOptions::BackupRefPtr::kDisabled,
+            PartitionOptions::BackupRefPtrZapping::kDisabled,
+            PartitionOptions::UseConfigurablePool::kNo,
+        });
     allocator_->root()->UncapEmptySlotSpanMemoryForTesting();
+    PartitionAllocGlobalInit(HandleOOM);
   }
 
-  void TearDown() override {
+  ~MemoryReclaimerTest() override {
+    // Since MemoryReclaimer::UnregisterPartition() checks whether
+    // the given partition is managed by MemoryReclaimer, need to
+    // destruct |allocator_| before ResetForTesting().
     allocator_ = nullptr;
-    MemoryReclaimer::Instance()->ResetForTesting();
     PartitionAllocGlobalUninitForTesting();
   }
 
@@ -69,7 +72,7 @@ class MemoryReclaimerTest : public ::testing::Test {
     allocator_->root()->Free(data);
   }
 
-  std::unique_ptr<PartitionAllocator> allocator_;
+  std::unique_ptr<PartitionAllocatorForTesting> allocator_;
 };
 
 TEST_F(MemoryReclaimerTest, FreesMemory) {

base/allocator/partition_allocator/partition_alloc.cc

@@ -104,16 +104,6 @@ void PartitionAllocGlobalUninitForTesting() {
 #if BUILDFLAG(ENABLE_PKEYS)
   internal::PartitionAddressSpace::UninitPkeyPoolForTesting();
 #endif
-#if BUILDFLAG(USE_STARSCAN)
-  internal::PCScan::UninitForTesting();  // IN-TEST
-#endif
-#if !BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
-#if BUILDFLAG(HAS_64_BIT_POINTERS)
-  internal::PartitionAddressSpace::UninitForTesting();
-#else
-  internal::AddressPoolManager::GetInstance().ResetForTesting();
-#endif  // BUILDFLAG(HAS_64_BIT_POINTERS)
-#endif  // !BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
   internal::g_oom_handling_function = nullptr;
 }
 

base/allocator/partition_allocator/partition_alloc_for_testing.h

@@ -0,0 +1,50 @@
+// Copyright 2023 The Chromium Authors
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef BASE_ALLOCATOR_PARTITION_ALLOCATOR_PARTITION_ALLOC_FOR_TESTING_H_
+#define BASE_ALLOCATOR_PARTITION_ALLOCATOR_PARTITION_ALLOC_FOR_TESTING_H_
+
+#include "base/allocator/partition_allocator/partition_alloc.h"
+
+namespace partition_alloc {
+namespace internal {
+
+constexpr bool AllowLeaks = true;
+constexpr bool DisallowLeaks = false;
+
+// A subclass of PartitionAllocator for testing. It will free all resources,
+// i.e. allocated memory, memory inside freelist, and so on, when destructing
+// it or when manually invoking reset().
+// If need to check if there are any memory allocated but not freed yet,
+// use allow_leaks=false. We will see CHECK failure inside reset() if any
+// leak is detected. Otherwise (e.g. intentional leaks), use allow_leaks=true.
+template <bool thread_safe, bool allow_leaks>
+struct PartitionAllocatorForTesting : public PartitionAllocator<thread_safe> {
+  PartitionAllocatorForTesting() : PartitionAllocator<thread_safe>() {}
+
+  explicit PartitionAllocatorForTesting(PartitionOptions opts)
+      : PartitionAllocator<thread_safe>() {
+    PartitionAllocator<thread_safe>::init(opts);
+  }
+
+  ~PartitionAllocatorForTesting() { reset(); }
+
+  PA_ALWAYS_INLINE void reset() {
+    PartitionAllocator<thread_safe>::root()->ResetForTesting(allow_leaks);
+  }
+};
+
+}  // namespace internal
+
+using PartitionAllocatorForTesting =
+    internal::PartitionAllocatorForTesting<internal::ThreadSafe,
+                                           internal::DisallowLeaks>;
+
+using PartitionAllocatorAllowLeaksForTesting =
+    internal::PartitionAllocatorForTesting<internal::ThreadSafe,
+                                           internal::AllowLeaks>;
+
+}  // namespace partition_alloc
+
+#endif  // BASE_ALLOCATOR_PARTITION_ALLOCATOR_PARTITION_ALLOC_FOR_TESTING_H_

base/allocator/partition_allocator/partition_alloc_unittest.cc

@@ -2,7 +2,7 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
-#include "base/allocator/partition_allocator/partition_alloc.h"
+#include "base/allocator/partition_allocator/partition_alloc_for_testing.h"
 
 #include <algorithm>
 #include <cstddef>
@@ -590,10 +590,10 @@ class PartitionAllocTest
   bool UsePkeyPool() const { return GetParam().use_pkey_pool; }
   bool UseBRPPool() const { return allocator.root()->brp_enabled(); }
 
-  PartitionAllocator<internal::ThreadSafe> allocator;
-  PartitionAllocator<internal::ThreadSafe> aligned_allocator;
+  partition_alloc::PartitionAllocatorForTesting allocator;
+  partition_alloc::PartitionAllocatorForTesting aligned_allocator;
 #if BUILDFLAG(ENABLE_PKEYS)
-  PartitionAllocator<internal::ThreadSafe> pkey_allocator;
+  partition_alloc::PartitionAllocatorForTesting pkey_allocator;
 #endif
   size_t test_bucket_index_;
 
@@ -1301,6 +1301,7 @@ TEST_P(PartitionAllocTest, AllocGetSizeAndStart) {
     }
   }
 #endif  // BUILDFLAG(ENABLE_BACKUP_REF_PTR_SUPPORT)
+  allocator.root()->Free(ptr);
 
   // Allocate the maximum allowed bucketed size.
   requested_size = kMaxBucketed - ExtraAllocSize(allocator);
@@ -2607,6 +2608,8 @@ TEST_P(PartitionAllocDeathTest, FreelistCorruption) {
   // Restore the freelist entry value, otherwise freelist corruption is detected
   // in TearDown(), crashing this process.
   uaf_data[0] = previous_uaf_data;
+
+  allocator.root()->Free(fake_freelist_entry);
 }
 
 // With BUILDFLAG(PA_DCHECK_IS_ON), cookie already handles off-by-one detection.
@@ -3780,6 +3783,8 @@ TEST_P(PartitionAllocTest, GetUsableSizeWithMac11MallocSizeHack) {
       PartitionRoot<ThreadSafe>::GetUsableSizeWithMac11MallocSizeHack(ptr);
   EXPECT_EQ(usable_size, internal::kMac11MallocSizeHackUsableSize);
   EXPECT_EQ(usable_size_with_hack, size);
+
+  allocator.root()->Free(ptr);
 }
 #endif  // PA_CONFIG(ENABLE_MAC11_MALLOC_SIZE_HACK)
 
@@ -4126,6 +4131,7 @@ void PartitionAllocTest::RunRefCountReallocSubtest(size_t orig_size,
     EXPECT_TRUE(ref_count2->IsAliveWithNoKnownRefs());
 
     EXPECT_TRUE(ref_count1->Release());
+    PartitionAllocFreeForRefCounting(allocator.root()->ObjectToSlotStart(ptr1));
   }
 
   allocator.root()->Free(ptr2);
@@ -4218,6 +4224,8 @@ TEST_P(UnretainedDanglingRawPtrTest, UnretainedDanglingPtrShouldReport) {
   ref_count->ReportIfDangling();
   EXPECT_EQ(g_unretained_dangling_raw_ptr_detected_count, 1);
   EXPECT_TRUE(ref_count->Release());
+
+  PartitionAllocFreeForRefCounting(allocator.root()->ObjectToSlotStart(ptr));
 }
 
 #if !BUILDFLAG(HAS_64_BIT_POINTERS)
@@ -5500,6 +5508,8 @@ TEST_P(PartitionAllocTest, FreeSlotBitmapMarkedAsUsedAfterAlloc) {
   void* ptr = allocator.root()->Alloc(kTestAllocSize, type_name);
   uintptr_t slot_start = allocator.root()->ObjectToSlotStart(ptr);
   EXPECT_TRUE(FreeSlotBitmapSlotIsUsed(slot_start));
+
+  allocator.root()->Free(ptr);
 }
 
 TEST_P(PartitionAllocTest, FreeSlotBitmapMarkedAsFreeAfterFree) {

base/allocator/partition_allocator/partition_root.cc

@@ -731,6 +731,17 @@ void PartitionRoot<thread_safe>::DestructForTesting() {
   // this function on PartitionRoots without a thread cache.
   PA_CHECK(!flags.with_thread_cache);
   auto pool_handle = ChoosePool();
+#if BUILDFLAG(ENABLE_PKEYS)
+  // The pages managed by pkey will be free-ed at UninitPKeyForTesting().
+  // Don't invoke FreePages() for the pages.
+  if (pool_handle == internal::kPkeyPoolHandle) {
+    return;
+  }
+  PA_DCHECK(pool_handle < internal::kNumPools);
+#else
+  PA_DCHECK(pool_handle <= internal::kNumPools);
+#endif
+
   auto* curr = first_extent;
   while (curr != nullptr) {
     auto* next = curr->next;
@@ -1494,6 +1505,73 @@ void PartitionRoot<thread_safe>::DeleteForTesting(
   delete partition_root;
 }
 
+template <bool thread_safe>
+void PartitionRoot<thread_safe>::ResetForTesting(bool allow_leaks) {
+  if (flags.with_thread_cache) {
+    ThreadCache::SwapForTesting(nullptr);
+    flags.with_thread_cache = false;
+  }
+
+  ::partition_alloc::internal::ScopedGuard guard(lock_);
+
+#if BUILDFLAG(PA_DCHECK_IS_ON)
+  if (!allow_leaks) {
+    unsigned num_allocated_slots = 0;
+    for (Bucket& bucket : buckets) {
+      if (bucket.active_slot_spans_head !=
+          internal::SlotSpanMetadata<thread_safe>::get_sentinel_slot_span()) {
+        for (internal::SlotSpanMetadata<thread_safe>* slot_span =
+                 bucket.active_slot_spans_head;
+             slot_span; slot_span = slot_span->next_slot_span) {
+          num_allocated_slots += slot_span->num_allocated_slots;
+        }
+      }
+      // Full slot spans are nowhere. Need to see bucket.num_full_slot_spans
+      // to count the number of full slot spans' slots.
+      if (bucket.num_full_slot_spans) {
+        num_allocated_slots +=
+            bucket.num_full_slot_spans * bucket.get_slots_per_span();
+      }
+    }
+    PA_DCHECK(num_allocated_slots == 0);
+
+    // Check for direct-mapped allocations.
+    PA_DCHECK(!direct_map_list);
+  }
+#endif
+
+  DestructForTesting();  // IN-TEST
+
+#if PA_CONFIG(USE_PARTITION_ROOT_ENUMERATOR)
+  if (initialized) {
+    internal::PartitionRootEnumerator::Instance().Unregister(this);
+  }
+#endif  // PA_CONFIG(USE_PARTITION_ROOT_ENUMERATOR)
+
+  for (Bucket& bucket : buckets) {
+    bucket.active_slot_spans_head =
+        SlotSpan::get_sentinel_slot_span_non_const();
+    bucket.empty_slot_spans_head = nullptr;
+    bucket.decommitted_slot_spans_head = nullptr;
+    bucket.num_full_slot_spans = 0;
+  }
+
+  next_super_page = 0;
+  next_partition_page = 0;
+  next_partition_page_end = 0;
+  current_extent = nullptr;
+  first_extent = nullptr;
+
+  direct_map_list = nullptr;
+  for (auto& entity : global_empty_slot_span_ring) {
+    entity = nullptr;
+  }
+
+  global_empty_slot_span_ring_index = 0;
+  global_empty_slot_span_ring_size = internal::kDefaultEmptySlotSpanRingSize;
+  initialized = false;
+}
+
 template <bool thread_safe>
 void PartitionRoot<thread_safe>::ResetBookkeepingForTesting() {
   ::partition_alloc::internal::ScopedGuard guard{lock_};

base/allocator/partition_allocator/partition_root.h

@@ -405,6 +405,8 @@ struct PA_ALIGNAS(64) PA_COMPONENT_EXPORT(PARTITION_ALLOC) PartitionRoot {
   PartitionRoot()
       : flags{QuarantineMode::kAlwaysDisabled, ScanMode::kDisabled} {}
   explicit PartitionRoot(PartitionOptions opts) : flags() { Init(opts); }
+  // TODO(tasak): remove ~PartitionRoot() after confirming all tests
+  // don't need ~PartitionRoot().
   ~PartitionRoot();
 
   // This will unreserve any space in the pool that the PartitionRoot is
@@ -584,6 +586,7 @@ struct PA_ALIGNAS(64) PA_COMPONENT_EXPORT(PARTITION_ALLOC) PartitionRoot {
                  PartitionStatsDumper* partition_stats_dumper);
 
   static void DeleteForTesting(PartitionRoot* partition_root);
+  void ResetForTesting(bool allow_leaks);
   void ResetBookkeepingForTesting();
 
   PA_ALWAYS_INLINE BucketDistribution GetBucketDistribution() const {

base/allocator/partition_allocator/starscan/metadata_allocator.cc

@@ -30,13 +30,14 @@ ThreadSafePartitionRoot& PCScanMetadataAllocator() {
   return *allocator;
 }
 
+// TODO(tasak): investigate whether PartitionAlloc tests really need this
+// function or not. If we found no tests need, remove it.
 void ReinitPCScanMetadataAllocatorForTesting() {
   // First, purge memory owned by PCScanMetadataAllocator.
   PCScanMetadataAllocator().PurgeMemory(PurgeFlags::kDecommitEmptySlotSpans |
                                         PurgeFlags::kDiscardUnusedSystemPages);
   // Then, reinit the allocator.
-  PCScanMetadataAllocator().~PartitionRoot();
-  memset(&PCScanMetadataAllocator(), 0, sizeof(PCScanMetadataAllocator()));
+  PCScanMetadataAllocator().ResetForTesting(true);  // IN-TEST
   PCScanMetadataAllocator().Init(kConfig);
 }
 

base/allocator/partition_allocator/starscan/pcscan_unittest.cc

@@ -9,13 +9,13 @@
 #include "base/allocator/partition_allocator/starscan/pcscan.h"
 
 #include "base/allocator/partition_allocator/partition_alloc-inl.h"
-#include "base/allocator/partition_allocator/partition_alloc.h"
 #include "base/allocator/partition_allocator/partition_alloc_base/compiler_specific.h"
 #include "base/allocator/partition_allocator/partition_alloc_base/cpu.h"
 #include "base/allocator/partition_allocator/partition_alloc_base/logging.h"
 #include "base/allocator/partition_allocator/partition_alloc_buildflags.h"
 #include "base/allocator/partition_allocator/partition_alloc_config.h"
 #include "base/allocator/partition_allocator/partition_alloc_constants.h"
+#include "base/allocator/partition_allocator/partition_alloc_for_testing.h"
 #include "base/allocator/partition_allocator/partition_root.h"
 #include "base/allocator/partition_allocator/starscan/stack/stack.h"
 #include "base/allocator/partition_allocator/tagging.h"
@@ -101,7 +101,7 @@ class PartitionAllocPCScanTestBase : public testing::Test {
 
  private:
   // Leverage the already-templated version outside `internal::`.
-  partition_alloc::PartitionAllocator allocator_;
+  partition_alloc::PartitionAllocatorAllowLeaksForTesting allocator_;
 };
 
 // The test that expects free() being quarantined only when tag overflow occurs.