[scudo] Scudo thread specific data refactor, part 2

[cryptoad]

Summary

Following D38139, we now consolidate the TSD definition, merging the shared
TSD definition with the exclusive TSD definition. We introduce a boolean set
at initializaton denoting the need for the TSD to be unlocked or not. This
adds some unused members to the exclusive TSD, but increases consistency and
reduces the definitions fragmentation.

We remove the fallback mechanism from scudo_allocator.cpp and add a fallback
TSD in the non-shared version. Since the shared version doesn't require one,
this makes overall more sense.

There are a couple of additional cosmetic changes: removing the header guards
from the remaining .inc files, added error string to a CHECK.

Question to reviewers: I thought about friending getTSDAndLock in ScudoTSD
so that the FallbackTSD could Mutex.Lock() directly instead of lock()
which involved zeroing out the Precedence, which is unused otherwise. Is it
worth doing?

lib/scudo/scudo_allocator.cpp

struct ScudoAllocator {
   typedef ReturnNullOrDieOnFailure FailureHandler;
 
   ScudoBackendAllocator BackendAllocator;
   ScudoQuarantine AllocatorQuarantine;
 
   StaticSpinMutex GlobalPrngMutex;
   ScudoPrng GlobalPrng;
 
-  // The fallback caches are used when the thread local caches have been
-  // 'detroyed' on thread tear-down. They are protected by a Mutex as they can
-  // be accessed by different threads.
-  StaticSpinMutex FallbackMutex;
-  AllocatorCache FallbackAllocatorCache;
-  ScudoQuarantineCache FallbackQuarantineCache;
-  ScudoPrng FallbackPrng;
-
   u32 QuarantineChunksUpToSize;
 
   bool DeallocationTypeMismatch;
   bool ZeroContents;
   bool DeleteSizeMismatch;
 
   explicit ScudoAllocator(LinkerInitialized)
-    : AllocatorQuarantine(LINKER_INITIALIZED),
-      FallbackQuarantineCache(LINKER_INITIALIZED) {}
+    : AllocatorQuarantine(LINKER_INITIALIZED) {}
 
