[scudo] Improve the scalability of the shared TSD model

[cryptoad]

Summary

The shared TSD model in its current form doesn't scale. Here is an example of
rpc2-benchmark (with default parameters, which is threading heavy) on a 72-core
machines (defaulting to a CompactSizeClassMap and no Quarantine):

• with tcmalloc: 337K reqs/sec, peak RSS of 338MB;
• with scudo (exclusive): 321K reqs/sec, peak RSS of 637MB;
• with scudo (shared): 241K reqs/sec, peak RSS of 324MB.

This isn't great, since the exclusive model uses a lot of memory, while the
shared model doesn't even come close to be competitive.

This is mostly due to the fact that we are consistently scanning the TSD pool
starting at index 0 for an available TSD, which can result in a lot of failed
lock attempts, and touching some memory that needs not be touched.

This CL attempts to make things better in most situations:

• first, use a thread local variable on Linux (intead of pthread APIs) to store the current TSD in the shared model;
• move the locking boolean out of the TSD: this allows the compiler to use a register and potentially optimize out a branch instead of reading it from the TSD everytime (we also save a tiny bit of memory per TSD);
• 64-bit atomic operations on 32-bit ARM platforms happen to be expensive: so store the Precedence in a uptr instead of a u64. We lose some nanoseconds of precision and we'll wrap around at some point, but the benefit is worth it;
• change a CHECK to a DCHECK: this should never happen, but if something is ever terribly wrong, we'll crash on a near null AV if the TSD happens to be null;
• based on an idea by dvyukov@, we are implementing a bound random scan for an available TSD. This requires computing the coprimes for the number of TSDs, and attempting to lock up to 4 TSDs in an random order before falling back to the current one. This is obviously slightly more expansive when we have just 2 TSDs (barely noticeable) but is otherwise beneficial. The Precedence still basically corresponds to the moment of the first contention on a TSD. To seed on random choice, we use the precedence of the current TSD since it is very likely to be non-zero (since we are in the slow path after a failed tryLock)

With those modifications, the benchmark yields to:

• with scudo (shared): 330K reqs/sec, peak RSS of 327MB.

So the shared model for this specific situation not only becomes competitive but
outperforms the exclusive model. I experimented with some values greater than 4
for the number of TSDs to attempt to lock and it yielded a decrease in QPS. Just
sticking with the current TSD is also a tad slower. Numbers on platforms with
less cores (eg: Android) remain similar.

[dvyukov]

lib/scudo/scudo_tsd_shared.cpp

71

The comment from the change description as to why we use precedence as rand state should be duplicated as a comment here. This is confusing.

73

Do you see getTSDAndLockSlow in profiles? Here we have 2 divisions. Divisions are expensive. We could replace them with multiplication+shift if they show up in profiles.

78

Have you tried dropping precedence entirely and then just going over all TSDs pseudo-randomly once (or maybe twice) and tryLock'ing them? If all tryLock's fail then we could lock the old one, or maybe a random one.
That would shave few cycles from malloc/free fast paths. The benefit of using precedence scheme is still unproved, right?
Also, if we have a hundred of cores/caches, then giving up after just 4 can lead to unnecessary blocking (there still can be dozens of available caches).

[cryptoad]

I will experiment with Dmitry's suggestions and report back with results.

[cryptoad]

Here are some answers to Dmitry's requests:

• Regarding getTSDAndLockSlow and the division: pprof shows no significant time spent in the function outside of the tryLock & lock so I think we are good here;
• Regarding the precedence: I tested a version where I dropped it entirely, results are mixed:
  • For Android's "improved" memory_replay: it is faster in all cases, but we only have 2 caches for that specific platform (due to memory constraints compared to the default allocator);
  • For rpc2-benchmark: mostly similar numbers;
  • For t-test1: the version with precedence shows better performances in almost all situations; this benchmark also demonstrates a slowdown with the number of TSDs scanned in the slowpath, eg: scanning 4 and slow locking if they all failed to tryLock performs better overall than scanning 32. And this can be a significant slowdown, for example with t-test1 800 40 800000 100000, it's 900s spent in allocation functions vs 1150s. The argument here is that this benchmark only does {de}allocations (& memset) and as such isn't very representative of "real" programs, but it's exercising the most contention on the caches.

