• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 // Copyright 2020 The Chromium Authors
2 // Use of this source code is governed by a BSD-style license that can be
3 // found in the LICENSE file.
4 
5 #include "base/allocator/partition_allocator/partition_root.h"
6 
7 #include <cstdint>
8 
9 #include "base/allocator/partition_allocator/freeslot_bitmap.h"
10 #include "base/allocator/partition_allocator/oom.h"
11 #include "base/allocator/partition_allocator/page_allocator.h"
12 #include "base/allocator/partition_allocator/partition_address_space.h"
13 #include "base/allocator/partition_allocator/partition_alloc_base/bits.h"
14 #include "base/allocator/partition_allocator/partition_alloc_base/compiler_specific.h"
15 #include "base/allocator/partition_allocator/partition_alloc_base/component_export.h"
16 #include "base/allocator/partition_allocator/partition_alloc_base/debug/debugging_buildflags.h"
17 #include "base/allocator/partition_allocator/partition_alloc_base/thread_annotations.h"
18 #include "base/allocator/partition_allocator/partition_alloc_buildflags.h"
19 #include "base/allocator/partition_allocator/partition_alloc_check.h"
20 #include "base/allocator/partition_allocator/partition_alloc_config.h"
21 #include "base/allocator/partition_allocator/partition_alloc_constants.h"
22 #include "base/allocator/partition_allocator/partition_bucket.h"
23 #include "base/allocator/partition_allocator/partition_cookie.h"
24 #include "base/allocator/partition_allocator/partition_oom.h"
25 #include "base/allocator/partition_allocator/partition_page.h"
26 #include "base/allocator/partition_allocator/partition_ref_count.h"
27 #include "base/allocator/partition_allocator/pkey.h"
28 #include "base/allocator/partition_allocator/reservation_offset_table.h"
29 #include "base/allocator/partition_allocator/tagging.h"
30 #include "build/build_config.h"
31 
32 #if PA_CONFIG(ENABLE_MAC11_MALLOC_SIZE_HACK) && BUILDFLAG(IS_APPLE)
33 #include "base/allocator/partition_allocator/partition_alloc_base/mac/mac_util.h"
34 #endif  // PA_CONFIG(ENABLE_MAC11_MALLOC_SIZE_HACK) && BUILDFLAG(IS_APPLE)
35 
36 #if BUILDFLAG(USE_STARSCAN)
37 #include "base/allocator/partition_allocator/starscan/pcscan.h"
38 #endif
39 
40 #if !BUILDFLAG(HAS_64_BIT_POINTERS)
41 #include "base/allocator/partition_allocator/address_pool_manager_bitmap.h"
42 #endif
43 
44 #if BUILDFLAG(IS_WIN)
45 #include <windows.h>
46 #include "wow64apiset.h"
47 #endif
48 
49 #if BUILDFLAG(IS_LINUX) || BUILDFLAG(IS_CHROMEOS)
50 #include <pthread.h>
51 #endif
52 
53 namespace partition_alloc::internal {
54 
55 #if BUILDFLAG(RECORD_ALLOC_INFO)
56 // Even if this is not hidden behind a BUILDFLAG, it should not use any memory
57 // when recording is disabled, since it ends up in the .bss section.
58 AllocInfo g_allocs = {};
59 
RecordAllocOrFree(uintptr_t addr,size_t size)60 void RecordAllocOrFree(uintptr_t addr, size_t size) {
61   g_allocs.allocs[g_allocs.index.fetch_add(1, std::memory_order_relaxed) %
62                   kAllocInfoSize] = {addr, size};
63 }
64 #endif  // BUILDFLAG(RECORD_ALLOC_INFO)
65 
66 #if BUILDFLAG(ENABLE_BACKUP_REF_PTR_SUPPORT)
IsPtrWithinSameAlloc(uintptr_t orig_address,uintptr_t test_address,size_t type_size)67 PtrPosWithinAlloc IsPtrWithinSameAlloc(uintptr_t orig_address,
68                                        uintptr_t test_address,
69                                        size_t type_size) {
70   // Required for pointers right past an allocation. See
71   // |PartitionAllocGetSlotStartInBRPPool()|.
72   uintptr_t adjusted_address =
73       orig_address - kPartitionPastAllocationAdjustment;
74   PA_DCHECK(IsManagedByNormalBucketsOrDirectMap(adjusted_address));
75   DCheckIfManagedByPartitionAllocBRPPool(adjusted_address);
76 
77   uintptr_t slot_start = PartitionAllocGetSlotStartInBRPPool(adjusted_address);
78   // Don't use |adjusted_address| beyond this point at all. It was needed to
79   // pick the right slot, but now we're dealing with very concrete addresses.
80   // Zero it just in case, to catch errors.
81   adjusted_address = 0;
82 
83   auto* slot_span = SlotSpanMetadata<ThreadSafe>::FromSlotStart(slot_start);
84   auto* root = PartitionRoot<ThreadSafe>::FromSlotSpan(slot_span);
85   // Double check that ref-count is indeed present.
86   PA_DCHECK(root->brp_enabled());
87 
88   uintptr_t object_addr = root->SlotStartToObjectAddr(slot_start);
89   uintptr_t object_end = object_addr + slot_span->GetUsableSize(root);
90   if (test_address < object_addr || object_end < test_address) {
91     return PtrPosWithinAlloc::kFarOOB;
92 #if BUILDFLAG(BACKUP_REF_PTR_POISON_OOB_PTR)
93   } else if (object_end - type_size < test_address) {
94     // Not even a single element of the type referenced by the pointer can fit
95     // between the pointer and the end of the object.
96     return PtrPosWithinAlloc::kAllocEnd;
97 #endif
98   } else {
99     return PtrPosWithinAlloc::kInBounds;
100   }
101 }
102 #endif  // BUILDFLAG(ENABLE_BACKUP_REF_PTR_SUPPORT)
103 
104 }  // namespace partition_alloc::internal
105 
106 namespace partition_alloc {
107 
108 #if PA_CONFIG(USE_PARTITION_ROOT_ENUMERATOR)
109 
110 namespace {
111 
112 internal::Lock g_root_enumerator_lock;
113 }
114 
115 template <bool thread_safe>
GetEnumeratorLock()116 internal::Lock& PartitionRoot<thread_safe>::GetEnumeratorLock() {
117   return g_root_enumerator_lock;
118 }
119 
120 namespace internal {
121 
122 class PartitionRootEnumerator {
123  public:
124   using EnumerateCallback = void (*)(ThreadSafePartitionRoot* root,
125                                      bool in_child);
126   enum EnumerateOrder {
127     kNormal,
128     kReverse,
129   };
130 
Instance()131   static PartitionRootEnumerator& Instance() {
132     static PartitionRootEnumerator instance;
133     return instance;
134   }
135 
Enumerate(EnumerateCallback callback,bool in_child,EnumerateOrder order)136   void Enumerate(EnumerateCallback callback,
137                  bool in_child,
138                  EnumerateOrder order) PA_NO_THREAD_SAFETY_ANALYSIS {
139     if (order == kNormal) {
140       ThreadSafePartitionRoot* root;
141       for (root = Head(partition_roots_); root != nullptr;
142            root = root->next_root) {
143         callback(root, in_child);
144       }
145     } else {
146       PA_DCHECK(order == kReverse);
147       ThreadSafePartitionRoot* root;
148       for (root = Tail(partition_roots_); root != nullptr;
149            root = root->prev_root) {
150         callback(root, in_child);
151       }
152     }
153   }
154 
Register(ThreadSafePartitionRoot * root)155   void Register(ThreadSafePartitionRoot* root) {
156     internal::ScopedGuard guard(ThreadSafePartitionRoot::GetEnumeratorLock());
157     root->next_root = partition_roots_;
158     root->prev_root = nullptr;
159     if (partition_roots_) {
160       partition_roots_->prev_root = root;
161     }
162     partition_roots_ = root;
163   }
164 
Unregister(ThreadSafePartitionRoot * root)165   void Unregister(ThreadSafePartitionRoot* root) {
166     internal::ScopedGuard guard(ThreadSafePartitionRoot::GetEnumeratorLock());
167     ThreadSafePartitionRoot* prev = root->prev_root;
168     ThreadSafePartitionRoot* next = root->next_root;
169     if (prev) {
170       PA_DCHECK(prev->next_root == root);
171       prev->next_root = next;
172     } else {
173       PA_DCHECK(partition_roots_ == root);
174       partition_roots_ = next;
175     }
176     if (next) {
177       PA_DCHECK(next->prev_root == root);
178       next->prev_root = prev;
179     }
180     root->next_root = nullptr;
181     root->prev_root = nullptr;
182   }
183 
184  private:
185   constexpr PartitionRootEnumerator() = default;
186 
Head(ThreadSafePartitionRoot * roots)187   ThreadSafePartitionRoot* Head(ThreadSafePartitionRoot* roots) {
188     return roots;
189   }
190 
Tail(ThreadSafePartitionRoot * roots)191   ThreadSafePartitionRoot* Tail(ThreadSafePartitionRoot* roots)
192       PA_NO_THREAD_SAFETY_ANALYSIS {
193     if (!roots) {
194       return nullptr;
195     }
196     ThreadSafePartitionRoot* node = roots;
197     for (; node->next_root != nullptr; node = node->next_root)
198       ;
199     return node;
200   }
201 
202   ThreadSafePartitionRoot* partition_roots_
203       PA_GUARDED_BY(ThreadSafePartitionRoot::GetEnumeratorLock()) = nullptr;
204 };
205 
206 }  // namespace internal
207 
208 #endif  // PA_USE_PARTITION_ROOT_ENUMERATOR
209 
210 #if BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
211 
212 namespace {
213 
214 #if PA_CONFIG(HAS_ATFORK_HANDLER)
215 
LockRoot(PartitionRoot<internal::ThreadSafe> * root,bool)216 void LockRoot(PartitionRoot<internal::ThreadSafe>* root,
217               bool) PA_NO_THREAD_SAFETY_ANALYSIS {
218   PA_DCHECK(root);
219   root->lock_.Acquire();
220 }
221 
222 // PA_NO_THREAD_SAFETY_ANALYSIS: acquires the lock and doesn't release it, by
223 // design.
BeforeForkInParent()224 void BeforeForkInParent() PA_NO_THREAD_SAFETY_ANALYSIS {
225   // ThreadSafePartitionRoot::GetLock() is private. So use
226   // g_root_enumerator_lock here.
227   g_root_enumerator_lock.Acquire();
228   internal::PartitionRootEnumerator::Instance().Enumerate(
229       LockRoot, false,
230       internal::PartitionRootEnumerator::EnumerateOrder::kNormal);
231 
232   ThreadCacheRegistry::GetLock().Acquire();
233 }
234 
235 template <typename T>
UnlockOrReinit(T & lock,bool in_child)236 void UnlockOrReinit(T& lock, bool in_child) PA_NO_THREAD_SAFETY_ANALYSIS {
237   // Only re-init the locks in the child process, in the parent can unlock
238   // normally.
239   if (in_child) {
240     lock.Reinit();
241   } else {
242     lock.Release();
243   }
244 }
245 
UnlockOrReinitRoot(PartitionRoot<internal::ThreadSafe> * root,bool in_child)246 void UnlockOrReinitRoot(PartitionRoot<internal::ThreadSafe>* root,
247                         bool in_child) PA_NO_THREAD_SAFETY_ANALYSIS {
248   UnlockOrReinit(root->lock_, in_child);
249 }
250 
ReleaseLocks(bool in_child)251 void ReleaseLocks(bool in_child) PA_NO_THREAD_SAFETY_ANALYSIS {
252   // In reverse order, even though there are no lock ordering dependencies.
253   UnlockOrReinit(ThreadCacheRegistry::GetLock(), in_child);
254   internal::PartitionRootEnumerator::Instance().Enumerate(
255       UnlockOrReinitRoot, in_child,
256       internal::PartitionRootEnumerator::EnumerateOrder::kReverse);
257 
258   // ThreadSafePartitionRoot::GetLock() is private. So use
259   // g_root_enumerator_lock here.
260   UnlockOrReinit(g_root_enumerator_lock, in_child);
261 }
262 
AfterForkInParent()263 void AfterForkInParent() {
264   ReleaseLocks(/* in_child = */ false);
265 }
266 
AfterForkInChild()267 void AfterForkInChild() {
268   ReleaseLocks(/* in_child = */ true);
269   // Unsafe, as noted in the name. This is fine here however, since at this
270   // point there is only one thread, this one (unless another post-fork()
271   // handler created a thread, but it would have needed to allocate, which would
272   // have deadlocked the process already).
273   //
274   // If we don't reclaim this memory, it is lost forever. Note that this is only
275   // really an issue if we fork() a multi-threaded process without calling
276   // exec() right away, which is discouraged.
277   ThreadCacheRegistry::Instance().ForcePurgeAllThreadAfterForkUnsafe();
278 }
279 #endif  // PA_CONFIG(HAS_ATFORK_HANDLER)
280 
281 std::atomic<bool> g_global_init_called;
PartitionAllocMallocInitOnce()282 void PartitionAllocMallocInitOnce() {
283   bool expected = false;
284   // No need to block execution for potential concurrent initialization, merely
285   // want to make sure this is only called once.
286   if (!g_global_init_called.compare_exchange_strong(expected, true)) {
287     return;
288   }
289 
290 #if BUILDFLAG(IS_LINUX) || BUILDFLAG(IS_CHROMEOS)
291   // When fork() is called, only the current thread continues to execute in the
292   // child process. If the lock is held, but *not* by this thread when fork() is
293   // called, we have a deadlock.
294   //
295   // The "solution" here is to acquire the lock on the forking thread before
296   // fork(), and keep it held until fork() is done, in the parent and the
297   // child. To clean up memory, we also must empty the thread caches in the
298   // child, which is easier, since no threads except for the current one are
299   // running right after the fork().
300   //
301   // This is not perfect though, since:
302   // - Multiple pre/post-fork() handlers can be registered, they are then run in
303   //   LIFO order for the pre-fork handler, and FIFO order for the post-fork
304   //   one. So unless we are the first to register a handler, if another handler
305   //   allocates, then we deterministically deadlock.
306   // - pthread handlers are *not* called when the application calls clone()
307   //   directly, which is what Chrome does to launch processes.
308   //
309   // However, no perfect solution really exists to make threads + fork()
310   // cooperate, but deadlocks are real (and fork() is used in DEATH_TEST()s),
311   // and other malloc() implementations use the same techniques.
312   int err =
313       pthread_atfork(BeforeForkInParent, AfterForkInParent, AfterForkInChild);
314   PA_CHECK(err == 0);
315 #endif  // BUILDFLAG(IS_LINUX) || BUILDFLAG(IS_CHROMEOS)
316 }
317 
318 }  // namespace
319 
320 #if BUILDFLAG(IS_APPLE)
PartitionAllocMallocHookOnBeforeForkInParent()321 void PartitionAllocMallocHookOnBeforeForkInParent() {
322   BeforeForkInParent();
323 }
324 
PartitionAllocMallocHookOnAfterForkInParent()325 void PartitionAllocMallocHookOnAfterForkInParent() {
326   AfterForkInParent();
327 }
328 
PartitionAllocMallocHookOnAfterForkInChild()329 void PartitionAllocMallocHookOnAfterForkInChild() {
330   AfterForkInChild();
331 }
332 #endif  // BUILDFLAG(IS_APPLE)
333 
334 #endif  // BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
335 
336 namespace internal {
337 
338 namespace {
339 // 64 was chosen arbitrarily, as it seems like a reasonable trade-off between
340 // performance and purging opportunity. Higher value (i.e. smaller slots)
341 // wouldn't necessarily increase chances of purging, but would result in
342 // more work and larger |slot_usage| array. Lower value would probably decrease
343 // chances of purging. Not empirically tested.
344 constexpr size_t kMaxPurgeableSlotsPerSystemPage = 64;
345 PA_ALWAYS_INLINE PAGE_ALLOCATOR_CONSTANTS_DECLARE_CONSTEXPR size_t
MinPurgeableSlotSize()346 MinPurgeableSlotSize() {
347   return SystemPageSize() / kMaxPurgeableSlotsPerSystemPage;
348 }
349 }  // namespace
350 
351 template <bool thread_safe>
PartitionPurgeSlotSpan(internal::SlotSpanMetadata<thread_safe> * slot_span,bool discard)352 static size_t PartitionPurgeSlotSpan(
353     internal::SlotSpanMetadata<thread_safe>* slot_span,
354     bool discard) {
355   auto* root = PartitionRoot<thread_safe>::FromSlotSpan(slot_span);
356   const internal::PartitionBucket<thread_safe>* bucket = slot_span->bucket;
357   size_t slot_size = bucket->slot_size;
358 
359   if (slot_size < MinPurgeableSlotSize() || !slot_span->num_allocated_slots) {
360     return 0;
361   }
362 
363   size_t bucket_num_slots = bucket->get_slots_per_span();
364   size_t discardable_bytes = 0;
365 
366   if (slot_span->CanStoreRawSize()) {
367     uint32_t utilized_slot_size = static_cast<uint32_t>(
368         RoundUpToSystemPage(slot_span->GetUtilizedSlotSize()));
369     discardable_bytes = bucket->slot_size - utilized_slot_size;
370     if (discardable_bytes && discard) {
371       uintptr_t slot_span_start =
372           internal::SlotSpanMetadata<thread_safe>::ToSlotSpanStart(slot_span);
373       uintptr_t committed_data_end = slot_span_start + utilized_slot_size;
374       ScopedSyscallTimer timer{root};
375       DiscardSystemPages(committed_data_end, discardable_bytes);
376     }
377     return discardable_bytes;
378   }
379 
380 #if defined(PAGE_ALLOCATOR_CONSTANTS_ARE_CONSTEXPR)
381   constexpr size_t kMaxSlotCount =
382       (PartitionPageSize() * kMaxPartitionPagesPerRegularSlotSpan) /
383       MinPurgeableSlotSize();
384 #elif BUILDFLAG(IS_APPLE) || (BUILDFLAG(IS_LINUX) && defined(ARCH_CPU_ARM64))
385   // It's better for slot_usage to be stack-allocated and fixed-size, which
386   // demands that its size be constexpr. On IS_APPLE and Linux on arm64,
387   // PartitionPageSize() is always SystemPageSize() << 2, so regardless of
388   // what the run time page size is, kMaxSlotCount can always be simplified
389   // to this expression.
390   constexpr size_t kMaxSlotCount =
391       4 * kMaxPurgeableSlotsPerSystemPage *
392       internal::kMaxPartitionPagesPerRegularSlotSpan;
393   PA_CHECK(kMaxSlotCount == (PartitionPageSize() *
394                              internal::kMaxPartitionPagesPerRegularSlotSpan) /
395                                 MinPurgeableSlotSize());
396 #endif
397   PA_DCHECK(bucket_num_slots <= kMaxSlotCount);
398   PA_DCHECK(slot_span->num_unprovisioned_slots < bucket_num_slots);
399   size_t num_slots = bucket_num_slots - slot_span->num_unprovisioned_slots;
400   char slot_usage[kMaxSlotCount];
401 #if !BUILDFLAG(IS_WIN)
402   // The last freelist entry should not be discarded when using OS_WIN.
403   // DiscardVirtualMemory makes the contents of discarded memory undefined.
404   size_t last_slot = static_cast<size_t>(-1);
405 #endif
406   memset(slot_usage, 1, num_slots);
407   uintptr_t slot_span_start =
408       SlotSpanMetadata<thread_safe>::ToSlotSpanStart(slot_span);
409   // First, walk the freelist for this slot span and make a bitmap of which
410   // slots are not in use.
411   for (PartitionFreelistEntry* entry = slot_span->get_freelist_head(); entry;
412        /**/) {
413     size_t slot_number =
414         bucket->GetSlotNumber(SlotStartPtr2Addr(entry) - slot_span_start);
415     PA_DCHECK(slot_number < num_slots);
416     slot_usage[slot_number] = 0;
417 #if !BUILDFLAG(IS_WIN)
418     // If we have a slot where the encoded next pointer is 0, we can actually
419     // discard that entry because touching a discarded page is guaranteed to
420     // return the original content or 0. (Note that this optimization won't be
421     // effective on big-endian machines because the masking function is
422     // negation.)
423     if (entry->IsEncodedNextPtrZero()) {
424       last_slot = slot_number;
425     }
426 #endif
427     entry = entry->GetNext(slot_size);
428   }
429 
430   // If the slot(s) at the end of the slot span are not in used, we can truncate
431   // them entirely and rewrite the freelist.
432   size_t truncated_slots = 0;
433   while (!slot_usage[num_slots - 1]) {
434     truncated_slots++;
435     num_slots--;
436     PA_DCHECK(num_slots);
437   }
438   // First, do the work of calculating the discardable bytes. Don't actually
439   // discard anything unless the discard flag was passed in.
440   if (truncated_slots) {
441     size_t unprovisioned_bytes = 0;
442     uintptr_t begin_addr = slot_span_start + (num_slots * slot_size);
443     uintptr_t end_addr = begin_addr + (slot_size * truncated_slots);
444 
445     // The slots that do not contain discarded pages should not be included to
446     // |truncated_slots|. Detects those slots and fixes |truncated_slots| and
447     // |num_slots| accordingly.
448     uintptr_t rounded_up_truncatation_begin_addr =
449         RoundUpToSystemPage(begin_addr);
450     while (begin_addr + slot_size <= rounded_up_truncatation_begin_addr) {
451       begin_addr += slot_size;
452       PA_DCHECK(truncated_slots);
453       --truncated_slots;
454       ++num_slots;
455     }
456     begin_addr = rounded_up_truncatation_begin_addr;
457 
458     // We round the end address here up and not down because we're at the end of
459     // a slot span, so we "own" all the way up the page boundary.
460     end_addr = RoundUpToSystemPage(end_addr);
461     PA_DCHECK(end_addr <= slot_span_start + bucket->get_bytes_per_span());
462     if (begin_addr < end_addr) {
463       unprovisioned_bytes = end_addr - begin_addr;
464       discardable_bytes += unprovisioned_bytes;
465     }
466     if (unprovisioned_bytes && discard) {
467       PA_DCHECK(truncated_slots > 0);
468       size_t new_unprovisioned_slots =
469           truncated_slots + slot_span->num_unprovisioned_slots;
470       PA_DCHECK(new_unprovisioned_slots <= bucket->get_slots_per_span());
471       slot_span->num_unprovisioned_slots = new_unprovisioned_slots;
472 
473       // Rewrite the freelist.
474       internal::PartitionFreelistEntry* head = nullptr;
475       internal::PartitionFreelistEntry* back = head;
476       size_t num_new_entries = 0;
477       for (size_t slot_index = 0; slot_index < num_slots; ++slot_index) {
478         if (slot_usage[slot_index]) {
479           continue;
480         }
481 
482         auto* entry = PartitionFreelistEntry::EmplaceAndInitNull(
483             slot_span_start + (slot_size * slot_index));
484         if (!head) {
485           head = entry;
486           back = entry;
487         } else {
488           back->SetNext(entry);
489           back = entry;
490         }
491         num_new_entries++;
492 #if !BUILDFLAG(IS_WIN)
493         last_slot = slot_index;
494 #endif
495       }
496 
497       slot_span->SetFreelistHead(head);
498 
499       PA_DCHECK(num_new_entries == num_slots - slot_span->num_allocated_slots);
500 
501 #if BUILDFLAG(USE_FREESLOT_BITMAP)
502       FreeSlotBitmapReset(slot_span_start + (slot_size * num_slots), end_addr,
503                           slot_size);
504 #endif
505 
506       // Discard the memory.
507       ScopedSyscallTimer timer{root};
508       DiscardSystemPages(begin_addr, unprovisioned_bytes);
509     }
510   }
511 
512   if (slot_size < SystemPageSize()) {
513     // Returns here because implementing the following steps for smaller slot
514     // size will need a complicated logic and make the code messy.
515     return discardable_bytes;
516   }
517 
518   // Next, walk the slots and for any not in use, consider which system pages
519   // are no longer needed. We can release any system pages back to the system as
520   // long as we don't interfere with a freelist pointer or an adjacent used
521   // slot.
522   for (size_t i = 0; i < num_slots; ++i) {
523     if (slot_usage[i]) {
524       continue;
525     }
526 
527     // The first address we can safely discard is just after the freelist
528     // pointer. There's one quirk: if the freelist pointer is actually nullptr,
529     // we can discard that pointer value too.
530     uintptr_t begin_addr = slot_span_start + (i * slot_size);
531     uintptr_t end_addr = begin_addr + slot_size;
532 
533     bool can_discard_free_list_pointer = false;
534 #if !BUILDFLAG(IS_WIN)
535     if (i != last_slot) {
536       begin_addr += sizeof(internal::PartitionFreelistEntry);
537     } else {
538       can_discard_free_list_pointer = true;
539     }
540 #else
541     begin_addr += sizeof(internal::PartitionFreelistEntry);
542 #endif
543 
544     uintptr_t rounded_up_begin_addr = RoundUpToSystemPage(begin_addr);
545     uintptr_t rounded_down_begin_addr = RoundDownToSystemPage(begin_addr);
546     end_addr = RoundDownToSystemPage(end_addr);
547 
548     // |rounded_up_begin_addr| could be greater than |end_addr| only if slot
549     // size was less than system page size, or if free list pointer crossed the
550     // page boundary. Neither is possible here.
551     PA_DCHECK(rounded_up_begin_addr <= end_addr);
552 
553     if (rounded_down_begin_addr < rounded_up_begin_addr && i != 0 &&
554         !slot_usage[i - 1] && can_discard_free_list_pointer) {
555       // This slot contains a partial page in the beginning. The rest of that
556       // page is contained in the slot[i-1], which is also discardable.
557       // Therefore we can discard this page.
558       begin_addr = rounded_down_begin_addr;
559     } else {
560       begin_addr = rounded_up_begin_addr;
561     }
562 
563     if (begin_addr < end_addr) {
564       size_t partial_slot_bytes = end_addr - begin_addr;
565       discardable_bytes += partial_slot_bytes;
566       if (discard) {
567         ScopedSyscallTimer timer{root};
568         DiscardSystemPages(begin_addr, partial_slot_bytes);
569       }
570     }
571   }
572 
573   return discardable_bytes;
574 }
575 
576 template <bool thread_safe>
PartitionPurgeBucket(internal::PartitionBucket<thread_safe> * bucket)577 static void PartitionPurgeBucket(
578     internal::PartitionBucket<thread_safe>* bucket) {
579   if (bucket->active_slot_spans_head !=
580       internal::SlotSpanMetadata<thread_safe>::get_sentinel_slot_span()) {
581     for (internal::SlotSpanMetadata<thread_safe>* slot_span =
582              bucket->active_slot_spans_head;
583          slot_span; slot_span = slot_span->next_slot_span) {
584       PA_DCHECK(
585           slot_span !=
586           internal::SlotSpanMetadata<thread_safe>::get_sentinel_slot_span());
587       PartitionPurgeSlotSpan(slot_span, true);
588     }
589   }
590 }
591 
592 template <bool thread_safe>
PartitionDumpSlotSpanStats(PartitionBucketMemoryStats * stats_out,internal::SlotSpanMetadata<thread_safe> * slot_span)593 static void PartitionDumpSlotSpanStats(
594     PartitionBucketMemoryStats* stats_out,
595     internal::SlotSpanMetadata<thread_safe>* slot_span) {
596   uint16_t bucket_num_slots = slot_span->bucket->get_slots_per_span();
597 
598   if (slot_span->is_decommitted()) {
599     ++stats_out->num_decommitted_slot_spans;
600     return;
601   }
602 
603   stats_out->discardable_bytes += PartitionPurgeSlotSpan(slot_span, false);
604 
605   if (slot_span->CanStoreRawSize()) {
606     stats_out->active_bytes += static_cast<uint32_t>(slot_span->GetRawSize());
607   } else {
608     stats_out->active_bytes +=
609         (slot_span->num_allocated_slots * stats_out->bucket_slot_size);
610   }
611   stats_out->active_count += slot_span->num_allocated_slots;
612 
613   size_t slot_span_bytes_resident = RoundUpToSystemPage(
614       (bucket_num_slots - slot_span->num_unprovisioned_slots) *
615       stats_out->bucket_slot_size);
616   stats_out->resident_bytes += slot_span_bytes_resident;
617   if (slot_span->is_empty()) {
618     stats_out->decommittable_bytes += slot_span_bytes_resident;
619     ++stats_out->num_empty_slot_spans;
620   } else if (slot_span->is_full()) {
621     ++stats_out->num_full_slot_spans;
622   } else {
623     PA_DCHECK(slot_span->is_active());
624     ++stats_out->num_active_slot_spans;
625   }
626 }
627 
628 template <bool thread_safe>
PartitionDumpBucketStats(PartitionBucketMemoryStats * stats_out,const internal::PartitionBucket<thread_safe> * bucket)629 static void PartitionDumpBucketStats(
630     PartitionBucketMemoryStats* stats_out,
631     const internal::PartitionBucket<thread_safe>* bucket) {
632   PA_DCHECK(!bucket->is_direct_mapped());
633   stats_out->is_valid = false;
634   // If the active slot span list is empty (==
635   // internal::SlotSpanMetadata::get_sentinel_slot_span()), the bucket might
636   // still need to be reported if it has a list of empty, decommitted or full
637   // slot spans.
638   if (bucket->active_slot_spans_head ==
639           internal::SlotSpanMetadata<thread_safe>::get_sentinel_slot_span() &&
640       !bucket->empty_slot_spans_head && !bucket->decommitted_slot_spans_head &&
641       !bucket->num_full_slot_spans) {
642     return;
643   }
644 
645   memset(stats_out, '\0', sizeof(*stats_out));
646   stats_out->is_valid = true;
647   stats_out->is_direct_map = false;
648   stats_out->num_full_slot_spans =
649       static_cast<size_t>(bucket->num_full_slot_spans);
650   stats_out->bucket_slot_size = bucket->slot_size;
651   uint16_t bucket_num_slots = bucket->get_slots_per_span();
652   size_t bucket_useful_storage = stats_out->bucket_slot_size * bucket_num_slots;
653   stats_out->allocated_slot_span_size = bucket->get_bytes_per_span();
654   stats_out->active_bytes = bucket->num_full_slot_spans * bucket_useful_storage;
655   stats_out->active_count = bucket->num_full_slot_spans * bucket_num_slots;
656   stats_out->resident_bytes =
657       bucket->num_full_slot_spans * stats_out->allocated_slot_span_size;
658 
659   for (internal::SlotSpanMetadata<thread_safe>* slot_span =
660            bucket->empty_slot_spans_head;
661        slot_span; slot_span = slot_span->next_slot_span) {
662     PA_DCHECK(slot_span->is_empty() || slot_span->is_decommitted());
663     PartitionDumpSlotSpanStats(stats_out, slot_span);
664   }
665   for (internal::SlotSpanMetadata<thread_safe>* slot_span =
666            bucket->decommitted_slot_spans_head;
667        slot_span; slot_span = slot_span->next_slot_span) {
668     PA_DCHECK(slot_span->is_decommitted());
669     PartitionDumpSlotSpanStats(stats_out, slot_span);
670   }
671 
672   if (bucket->active_slot_spans_head !=
673       internal::SlotSpanMetadata<thread_safe>::get_sentinel_slot_span()) {
674     for (internal::SlotSpanMetadata<thread_safe>* slot_span =
675              bucket->active_slot_spans_head;
676          slot_span; slot_span = slot_span->next_slot_span) {
677       PA_DCHECK(
678           slot_span !=
679           internal::SlotSpanMetadata<thread_safe>::get_sentinel_slot_span());
680       PartitionDumpSlotSpanStats(stats_out, slot_span);
681     }
682   }
683 }
684 
685 #if BUILDFLAG(PA_DCHECK_IS_ON)
DCheckIfManagedByPartitionAllocBRPPool(uintptr_t address)686 void DCheckIfManagedByPartitionAllocBRPPool(uintptr_t address) {
687   PA_DCHECK(IsManagedByPartitionAllocBRPPool(address));
688 }
689 #endif
690 
691 #if BUILDFLAG(ENABLE_PKEYS)
PartitionAllocPkeyInit(int pkey)692 void PartitionAllocPkeyInit(int pkey) {
693   PkeySettings::settings.enabled = true;
694   PartitionAddressSpace::InitPkeyPool(pkey);
695   // Call TagGlobalsWithPkey last since we might not have write permissions to
696   // to memory tagged with `pkey` at this point.
697   TagGlobalsWithPkey(pkey);
698 }
699 #endif  // BUILDFLAG(ENABLE_PKEYS)
700 
701 }  // namespace internal
702 
703 template <bool thread_safe>
OutOfMemory(size_t size)704 [[noreturn]] PA_NOINLINE void PartitionRoot<thread_safe>::OutOfMemory(
705     size_t size) {
706   const size_t virtual_address_space_size =
707       total_size_of_super_pages.load(std::memory_order_relaxed) +
708       total_size_of_direct_mapped_pages.load(std::memory_order_relaxed);
709 #if !defined(ARCH_CPU_64_BITS)
710   const size_t uncommitted_size =
711       virtual_address_space_size -
712       total_size_of_committed_pages.load(std::memory_order_relaxed);
713 
714   // Check whether this OOM is due to a lot of super pages that are allocated
715   // but not committed, probably due to http://crbug.com/421387.
716   if (uncommitted_size > internal::kReasonableSizeOfUnusedPages) {
717     internal::PartitionOutOfMemoryWithLotsOfUncommitedPages(size);
718   }
719 
720 #if BUILDFLAG(IS_WIN)
721   // If true then we are running on 64-bit Windows.
722   BOOL is_wow_64 = FALSE;
723   // Intentionally ignoring failures.
724   IsWow64Process(GetCurrentProcess(), &is_wow_64);
725   // 32-bit address space on Windows is typically either 2 GiB (on 32-bit
726   // Windows) or 4 GiB (on 64-bit Windows). 2.8 and 1.0 GiB are just rough
727   // guesses as to how much address space PA can consume (note that code,
728   // stacks, and other allocators will also consume address space).
729   const size_t kReasonableVirtualSize = (is_wow_64 ? 2800 : 1024) * 1024 * 1024;
730   // Make it obvious whether we are running on 64-bit Windows.
731   PA_DEBUG_DATA_ON_STACK("is_wow_64", static_cast<size_t>(is_wow_64));
732 #else
733   constexpr size_t kReasonableVirtualSize =
734       // 1.5GiB elsewhere, since address space is typically 3GiB.
735       (1024 + 512) * 1024 * 1024;
736 #endif
737   if (virtual_address_space_size > kReasonableVirtualSize) {
738     internal::PartitionOutOfMemoryWithLargeVirtualSize(
739         virtual_address_space_size);
740   }
741 #endif  // #if !defined(ARCH_CPU_64_BITS)
742 
743   // Out of memory can be due to multiple causes, such as:
744   // - Out of virtual address space in the desired pool
745   // - Out of commit due to either our process, or another one
746   // - Excessive allocations in the current process
747   //
748   // Saving these values make it easier to distinguish between these. See the
749   // documentation in PA_CONFIG(DEBUG_DATA_ON_STACK) on how to get these from
750   // minidumps.
751   PA_DEBUG_DATA_ON_STACK("va_size", virtual_address_space_size);
752   PA_DEBUG_DATA_ON_STACK("alloc", get_total_size_of_allocated_bytes());
753   PA_DEBUG_DATA_ON_STACK("commit", get_total_size_of_committed_pages());
754   PA_DEBUG_DATA_ON_STACK("size", size);
755 
756   if (internal::g_oom_handling_function) {
757     (*internal::g_oom_handling_function)(size);
758   }
759   OOM_CRASH(size);
760 }
761 
762 template <bool thread_safe>
DecommitEmptySlotSpans()763 void PartitionRoot<thread_safe>::DecommitEmptySlotSpans() {
764   ShrinkEmptySlotSpansRing(0);
765   // Just decommitted everything, and holding the lock, should be exactly 0.
766   PA_DCHECK(empty_slot_spans_dirty_bytes == 0);
767 }
768 
769 template <bool thread_safe>
DestructForTesting()770 void PartitionRoot<thread_safe>::DestructForTesting() {
771   // We need to destruct the thread cache before we unreserve any of the super
772   // pages below, which we currently are not doing. So, we should only call
773   // this function on PartitionRoots without a thread cache.
774   PA_CHECK(!flags.with_thread_cache);
775   auto pool_handle = ChoosePool();
776 #if BUILDFLAG(ENABLE_PKEYS)
777   // The pages managed by pkey will be free-ed at UninitPKeyForTesting().
778   // Don't invoke FreePages() for the pages.
779   if (pool_handle == internal::kPkeyPoolHandle) {
780     return;
781   }
782   PA_DCHECK(pool_handle < internal::kNumPools);
783 #else
784   PA_DCHECK(pool_handle <= internal::kNumPools);
785 #endif
786 
787   auto* curr = first_extent;
788   while (curr != nullptr) {
789     auto* next = curr->next;
790     uintptr_t address = SuperPagesBeginFromExtent(curr);
791     size_t size =
792         internal::kSuperPageSize * curr->number_of_consecutive_super_pages;
793 #if !BUILDFLAG(HAS_64_BIT_POINTERS)
794     internal::AddressPoolManager::GetInstance().MarkUnused(pool_handle, address,
795                                                            size);
796 #endif
797     internal::AddressPoolManager::GetInstance().UnreserveAndDecommit(
798         pool_handle, address, size);
799     curr = next;
800   }
801 }
802 
803 #if PA_CONFIG(ENABLE_MAC11_MALLOC_SIZE_HACK)
804 template <bool thread_safe>
EnableMac11MallocSizeHackForTesting()805 void PartitionRoot<thread_safe>::EnableMac11MallocSizeHackForTesting() {
806   flags.mac11_malloc_size_hack_enabled_ = true;
807 }
808 #endif  // PA_CONFIG(ENABLE_MAC11_MALLOC_SIZE_HACK)
809 
810 #if BUILDFLAG(ENABLE_BACKUP_REF_PTR_SUPPORT) && !BUILDFLAG(HAS_64_BIT_POINTERS)
811 namespace {
812 std::atomic<bool> g_reserve_brp_guard_region_called;
813 // An address constructed by repeating `kQuarantinedByte` shouldn't never point
814 // to valid memory. Preemptively reserve a memory region around that address and
815 // make it inaccessible. Not needed for 64-bit platforms where the address is
816 // guaranteed to be non-canonical. Safe to call multiple times.
ReserveBackupRefPtrGuardRegionIfNeeded()817 void ReserveBackupRefPtrGuardRegionIfNeeded() {
818   bool expected = false;
819   // No need to block execution for potential concurrent initialization, merely
820   // want to make sure this is only called once.
821   if (!g_reserve_brp_guard_region_called.compare_exchange_strong(expected,
822                                                                  true)) {
823     return;
824   }
825 
826   size_t alignment = internal::PageAllocationGranularity();
827   uintptr_t requested_address;
828   memset(&requested_address, internal::kQuarantinedByte,
829          sizeof(requested_address));
830   requested_address = RoundDownToPageAllocationGranularity(requested_address);
831 
832   // Request several pages so that even unreasonably large C++ objects stay
833   // within the inaccessible region. If some of the pages can't be reserved,
834   // it's still preferable to try and reserve the rest.
835   for (size_t i = 0; i < 4; ++i) {
836     [[maybe_unused]] uintptr_t allocated_address =
837         AllocPages(requested_address, alignment, alignment,
838                    PageAccessibilityConfiguration(
839                        PageAccessibilityConfiguration::kInaccessible),
840                    PageTag::kPartitionAlloc);
841     requested_address += alignment;
842   }
843 }
844 }  // namespace
845 #endif  // BUILDFLAG(ENABLE_BACKUP_REF_PTR_SUPPORT) &&
846         // !BUILDFLAG(HAS_64_BIT_POINTERS)
847 
848 template <bool thread_safe>
Init(PartitionOptions opts)849 void PartitionRoot<thread_safe>::Init(PartitionOptions opts) {
850   {
851 #if BUILDFLAG(IS_APPLE)
852     // Needed to statically bound page size, which is a runtime constant on
853     // apple OSes.
854     PA_CHECK((internal::SystemPageSize() == (size_t{1} << 12)) ||
855              (internal::SystemPageSize() == (size_t{1} << 14)));
856 #elif BUILDFLAG(IS_LINUX) && defined(ARCH_CPU_ARM64)
857     // Check runtime pagesize. Though the code is currently the same, it is
858     // not merged with the IS_APPLE case above as a 1 << 16 case needs to be
859     // added here in the future, to allow 64 kiB pagesize. That is only
860     // supported on Linux on arm64, not on IS_APPLE, but not yet present here
861     // as the rest of the partition allocator does not currently support it.
862     PA_CHECK((internal::SystemPageSize() == (size_t{1} << 12)) ||
863              (internal::SystemPageSize() == (size_t{1} << 14)));
864 #endif
865 
866     ::partition_alloc::internal::ScopedGuard guard{lock_};
867     if (initialized) {
868       return;
869     }
870 
871     // Swaps out the active no-op tagging intrinsics with MTE-capable ones, if
872     // running on the right hardware.
873     ::partition_alloc::internal::InitializeMTESupportIfNeeded();
874 
875 #if BUILDFLAG(HAS_64_BIT_POINTERS)
876     // Reserve address space for partition alloc.
877     internal::PartitionAddressSpace::Init();
878 #endif
879 
880 #if BUILDFLAG(ENABLE_BACKUP_REF_PTR_SUPPORT) && !BUILDFLAG(HAS_64_BIT_POINTERS)
881     ReserveBackupRefPtrGuardRegionIfNeeded();
882 #endif
883 
884     flags.allow_aligned_alloc =
885         opts.aligned_alloc == PartitionOptions::AlignedAlloc::kAllowed;
886     flags.allow_cookie = opts.cookie == PartitionOptions::Cookie::kAllowed;
887 #if BUILDFLAG(ENABLE_BACKUP_REF_PTR_SUPPORT)
888     flags.brp_enabled_ =
889         opts.backup_ref_ptr == PartitionOptions::BackupRefPtr::kEnabled;
890     flags.brp_zapping_enabled_ =
891         opts.backup_ref_ptr_zapping ==
892         PartitionOptions::BackupRefPtrZapping::kEnabled;
893     PA_CHECK(!flags.brp_zapping_enabled_ || flags.brp_enabled_);
894 #if PA_CONFIG(ENABLE_MAC11_MALLOC_SIZE_HACK) && BUILDFLAG(IS_APPLE)
895     flags.mac11_malloc_size_hack_enabled_ =
896         flags.brp_enabled_ && internal::base::mac::IsOS11();
897 #endif  // PA_CONFIG(ENABLE_MAC11_MALLOC_SIZE_HACK) && BUILDFLAG(IS_APPLE)
898 #else   // BUILDFLAG(ENABLE_BACKUP_REF_PTR_SUPPORT)
899     PA_CHECK(opts.backup_ref_ptr == PartitionOptions::BackupRefPtr::kDisabled);
900 #endif  // BUILDFLAG(ENABLE_BACKUP_REF_PTR_SUPPORT)
901     flags.use_configurable_pool =
902         (opts.use_configurable_pool ==
903          PartitionOptions::UseConfigurablePool::kIfAvailable) &&
904         IsConfigurablePoolAvailable();
905     PA_DCHECK(!flags.use_configurable_pool || IsConfigurablePoolAvailable());
906 
907     // brp_enabled() is not supported in the configurable pool because
908     // BRP requires objects to be in a different Pool.
909     PA_CHECK(!(flags.use_configurable_pool && brp_enabled()));
910 
911 #if BUILDFLAG(ENABLE_PKEYS)
912     // BRP and pkey mode use different pools, so they can't be enabled at the
913     // same time.
914     PA_CHECK(opts.pkey == internal::kDefaultPkey ||
915              opts.backup_ref_ptr == PartitionOptions::BackupRefPtr::kDisabled);
916     flags.pkey = opts.pkey;
917 #endif
918 
919     // Ref-count messes up alignment needed for AlignedAlloc, making this
920     // option incompatible. However, except in the
921     // PUT_REF_COUNT_IN_PREVIOUS_SLOT case.
922 #if BUILDFLAG(ENABLE_BACKUP_REF_PTR_SUPPORT) && \
923     !BUILDFLAG(PUT_REF_COUNT_IN_PREVIOUS_SLOT)
924     PA_CHECK(!flags.allow_aligned_alloc || !flags.brp_enabled_);
925 #endif
926 
927 #if PA_CONFIG(EXTRAS_REQUIRED)
928     flags.extras_size = 0;
929     flags.extras_offset = 0;
930 
931     if (flags.allow_cookie) {
932       flags.extras_size += internal::kPartitionCookieSizeAdjustment;
933     }
934 
935     if (brp_enabled()) {
936       // TODO(tasak): In the PUT_REF_COUNT_IN_PREVIOUS_SLOT case, ref-count is
937       // stored out-of-line for single-slot slot spans, so no need to
938       // add/subtract its size in this case.
939       flags.extras_size += internal::kPartitionRefCountSizeAdjustment;
940       flags.extras_offset += internal::kPartitionRefCountOffsetAdjustment;
941     }
942     if (opts.add_dummy_ref_count ==
943         PartitionOptions::AddDummyRefCount::kEnabled) {
944       // AddDummyRefCount will increase the size to simulate adding
945       // PartitionRefCount, but non of the BRP logic will run.
946       PA_CHECK(!brp_enabled());
947       flags.extras_size += internal::kPartitionRefCountSizeAdjustment;
948     }
949 #endif  // PA_CONFIG(EXTRAS_REQUIRED)
950 
951     // Re-confirm the above PA_CHECKs, by making sure there are no
952     // pre-allocation extras when AlignedAlloc is allowed. Post-allocation
953     // extras are ok.
954     PA_CHECK(!flags.allow_aligned_alloc || !flags.extras_offset);
955 
956     flags.quarantine_mode =
957 #if BUILDFLAG(USE_STARSCAN)
958         (opts.quarantine == PartitionOptions::Quarantine::kDisallowed
959              ? QuarantineMode::kAlwaysDisabled
960              : QuarantineMode::kDisabledByDefault);
961 #else
962         QuarantineMode::kAlwaysDisabled;
963 #endif  // BUILDFLAG(USE_STARSCAN)
964 
965     // We mark the sentinel slot span as free to make sure it is skipped by our
966     // logic to find a new active slot span.
967     memset(&sentinel_bucket, 0, sizeof(sentinel_bucket));
968     sentinel_bucket.active_slot_spans_head =
969         SlotSpan::get_sentinel_slot_span_non_const();
970 
971     // This is a "magic" value so we can test if a root pointer is valid.
972     inverted_self = ~reinterpret_cast<uintptr_t>(this);
973 
974     // Set up the actual usable buckets first.
975     constexpr internal::BucketIndexLookup lookup{};
976     size_t bucket_index = 0;
977     while (lookup.bucket_sizes()[bucket_index] !=
978            internal::kInvalidBucketSize) {
979       buckets[bucket_index].Init(lookup.bucket_sizes()[bucket_index]);
980       bucket_index++;
981     }
982     PA_DCHECK(bucket_index < internal::kNumBuckets);
983 
984     // Remaining buckets are not usable, and not real.
985     for (size_t index = bucket_index; index < internal::kNumBuckets; index++) {
986       // Cannot init with size 0 since it computes 1 / size, but make sure the
987       // bucket is invalid.
988       buckets[index].Init(internal::kInvalidBucketSize);
989       buckets[index].active_slot_spans_head = nullptr;
990       PA_DCHECK(!buckets[index].is_valid());
991     }
992 
993 #if !PA_CONFIG(THREAD_CACHE_SUPPORTED)
994     // TLS in ThreadCache not supported on other OSes.
995     flags.with_thread_cache = false;
996 #else
997     ThreadCache::EnsureThreadSpecificDataInitialized();
998     flags.with_thread_cache =
999         (opts.thread_cache == PartitionOptions::ThreadCache::kEnabled);
1000 
1001     if (flags.with_thread_cache) {
1002       ThreadCache::Init(this);
1003     }
1004 #endif  // !PA_CONFIG(THREAD_CACHE_SUPPORTED)
1005 
1006 #if PA_CONFIG(USE_PARTITION_ROOT_ENUMERATOR)
1007     internal::PartitionRootEnumerator::Instance().Register(this);
1008 #endif
1009 
1010     initialized = true;
1011   }
1012 
1013   // Called without the lock, might allocate.
1014 #if BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
1015   PartitionAllocMallocInitOnce();
1016 #endif
1017 
1018 #if BUILDFLAG(ENABLE_PKEYS)
1019   if (flags.pkey != internal::kDefaultPkey) {
1020     internal::PartitionAllocPkeyInit(flags.pkey);
1021   }
1022 #endif
1023 }
1024 
1025 template <bool thread_safe>
~PartitionRoot()1026 PartitionRoot<thread_safe>::~PartitionRoot() {
1027 #if BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
1028   PA_CHECK(!flags.with_thread_cache)
1029       << "Must not destroy a partition with a thread cache";
1030 #endif  // BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
1031 
1032 #if PA_CONFIG(USE_PARTITION_ROOT_ENUMERATOR)
1033   if (initialized) {
1034     internal::PartitionRootEnumerator::Instance().Unregister(this);
1035   }
1036 #endif  // PA_CONFIG(USE_PARTITION_ALLOC_ENUMERATOR)
1037 }
1038 
1039 template <bool thread_safe>
EnableThreadCacheIfSupported()1040 void PartitionRoot<thread_safe>::EnableThreadCacheIfSupported() {
1041 #if PA_CONFIG(THREAD_CACHE_SUPPORTED)
1042   ::partition_alloc::internal::ScopedGuard guard{lock_};
1043   PA_CHECK(!flags.with_thread_cache);
1044   // By the time we get there, there may be multiple threads created in the
1045   // process. Since `with_thread_cache` is accessed without a lock, it can
1046   // become visible to another thread before the effects of
1047   // `internal::ThreadCacheInit()` are visible. To prevent that, we fake thread
1048   // cache creation being in-progress while this is running.
1049   //
1050   // This synchronizes with the acquire load in `MaybeInitThreadCacheAndAlloc()`
1051   // to ensure that we don't create (and thus use) a ThreadCache before
1052   // ThreadCache::Init()'s effects are visible.
1053   int before =
1054       thread_caches_being_constructed_.fetch_add(1, std::memory_order_acquire);
1055   PA_CHECK(before == 0);
1056   ThreadCache::Init(this);
1057   thread_caches_being_constructed_.fetch_sub(1, std::memory_order_release);
1058   flags.with_thread_cache = true;
1059 #endif  // PA_CONFIG(THREAD_CACHE_SUPPORTED)
1060 }
1061 
1062 template <bool thread_safe>
TryReallocInPlaceForDirectMap(internal::SlotSpanMetadata<thread_safe> * slot_span,size_t requested_size)1063 bool PartitionRoot<thread_safe>::TryReallocInPlaceForDirectMap(
1064     internal::SlotSpanMetadata<thread_safe>* slot_span,
1065     size_t requested_size) {
1066   PA_DCHECK(slot_span->bucket->is_direct_mapped());
1067   // Slot-span metadata isn't MTE-tagged.
1068   PA_DCHECK(
1069       internal::IsManagedByDirectMap(reinterpret_cast<uintptr_t>(slot_span)));
1070 
1071   size_t raw_size = AdjustSizeForExtrasAdd(requested_size);
1072   auto* extent = DirectMapExtent::FromSlotSpan(slot_span);
1073   size_t current_reservation_size = extent->reservation_size;
1074   // Calculate the new reservation size the way PartitionDirectMap() would, but
1075   // skip the alignment, because this call isn't requesting it.
1076   size_t new_reservation_size = GetDirectMapReservationSize(raw_size);
1077 
1078   // If new reservation would be larger, there is nothing we can do to
1079   // reallocate in-place.
1080   if (new_reservation_size > current_reservation_size) {
1081     return false;
1082   }
1083 
1084   // Don't reallocate in-place if new reservation size would be less than 80 %
1085   // of the current one, to avoid holding on to too much unused address space.
1086   // Make this check before comparing slot sizes, as even with equal or similar
1087   // slot sizes we can save a lot if the original allocation was heavily padded
1088   // for alignment.
1089   if ((new_reservation_size >> internal::SystemPageShift()) * 5 <
1090       (current_reservation_size >> internal::SystemPageShift()) * 4) {
1091     return false;
1092   }
1093 
1094   // Note that the new size isn't a bucketed size; this function is called
1095   // whenever we're reallocating a direct mapped allocation, so calculate it
1096   // the way PartitionDirectMap() would.
1097   size_t new_slot_size = GetDirectMapSlotSize(raw_size);
1098   if (new_slot_size < internal::kMinDirectMappedDownsize) {
1099     return false;
1100   }
1101 
1102   // Past this point, we decided we'll attempt to reallocate without relocating,
1103   // so we have to honor the padding for alignment in front of the original
1104   // allocation, even though this function isn't requesting any alignment.
1105 
1106   // bucket->slot_size is the currently committed size of the allocation.
1107   size_t current_slot_size = slot_span->bucket->slot_size;
1108   size_t current_usable_size = slot_span->GetUsableSize(this);
1109   uintptr_t slot_start = SlotSpan::ToSlotSpanStart(slot_span);
1110   // This is the available part of the reservation up to which the new
1111   // allocation can grow.
1112   size_t available_reservation_size =
1113       current_reservation_size - extent->padding_for_alignment -
1114       PartitionRoot<thread_safe>::GetDirectMapMetadataAndGuardPagesSize();
1115 #if BUILDFLAG(PA_DCHECK_IS_ON)
1116   uintptr_t reservation_start = slot_start & internal::kSuperPageBaseMask;
1117   PA_DCHECK(internal::IsReservationStart(reservation_start));
1118   PA_DCHECK(slot_start + available_reservation_size ==
1119             reservation_start + current_reservation_size -
1120                 GetDirectMapMetadataAndGuardPagesSize() +
1121                 internal::PartitionPageSize());
1122 #endif
1123 
1124   if (new_slot_size == current_slot_size) {
1125     // No need to move any memory around, but update size and cookie below.
1126     // That's because raw_size may have changed.
1127   } else if (new_slot_size < current_slot_size) {
1128     // Shrink by decommitting unneeded pages and making them inaccessible.
1129     size_t decommit_size = current_slot_size - new_slot_size;
1130     DecommitSystemPagesForData(slot_start + new_slot_size, decommit_size,
1131                                PageAccessibilityDisposition::kRequireUpdate);
1132     // Since the decommited system pages are still reserved, we don't need to
1133     // change the entries for decommitted pages in the reservation offset table.
1134   } else if (new_slot_size <= available_reservation_size) {
1135     // Grow within the actually reserved address space. Just need to make the
1136     // pages accessible again.
1137     size_t recommit_slot_size_growth = new_slot_size - current_slot_size;
1138     RecommitSystemPagesForData(slot_start + current_slot_size,
1139                                recommit_slot_size_growth,
1140                                PageAccessibilityDisposition::kRequireUpdate);
1141     // The recommited system pages had been already reserved and all the
1142     // entries in the reservation offset table (for entire reservation_size
1143     // region) have been already initialized.
1144 
1145 #if BUILDFLAG(PA_DCHECK_IS_ON)
1146     memset(reinterpret_cast<void*>(slot_start + current_slot_size),
1147            internal::kUninitializedByte, recommit_slot_size_growth);
1148 #endif
1149   } else {
1150     // We can't perform the realloc in-place.
1151     // TODO: support this too when possible.
1152     return false;
1153   }
1154 
1155   DecreaseTotalSizeOfAllocatedBytes(reinterpret_cast<uintptr_t>(slot_span),
1156                                     slot_span->bucket->slot_size);
1157   slot_span->SetRawSize(raw_size);
1158   slot_span->bucket->slot_size = new_slot_size;
1159   IncreaseTotalSizeOfAllocatedBytes(reinterpret_cast<uintptr_t>(slot_span),
1160                                     slot_span->bucket->slot_size, raw_size);
1161 
1162   // Always record in-place realloc() as free()+malloc() pair.
1163   //
1164   // The early returns above (`return false`) will fall back to free()+malloc(),
1165   // so this is consistent.
1166   auto* thread_cache = GetOrCreateThreadCache();
1167   if (ThreadCache::IsValid(thread_cache)) {
1168     thread_cache->RecordDeallocation(current_usable_size);
1169     thread_cache->RecordAllocation(slot_span->GetUsableSize(this));
1170   }
1171 
1172 #if BUILDFLAG(PA_DCHECK_IS_ON)
1173   // Write a new trailing cookie.
1174   if (flags.allow_cookie) {
1175     auto* object = static_cast<unsigned char*>(SlotStartToObject(slot_start));
1176     internal::PartitionCookieWriteValue(object +
1177                                         slot_span->GetUsableSize(this));
1178   }
1179 #endif
1180 
1181   return true;
1182 }
1183 
1184 template <bool thread_safe>
TryReallocInPlaceForNormalBuckets(void * object,SlotSpan * slot_span,size_t new_size)1185 bool PartitionRoot<thread_safe>::TryReallocInPlaceForNormalBuckets(
1186     void* object,
1187     SlotSpan* slot_span,
1188     size_t new_size) {
1189   uintptr_t slot_start = ObjectToSlotStart(object);
1190   PA_DCHECK(internal::IsManagedByNormalBuckets(slot_start));
1191 
1192   // TODO: note that tcmalloc will "ignore" a downsizing realloc() unless the
1193   // new size is a significant percentage smaller. We could do the same if we
1194   // determine it is a win.
1195   if (AllocationCapacityFromRequestedSize(new_size) !=
1196       AllocationCapacityFromSlotStart(slot_start)) {
1197     return false;
1198   }
1199   size_t current_usable_size = slot_span->GetUsableSize(this);
1200 
1201   // Trying to allocate |new_size| would use the same amount of underlying
1202   // memory as we're already using, so re-use the allocation after updating
1203   // statistics (and cookie, if present).
1204   if (slot_span->CanStoreRawSize()) {
1205 #if BUILDFLAG(PUT_REF_COUNT_IN_PREVIOUS_SLOT) && BUILDFLAG(PA_DCHECK_IS_ON)
1206     internal::PartitionRefCount* old_ref_count;
1207     if (brp_enabled()) {
1208       old_ref_count = internal::PartitionRefCountPointer(slot_start);
1209     }
1210 #endif  // BUILDFLAG(PUT_REF_COUNT_IN_PREVIOUS_SLOT) &&
1211         // BUILDFLAG(PA_DCHECK_IS_ON)
1212     size_t new_raw_size = AdjustSizeForExtrasAdd(new_size);
1213     slot_span->SetRawSize(new_raw_size);
1214 #if BUILDFLAG(PUT_REF_COUNT_IN_PREVIOUS_SLOT) && BUILDFLAG(PA_DCHECK_IS_ON)
1215     if (brp_enabled()) {
1216       internal::PartitionRefCount* new_ref_count =
1217           internal::PartitionRefCountPointer(slot_start);
1218       PA_DCHECK(new_ref_count == old_ref_count);
1219     }
1220 #endif  // BUILDFLAG(PUT_REF_COUNT_IN_PREVIOUS_SLOT) &&
1221         // BUILDFLAG(PA_DCHECK_IS_ON)
1222 #if BUILDFLAG(PA_DCHECK_IS_ON)
1223     // Write a new trailing cookie only when it is possible to keep track
1224     // raw size (otherwise we wouldn't know where to look for it later).
1225     if (flags.allow_cookie) {
1226       internal::PartitionCookieWriteValue(static_cast<unsigned char*>(object) +
1227                                           slot_span->GetUsableSize(this));
1228     }
1229 #endif  // BUILDFLAG(PA_DCHECK_IS_ON)
1230   }
1231 
1232   // Always record a realloc() as a free() + malloc(), even if it's in
1233   // place. When we cannot do it in place (`return false` above), the allocator
1234   // falls back to free()+malloc(), so this is consistent.
1235   ThreadCache* thread_cache = GetOrCreateThreadCache();
1236   if (PA_LIKELY(ThreadCache::IsValid(thread_cache))) {
1237     thread_cache->RecordDeallocation(current_usable_size);
1238     thread_cache->RecordAllocation(slot_span->GetUsableSize(this));
1239   }
1240 
1241   return object;
1242 }
1243 
1244 template <bool thread_safe>
ReallocWithFlags(unsigned int flags,void * ptr,size_t new_size,const char * type_name)1245 void* PartitionRoot<thread_safe>::ReallocWithFlags(unsigned int flags,
1246                                                    void* ptr,
1247                                                    size_t new_size,
1248                                                    const char* type_name) {
1249 #if defined(MEMORY_TOOL_REPLACES_ALLOCATOR)
1250   CHECK_MAX_SIZE_OR_RETURN_NULLPTR(new_size, flags);
1251   void* result = realloc(ptr, new_size);
1252   PA_CHECK(result || flags & AllocFlags::kReturnNull);
1253   return result;
1254 #else
1255   bool no_hooks = flags & AllocFlags::kNoHooks;
1256   if (PA_UNLIKELY(!ptr)) {
1257     return no_hooks
1258                ? AllocWithFlagsNoHooks(flags, new_size,
1259                                        internal::PartitionPageSize())
1260                : AllocWithFlagsInternal(
1261                      flags, new_size, internal::PartitionPageSize(), type_name);
1262   }
1263 
1264   if (PA_UNLIKELY(!new_size)) {
1265     Free(ptr);
1266     return nullptr;
1267   }
1268 
1269   if (new_size > internal::MaxDirectMapped()) {
1270     if (flags & AllocFlags::kReturnNull) {
1271       return nullptr;
1272     }
1273     internal::PartitionExcessiveAllocationSize(new_size);
1274   }
1275 
1276   const bool hooks_enabled = PartitionAllocHooks::AreHooksEnabled();
1277   bool overridden = false;
1278   size_t old_usable_size;
1279   if (PA_UNLIKELY(!no_hooks && hooks_enabled)) {
1280     overridden = PartitionAllocHooks::ReallocOverrideHookIfEnabled(
1281         &old_usable_size, ptr);
1282   }
1283   if (PA_LIKELY(!overridden)) {
1284     // |ptr| may have been allocated in another root.
1285     SlotSpan* slot_span = SlotSpan::FromObject(ptr);
1286     auto* old_root = PartitionRoot::FromSlotSpan(slot_span);
1287     bool success = false;
1288     bool tried_in_place_for_direct_map = false;
1289     {
1290       ::partition_alloc::internal::ScopedGuard guard{old_root->lock_};
1291       // TODO(crbug.com/1257655): See if we can afford to make this a CHECK.
1292       PA_DCHECK(IsValidSlotSpan(slot_span));
1293       old_usable_size = slot_span->GetUsableSize(old_root);
1294 
1295       if (PA_UNLIKELY(slot_span->bucket->is_direct_mapped())) {
1296         tried_in_place_for_direct_map = true;
1297         // We may be able to perform the realloc in place by changing the
1298         // accessibility of memory pages and, if reducing the size, decommitting
1299         // them.
1300         success = old_root->TryReallocInPlaceForDirectMap(slot_span, new_size);
1301       }
1302     }
1303     if (success) {
1304       if (PA_UNLIKELY(!no_hooks && hooks_enabled)) {
1305         PartitionAllocHooks::ReallocObserverHookIfEnabled(ptr, ptr, new_size,
1306                                                           type_name);
1307       }
1308       return ptr;
1309     }
1310 
1311     if (PA_LIKELY(!tried_in_place_for_direct_map)) {
1312       if (old_root->TryReallocInPlaceForNormalBuckets(ptr, slot_span,
1313                                                       new_size)) {
1314         return ptr;
1315       }
1316     }
1317   }
1318 
1319   // This realloc cannot be resized in-place. Sadness.
1320   void* ret =
1321       no_hooks ? AllocWithFlagsNoHooks(flags, new_size,
1322                                        internal::PartitionPageSize())
1323                : AllocWithFlagsInternal(
1324                      flags, new_size, internal::PartitionPageSize(), type_name);
1325   if (!ret) {
1326     if (flags & AllocFlags::kReturnNull) {
1327       return nullptr;
1328     }
1329     internal::PartitionExcessiveAllocationSize(new_size);
1330   }
1331 
1332   memcpy(ret, ptr, std::min(old_usable_size, new_size));
1333   Free(ptr);  // Implicitly protects the old ptr on MTE systems.
1334   return ret;
1335 #endif
1336 }
1337 
1338 template <bool thread_safe>
PurgeMemory(int flags)1339 void PartitionRoot<thread_safe>::PurgeMemory(int flags) {
1340   {
1341     ::partition_alloc::internal::ScopedGuard guard{lock_};
1342 #if BUILDFLAG(USE_STARSCAN)
1343     // Avoid purging if there is PCScan task currently scheduled. Since pcscan
1344     // takes snapshot of all allocated pages, decommitting pages here (even
1345     // under the lock) is racy.
1346     // TODO(bikineev): Consider rescheduling the purging after PCScan.
1347     if (PCScan::IsInProgress()) {
1348       return;
1349     }
1350 #endif  // BUILDFLAG(USE_STARSCAN)
1351 
1352     if (flags & PurgeFlags::kDecommitEmptySlotSpans) {
1353       DecommitEmptySlotSpans();
1354     }
1355     if (flags & PurgeFlags::kDiscardUnusedSystemPages) {
1356       for (Bucket& bucket : buckets) {
1357         if (bucket.slot_size == internal::kInvalidBucketSize) {
1358           continue;
1359         }
1360 
1361         if (bucket.slot_size >= internal::MinPurgeableSlotSize()) {
1362           internal::PartitionPurgeBucket(&bucket);
1363         } else {
1364           bucket.SortSlotSpanFreelists();
1365         }
1366 
1367         // Do it at the end, as the actions above change the status of slot
1368         // spans (e.g. empty -> decommitted).
1369         bucket.MaintainActiveList();
1370 
1371         if (sort_active_slot_spans_) {
1372           bucket.SortActiveSlotSpans();
1373         }
1374       }
1375     }
1376   }
1377 }
1378 
1379 template <bool thread_safe>
ShrinkEmptySlotSpansRing(size_t limit)1380 void PartitionRoot<thread_safe>::ShrinkEmptySlotSpansRing(size_t limit) {
1381   int16_t index = global_empty_slot_span_ring_index;
1382   int16_t starting_index = index;
1383   while (empty_slot_spans_dirty_bytes > limit) {
1384     SlotSpan* slot_span = global_empty_slot_span_ring[index];
1385     // The ring is not always full, may be nullptr.
1386     if (slot_span) {
1387       slot_span->DecommitIfPossible(this);
1388       global_empty_slot_span_ring[index] = nullptr;
1389     }
1390     index += 1;
1391     // Walk through the entirety of possible slots, even though the last ones
1392     // are unused, if global_empty_slot_span_ring_size is smaller than
1393     // kMaxFreeableSpans. It's simpler, and does not cost anything, since all
1394     // the pointers are going to be nullptr.
1395     if (index == internal::kMaxFreeableSpans) {
1396       index = 0;
1397     }
1398 
1399     // Went around the whole ring, since this is locked,
1400     // empty_slot_spans_dirty_bytes should be exactly 0.
1401     if (index == starting_index) {
1402       PA_DCHECK(empty_slot_spans_dirty_bytes == 0);
1403       // Metrics issue, don't crash, return.
1404       break;
1405     }
1406   }
1407 }
1408 
1409 template <bool thread_safe>
DumpStats(const char * partition_name,bool is_light_dump,PartitionStatsDumper * dumper)1410 void PartitionRoot<thread_safe>::DumpStats(const char* partition_name,
1411                                            bool is_light_dump,
1412                                            PartitionStatsDumper* dumper) {
1413   static const size_t kMaxReportableDirectMaps = 4096;
1414   // Allocate on the heap rather than on the stack to avoid stack overflow
1415   // skirmishes (on Windows, in particular). Allocate before locking below,
1416   // otherwise when PartitionAlloc is malloc() we get reentrancy issues. This
1417   // inflates reported values a bit for detailed dumps though, by 16kiB.
1418   std::unique_ptr<uint32_t[]> direct_map_lengths;
1419   if (!is_light_dump) {
1420     direct_map_lengths =
1421         std::unique_ptr<uint32_t[]>(new uint32_t[kMaxReportableDirectMaps]);
1422   }
1423   PartitionBucketMemoryStats bucket_stats[internal::kNumBuckets];
1424   size_t num_direct_mapped_allocations = 0;
1425   PartitionMemoryStats stats = {0};
1426 
1427   stats.syscall_count = syscall_count.load(std::memory_order_relaxed);
1428   stats.syscall_total_time_ns =
1429       syscall_total_time_ns.load(std::memory_order_relaxed);
1430 
1431   // Collect data with the lock held, cannot allocate or call third-party code
1432   // below.
1433   {
1434     ::partition_alloc::internal::ScopedGuard guard{lock_};
1435     PA_DCHECK(total_size_of_allocated_bytes <= max_size_of_allocated_bytes);
1436 
1437     stats.total_mmapped_bytes =
1438         total_size_of_super_pages.load(std::memory_order_relaxed) +
1439         total_size_of_direct_mapped_pages.load(std::memory_order_relaxed);
1440     stats.total_committed_bytes =
1441         total_size_of_committed_pages.load(std::memory_order_relaxed);
1442     stats.max_committed_bytes =
1443         max_size_of_committed_pages.load(std::memory_order_relaxed);
1444     stats.total_allocated_bytes = total_size_of_allocated_bytes;
1445     stats.max_allocated_bytes = max_size_of_allocated_bytes;
1446 #if BUILDFLAG(ENABLE_BACKUP_REF_PTR_SUPPORT)
1447     stats.total_brp_quarantined_bytes =
1448         total_size_of_brp_quarantined_bytes.load(std::memory_order_relaxed);
1449     stats.total_brp_quarantined_count =
1450         total_count_of_brp_quarantined_slots.load(std::memory_order_relaxed);
1451     stats.cumulative_brp_quarantined_bytes =
1452         cumulative_size_of_brp_quarantined_bytes.load(
1453             std::memory_order_relaxed);
1454     stats.cumulative_brp_quarantined_count =
1455         cumulative_count_of_brp_quarantined_slots.load(
1456             std::memory_order_relaxed);
1457 #endif
1458 
1459     size_t direct_mapped_allocations_total_size = 0;
1460     for (size_t i = 0; i < internal::kNumBuckets; ++i) {
1461       const Bucket* bucket = &bucket_at(i);
1462       // Don't report the pseudo buckets that the generic allocator sets up in
1463       // order to preserve a fast size->bucket map (see
1464       // PartitionRoot::Init() for details).
1465       if (!bucket->is_valid()) {
1466         bucket_stats[i].is_valid = false;
1467       } else {
1468         internal::PartitionDumpBucketStats(&bucket_stats[i], bucket);
1469       }
1470       if (bucket_stats[i].is_valid) {
1471         stats.total_resident_bytes += bucket_stats[i].resident_bytes;
1472         stats.total_active_bytes += bucket_stats[i].active_bytes;
1473         stats.total_active_count += bucket_stats[i].active_count;
1474         stats.total_decommittable_bytes += bucket_stats[i].decommittable_bytes;
1475         stats.total_discardable_bytes += bucket_stats[i].discardable_bytes;
1476       }
1477     }
1478 
1479     for (DirectMapExtent* extent = direct_map_list;
1480          extent && num_direct_mapped_allocations < kMaxReportableDirectMaps;
1481          extent = extent->next_extent, ++num_direct_mapped_allocations) {
1482       PA_DCHECK(!extent->next_extent ||
1483                 extent->next_extent->prev_extent == extent);
1484       size_t slot_size = extent->bucket->slot_size;
1485       direct_mapped_allocations_total_size += slot_size;
1486       if (is_light_dump) {
1487         continue;
1488       }
1489       direct_map_lengths[num_direct_mapped_allocations] = slot_size;
1490     }
1491 
1492     stats.total_resident_bytes += direct_mapped_allocations_total_size;
1493     stats.total_active_bytes += direct_mapped_allocations_total_size;
1494     stats.total_active_count += num_direct_mapped_allocations;
1495 
1496     stats.has_thread_cache = flags.with_thread_cache;
1497     if (stats.has_thread_cache) {
1498       ThreadCacheRegistry::Instance().DumpStats(
1499           true, &stats.current_thread_cache_stats);
1500       ThreadCacheRegistry::Instance().DumpStats(false,
1501                                                 &stats.all_thread_caches_stats);
1502     }
1503   }
1504 
1505   // Do not hold the lock when calling |dumper|, as it may allocate.
1506   if (!is_light_dump) {
1507     for (auto& stat : bucket_stats) {
1508       if (stat.is_valid) {
1509         dumper->PartitionsDumpBucketStats(partition_name, &stat);
1510       }
1511     }
1512 
1513     for (size_t i = 0; i < num_direct_mapped_allocations; ++i) {
1514       uint32_t size = direct_map_lengths[i];
1515 
1516       PartitionBucketMemoryStats mapped_stats = {};
1517       mapped_stats.is_valid = true;
1518       mapped_stats.is_direct_map = true;
1519       mapped_stats.num_full_slot_spans = 1;
1520       mapped_stats.allocated_slot_span_size = size;
1521       mapped_stats.bucket_slot_size = size;
1522       mapped_stats.active_bytes = size;
1523       mapped_stats.active_count = 1;
1524       mapped_stats.resident_bytes = size;
1525       dumper->PartitionsDumpBucketStats(partition_name, &mapped_stats);
1526     }
1527   }
1528   dumper->PartitionDumpTotals(partition_name, &stats);
1529 }
1530 
1531 // static
1532 template <bool thread_safe>
DeleteForTesting(PartitionRoot * partition_root)1533 void PartitionRoot<thread_safe>::DeleteForTesting(
1534     PartitionRoot* partition_root) {
1535   if (partition_root->flags.with_thread_cache) {
1536     ThreadCache::SwapForTesting(nullptr);
1537     partition_root->flags.with_thread_cache = false;
1538   }
1539 
1540   partition_root->DestructForTesting();  // IN-TEST
1541 
1542   delete partition_root;
1543 }
1544 
1545 template <bool thread_safe>
ResetForTesting(bool allow_leaks)1546 void PartitionRoot<thread_safe>::ResetForTesting(bool allow_leaks) {
1547   if (flags.with_thread_cache) {
1548     ThreadCache::SwapForTesting(nullptr);
1549     flags.with_thread_cache = false;
1550   }
1551 
1552   ::partition_alloc::internal::ScopedGuard guard(lock_);
1553 
1554 #if BUILDFLAG(PA_DCHECK_IS_ON)
1555   if (!allow_leaks) {
1556     unsigned num_allocated_slots = 0;
1557     for (Bucket& bucket : buckets) {
1558       if (bucket.active_slot_spans_head !=
1559           internal::SlotSpanMetadata<thread_safe>::get_sentinel_slot_span()) {
1560         for (internal::SlotSpanMetadata<thread_safe>* slot_span =
1561                  bucket.active_slot_spans_head;
1562              slot_span; slot_span = slot_span->next_slot_span) {
1563           num_allocated_slots += slot_span->num_allocated_slots;
1564         }
1565       }
1566       // Full slot spans are nowhere. Need to see bucket.num_full_slot_spans
1567       // to count the number of full slot spans' slots.
1568       if (bucket.num_full_slot_spans) {
1569         num_allocated_slots +=
1570             bucket.num_full_slot_spans * bucket.get_slots_per_span();
1571       }
1572     }
1573     PA_DCHECK(num_allocated_slots == 0);
1574 
1575     // Check for direct-mapped allocations.
1576     PA_DCHECK(!direct_map_list);
1577   }
1578 #endif
1579 
1580   DestructForTesting();  // IN-TEST
1581 
1582 #if PA_CONFIG(USE_PARTITION_ROOT_ENUMERATOR)
1583   if (initialized) {
1584     internal::PartitionRootEnumerator::Instance().Unregister(this);
1585   }
1586 #endif  // PA_CONFIG(USE_PARTITION_ROOT_ENUMERATOR)
1587 
1588   for (Bucket& bucket : buckets) {
1589     bucket.active_slot_spans_head =
1590         SlotSpan::get_sentinel_slot_span_non_const();
1591     bucket.empty_slot_spans_head = nullptr;
1592     bucket.decommitted_slot_spans_head = nullptr;
1593     bucket.num_full_slot_spans = 0;
1594   }
1595 
1596   next_super_page = 0;
1597   next_partition_page = 0;
1598   next_partition_page_end = 0;
1599   current_extent = nullptr;
1600   first_extent = nullptr;
1601 
1602   direct_map_list = nullptr;
1603   for (auto& entity : global_empty_slot_span_ring) {
1604     entity = nullptr;
1605   }
1606 
1607   global_empty_slot_span_ring_index = 0;
1608   global_empty_slot_span_ring_size = internal::kDefaultEmptySlotSpanRingSize;
1609   initialized = false;
1610 }
1611 
1612 template <bool thread_safe>
ResetBookkeepingForTesting()1613 void PartitionRoot<thread_safe>::ResetBookkeepingForTesting() {
1614   ::partition_alloc::internal::ScopedGuard guard{lock_};
1615   max_size_of_allocated_bytes = total_size_of_allocated_bytes;
1616   max_size_of_committed_pages.store(total_size_of_committed_pages);
1617 }
1618 
1619 template <>
MaybeInitThreadCache()1620 ThreadCache* PartitionRoot<internal::ThreadSafe>::MaybeInitThreadCache() {
1621   auto* tcache = ThreadCache::Get();
1622   // See comment in `EnableThreadCacheIfSupport()` for why this is an acquire
1623   // load.
1624   if (ThreadCache::IsTombstone(tcache) ||
1625       thread_caches_being_constructed_.load(std::memory_order_acquire)) {
1626     // Two cases:
1627     // 1. Thread is being terminated, don't try to use the thread cache, and
1628     //    don't try to resurrect it.
1629     // 2. Someone, somewhere is currently allocating a thread cache. This may
1630     //    be us, in which case we are re-entering and should not create a thread
1631     //    cache. If it is not us, then this merely delays thread cache
1632     //    construction a bit, which is not an issue.
1633     return nullptr;
1634   }
1635 
1636   // There is no per-thread ThreadCache allocated here yet, and this partition
1637   // has a thread cache, allocate a new one.
1638   //
1639   // The thread cache allocation itself will not reenter here, as it sidesteps
1640   // the thread cache by using placement new and |RawAlloc()|. However,
1641   // internally to libc, allocations may happen to create a new TLS
1642   // variable. This would end up here again, which is not what we want (and
1643   // likely is not supported by libc).
1644   //
1645   // To avoid this sort of reentrancy, increase the count of thread caches that
1646   // are currently allocating a thread cache.
1647   //
1648   // Note that there is no deadlock or data inconsistency concern, since we do
1649   // not hold the lock, and has such haven't touched any internal data.
1650   int before =
1651       thread_caches_being_constructed_.fetch_add(1, std::memory_order_relaxed);
1652   PA_CHECK(before < std::numeric_limits<int>::max());
1653   tcache = ThreadCache::Create(this);
1654   thread_caches_being_constructed_.fetch_sub(1, std::memory_order_relaxed);
1655 
1656   return tcache;
1657 }
1658 
1659 template <>
EnableSortActiveSlotSpans()1660 void PartitionRoot<internal::ThreadSafe>::EnableSortActiveSlotSpans() {
1661   sort_active_slot_spans_ = true;
1662 }
1663 
1664 template struct PA_COMPONENT_EXPORT(PARTITION_ALLOC)
1665     PartitionRoot<internal::ThreadSafe>;
1666 
1667 static_assert(offsetof(PartitionRoot<internal::ThreadSafe>, sentinel_bucket) ==
1668                   offsetof(PartitionRoot<internal::ThreadSafe>, buckets) +
1669                       internal::kNumBuckets *
1670                           sizeof(PartitionRoot<internal::ThreadSafe>::Bucket),
1671               "sentinel_bucket must be just after the regular buckets.");
1672 
1673 static_assert(
1674     offsetof(PartitionRoot<internal::ThreadSafe>, lock_) >= 64,
1675     "The lock should not be on the same cacheline as the read-mostly flags");
1676 
1677 }  // namespace partition_alloc
1678