1//===-- tsd_shared.h --------------------------------------------*- C++ -*-===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8
9#ifndef SCUDO_TSD_SHARED_H_
10#define SCUDO_TSD_SHARED_H_
11
12#include "tsd.h"
13
14#include "string_utils.h"
15
16#if SCUDO_HAS_PLATFORM_TLS_SLOT
17// This is a platform-provided header that needs to be on the include path when
18// Scudo is compiled. It must declare a function with the prototype:
19// uintptr_t *getPlatformAllocatorTlsSlot()
20// that returns the address of a thread-local word of storage reserved for
21// Scudo, that must be zero-initialized in newly created threads.
22#include "scudo_platform_tls_slot.h"
23#endif
24
25namespace scudo {
26
27template <class Allocator, u32 TSDsArraySize, u32 DefaultTSDCount>
28struct TSDRegistrySharedT {
29 using ThisT = TSDRegistrySharedT<Allocator, TSDsArraySize, DefaultTSDCount>;
30
31 struct ScopedTSD {
32 ALWAYS_INLINE ScopedTSD(ThisT &TSDRegistry) {
33 CurrentTSD = TSDRegistry.getTSDAndLock();
34 DCHECK_NE(CurrentTSD, nullptr);
35 }
36
37 ~ScopedTSD() { CurrentTSD->unlock(); }
38
39 TSD<Allocator> &operator*() { return *CurrentTSD; }
40
41 TSD<Allocator> *operator->() {
42 CurrentTSD->assertLocked(/*BypassCheck=*/false);
43 return CurrentTSD;
44 }
45
46 private:
47 TSD<Allocator> *CurrentTSD;
48 };
49
50 void init(Allocator *Instance) EXCLUDES(Mutex) {
51 ScopedLock L(Mutex);
52 // If more than one thread is initializing at the exact same moment, the
53 // threads that lose don't need to do anything.
54 if (UNLIKELY(atomic_load_relaxed(&Initialized) != 0))
55 return;
56
57 Instance->init();
58 for (u32 I = 0; I < TSDsArraySize; I++)
59 TSDs[I].init(Instance);
60 const u32 NumberOfCPUs = getNumberOfCPUs();
61 setNumberOfTSDs((NumberOfCPUs == 0) ? DefaultTSDCount
62 : Min(A: NumberOfCPUs, B: DefaultTSDCount));
63 atomic_store_relaxed(A: &Initialized, V: 1);
64 }
65
66 void initOnceMaybe(Allocator *Instance) {
67 if (LIKELY(atomic_load_relaxed(&Initialized) != 0))
68 return;
69 init(Instance); // Sets Initialized.
70 }
71
72 void unmapTestOnly(Allocator *Instance) EXCLUDES(Mutex) {
73 for (u32 I = 0; I < TSDsArraySize; I++) {
74 TSDs[I].commitBack(Instance);
75 TSDs[I] = {};
76 }
77 setCurrentTSD(nullptr);
78 ScopedLock L(Mutex);
79 atomic_store_relaxed(A: &Initialized, V: 0);
80 }
81
82 void drainCaches(Allocator *Instance) {
83 ScopedLock L(Mutex);
84 for (uptr I = 0; I < NumberOfTSDs; ++I) {
85 TSDs[I].lock();
86 Instance->drainCache(&TSDs[I]);
87 TSDs[I].unlock();
88 }
89 }
90
91 ALWAYS_INLINE void initThreadMaybe(Allocator *Instance,
92 UNUSED bool MinimalInit) {
93 if (LIKELY(getCurrentTSD()))
94 return;
95 initThread(Instance);
96 }
97
98 void disable() NO_THREAD_SAFETY_ANALYSIS {
99 Mutex.lock();
100 for (u32 I = 0; I < TSDsArraySize; I++)
101 TSDs[I].lock();
102 }
103
104 void enable() NO_THREAD_SAFETY_ANALYSIS {
105 for (s32 I = static_cast<s32>(TSDsArraySize - 1); I >= 0; I--)
106 TSDs[I].unlock();
107 Mutex.unlock();
108 }
109
110 bool setOption(Option O, sptr Value) {
111 if (O == Option::MaxTSDsCount) {
112 ScopedLock L(Mutex);
113 return setNumberOfTSDs(static_cast<u32>(Value));
114 }
115 if (O == Option::ThreadDisableMemInit)
116 setDisableMemInit(Value);
117 // Not supported by the TSD Registry, but not an error either.
118 return true;
119 }
120
121 bool getDisableMemInit() const { return *getTlsPtr() & 1; }
122
123 void getStats(ScopedString *Str) EXCLUDES(Mutex) {
124 ScopedLock L(Mutex);
125
126 Str->append(Format: "Stats: SharedTSDs: %u available; total %u\n", NumberOfTSDs,
127 TSDsArraySize);
128 for (uptr I = 0; I < NumberOfTSDs; ++I) {
129 TSDs[I].lock();
130 // Theoretically, we want to mark TSD::lock()/TSD::unlock() with proper
131 // thread annotations. However, given the TSD is only locked on shared
132 // path, do the assertion in a separate path to avoid confusing the
133 // analyzer.
134 TSDs[I].assertLocked(/*BypassCheck=*/true);
135 Str->append(Format: " Shared TSD[%zu]:\n", I);
136 TSDs[I].getSizeClassAllocator().getStats(Str);
137 TSDs[I].unlock();
138 }
139 }
140
141private:
142 ALWAYS_INLINE TSD<Allocator> *getTSDAndLock() NO_THREAD_SAFETY_ANALYSIS {
143 TSD<Allocator> *TSD = getCurrentTSD();
144 DCHECK(TSD);
145 // Try to lock the currently associated context.
146 if (TSD->tryLock())
147 return TSD;
148 // If that fails, go down the slow path.
149 if (TSDsArraySize == 1U) {
150 // Only 1 TSD, not need to go any further.
151 // The compiler will optimize this one way or the other.
152 TSD->lock();
153 return TSD;
154 }
155 return getTSDAndLockSlow(CurrentTSD: TSD);
156 }
157
158 ALWAYS_INLINE uptr *getTlsPtr() const {
159#if SCUDO_HAS_PLATFORM_TLS_SLOT
160 return reinterpret_cast<uptr *>(getPlatformAllocatorTlsSlot());
161#else
162 static thread_local uptr ThreadTSD;
163 return &ThreadTSD;
164#endif
165 }
166
167 static_assert(alignof(TSD<Allocator>) >= 2, "");
168
169 ALWAYS_INLINE void setCurrentTSD(TSD<Allocator> *CurrentTSD) {
170 *getTlsPtr() &= 1;
171 *getTlsPtr() |= reinterpret_cast<uptr>(CurrentTSD);
172 }
173
174 ALWAYS_INLINE TSD<Allocator> *getCurrentTSD() {
175 return reinterpret_cast<TSD<Allocator> *>(*getTlsPtr() & ~1ULL);
176 }
177
178 bool setNumberOfTSDs(u32 N) REQUIRES(Mutex) {
179 if (N < NumberOfTSDs)
180 return false;
181 if (N > TSDsArraySize)
182 N = TSDsArraySize;
183 NumberOfTSDs = N;
184 NumberOfCoPrimes = 0;
185 // Compute all the coprimes of NumberOfTSDs. This will be used to walk the
186 // array of TSDs in a random order. For details, see:
187 // https://lemire.me/blog/2017/09/18/visiting-all-values-in-an-array-exactly-once-in-random-order/
188 for (u32 I = 0; I < N; I++) {
189 u32 A = I + 1;
190 u32 B = N;
191 // Find the GCD between I + 1 and N. If 1, they are coprimes.
192 while (B != 0) {
193 const u32 T = A;
194 A = B;
195 B = T % B;
196 }
197 if (A == 1)
198 CoPrimes[NumberOfCoPrimes++] = I + 1;
199 }
200 return true;
201 }
202
203 void setDisableMemInit(bool B) {
204 *getTlsPtr() &= ~1ULL;
205 *getTlsPtr() |= B;
206 }
207
208 NOINLINE void initThread(Allocator *Instance) NO_THREAD_SAFETY_ANALYSIS {
209 initOnceMaybe(Instance);
210 // Initial context assignment is done in a plain round-robin fashion.
211 const u32 Index = atomic_fetch_add(A: &CurrentIndex, V: 1U, MO: memory_order_relaxed);
212 setCurrentTSD(&TSDs[Index % NumberOfTSDs]);
213 Instance->callPostInitCallback();
214 }
215
216 // TSDs is an array of locks which is not supported for marking thread-safety
217 // capability.
218 NOINLINE TSD<Allocator> *getTSDAndLockSlow(TSD<Allocator> *CurrentTSD)
219 EXCLUDES(Mutex) {
220 // Use the Precedence of the current TSD as our random seed. Since we are
221 // in the slow path, it means that tryLock failed, and as a result it's
222 // very likely that said Precedence is non-zero.
223 const u32 R = static_cast<u32>(CurrentTSD->getPrecedence());
224 u32 N, Inc;
225 {
226 ScopedLock L(Mutex);
227 N = NumberOfTSDs;
228 DCHECK_NE(NumberOfCoPrimes, 0U);
229 Inc = CoPrimes[R % NumberOfCoPrimes];
230 }
231 if (N > 1U) {
232 u32 Index = R % N;
233 uptr LowestPrecedence = UINTPTR_MAX;
234 TSD<Allocator> *CandidateTSD = nullptr;
235 // Go randomly through at most 4 contexts and find a candidate.
236 for (u32 I = 0; I < Min(A: 4U, B: N); I++) {
237 if (TSDs[Index].tryLock()) {
238 setCurrentTSD(&TSDs[Index]);
239 return &TSDs[Index];
240 }
241 const uptr Precedence = TSDs[Index].getPrecedence();
242 // A 0 precedence here means another thread just locked this TSD.
243 if (Precedence && Precedence < LowestPrecedence) {
244 CandidateTSD = &TSDs[Index];
245 LowestPrecedence = Precedence;
246 }
247 Index += Inc;
248 if (Index >= N)
249 Index -= N;
250 }
251 if (CandidateTSD) {
252 CandidateTSD->lock();
253 setCurrentTSD(CandidateTSD);
254 return CandidateTSD;
255 }
256 }
257 // Last resort, stick with the current one.
258 CurrentTSD->lock();
259 return CurrentTSD;
260 }
261
262 atomic_u32 CurrentIndex = {};
263 u32 NumberOfTSDs GUARDED_BY(Mutex) = 0;
264 u32 NumberOfCoPrimes GUARDED_BY(Mutex) = 0;
265 u32 CoPrimes[TSDsArraySize] GUARDED_BY(Mutex) = {};
266 atomic_u8 Initialized = {};
267 // Used for global initialization and TSDs access.
268 // Acquiring the global initialization should only lock once in normal
269 // operation, which is why using it for TSDs access should not cause
270 // any interference.
271 HybridMutex Mutex;
272 TSD<Allocator> TSDs[TSDsArraySize];
273};
274
275} // namespace scudo
276
277#endif // SCUDO_TSD_SHARED_H_
278