I can't seem to get a definitive answer overall as with or without precedence have both win & lose situations.
The only sure thing so far is that both are better than the current version.

I am open to suggestion or potential improvements, otherwise I'd keep the current version of the CL (and will address the review comments).

[dvyukov]

What exactly versions did you test? Current vs removing precedence and trylocking all caches? If yes, then it's unclear what factor affects performance as we change 2 things at the same time. So I would also try:

1. Remove precedence, but scan only 4 random caches, then lock the current cache.
2. Remove precedence, but scan only 4 random caches, then lock a random cache.
3. Keep precedence but scan all caches.

I would not expect 3 to be faster, but who knows.
Looking at your results I think it mostly needs to be tested on t-test1, because it should not affect android with 2/4 caches.

[cryptoad]

Here are more detail numbers for t-test1.

The machine has 72 cores. We are using the shared TSD version with 32 caches (to exercise some contention).
The numbers are the total time (averaged and rounded over 3 consecutive runs) spent in allocation functions only with 40, then 80 concurrent threads:

• current upstream: 960s, 3315s
• with precedence, max 4 caches scanned, lock current: 810s, 3200s (current CL proposed)
• with precedence, max 4 caches scanned, lock random: 815s, 3125s
• with precedence, all caches scanned, lock current: 880s, 3940s
• with precedence, all caches scanned, lock random: 890s, 3755s
• no precedence, max 4 caches scanned, lock current: 900s, 3365s
• no precedence, max 4 caches scanned, lock random: 840s, 3300s
• no precedence, all caches scanned, lock current: 1025s, 3600s
• no precedence, all caches scanned, lock random: 890s, 3785s

Locking a random cache in the event of heavier contention seems to be beneficial, but not necessarily with lesser contention.
Since I am more interested in striking a middle ground rather than aiming for contentious applications, it looks like the precedence matters, as well as not scanning all the caches but limiting ourselves to 4.

[dvyukov]

So your current version is actually the best one.
Thanks for bearing with me.

[cryptoad]

Adding a comment to explain why we use the precedence as random seed.
Pass the current TSD to getTSDAndLockSlow to avoid an extra getCurrentTSD
call.

[cryptoad]

Remove the continue clause from the loop when the TSD is the current TSD.
This is greatly beneficial to scenarios (Android for example) where we only
have a small number of TSDs. This doesn't impact larger pools of TSDs.

[cryptoad]

Thanks for the review Aleksey &Dmitry.
If I could get a final LGTM after the last couple of changes that would be great.

[alekseyshl]

LGTM

lib/scudo/scudo_allocator.cpp

