xref: /aosp_15_r20/external/cronet/base/allocator/partition_allocator/src/partition_alloc/extended_api.cc (revision 6777b5387eb2ff775bb5750e3f5d96f37fb7352b) 
1 // Copyright 2021 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 "partition_alloc/extended_api.h"
6 
7 #include "partition_alloc/partition_alloc_buildflags.h"
8 #include "partition_alloc/partition_alloc_config.h"
9 #include "partition_alloc/shim/allocator_shim_default_dispatch_to_partition_alloc.h"
10 #include "partition_alloc/thread_cache.h"
11 
12 namespace partition_alloc::internal {
13 
14 #if PA_CONFIG(THREAD_CACHE_SUPPORTED)
15 
16 namespace {
17 
DisableThreadCacheForRootIfEnabled(PartitionRoot * root)18 void DisableThreadCacheForRootIfEnabled(PartitionRoot* root) {
19   // Some platforms don't have a thread cache, or it could already have been
20   // disabled.
21   if (!root || !root->settings.with_thread_cache) {
22     return;
23   }
24 
25   ThreadCacheRegistry::Instance().PurgeAll();
26   root->settings.with_thread_cache = false;
27   // Doesn't destroy the thread cache object(s). For background threads, they
28   // will be collected (and free cached memory) at thread destruction
29   // time. For the main thread, we leak it.
30 }
31 
EnablePartitionAllocThreadCacheForRootIfDisabled(PartitionRoot * root)32 void EnablePartitionAllocThreadCacheForRootIfDisabled(PartitionRoot* root) {
33   if (!root) {
34     return;
35   }
36   root->settings.with_thread_cache = true;
37 }
38 
39 #if BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
DisablePartitionAllocThreadCacheForProcess()40 void DisablePartitionAllocThreadCacheForProcess() {
41   PA_CHECK(allocator_shim::internal::PartitionAllocMalloc::
42                AllocatorConfigurationFinalized());
43   DisableThreadCacheForRootIfEnabled(
44       allocator_shim::internal::PartitionAllocMalloc::Allocator());
45   DisableThreadCacheForRootIfEnabled(
46       allocator_shim::internal::PartitionAllocMalloc::OriginalAllocator());
47 }
48 #endif  // BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
49 
50 }  // namespace
51 
52 #endif  // PA_CONFIG(THREAD_CACHE_SUPPORTED)
53 
GetAllocStatsForCurrentThread()54 ThreadAllocStats GetAllocStatsForCurrentThread() {
55   ThreadCache* thread_cache = ThreadCache::Get();
56   if (ThreadCache::IsValid(thread_cache)) {
57     return thread_cache->thread_alloc_stats();
58   }
59   return {};
60 }
61 
62 #if PA_CONFIG(THREAD_CACHE_SUPPORTED)
ThreadCacheProcessScopeForTesting(PartitionRoot * root)63 ThreadCacheProcessScopeForTesting::ThreadCacheProcessScopeForTesting(
64     PartitionRoot* root)
65     : root_(root) {
66 #if BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
67   auto* regular_allocator =
68       allocator_shim::internal::PartitionAllocMalloc::Allocator();
69   regular_was_enabled_ =
70       regular_allocator && regular_allocator->settings.with_thread_cache;
71 
72   if (root_ != regular_allocator) {
73     // Another |root| is ThreadCache's PartitionRoot. Need to disable
74     // thread cache for the process.
75     DisablePartitionAllocThreadCacheForProcess();
76     EnablePartitionAllocThreadCacheForRootIfDisabled(root_);
77     // Replace ThreadCache's PartitionRoot.
78     ThreadCache::SwapForTesting(root_);
79   } else {
80     if (!regular_was_enabled_) {
81       EnablePartitionAllocThreadCacheForRootIfDisabled(root_);
82       ThreadCache::SwapForTesting(root_);
83     }
84   }
85 #else
86   PA_CHECK(!ThreadCache::IsValid(ThreadCache::Get()));
87   EnablePartitionAllocThreadCacheForRootIfDisabled(root_);
88   ThreadCache::SwapForTesting(root_);
89 #endif  // BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
90 
91   PA_CHECK(ThreadCache::Get());
92 }
93 
~ThreadCacheProcessScopeForTesting()94 ThreadCacheProcessScopeForTesting::~ThreadCacheProcessScopeForTesting() {
95 #if BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
96   auto* regular_allocator =
97       allocator_shim::internal::PartitionAllocMalloc::Allocator();
98   bool regular_enabled =
99       regular_allocator && regular_allocator->settings.with_thread_cache;
100 
101   if (regular_was_enabled_) {
102     if (!regular_enabled) {
103       // Need to re-enable ThreadCache for the process.
104       EnablePartitionAllocThreadCacheForRootIfDisabled(regular_allocator);
105       // In the case, |regular_allocator| must be ThreadCache's root.
106       ThreadCache::SwapForTesting(regular_allocator);
107     } else {
108       // ThreadCache is enabled for the process, but we need to be
109       // careful about ThreadCache's PartitionRoot. If it is different from
110       // |regular_allocator|, we need to invoke SwapForTesting().
111       if (regular_allocator != root_) {
112         ThreadCache::SwapForTesting(regular_allocator);
113       }
114     }
115   } else {
116     // ThreadCache for all processes was disabled.
117     DisableThreadCacheForRootIfEnabled(regular_allocator);
118     ThreadCache::SwapForTesting(nullptr);
119   }
120 #else
121   // First, disable the test thread cache we have.
122   DisableThreadCacheForRootIfEnabled(root_);
123 
124   ThreadCache::SwapForTesting(nullptr);
125 #endif  // BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
126 }
127 #endif  // PA_CONFIG(THREAD_CACHE_SUPPORTED)
128 
129 }  // namespace partition_alloc::internal
130