   void init(const AllocatorOptions &Options) {
     // Verify that the header offset field can hold the maximum offset. In the
     // case of the Secondary allocator, it takes care of alignment and the
     // offset will always be 0. In the case of the Primary, the worst case
     // scenario happens in the last size class, when the backend allocation
     // would already be aligned on the requested alignment, which would happen
     // to be the maximum alignment that would fit in that size class. As a
void init(const AllocatorOptions &Options) {
     SetAllocatorMayReturnNull(Options.MayReturnNull);
     BackendAllocator.init(Options.ReleaseToOSIntervalMs);
     AllocatorQuarantine.Init(
         static_cast<uptr>(Options.QuarantineSizeKb) << 10,
         static_cast<uptr>(Options.ThreadLocalQuarantineSizeKb) << 10);
     QuarantineChunksUpToSize = Options.QuarantineChunksUpToSize;
     GlobalPrng.init();
     Cookie = GlobalPrng.getU64();
-    BackendAllocator.initCache(&FallbackAllocatorCache);
-    FallbackPrng.init();
   }
 
   // Helper function that checks for a valid Scudo chunk. nullptr isn't.
   bool isValidPointer(const void *UserPtr) {
     initThreadMaybe();
     if (UNLIKELY(!UserPtr))
       return false;
     uptr UserBeg = reinterpret_cast<uptr>(UserPtr);
void *allocate(uptr Size, uptr Alignment, AllocType Type,
     bool FromPrimary = PrimaryAllocator::CanAllocate(AlignedSize, MinAlignment);
 
     void *Ptr;
     u8 Salt;
     uptr AllocSize;
     if (FromPrimary) {
       AllocSize = AlignedSize;
       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 {
       {
         SpinMutexLock l(&GlobalPrngMutex);
         Salt = GlobalPrng.getU8();
       }
       AllocSize = NeededSize;
       Ptr = BackendAllocator.allocateSecondary(AllocSize, Alignment);
     }
     if (UNLIKELY(!Ptr))
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) {
         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);
       ScudoTSD *TSD = getTSDAndLock();
-      if (LIKELY(TSD)) {
-        AllocatorQuarantine.Put(getQuarantineCache(TSD),
-                                QuarantineCallback(&TSD->Cache),
-                                Chunk, EstimatedSize);
-        TSD->unlock();
-      } else {
-        SpinMutexLock l(&FallbackMutex);
-        AllocatorQuarantine.Put(&FallbackQuarantineCache,
-                                QuarantineCallback(&FallbackAllocatorCache),
-                                Chunk, EstimatedSize);
-      }
     }
   }
 
   // Deallocates a Chunk, which means adding it to the delayed free list (or
   // Quarantine).
   void deallocate(void *UserPtr, uptr DeleteSize, AllocType Type) {
     // For a deallocation, we only ensure minimal initialization, meaning thread
     // local data will be left uninitialized for now (when using ELF TLS). The

 static ScudoBackendAllocator &getBackendAllocator() {
   return Instance.BackendAllocator;
 }
 
 static void initScudoInternal(const AllocatorOptions &Options) {
   Instance.init(Options);
 }
 
-void ScudoTSD::init() {
+void ScudoTSD::init(bool Shared) {
+  UnlockRequired = Shared;
   getBackendAllocator().initCache(&Cache);
   Prng.init();
   memset(QuarantineCachePlaceHolder, 0, sizeof(QuarantineCachePlaceHolder));
 }
 
 void ScudoTSD::commitBack() {
   Instance.commitBack(this);
 }

lib/scudo/scudo_tls.h

 #include "scudo_allocator.h"
 #include "scudo_utils.h"
 
 #include "sanitizer_common/sanitizer_linux.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) ScudoTSD : public ScudoTSDPlatform {
+struct ALIGNED(64) ScudoTSD {
   AllocatorCache Cache;
   ScudoPrng Prng;
   uptr QuarantineCachePlaceHolder[4];
-  void init();
+
+  void init(bool Shared);
   void commitBack();
+
+  INLINE bool tryLock() {
+    if (Mutex.TryLock()) {
+      atomic_store_relaxed(&Precedence, 0);
+      return true;
+    }
+    if (atomic_load_relaxed(&Precedence) == 0)
+      atomic_store_relaxed(&Precedence, NanoTime());
+    return false;
+  }
+
+  INLINE void lock() {
+    Mutex.Lock();
+    atomic_store_relaxed(&Precedence, 0);
+  }
+
+  INLINE void unlock() {
+    if (!UnlockRequired)
+      return;
+    Mutex.Unlock();
+  }
+
+  INLINE u64 getPrecedence() {
+    return atomic_load_relaxed(&Precedence);
+  }
+
+ private:
+  bool UnlockRequired;
+  StaticSpinMutex Mutex;
+  atomic_uint64_t Precedence;
 };
 
 void initThread(bool MinimalInit);
 
 // 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

static void initOnce() {
   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();
+    TSDs[i].init(/*Shared=*/true);
 }
 
 void initThread(bool MinimalInit) {
   pthread_once(&GlobalInitialized, initOnce);
   // Initial context assignment is done in a plain round-robin fashion.
   u32 Index = atomic_fetch_add(&CurrentIndex, 1, memory_order_relaxed);
   ScudoTSD *TSD = &TSDs[Index % NumberOfTSDs];
   *get_android_tls_ptr() = reinterpret_cast<uptr>(TSD);

lib/scudo/scudo_tls_android.inc

 //===-- scudo_tls_android.inc -----------------------------------*- C++ -*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 ///
 /// Scudo thread local structure fastpath functions implementation for Android.
 ///
 //===----------------------------------------------------------------------===//
 
-#ifndef SCUDO_TLS_ANDROID_H_
-#define SCUDO_TLS_ANDROID_H_
-
 #ifndef SCUDO_TLS_H_
 # error "This file must be included inside scudo_tls.h."
 #endif  // SCUDO_TLS_H_
 
 #if SANITIZER_LINUX && SANITIZER_ANDROID
 
 ALWAYS_INLINE void initThreadMaybe(bool MinimalInit = false) {
   if (LIKELY(*get_android_tls_ptr()))
     return;
   initThread(MinimalInit);
 }
 
 ScudoTSD *getTSDAndLockSlow();
 
 ALWAYS_INLINE ScudoTSD *getTSDAndLock() {
   ScudoTSD *TSD = reinterpret_cast<ScudoTSD *>(*get_android_tls_ptr());
-  CHECK(TSD);
+  CHECK(TSD && "No TSD associated with the current thread!");
   // Try to lock the currently associated context.
   if (TSD->tryLock())
     return TSD;
   // If it failed, go the slow path.
   return getTSDAndLockSlow();
 }
 
 #endif  // SANITIZER_LINUX && SANITIZER_ANDROID
-
-#endif  // SCUDO_TLS_ANDROID_H_

lib/scudo/scudo_tls_context_android.inc

-//===-- scudo_tls_context_android.inc ---------------------------*- C++ -*-===//
-//
-//                     The LLVM Compiler Infrastructure
-//
-// This file is distributed under the University of Illinois Open Source
-// License. See LICENSE.TXT for details.
-//
-//===----------------------------------------------------------------------===//
-///
-/// Android specific base thread context definition.
-///
-//===----------------------------------------------------------------------===//
-
-#ifndef 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 ScudoTSDPlatform {
-  INLINE bool tryLock() {
-    if (Mutex.TryLock()) {
-      atomic_store_relaxed(&Precedence, 0);
-      return true;
-    }
-    if (atomic_load_relaxed(&Precedence) == 0)
-      atomic_store_relaxed(&Precedence, NanoTime());
-    return false;
-  }
-
-  INLINE void lock() {
-    Mutex.Lock();
-    atomic_store_relaxed(&Precedence, 0);
-  }
-
-  INLINE void unlock() {
-    Mutex.Unlock();
-  }
-
-  INLINE u64 getPrecedence() {
-    return atomic_load_relaxed(&Precedence);
-  }
-
- private:
-  StaticSpinMutex Mutex;
-  atomic_uint64_t Precedence;
-};
-
-#endif  // SANITIZER_LINUX && SANITIZER_ANDROID
-
-#endif  // SCUDO_TLS_CONTEXT_ANDROID_INC_

lib/scudo/scudo_tls_context_linux.inc

-//===-- scudo_tls_context_linux.inc -----------------------------*- C++ -*-===//
-//
-//                     The LLVM Compiler Infrastructure
-//
-// This file is distributed under the University of Illinois Open Source
-// License. See LICENSE.TXT for details.
-//
-//===----------------------------------------------------------------------===//
-///
-/// Linux specific base thread context definition.
-///
-//===----------------------------------------------------------------------===//
-
-#ifndef 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 ScudoTSDPlatform {
-  ALWAYS_INLINE void unlock() {}
-};
-
-#endif  // SANITIZER_LINUX && !SANITIZER_ANDROID
-
-#endif  // SCUDO_TLS_CONTEXT_LINUX_INC_

lib/scudo/scudo_tls_linux.cpp

 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 ScudoTSD TSD;
 
+// Fallback TSD for when the thread isn't initialized yet or is torn down. It
+// can be shared between multiple threads and as such must be locked.
+ScudoTSD FallbackTSD;
+
 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;
   }
   TSD.commitBack();
   ScudoThreadState = ThreadTornDown;
 }
 
 
 static void initOnce() {
   CHECK_EQ(pthread_key_create(&PThreadKey, teardownThread), 0);
   initScudo();
+  FallbackTSD.init(/*Shared=*/true);
 }
 
 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);