void *allocate(uptr Size, uptr Alignment, AllocType Type,
     // deal with alignment requirements of Primary serviced allocations here,
     // but the Secondary will take care of its own alignment needs.
     void *BackendPtr;
     uptr BackendSize;
     u8 ClassId;
     if (PrimaryAllocator::CanAllocate(AlignedSize, MinAlignment)) {
       BackendSize = AlignedSize;
       ClassId = SizeClassMap::ClassID(BackendSize);
-      ScudoTSD *TSD = getTSDAndLock();
+      bool UnlockRequired;
+      ScudoTSD *TSD = getTSDAndLock(&UnlockRequired);
       BackendPtr = BackendAllocator.allocatePrimary(&TSD->Cache, ClassId);
+      if (UnlockRequired)
         TSD->unlock();
     } else {
       BackendSize = NeededSize;
       ClassId = 0;
       BackendPtr = BackendAllocator.allocateSecondary(BackendSize, Alignment);
     }
     if (UNLIKELY(!BackendPtr))
       return FailureHandler::OnOOM();
 
struct ScudoAllocator {
   void quarantineOrDeallocateChunk(void *Ptr, UnpackedHeader *Header,
                                    uptr Size) {
     const bool BypassQuarantine = (AllocatorQuarantine.GetCacheSize() == 0) ||
         (Size > QuarantineChunksUpToSize);
     if (BypassQuarantine) {
       Chunk::eraseHeader(Ptr);
       void *BackendPtr = Chunk::getBackendPtr(Ptr, Header);
       if (Header->ClassId) {
-        ScudoTSD *TSD = getTSDAndLock();
+        bool UnlockRequired;
+        ScudoTSD *TSD = getTSDAndLock(&UnlockRequired);
         getBackendAllocator().deallocatePrimary(&TSD->Cache, BackendPtr,
                                                 Header->ClassId);
+        if (UnlockRequired)
           TSD->unlock();
       } else {
         getBackendAllocator().deallocateSecondary(BackendPtr);
       }
     } 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.
       const uptr EstimatedSize = Size + (Header->Offset << MinAlignmentLog);
       UnpackedHeader NewHeader = *Header;
       NewHeader.State = ChunkQuarantine;
       Chunk::compareExchangeHeader(Ptr, &NewHeader, Header);
-      ScudoTSD *TSD = getTSDAndLock();
+      bool UnlockRequired;
+      ScudoTSD *TSD = getTSDAndLock(&UnlockRequired);
       AllocatorQuarantine.Put(getQuarantineCache(TSD),
                               QuarantineCallback(&TSD->Cache), Ptr,
                               EstimatedSize);
+      if (UnlockRequired)
         TSD->unlock();
     }
   }
 
   // Deallocates a Chunk, which means either adding it to the quarantine or
   // directly returning it to the backend if criteria are met.
   void deallocate(void *Ptr, 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;
 }
 
 void initScudo() {
   Instance.init();
 }
 
-void ScudoTSD::init(bool Shared) {
-  UnlockRequired = Shared;
+void ScudoTSD::init() {
   getBackendAllocator().initCache(&Cache);
   memset(QuarantineCachePlaceHolder, 0, sizeof(QuarantineCachePlaceHolder));
 }
 
 void ScudoTSD::commitBack() {
   Instance.commitBack(this);
 }
 

lib/scudo/scudo_tsd.h

 
 #include "scudo_allocator.h"
 #include "scudo_utils.h"
 
 #include <pthread.h>
 
 namespace __scudo {
 
-struct ALIGNED(64) ScudoTSD {
+struct ALIGNED(SANITIZER_CACHE_LINE_SIZE) ScudoTSD {
   AllocatorCache Cache;
   uptr QuarantineCachePlaceHolder[4];
 
-  void init(bool Shared);
+  void init();
   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, MonotonicNanoTime());
+      atomic_store_relaxed(&Precedence, static_cast<uptr>(
+          MonotonicNanoTime() >> FIRST_32_SECOND_64(16, 0)));
     return false;
   }
 
   INLINE void lock() {
-    Mutex.Lock();
     atomic_store_relaxed(&Precedence, 0);
+    Mutex.Lock();
   }
 
-  INLINE void unlock() {
-    if (!UnlockRequired)
-      return;
-    Mutex.Unlock();
-  }
+  INLINE void unlock() { Mutex.Unlock(); }
 
-  INLINE u64 getPrecedence() {
-    return atomic_load_relaxed(&Precedence);
-  }
+  INLINE uptr getPrecedence() { return atomic_load_relaxed(&Precedence); }
 
  private:
-  bool UnlockRequired;
   StaticSpinMutex Mutex;
-  atomic_uint64_t Precedence;
+  atomic_uintptr_t Precedence;
 };
 
 void initThread(bool MinimalInit);
 
 // TSD model specific fastpath functions definitions.
 #include "scudo_tsd_exclusive.inc"
 #include "scudo_tsd_shared.inc"
 
 }  // namespace __scudo
 
 #endif  // SCUDO_TSD_H_

lib/scudo/scudo_tsd_exclusive.cpp

static void teardownThread(void *Ptr) {
   TSD.commitBack();
   ScudoThreadState = ThreadTornDown;
 }
 
 
 static void initOnce() {
   CHECK_EQ(pthread_key_create(&PThreadKey, teardownThread), 0);
   initScudo();
-  FallbackTSD.init(/*Shared=*/true);
+  FallbackTSD.init();
 }
 
 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(/*Shared=*/false);
+  TSD.init();
   ScudoThreadState = ThreadInitialized;
 }
 
 }  // namespace __scudo
 
 #endif  // SCUDO_TSD_EXCLUSIVE

