1 // Copyright 2023 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 #ifdef PARTITION_ALLOC_SHIM_ALLOCATOR_SHIM_FUNCTIONS_H_
6 #error This header is meant to be included only once by allocator_shim*.cc except allocator_shim_win_component.cc
7 #endif
8
9 #ifndef PARTITION_ALLOC_SHIM_ALLOCATOR_SHIM_FUNCTIONS_H_
10 #define PARTITION_ALLOC_SHIM_ALLOCATOR_SHIM_FUNCTIONS_H_
11
12 #include <atomic>
13 #include <cstddef>
14 #include <new>
15
16 #include "build/build_config.h"
17 #include "partition_alloc/partition_alloc_base/compiler_specific.h"
18 #include "partition_alloc/partition_alloc_check.h"
19 #include "partition_alloc/shim/allocator_dispatch.h"
20 #include "partition_alloc/shim/allocator_shim_internals.h"
21
22 #if BUILDFLAG(IS_WIN)
23 #include "partition_alloc/shim/winheap_stubs_win.h"
24 #endif
25
26 namespace allocator_shim {
27 namespace internal {
28
29 std::atomic<const allocator_shim::AllocatorDispatch*> g_chain_head{
30 &allocator_shim::AllocatorDispatch::default_dispatch};
31
32 bool g_call_new_handler_on_malloc_failure = false;
33
34 // Calls the std::new handler thread-safely. Returns true if a new_handler was
35 // set and called, false if no new_handler was set.
CallNewHandler(size_t size)36 bool CallNewHandler(size_t size) {
37 #if BUILDFLAG(IS_WIN)
38 return allocator_shim::WinCallNewHandler(size);
39 #else
40 std::new_handler nh = std::get_new_handler();
41 if (!nh) {
42 return false;
43 }
44 (*nh)();
45 // Assume the new_handler will abort if it fails. Exception are disabled and
46 // we don't support the case of a new_handler throwing std::bad_balloc.
47 return true;
48 #endif
49 }
50
51 #if !(BUILDFLAG(IS_WIN) && defined(COMPONENT_BUILD))
52 PA_ALWAYS_INLINE
53 #endif
GetChainHead()54 const allocator_shim::AllocatorDispatch* GetChainHead() {
55 return internal::g_chain_head.load(std::memory_order_relaxed);
56 }
57
58 } // namespace internal
59
SetCallNewHandlerOnMallocFailure(bool value)60 void SetCallNewHandlerOnMallocFailure(bool value) {
61 internal::g_call_new_handler_on_malloc_failure = value;
62 }
63
UncheckedAlloc(size_t size)64 void* UncheckedAlloc(size_t size) {
65 const AllocatorDispatch* const chain_head = internal::GetChainHead();
66 return chain_head->alloc_unchecked_function(chain_head, size, nullptr);
67 }
68
UncheckedFree(void * ptr)69 void UncheckedFree(void* ptr) {
70 const AllocatorDispatch* const chain_head = internal::GetChainHead();
71 return chain_head->free_function(chain_head, ptr, nullptr);
72 }
73
InsertAllocatorDispatch(AllocatorDispatch * dispatch)74 void InsertAllocatorDispatch(AllocatorDispatch* dispatch) {
75 // Loop in case of (an unlikely) race on setting the list head.
76 size_t kMaxRetries = 7;
77 for (size_t i = 0; i < kMaxRetries; ++i) {
78 const AllocatorDispatch* chain_head = internal::GetChainHead();
79 dispatch->next = chain_head;
80
81 // This function guarantees to be thread-safe w.r.t. concurrent
82 // insertions. It also has to guarantee that all the threads always
83 // see a consistent chain, hence the atomic_thread_fence() below.
84 // InsertAllocatorDispatch() is NOT a fastpath, as opposite to malloc(), so
85 // we don't really want this to be a release-store with a corresponding
86 // acquire-load during malloc().
87 std::atomic_thread_fence(std::memory_order_seq_cst);
88 // Set the chain head to the new dispatch atomically. If we lose the race,
89 // retry.
90 if (internal::g_chain_head.compare_exchange_strong(
91 chain_head, dispatch, std::memory_order_relaxed,
92 std::memory_order_relaxed)) {
93 // Success.
94 return;
95 }
96 }
97
98 PA_CHECK(false); // Too many retries, this shouldn't happen.
99 }
100
RemoveAllocatorDispatchForTesting(AllocatorDispatch * dispatch)101 void RemoveAllocatorDispatchForTesting(AllocatorDispatch* dispatch) {
102 PA_DCHECK(internal::GetChainHead() == dispatch);
103 internal::g_chain_head.store(dispatch->next, std::memory_order_relaxed);
104 }
105
106 } // namespace allocator_shim
107
108 #endif // PARTITION_ALLOC_SHIM_ALLOCATOR_SHIM_FUNCTIONS_H_
109