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