lib/scudo/scudo_tsd_exclusive.inc

 extern ScudoTSD FallbackTSD;
 
 ALWAYS_INLINE void initThreadMaybe(bool MinimalInit = false) {
   if (LIKELY(ScudoThreadState != ThreadNotInitialized))
     return;
   initThread(MinimalInit);
 }
 
-ALWAYS_INLINE ScudoTSD *getTSDAndLock() {
+ALWAYS_INLINE ScudoTSD *getTSDAndLock(bool *UnlockRequired) {
   if (UNLIKELY(ScudoThreadState != ThreadInitialized)) {
     FallbackTSD.lock();
+    *UnlockRequired = true;
     return &FallbackTSD;
   }
+  *UnlockRequired = false;
   return &TSD;
 }
 
 #endif  // SCUDO_TSD_EXCLUSIVE

lib/scudo/scudo_tsd_shared.cpp

 namespace __scudo {
 
 static pthread_once_t GlobalInitialized = PTHREAD_ONCE_INIT;
 pthread_key_t PThreadKey;
 
 static atomic_uint32_t CurrentIndex;
 static ScudoTSD *TSDs;
 static u32 NumberOfTSDs;
+static u32 CoPrimes[SCUDO_SHARED_TSD_POOL_SIZE];
+static u32 NumberOfCoPrimes = 0;
+
+#if SANITIZER_LINUX && !SANITIZER_ANDROID
+__attribute__((tls_model("initial-exec")))
+THREADLOCAL ScudoTSD *CurrentTSD;
+#endif
 
 static void initOnce() {
   CHECK_EQ(pthread_key_create(&PThreadKey, NULL), 0);
   initScudo();
   NumberOfTSDs = Min(Max(1U, GetNumberOfCPUsCached()),
                      static_cast<u32>(SCUDO_SHARED_TSD_POOL_SIZE));
   TSDs = reinterpret_cast<ScudoTSD *>(
       MmapOrDie(sizeof(ScudoTSD) * NumberOfTSDs, "ScudoTSDs"));
-  for (u32 i = 0; i < NumberOfTSDs; i++)
-    TSDs[i].init(/*Shared=*/true);
+  for (u32 I = 0; I < NumberOfTSDs; I++) {
+    TSDs[I].init();
+    u32 A = I + 1;
+    u32 B = NumberOfTSDs;
+    while (B != 0) { const u32 T = A; A = B; B = T % B; }
+    if (A == 1)
+      CoPrimes[NumberOfCoPrimes++] = I + 1;
+  }
 }
 
 ALWAYS_INLINE void setCurrentTSD(ScudoTSD *TSD) {
 #if SANITIZER_ANDROID
   *get_android_tls_ptr() = reinterpret_cast<uptr>(TSD);
+#elif SANITIZER_LINUX
+  CurrentTSD = TSD;
 #else
   CHECK_EQ(pthread_setspecific(PThreadKey, reinterpret_cast<void *>(TSD)), 0);
 #endif  // SANITIZER_ANDROID
 }
 
 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);
   setCurrentTSD(&TSDs[Index % NumberOfTSDs]);
 }
 
 ScudoTSD *getTSDAndLockSlow() {
-  ScudoTSD *TSD;
+  ScudoTSD *TSD = getCurrentTSD();
   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()) {
-        setCurrentTSD(TSD);
-        return TSD;
-      }
-    }
-    // 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];
+    u32 RandState = static_cast<u32>(TSD->getPrecedence());
+    const u32 R = Rand(&RandState);
+    const u32 Inc = CoPrimes[R % NumberOfCoPrimes];
+    u32 Index = R % NumberOfTSDs;
+    uptr LowestPrecedence = UINTPTR_MAX;
+    ScudoTSD *CandidateTSD = nullptr;
+    // Go randomly through at most 4 contexts and find a candidate.
+    for (u32 I = 0; I < Min(4U, NumberOfTSDs); I++) {
+      if (&TSDs[Index] == TSD)
+        continue;
+      if (TSDs[Index].tryLock()) {
+        setCurrentTSD(&TSDs[Index]);
+        return &TSDs[Index];
+      }
+      const uptr Precedence = TSDs[Index].getPrecedence();
+      // A 0 precedence here means another thread just locked this TSD.
+      if (UNLIKELY(Precedence == 0))
+        continue;
+      if (Precedence < LowestPrecedence) {
+        CandidateTSD = &TSDs[Index];
         LowestPrecedence = Precedence;
       }
-    }
-    if (LIKELY(LowestPrecedence != UINT64_MAX)) {
-      TSD->lock();
-      setCurrentTSD(TSD);
-      return TSD;
+      Index += Inc;
+      if (Index >= NumberOfTSDs)
+        Index -= NumberOfTSDs;
+    }
+    if (CandidateTSD) {
+      CandidateTSD->lock();
+      setCurrentTSD(CandidateTSD);
+      return CandidateTSD;
     }
   }
   // Last resort, stick with the current one.
