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)18void 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)32void 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()40void DisablePartitionAllocThreadCacheForProcess() { 41 PA_CHECK(allocator_shim::internal::PartitionAllocMalloc:: 42 AllocatorConfigurationFinalized()); 43 auto* regular_allocator = 44 allocator_shim::internal::PartitionAllocMalloc::Allocator(); 45 auto* aligned_allocator = 46 allocator_shim::internal::PartitionAllocMalloc::AlignedAllocator(); 47 DisableThreadCacheForRootIfEnabled(regular_allocator); 48 if (aligned_allocator != regular_allocator) { 49 DisableThreadCacheForRootIfEnabled(aligned_allocator); 50 } 51 DisableThreadCacheForRootIfEnabled( 52 allocator_shim::internal::PartitionAllocMalloc::OriginalAllocator()); 53 } 54 #endif // BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC) 55 56 } // namespace 57 58 #endif // PA_CONFIG(THREAD_CACHE_SUPPORTED) 59 GetAllocStatsForCurrentThread()60ThreadAllocStats GetAllocStatsForCurrentThread() { 61 ThreadCache* thread_cache = ThreadCache::Get(); 62 if (ThreadCache::IsValid(thread_cache)) { 63 return thread_cache->thread_alloc_stats(); 64 } 65 return {}; 66 } 67 68 #if PA_CONFIG(THREAD_CACHE_SUPPORTED) ThreadCacheProcessScopeForTesting(PartitionRoot * root)69ThreadCacheProcessScopeForTesting::ThreadCacheProcessScopeForTesting( 70 PartitionRoot* root) 71 : root_(root) { 72 #if BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC) 73 auto* regular_allocator = 74 allocator_shim::internal::PartitionAllocMalloc::Allocator(); 75 regular_was_enabled_ = 76 regular_allocator && regular_allocator->settings.with_thread_cache; 77 78 if (root_ != regular_allocator) { 79 // Another |root| is ThreadCache's PartitionRoot. Need to disable 80 // thread cache for the process. 81 DisablePartitionAllocThreadCacheForProcess(); 82 EnablePartitionAllocThreadCacheForRootIfDisabled(root_); 83 // Replace ThreadCache's PartitionRoot. 84 ThreadCache::SwapForTesting(root_); 85 } else { 86 if (!regular_was_enabled_) { 87 EnablePartitionAllocThreadCacheForRootIfDisabled(root_); 88 ThreadCache::SwapForTesting(root_); 89 } 90 } 91 #else 92 PA_CHECK(!ThreadCache::IsValid(ThreadCache::Get())); 93 EnablePartitionAllocThreadCacheForRootIfDisabled(root_); 94 ThreadCache::SwapForTesting(root_); 95 #endif // BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC) 96 97 PA_CHECK(ThreadCache::Get()); 98 } 99 ~ThreadCacheProcessScopeForTesting()100ThreadCacheProcessScopeForTesting::~ThreadCacheProcessScopeForTesting() { 101 #if BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC) 102 auto* regular_allocator = 103 allocator_shim::internal::PartitionAllocMalloc::Allocator(); 104 bool regular_enabled = 105 regular_allocator && regular_allocator->settings.with_thread_cache; 106 107 if (regular_was_enabled_) { 108 if (!regular_enabled) { 109 // Need to re-enable ThreadCache for the process. 110 EnablePartitionAllocThreadCacheForRootIfDisabled(regular_allocator); 111 // In the case, |regular_allocator| must be ThreadCache's root. 112 ThreadCache::SwapForTesting(regular_allocator); 113 } else { 114 // ThreadCache is enabled for the process, but we need to be 115 // careful about ThreadCache's PartitionRoot. If it is different from 116 // |regular_allocator|, we need to invoke SwapForTesting(). 117 if (regular_allocator != root_) { 118 ThreadCache::SwapForTesting(regular_allocator); 119 } 120 } 121 } else { 122 // ThreadCache for all processes was disabled. 123 DisableThreadCacheForRootIfEnabled(regular_allocator); 124 ThreadCache::SwapForTesting(nullptr); 125 } 126 #else 127 // First, disable the test thread cache we have. 128 DisableThreadCacheForRootIfEnabled(root_); 129 130 ThreadCache::SwapForTesting(nullptr); 131 #endif // BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC) 132 } 133 #endif // PA_CONFIG(THREAD_CACHE_SUPPORTED) 134 135 } // namespace partition_alloc::internal 136