[scudo] Scudo thread specific data refactor, part 1

[cryptoad]

Summary

We are going through an overhaul of Scudo's TSD, to allow for new platforms
to be integrated more easily, and make the code more sound.

This first part is mostly renaming, preferring some shorter names, correcting
some comments. I removed getPrng and getAllocatorCache to directly access
the members, there was not really any benefit to them (and it was suggested by
Dmitry in D37590).

The only functional change is in scudo_tls_android.cpp: we enforce bounds to
the NumberOfTSDs and most of the logic in getTSDAndLockSlow is skipped if we
only have 1 TSD.

lib/scudo/scudo_allocator.cpp

struct QuarantineCallback {
   }
 
   AllocatorCache *Cache_;
 };
 
 typedef Quarantine<QuarantineCallback, ScudoChunk> ScudoQuarantine;
 typedef ScudoQuarantine::Cache ScudoQuarantineCache;
 COMPILER_CHECK(sizeof(ScudoQuarantineCache) <=
-               sizeof(ScudoThreadContext::QuarantineCachePlaceHolder));
+               sizeof(ScudoTSD::QuarantineCachePlaceHolder));
 
-AllocatorCache *getAllocatorCache(ScudoThreadContext *ThreadContext) {
-  return &ThreadContext->Cache;
-}
-
-ScudoQuarantineCache *getQuarantineCache(ScudoThreadContext *ThreadContext) {
-  return reinterpret_cast<
-      ScudoQuarantineCache *>(ThreadContext->QuarantineCachePlaceHolder);
-}
-
-ScudoPrng *getPrng(ScudoThreadContext *ThreadContext) {
-  return &ThreadContext->Prng;
+ScudoQuarantineCache *getQuarantineCache(ScudoTSD *TSD) {
+  return reinterpret_cast<ScudoQuarantineCache *>(
+      TSD->QuarantineCachePlaceHolder);
 }
 
 struct ScudoAllocator {
   static const uptr MaxAllowedMallocSize =
       FIRST_32_SECOND_64(2UL << 30, 1ULL << 40);
 
   typedef ReturnNullOrDieOnFailure FailureHandler;
 
void *allocate(uptr Size, uptr Alignment, AllocType Type,
     // but the Secondary will take care of its own alignment needs.
     bool FromPrimary = PrimaryAllocator::CanAllocate(AlignedSize, MinAlignment);
 
     void *Ptr;
     u8 Salt;
     uptr AllocSize;
     if (FromPrimary) {
       AllocSize = AlignedSize;
-      ScudoThreadContext *ThreadContext = getThreadContextAndLock();
-      if (LIKELY(ThreadContext)) {
-        Salt = getPrng(ThreadContext)->getU8();
-        Ptr = BackendAllocator.allocatePrimary(getAllocatorCache(ThreadContext),
-                                               AllocSize);
-        ThreadContext->unlock();
+      ScudoTSD *TSD = getTSDAndLock();
+      if (LIKELY(TSD)) {
+        Salt = TSD->Prng.getU8();
+        Ptr = BackendAllocator.allocatePrimary(&TSD->Cache, AllocSize);
+        TSD->unlock();
       } else {
         SpinMutexLock l(&FallbackMutex);
         Salt = FallbackPrng.getU8();
         Ptr = BackendAllocator.allocatePrimary(&FallbackAllocatorCache,
                                                AllocSize);
       }
     } else {
       {
struct ScudoAllocator {
   void quarantineOrDeallocateChunk(ScudoChunk *Chunk, UnpackedHeader *Header,
                                    uptr Size) {
     const bool BypassQuarantine = (AllocatorQuarantine.GetCacheSize() == 0) ||
         (Size > QuarantineChunksUpToSize);
     if (BypassQuarantine) {
       Chunk->eraseHeader();
       void *Ptr = Chunk->getAllocBeg(Header);
       if (Header->FromPrimary) {
-        ScudoThreadContext *ThreadContext = getThreadContextAndLock();
-        if (LIKELY(ThreadContext)) {
-          getBackendAllocator().deallocatePrimary(
-              getAllocatorCache(ThreadContext), Ptr);
-          ThreadContext->unlock();
+        ScudoTSD *TSD = getTSDAndLock();
+        if (LIKELY(TSD)) {
+          getBackendAllocator().deallocatePrimary(&TSD->Cache, Ptr);
+          TSD->unlock();
         } else {
           SpinMutexLock Lock(&FallbackMutex);
           getBackendAllocator().deallocatePrimary(&FallbackAllocatorCache, Ptr);
         }
       } else {
         getBackendAllocator().deallocateSecondary(Ptr);
       }
     } else {
       // If a small memory amount was allocated with a larger alignment, we want
       // to take that into account. Otherwise the Quarantine would be filled
       // with tiny chunks, taking a lot of VA memory. This is an approximation
       // of the usable size, that allows us to not call
       // GetActuallyAllocatedSize.
       uptr EstimatedSize = Size + (Header->Offset << MinAlignmentLog);
       UnpackedHeader NewHeader = *Header;
       NewHeader.State = ChunkQuarantine;
       Chunk->compareExchangeHeader(&NewHeader, Header);
-      ScudoThreadContext *ThreadContext = getThreadContextAndLock();
-      if (LIKELY(ThreadContext)) {
-        AllocatorQuarantine.Put(getQuarantineCache(ThreadContext),
-                                QuarantineCallback(
-                                    getAllocatorCache(ThreadContext)),
+      ScudoTSD *TSD = getTSDAndLock();
+      if (LIKELY(TSD)) {
+        AllocatorQuarantine.Put(getQuarantineCache(TSD),
+                                QuarantineCallback(&TSD->Cache),
                                 Chunk, EstimatedSize);
-        ThreadContext->unlock();
+        TSD->unlock();
       } else {
         SpinMutexLock l(&FallbackMutex);
         AllocatorQuarantine.Put(&FallbackQuarantineCache,
                                 QuarantineCallback(&FallbackAllocatorCache),
                                 Chunk, EstimatedSize);
       }
     }
   }
struct ScudoAllocator {
 
   void *calloc(uptr NMemB, uptr Size) {
     initThreadMaybe();
     if (UNLIKELY(CheckForCallocOverflow(NMemB, Size)))
       return FailureHandler::OnBadRequest();
     return allocate(NMemB * Size, MinAlignment, FromMalloc, true);
   }
 
-  void commitBack(ScudoThreadContext *ThreadContext) {
-    AllocatorCache *Cache = getAllocatorCache(ThreadContext);
-    AllocatorQuarantine.Drain(getQuarantineCache(ThreadContext),
-                              QuarantineCallback(Cache));
-    BackendAllocator.destroyCache(Cache);
+  void commitBack(ScudoTSD *TSD) {
+    AllocatorQuarantine.Drain(getQuarantineCache(TSD),
+                              QuarantineCallback(&TSD->Cache));
+    BackendAllocator.destroyCache(&TSD->Cache);
   }
 
   uptr getStats(AllocatorStat StatType) {
     initThreadMaybe();
     uptr stats[AllocatorStatCount];
     BackendAllocator.getStats(stats);
     return stats[StatType];
   }

 static ScudoBackendAllocator &getBackendAllocator() {
   return Instance.BackendAllocator;
 }
 
 static void initScudoInternal(const AllocatorOptions &Options) {
   Instance.init(Options);
 }
 
-void ScudoThreadContext::init() {
+void ScudoTSD::init() {
   getBackendAllocator().initCache(&Cache);
   Prng.init();
   memset(QuarantineCachePlaceHolder, 0, sizeof(QuarantineCachePlaceHolder));
 }
 
-void ScudoThreadContext::commitBack() {
+void ScudoTSD::commitBack() {
   Instance.commitBack(this);
 }
 
 void *scudoMalloc(uptr Size, AllocType Type) {
   return SetErrnoOnNull(Instance.allocate(Size, MinAlignment, Type));
 }
 
 void scudoFree(void *Ptr, AllocType Type) {

lib/scudo/scudo_tls.h

 #include "sanitizer_common/sanitizer_platform.h"
 
 namespace __scudo {
 
 // Platform specific base thread context definitions.
 #include "scudo_tls_context_android.inc"
 #include "scudo_tls_context_linux.inc"
 
-struct ALIGNED(64) ScudoThreadContext : public ScudoThreadContextPlatform {
+struct ALIGNED(64) ScudoTSD : public ScudoTSDPlatform {
   AllocatorCache Cache;
   ScudoPrng Prng;
   uptr QuarantineCachePlaceHolder[4];
   void init();
   void commitBack();
 };
 
 void initThread(bool MinimalInit);
 
-// Platform specific dastpath functions definitions.
+// Platform specific fastpath functions definitions.
 #include "scudo_tls_android.inc"
 #include "scudo_tls_linux.inc"
 
 }  // namespace __scudo
 
 #endif  // SCUDO_TLS_H_

lib/scudo/scudo_tls_android.cpp

 
 #include <pthread.h>
 
 namespace __scudo {
 
 static pthread_once_t GlobalInitialized = PTHREAD_ONCE_INIT;
 static pthread_key_t PThreadKey;
 
-static atomic_uint32_t ThreadContextCurrentIndex;
-static ScudoThreadContext *ThreadContexts;
-static uptr NumberOfContexts;
+static atomic_uint32_t CurrentIndex;
+static ScudoTSD *TSDs;
+static u32 NumberOfTSDs;
 
 // sysconf(_SC_NPROCESSORS_{CONF,ONLN}) cannot be used as they allocate memory.
 static uptr getNumberOfCPUs() {
   cpu_set_t CPUs;
   CHECK_EQ(sched_getaffinity(0, sizeof(cpu_set_t), &CPUs), 0);
   return CPU_COUNT(&CPUs);
 }
 
 static void initOnce() {
   // Hack: TLS_SLOT_TSAN was introduced in N. To be able to use it on M for
   // testing, we create an unused key. Since the key_data array follows the tls
   // array, it basically gives us the extra entry we need.
   // TODO(kostyak): remove and restrict to N and above.
   CHECK_EQ(pthread_key_create(&PThreadKey, NULL), 0);
   initScudo();
-  NumberOfContexts = getNumberOfCPUs();
-  ThreadContexts = reinterpret_cast<ScudoThreadContext *>(
-      MmapOrDie(sizeof(ScudoThreadContext) * NumberOfContexts, __func__));
-  for (uptr i = 0; i < NumberOfContexts; i++)
-    ThreadContexts[i].init();
+  NumberOfTSDs = getNumberOfCPUs();
+  if (NumberOfTSDs == 0)
+    NumberOfTSDs = 1;
+  if (NumberOfTSDs > 32)
+    NumberOfTSDs = 32;
+  TSDs = reinterpret_cast<ScudoTSD *>(
+      MmapOrDie(sizeof(ScudoTSD) * NumberOfTSDs, "ScudoTSDs"));
+  for (u32 i = 0; i < NumberOfTSDs; i++)
+    TSDs[i].init();
 }
 
 void initThread(bool MinimalInit) {
   pthread_once(&GlobalInitialized, initOnce);
   // Initial context assignment is done in a plain round-robin fashion.
-  u32 Index = atomic_fetch_add(&ThreadContextCurrentIndex, 1,
-                               memory_order_relaxed);
-  ScudoThreadContext *ThreadContext =
-      &ThreadContexts[Index % NumberOfContexts];
-  *get_android_tls_ptr() = reinterpret_cast<uptr>(ThreadContext);
+  u32 Index = atomic_fetch_add(&CurrentIndex, 1, memory_order_relaxed);
+  ScudoTSD *TSD = &TSDs[Index % NumberOfTSDs];
+  *get_android_tls_ptr() = reinterpret_cast<uptr>(TSD);
 }
 
-ScudoThreadContext *getThreadContextAndLockSlow() {
-  ScudoThreadContext *ThreadContext;
-  // Go through all the contexts and find the first unlocked one. 
-  for (u32 i = 0; i < NumberOfContexts; i++) {
-    ThreadContext = &ThreadContexts[i];
-    if (ThreadContext->tryLock()) {
-      *get_android_tls_ptr() = reinterpret_cast<uptr>(ThreadContext);
-      return ThreadContext;
+ScudoTSD *getTSDAndLockSlow() {
+  ScudoTSD *TSD;
+  if (NumberOfTSDs > 1) {
+    // Go through all the contexts and find the first unlocked one.
+    for (u32 i = 0; i < NumberOfTSDs; i++) {
+      TSD = &TSDs[i];
+      if (TSD->tryLock()) {
+        *get_android_tls_ptr() = reinterpret_cast<uptr>(TSD);
+        return TSD;
       }
     }
-  // No luck, find the one with the lowest precedence, and slow lock it.
-  u64 Precedence = UINT64_MAX;
-  for (u32 i = 0; i < NumberOfContexts; i++) {
-    u64 SlowLockPrecedence = ThreadContexts[i].getSlowLockPrecedence();
-    if (SlowLockPrecedence && SlowLockPrecedence < Precedence) {
-      ThreadContext = &ThreadContexts[i];
-      Precedence = SlowLockPrecedence;
+    // No luck, find the one with the lowest Precedence, and slow lock it.
+    u64 LowestPrecedence = UINT64_MAX;
+    for (u32 i = 0; i < NumberOfTSDs; i++) {
+      u64 Precedence = TSDs[i].getPrecedence();
+      if (Precedence && Precedence < LowestPrecedence) {
+        TSD = &TSDs[i];
+        LowestPrecedence = Precedence;
       }
     }
-  if (LIKELY(Precedence != UINT64_MAX)) {
-    ThreadContext->lock();
-    *get_android_tls_ptr() = reinterpret_cast<uptr>(ThreadContext);
-    return ThreadContext;
+    if (LIKELY(LowestPrecedence != UINT64_MAX)) {
+      TSD->lock();
+      *get_android_tls_ptr() = reinterpret_cast<uptr>(TSD);
+      return TSD;
     }
-  // Last resort (can this happen?), stick with the current one.
-  ThreadContext =
-      reinterpret_cast<ScudoThreadContext *>(*get_android_tls_ptr());
-  ThreadContext->lock();
-  return ThreadContext;
+  }
+  // Last resort, stick with the current one.
+  TSD = reinterpret_cast<ScudoTSD *>(*get_android_tls_ptr());
+  TSD->lock();
+  return TSD;
 }
 
 }  // namespace __scudo
 
 #endif  // SANITIZER_LINUX && SANITIZER_ANDROID

lib/scudo/scudo_tls_android.inc

 #if SANITIZER_LINUX && SANITIZER_ANDROID
 
 ALWAYS_INLINE void initThreadMaybe(bool MinimalInit = false) {
   if (LIKELY(*get_android_tls_ptr()))
     return;
   initThread(MinimalInit);
 }
 
-ScudoThreadContext *getThreadContextAndLockSlow();
+ScudoTSD *getTSDAndLockSlow();
 
-ALWAYS_INLINE ScudoThreadContext *getThreadContextAndLock() {
-  ScudoThreadContext *ThreadContext =
-      reinterpret_cast<ScudoThreadContext *>(*get_android_tls_ptr());
-  CHECK(ThreadContext);
+ALWAYS_INLINE ScudoTSD *getTSDAndLock() {
+  ScudoTSD *TSD = reinterpret_cast<ScudoTSD *>(*get_android_tls_ptr());
+  CHECK(TSD);
   // Try to lock the currently associated context.
-  if (ThreadContext->tryLock())
-    return ThreadContext;
+  if (TSD->tryLock())
+    return TSD;
   // If it failed, go the slow path.
-  return getThreadContextAndLockSlow();
+  return getTSDAndLockSlow();
 }
 
 #endif  // SANITIZER_LINUX && SANITIZER_ANDROID
 
 #endif  // SCUDO_TLS_ANDROID_H_

lib/scudo/scudo_tls_context_android.inc

 #define SCUDO_TLS_CONTEXT_ANDROID_INC_
 
 #ifndef SCUDO_TLS_H_
 # error "This file must be included inside scudo_tls.h."
 #endif  // SCUDO_TLS_H_
 
 #if SANITIZER_LINUX && SANITIZER_ANDROID
 
-struct ScudoThreadContextPlatform {
+struct ScudoTSDPlatform {
   INLINE bool tryLock() {
     if (Mutex.TryLock()) {
-      atomic_store_relaxed(&SlowLockPrecedence, 0);
+      atomic_store_relaxed(&Precedence, 0);
       return true;
     }
-    if (atomic_load_relaxed(&SlowLockPrecedence) == 0)
-      atomic_store_relaxed(&SlowLockPrecedence, NanoTime());
+    if (atomic_load_relaxed(&Precedence) == 0)
+      atomic_store_relaxed(&Precedence, NanoTime());
     return false;
   }
 
   INLINE void lock() {
     Mutex.Lock();
-    atomic_store_relaxed(&SlowLockPrecedence, 0);
+    atomic_store_relaxed(&Precedence, 0);
   }
 
   INLINE void unlock() {
     Mutex.Unlock();
   }
 
-  INLINE u64 getSlowLockPrecedence() {
-    return atomic_load_relaxed(&SlowLockPrecedence);
+  INLINE u64 getPrecedence() {
+    return atomic_load_relaxed(&Precedence);
   }
 
  private:
   StaticSpinMutex Mutex;
-  atomic_uint64_t SlowLockPrecedence;
+  atomic_uint64_t Precedence;
 };
 
 #endif  // SANITIZER_LINUX && SANITIZER_ANDROID
 
 #endif  // SCUDO_TLS_CONTEXT_ANDROID_INC_

lib/scudo/scudo_tls_context_linux.inc

 #define SCUDO_TLS_CONTEXT_LINUX_INC_
 
 #ifndef SCUDO_TLS_H_
 # error "This file must be included inside scudo_tls.h."
 #endif  // SCUDO_TLS_H_
 
 #if SANITIZER_LINUX && !SANITIZER_ANDROID
 
-struct ScudoThreadContextPlatform {
+struct ScudoTSDPlatform {
   ALWAYS_INLINE void unlock() {}
 };
 
 #endif  // SANITIZER_LINUX && !SANITIZER_ANDROID
 
 #endif  // SCUDO_TLS_CONTEXT_LINUX_INC_

lib/scudo/scudo_tls_linux.cpp

 namespace __scudo {
 
 static pthread_once_t GlobalInitialized = PTHREAD_ONCE_INIT;
 static pthread_key_t PThreadKey;
 
 __attribute__((tls_model("initial-exec")))
 THREADLOCAL ThreadState ScudoThreadState = ThreadNotInitialized;
 __attribute__((tls_model("initial-exec")))
-THREADLOCAL ScudoThreadContext ThreadLocalContext;
+THREADLOCAL ScudoTSD TSD;
 
 static void teardownThread(void *Ptr) {
   uptr I = reinterpret_cast<uptr>(Ptr);
   // The glibc POSIX thread-local-storage deallocation routine calls user
   // provided destructors in a loop of PTHREAD_DESTRUCTOR_ITERATIONS.
   // We want to be called last since other destructors might call free and the
   // like, so we wait until PTHREAD_DESTRUCTOR_ITERATIONS before draining the
   // quarantine and swallowing the cache.
   if (I > 1) {
     // If pthread_setspecific fails, we will go ahead with the teardown.
     if (LIKELY(pthread_setspecific(PThreadKey,
                                    reinterpret_cast<void *>(I - 1)) == 0))
       return;
   }
-  ThreadLocalContext.commitBack();
+  TSD.commitBack();
   ScudoThreadState = ThreadTornDown;
 }
 
 
 static void initOnce() {
   CHECK_EQ(pthread_key_create(&PThreadKey, teardownThread), 0);
   initScudo();
 }
 
 void initThread(bool MinimalInit) {
   CHECK_EQ(pthread_once(&GlobalInitialized, initOnce), 0);
   if (UNLIKELY(MinimalInit))
     return;
   CHECK_EQ(pthread_setspecific(PThreadKey, reinterpret_cast<void *>(
       GetPthreadDestructorIterations())), 0);
-  ThreadLocalContext.init();
+  TSD.init();
   ScudoThreadState = ThreadInitialized;
 }
 
 }  // namespace __scudo
 
 #endif  // SANITIZER_LINUX && !SANITIZER_ANDROID

lib/scudo/scudo_tls_linux.inc

 enum ThreadState : u8 {
   ThreadNotInitialized = 0,
   ThreadInitialized,
   ThreadTornDown,
 };
 __attribute__((tls_model("initial-exec")))
 extern THREADLOCAL ThreadState ScudoThreadState;
 __attribute__((tls_model("initial-exec")))
-extern THREADLOCAL ScudoThreadContext ThreadLocalContext;
+extern THREADLOCAL ScudoTSD TSD;
 
 ALWAYS_INLINE void initThreadMaybe(bool MinimalInit = false) {
   if (LIKELY(ScudoThreadState != ThreadNotInitialized))
     return;
   initThread(MinimalInit);
 }
 
-ALWAYS_INLINE ScudoThreadContext *getThreadContextAndLock() {
+ALWAYS_INLINE ScudoTSD *getTSDAndLock() {
   if (UNLIKELY(ScudoThreadState != ThreadInitialized))
     return nullptr;
-  return &ThreadLocalContext;
+  return &TSD;
 }
 
 #endif  // SANITIZER_LINUX && !SANITIZER_ANDROID
 
 #endif  // SCUDO_TLS_LINUX_H_