-  TSD = getCurrentTSD();
   TSD->lock();
   return TSD;
 }
 
 }  // namespace __scudo
 
 #endif  // !SCUDO_TSD_EXCLUSIVE

lib/scudo/scudo_tsd_shared.inc

 #ifndef SCUDO_TSD_H_
 # error "This file must be included inside scudo_tsd.h."
 #endif  // SCUDO_TSD_H_
 
 #if !SCUDO_TSD_EXCLUSIVE
 
 extern pthread_key_t PThreadKey;
 
+#if SANITIZER_LINUX && !SANITIZER_ANDROID
+__attribute__((tls_model("initial-exec")))
+extern THREADLOCAL ScudoTSD *CurrentTSD;
+#endif
+
 ALWAYS_INLINE ScudoTSD* getCurrentTSD() {
 #if SANITIZER_ANDROID
   return reinterpret_cast<ScudoTSD *>(*get_android_tls_ptr());
+#elif SANITIZER_LINUX
+  return CurrentTSD;
 #else
   return reinterpret_cast<ScudoTSD *>(pthread_getspecific(PThreadKey));
 #endif  // SANITIZER_ANDROID
 }
 
 ALWAYS_INLINE void initThreadMaybe(bool MinimalInit = false) {
   if (LIKELY(getCurrentTSD()))
     return;
   initThread(MinimalInit);
 }
 
 ScudoTSD *getTSDAndLockSlow();
 
-ALWAYS_INLINE ScudoTSD *getTSDAndLock() {
+ALWAYS_INLINE ScudoTSD *getTSDAndLock(bool *UnlockRequired) {
   ScudoTSD *TSD = getCurrentTSD();
-  CHECK(TSD && "No TSD associated with the current thread!");
+  DCHECK(TSD && "No TSD associated with the current thread!");
+  *UnlockRequired = true;
   // Try to lock the currently associated context.
   if (TSD->tryLock())
     return TSD;
   // If it failed, go the slow path.
   return getTSDAndLockSlow();
 }
 
 #endif  // !SCUDO_TSD_EXCLUSIVE