tsd_shared.h revision 360784 
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 "linux.h" // for getAndroidTlsPtr()
13#include "tsd.h"
14
15namespace scudo {
16
17template <class Allocator, u32 MaxTSDCount> struct TSDRegistrySharedT {
18  void initLinkerInitialized(Allocator *Instance) {
19    Instance->initLinkerInitialized();
20    CHECK_EQ(pthread_key_create(&PThreadKey, nullptr), 0); // For non-TLS
21    NumberOfTSDs = Min(Max(1U, getNumberOfCPUs()), MaxTSDCount);
22    TSDs = reinterpret_cast<TSD<Allocator> *>(
23        map(nullptr, sizeof(TSD<Allocator>) * NumberOfTSDs, "scudo:tsd"));
24    for (u32 I = 0; I < NumberOfTSDs; I++)
25      TSDs[I].initLinkerInitialized(Instance);
26    // Compute all the coprimes of NumberOfTSDs. This will be used to walk the
27    // array of TSDs in a random order. For details, see:
28    // https://lemire.me/blog/2017/09/18/visiting-all-values-in-an-array-exactly-once-in-random-order/
29    for (u32 I = 0; I < NumberOfTSDs; I++) {
30      u32 A = I + 1;
31      u32 B = NumberOfTSDs;
32      // Find the GCD between I + 1 and NumberOfTSDs. If 1, they are coprimes.
33      while (B != 0) {
34        const u32 T = A;
35        A = B;
36        B = T % B;
37      }
38      if (A == 1)
39        CoPrimes[NumberOfCoPrimes++] = I + 1;
40    }
41    Initialized = true;
42  }
43  void init(Allocator *Instance) {
44    memset(this, 0, sizeof(*this));
45    initLinkerInitialized(Instance);
46  }
47
48  void unmapTestOnly() {
49    unmap(reinterpret_cast<void *>(TSDs),
50          sizeof(TSD<Allocator>) * NumberOfTSDs);
51    setCurrentTSD(nullptr);
52    pthread_key_delete(PThreadKey);
53  }
54
55  ALWAYS_INLINE void initThreadMaybe(Allocator *Instance,
56                                     UNUSED bool MinimalInit) {
57    if (LIKELY(getCurrentTSD()))
58      return;
59    initThread(Instance);
60  }
61
62  ALWAYS_INLINE TSD<Allocator> *getTSDAndLock(bool *UnlockRequired) {
63    TSD<Allocator> *TSD = getCurrentTSD();
64    DCHECK(TSD);
65    *UnlockRequired = true;
66    // Try to lock the currently associated context.
67    if (TSD->tryLock())
68      return TSD;
69    // If that fails, go down the slow path.
70    return getTSDAndLockSlow(TSD);
71  }
72
73  void disable() {
74    Mutex.lock();
75    for (u32 I = 0; I < NumberOfTSDs; I++)
76      TSDs[I].lock();
77  }
78
79  void enable() {
80    for (s32 I = NumberOfTSDs - 1; I >= 0; I--)
81      TSDs[I].unlock();
82    Mutex.unlock();
83  }
84
85private:
86  ALWAYS_INLINE void setCurrentTSD(TSD<Allocator> *CurrentTSD) {
87#if _BIONIC
88    *getAndroidTlsPtr() = reinterpret_cast<uptr>(CurrentTSD);
89#elif SCUDO_LINUX
90    ThreadTSD = CurrentTSD;
91#else
92    CHECK_EQ(
93        pthread_setspecific(PThreadKey, reinterpret_cast<void *>(CurrentTSD)),
94        0);
95#endif
96  }
97
98  ALWAYS_INLINE TSD<Allocator> *getCurrentTSD() {
99#if _BIONIC
100    return reinterpret_cast<TSD<Allocator> *>(*getAndroidTlsPtr());
101#elif SCUDO_LINUX
102    return ThreadTSD;
103#else
104    return reinterpret_cast<TSD<Allocator> *>(pthread_getspecific(PThreadKey));
105#endif
106  }
107
108  void initOnceMaybe(Allocator *Instance) {
109    ScopedLock L(Mutex);
110    if (LIKELY(Initialized))
111      return;
112    initLinkerInitialized(Instance); // Sets Initialized.
113  }
114
115  NOINLINE void initThread(Allocator *Instance) {
116    initOnceMaybe(Instance);
117    // Initial context assignment is done in a plain round-robin fashion.
118    const u32 Index = atomic_fetch_add(&CurrentIndex, 1U, memory_order_relaxed);
119    setCurrentTSD(&TSDs[Index % NumberOfTSDs]);
120    Instance->callPostInitCallback();
121  }
122
123  NOINLINE TSD<Allocator> *getTSDAndLockSlow(TSD<Allocator> *CurrentTSD) {
124    if (MaxTSDCount > 1U && NumberOfTSDs > 1U) {
125      // Use the Precedence of the current TSD as our random seed. Since we are
126      // in the slow path, it means that tryLock failed, and as a result it's
127      // very likely that said Precedence is non-zero.
128      const u32 R = static_cast<u32>(CurrentTSD->getPrecedence());
129      const u32 Inc = CoPrimes[R % NumberOfCoPrimes];
130      u32 Index = R % NumberOfTSDs;
131      uptr LowestPrecedence = UINTPTR_MAX;
132      TSD<Allocator> *CandidateTSD = nullptr;
133      // Go randomly through at most 4 contexts and find a candidate.
134      for (u32 I = 0; I < Min(4U, NumberOfTSDs); I++) {
135        if (TSDs[Index].tryLock()) {
136          setCurrentTSD(&TSDs[Index]);
137          return &TSDs[Index];
138        }
139        const uptr Precedence = TSDs[Index].getPrecedence();
140        // A 0 precedence here means another thread just locked this TSD.
141        if (Precedence && Precedence < LowestPrecedence) {
142          CandidateTSD = &TSDs[Index];
143          LowestPrecedence = Precedence;
144        }
145        Index += Inc;
146        if (Index >= NumberOfTSDs)
147          Index -= NumberOfTSDs;
148      }
149      if (CandidateTSD) {
150        CandidateTSD->lock();
151        setCurrentTSD(CandidateTSD);
152        return CandidateTSD;
153      }
154    }
155    // Last resort, stick with the current one.
156    CurrentTSD->lock();
157    return CurrentTSD;
158  }
159
160  pthread_key_t PThreadKey;
161  atomic_u32 CurrentIndex;
162  u32 NumberOfTSDs;
163  TSD<Allocator> *TSDs;
164  u32 NumberOfCoPrimes;
165  u32 CoPrimes[MaxTSDCount];
166  bool Initialized;
167  HybridMutex Mutex;
168#if SCUDO_LINUX && !_BIONIC
169  static THREADLOCAL TSD<Allocator> *ThreadTSD;
170#endif
171};
172
173#if SCUDO_LINUX && !_BIONIC
174template <class Allocator, u32 MaxTSDCount>
175THREADLOCAL TSD<Allocator>
176    *TSDRegistrySharedT<Allocator, MaxTSDCount>::ThreadTSD;
177#endif
178
179} // namespace scudo
180
181#endif // SCUDO_TSD_SHARED_H_
182