-  TSD.init();
+  TSD.init(/*Shared=*/false);
   ScudoThreadState = ThreadInitialized;
 }
 
 }  // namespace __scudo
 
 #endif  // SANITIZER_LINUX && !SANITIZER_ANDROID

lib/scudo/scudo_tls_linux.inc

 //===-- scudo_tls_linux.inc -------------------------------------*- C++ -*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 ///
 /// Scudo thread local structure fastpath functions implementation for platforms
 /// supporting thread_local.
 ///
 //===----------------------------------------------------------------------===//
 
-#ifndef SCUDO_TLS_LINUX_H_
-#define SCUDO_TLS_LINUX_H_
-
 #ifndef SCUDO_TLS_H_
 # error "This file must be included inside scudo_tls.h."
 #endif  // SCUDO_TLS_H_
 
 #if SANITIZER_LINUX && !SANITIZER_ANDROID
 
 enum ThreadState : u8 {
   ThreadNotInitialized = 0,
   ThreadInitialized,
   ThreadTornDown,
 };
 __attribute__((tls_model("initial-exec")))
 extern THREADLOCAL ThreadState ScudoThreadState;
 __attribute__((tls_model("initial-exec")))
 extern THREADLOCAL ScudoTSD TSD;
 
+extern ScudoTSD FallbackTSD;
+
 ALWAYS_INLINE void initThreadMaybe(bool MinimalInit = false) {
   if (LIKELY(ScudoThreadState != ThreadNotInitialized))
     return;
   initThread(MinimalInit);
 }
 
 ALWAYS_INLINE ScudoTSD *getTSDAndLock() {
-  if (UNLIKELY(ScudoThreadState != ThreadInitialized))
-    return nullptr;
+  if (UNLIKELY(ScudoThreadState != ThreadInitialized)) {
+    FallbackTSD.lock();
+    return &FallbackTSD;
+  }
   return &TSD;
 }
 
 #endif  // SANITIZER_LINUX && !SANITIZER_ANDROID
-
-#endif  // SCUDO_TLS_LINUX_H_