[scudo][standalone] Introduce the thread specific data structures

[cryptoad]

Summary

This CL adds the structures dealing with thread specific data for the
allocator. This includes the thread specific data structure itself and
two registries for said structures: an exclusive one, where each thread
will have its own TSD struct, and a shared one, where a pool of TSD
structs will be shared by all threads, with dynamic reassignment at
runtime based on contention.

This departs from the current Scudo implementation: we intend to make
the Registry a template parameter of the allocator (as opposed to a
single global entity), allowing various allocators to coexist with
different TSD registry models. As a result, TSD registry and Allocator
are tightly coupled.

This also corrects a couple of things in other files that I noticed
while adding this.

[cryptoad]

Correct a test and some formatting.

[morehouse]

lib/scudo/standalone/tests/tsd_test.cc

34

I assume allocators will always be linker-initialized, which is why we have reset for testing instead of an init?

49

Allocator-> is cleaner than Allocator.get()->, unless there's a good reason to get the raw pointer.

lib/scudo/standalone/tsd.h

27

Mutex.initLinkerInitialized()

lib/scudo/standalone/tsd_exclusive.h

28

Shouldn't this also call initLinkerInitialized?

60

Will threads have different allocators?

64

What is the use-case for MinimalInit initialization?

74

Any reason this needs to be a pointer?

94

Assigning to N seems pointless. Can we simplify to:

if (TSDRegistryT::ThreadTSD.DestructorIterations > 1) {
  TSDRegistryT::ThreadTSD.DestructorIterations--;
 ...
}

lib/scudo/standalone/tsd_shared.h

21

Shouldn't this also call initLinkerInitialized?

23

Should Allocator initialization just be done in initLinkerInitialized?

61

What are the CoPrimes used for, and what is the strange calculation above doing?

74

Do we really need three ways to do TLS? Can we just use pthreads for all?

97

Is it sufficient to just check NumberOfTSDs?

123

Index %= NumberOfTSDs

144

Why do we need these guards here but not in tsd_exclusive.h?

[cryptoad]

Matt, thank you for all the reviews you are doing. Very insightful points.

lib/scudo/standalone/tests/tsd_test.cc

34

There is indeed a subtlety here, the initOnce function of a TSD will call the initLinkerInitialized function of the allocator.
Calling this init would (I feel), indicate the usual construct of init calling initLinkedInitialized, which shouldn't be the case here, we just want a nulled out structure but not call the initLinkedInitialized wanna-be-constructor.
Hence using reset.

lib/scudo/standalone/tsd_exclusive.h

60

Multiple allocators can coexist in a process, hence a need to know which allocator a registry belongs to.

64

Right, this is documented in combined.h that hasn't landed yet, here is the comment:

// For a deallocation, we only ensure minimal initialization, meaning thread
// local data will be left uninitialized for now (when using ELF TLS). The
// fallback cache will be used instead. This is a workaround for a situation
// where the only heap operation performed in a thread would be a free past
// the TLS destructors, ending up in initialized thread specific data never
// being destroyed properly. Any other heap operation will do a full init.

There is test case in the current Scudo for that situation that will be in as well.

74

The main reason is for an uninitialized allocator to take as little memory space as possible.
A TSD is usually around 8kB (varying based on the number of classes and pointers cached).
It costs an extra map() on init, but allows several allocators to coexist with using extra memory (the code footprint cost still being there).

lib/scudo/standalone/tsd_shared.h

23

This question made me realize that something is a bit sketchy.
initOnce is initializing the allocator, which in turn is calling initLinkerInitializedof the registry: at this point we are in the spinmutex, and call the initLinkerInitialized method of that same mutex. It all works out because 1) everything is zeroinitialized 2) the initLinkerInitialized of the spinmutex is a no-op.
But it is definitely logically skewed.
The original version was using pthread_once which was alleviating the need for the mutex, but I needed the Instance parameter for initOnce (pthread_one doesn't allow parameters to the once function).
I am going to have to rethink that.

61

I added a comment. The original idea was from Dmitry, but there are online reference such as:
https://lemire.me/blog/2017/09/18/visiting-all-values-in-an-array-exactly-once-in-random-order/

74

The ELF TLS would be the fastest as it's only a couple of instructions to access the data relative to %fs.
The pthread implementations vary widely. They require a call, which is not ideal, but also they range in terms of efficiency.
The getAndroidTlsPtr ends up being a few asm instructions as well.

Now with the introduction of ELF TLS in Android (beginning of this year), we might be able to get rid of the Android TLS ptr trick.
I haven't tested it yet. Also I am not sure the ELF TLS works from within the libc (it doesn't on Fuchsia for example).

97

This is to help the compiler optimize the block out.
If MaxTSDCount is 1 (likely the svelte Android config case), the compiler doesn't know that NumberOfTSDs is 1 and leaves the block in.
The MaxTSDCount check will be optimized at compile time, and for the 1 case, the whole block goes away.

123

Here the division is potentially costly, hence going for a comparison and subtraction.
I am not sure the performance difference is significant in this situation (slow path, 4 iterations at most) but I tried to avoid divisions everywhere as a general rule of thumb.

144

The exclusive version only uses ELF TLS. So if opting for an exclusive TSD it assumes the platform supports it.
For the shared version, it can work without ELF TLS, hence having THREADLOCAL variables within defines.
A couple of things to consider:

• Fuchsia supports THREADLOCAL, but not within the libc, and we are in the libc.
• Android recently got support for ELF TLS, but 1) I haven't tested it yet 2) I don't know if THREADLOCAL works from within the C library

I will revisit this when things settle, but as of now, I am sure this version works.

[cryptoad]

Addressing several of Matt's review points:

• adding comments to obscure code snippets
• calling initLinkerInitialized when needed
• changing unique_ptr's get()-> to ->
• simplifying some code constructs

[cryptoad]

Switching from a spin mutex to a blocking mutex for initOnce.

[morehouse]

lib/scudo/standalone/tsd_exclusive.h

28

I think switching to BlockingMutex introduced a new bug. BlockingMutex has a default constructor that will run at global ctor time. So the following sequence could happen:

1. TSDRegistryExT::initLInkerInitialized()
2. Mutex.lock()
3. TSDRegistryExT() implicit constructor runs, calling Mutex's ctor.

So now Mutex is unlocked even though unlock was never called.

(also applies to the other registry)

60

Ok, but will a single registry ever have multiple allocators? If not, we don't need to pass the allocator to the therad-specific initThread and should instead pass it to the global init.

64

Thanks for the explanation. Would it make sense to put that comment here, or is the combined a better place?

[cryptoad]

lib/scudo/standalone/tsd_exclusive.h

28

Damn, I failed.
The other option was to implement a call_once type function, that would be specific to the registry, using an atomic cas (like llvm::call_once) and a little spin, and not have a mutex at all.
I can't use std::call_once and pthread_once, but if you have another idea, I am happy to oblige.

60

Ah, I see what you mean. That is indeed the case, I'll reorganize that.

64

I'll add a short version of it here.

[morehouse]

lib/scudo/standalone/tsd_exclusive.h

28

Is there a reason we need the init loop to begin with? It would be nice to have a one-way initialization instead.

lib/scudo/standalone/tsd_shared.h

23

Seems this is a consequence of the tight coupling between registry and allocator.

I think ideally we'd have a one-way dependency so either allocator-has-a-registry, or registry-has-an-allocator. Then the initialization only ever makes sense one-way.

Loose coupling in general would also make it easier to understand each piece in isolation, which makes future code maintenance much easier.

[cryptoad]

lib/scudo/standalone/tsd_exclusive.h

28

The loop isn't required, it's sort of a residue at my attempt at lazy initializing everything, while keeping the init/initLinkerInitialized construct.
The initialization flow should be as follows:

1. we have a zero-initialized Combined structure
2. someone calls allocate() on that Combined
3. this does initThread, which calls initThread in the registry (since it's really a registry thing)
4. thread isn't initialized, neither is the allocator, so call the combined's initOnce
5. parse the flag, initialize the internal structures

[cryptoad]

This, hopefully, detangles a bit the initialization process for
the registry. In order to do that, we get rid of the mutex in favor
of our own call_once type construct (losely inspired by
llvm::call_once). We move the initialization code into
initLinkerInitialized.

This gets rid of the potential init loop, and hopefully makes the
code structure more coherent.

Additional, I had some local flakes while testing this, and realized
that I messed up some region log sizes in the tests. Since we are
modifyin that file here, unflake the test as well.

[cryptoad]

Add a comment for MinimalInit.

[cryptoad]

Change the Primary test again to be more forgiving to OOM.

[cryptoad]

Disable a test on 32-bit for now: we are running out of address space.
This will be fixed in a subsequent CL, but affects the tests of this one.

[cryptoad]

Ping pretty please! I think all the comments were addressed one way or another. The new structure doesn't "cycle" anymore.

[morehouse]

lib/scudo/standalone/tests/tsd_test.cc

30

Don't we need TSDRegistry.initLinkerInitialized()?

lib/scudo/standalone/tsd_exclusive.h

63

With this change, we now expect initialization to only happen through initThreadMaybe, right? Should we even keep initLinkerInitialized around then?

68

Is the fence necessary? Not an atomic expert, but doesn't the load achieve the same thing?

78

I think things would be much simpler if we can use a StaticSpinMutex. We're assuming we're zero-initialized anyway, so the mutex will be initialized before use as well.

[cryptoad]

lib/scudo/standalone/tsd_exclusive.h

63

If I did things correctly, the Registry can also now be initialized with initLinkerInitialized, which will carry out the "once" initialization. Then initThreadMaybe will do the thread initialization skipping the "once".
My current plan is to have everything lazy initialized the first time someone calls malloc (or whatever else), meaning the first initThreadMaybe will carry the "once" initialization, but the alternative of pre-initializing is now offered.

68

Not an atomic expert either, it's inspired from from llvm::call_once which does this.
I figured it had a purpose and that I should keep it.

78

I agree. But the issue that would come up using StaticSpinMutex, is that we would land in initLinkerInitializedwhere we should call the Mutex's initLinkerInitialized (to be consistent), while holding it. It works because it's a noop, but doesn't make sense from a logical perspective.
Implementing it this way (with our wannabe call_once) allows us to not have such a loophole.

[morehouse]

lib/scudo/standalone/tsd_exclusive.h

78

Can we simply remove initLinkerInitialized from StaticSpinMutex, as we have in sanitizer_common? We could possibly add a comment that calls to init are only necessary if the mutex is not linker initialized.

[cryptoad]

lib/scudo/standalone/tsd_exclusive.h

78

Definitely can!

[cryptoad]

As discussed through comments, re-introduce a StaticSpinMutex to the
"once" initialization of the TSD registry, and remove its no-op
initLinkerInitialized in the various places it was used.
Also adds a test to exercise the path of "direct" initialization
via calling initLinkerInitialized directly on the registry.

[morehouse]

lib/scudo/standalone/tsd_exclusive.h

70

This is unguarded by the mutex. I think what we need to do is call initOnce unconditionally.

lib/scudo/standalone/tsd_shared.h

100

This is also unguarded by mutex.

[cryptoad]

lib/scudo/standalone/tsd_exclusive.h

70

My understanding is that it's not needed.
If it's true then initialization has been carried, if false, then we lock the mutex in initOnce.
It's assuming that the bool can't be partially modified, but I guess I should enforce that with an atomic_u8 instead.

[cryptoad]

Use an atomic_u8 for Initialized instead of a bool.

[cryptoad]

clang-format'ing the code.

[morehouse]

lib/scudo/standalone/tsd_exclusive.h

70

I think another issue is that the compiler is allowed to reorder instructions within a locking context. So it is possible that Initialized = true but initialization hasn't fully finished. If we always acquire the lock before accessing Initialized, we guarantee that we're fully initialized since compiler can't reorder instructions outside of the locking context.

[cryptoad]

Always call initOnce, and rename it initOnceMaybe to reflect that
initialization might not necessarily occur if it already happened.

lib/scudo/standalone/CMakeLists.txt

set(SCUDO_HEADERS
   primary64.h
   quarantine.h
   release.h
   report.h
   secondary.h
   size_class_map.h
   stats.h
   string_utils.h
+  tsd.h
+  tsd_exclusive.h
+  tsd_shared.h
   vector.h)
 
 if(COMPILER_RT_HAS_SCUDO_STANDALONE)
   add_compiler_rt_object_libraries(RTScudoStandalone
     ARCHS ${SCUDO_STANDALONE_SUPPORTED_ARCH}
     SOURCES ${SCUDO_SOURCES}
     ADDITIONAL_HEADERS ${SCUDO_HEADERS}
     CFLAGS ${SCUDO_CFLAGS})

lib/scudo/standalone/internal_defs.h

 #include "platform.h"
 
 #include <stdint.h>
 
 #ifndef SCUDO_DEBUG
 #define SCUDO_DEBUG 0
 #endif
 
-#define ARRAY_SIZE(a) (sizeof(a) / sizeof((a)[0]))
+#define ARRAY_SIZE(A) (sizeof(A) / sizeof((A)[0]))
 
 // String related macros.
 
 #define STRINGIFY_(S) #S
 #define STRINGIFY(S) STRINGIFY_(S)
 #define CONCATENATE_(S, C) S##C
 #define CONCATENATE(S, C) CONCATENATE_(S, C)
 

lib/scudo/standalone/tests/CMakeLists.txt

set(SCUDO_UNIT_TEST_SOURCES
   primary_test.cc
   quarantine_test.cc
   release_test.cc
   report_test.cc
   secondary_test.cc
   size_class_map_test.cc
   stats_test.cc
   strings_test.cc
+  tsd_test.cc
   vector_test.cc
   scudo_unit_test_main.cc)
 
 add_scudo_unittest(ScudoUnitTest
   SOURCES ${SCUDO_UNIT_TEST_SOURCES})

lib/scudo/standalone/tests/primary_test.cc

while (!V.empty()) {
     V.pop_back();
   }
   Cache.destroy(nullptr);
 }
 
 template <typename Primary> static void testPrimaryThreaded() {
   std::unique_ptr<Primary> Allocator(new Primary);
   Allocator->init(/*ReleaseToOsInterval=*/-1);
-  std::thread Threads[10];
-  for (scudo::uptr I = 0; I < 10U; I++)
+  std::thread Threads[32];
+  for (scudo::uptr I = 0; I < ARRAY_SIZE(Threads); I++)
     Threads[I] = std::thread(performAllocations<Primary>, Allocator.get());
   {
     std::unique_lock<std::mutex> Lock(Mutex);
     Ready = true;
     Cv.notify_all();
   }
   for (auto &T : Threads)
     T.join();

lib/scudo/standalone/tests/tsd_test.cc

+//===-- tsd_test.cc ---------------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "tsd_exclusive.h"
+#include "tsd_shared.h"
+
+#include "gtest/gtest.h"
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+// We mock out an allocator with a TSD registry, mostly using empty stubs. The
+// cache contains a single volatile uptr, to be able to test that several
+// concurrent threads will not access or modify the same cache at the same time.
+template <class Config> class MockAllocator {
+public:
+  using ThisT = MockAllocator<Config>;
+  using TSDRegistryT = typename Config::template TSDRegistryT<ThisT>;
+  using CacheT = struct MockCache { volatile scudo::uptr Canary; };
+  using QuarantineCacheT = struct MockQuarantine {};
+
+  void initLinkerInitialized() {
+    TSDRegistry.initLinkerInitialized();
+    // This should only be called once by the registry.
+    EXPECT_FALSE(Initialized);
+    Initialized = true;
+  }
+  void reset() { memset(this, 0, sizeof(*this)); }
+
+  void initCache(CacheT *Cache) { memset(Cache, 0, sizeof(*Cache)); }
+  void commitBack(scudo::TSD<MockAllocator> *TSD) {}
+  TSDRegistryT *getTSDRegistry() { return &TSDRegistry; }
+
+  bool isInitialized() { return Initialized; }
+
+private:
+  bool Initialized;
+  TSDRegistryT TSDRegistry;
+};
+
+template <class AllocatorT> static void testRegistry() {
+  std::unique_ptr<AllocatorT> Allocator(new AllocatorT);
+  Allocator.get()->reset();
+  EXPECT_FALSE(Allocator.get()->isInitialized());
+
+  auto Registry = Allocator.get()->getTSDRegistry();
+  Registry->initThreadMaybe(Allocator.get(), /*MinimalInit=*/true);
+  EXPECT_TRUE(Allocator.get()->isInitialized());
+
+  bool UnlockRequired;
+  auto TSD = Registry->getTSDAndLock(&UnlockRequired);
+  EXPECT_NE(TSD, nullptr);
+  EXPECT_EQ(TSD->Cache.Canary, 0U);
+  if (UnlockRequired)
+    TSD->unlock();
+
+  Registry->initThreadMaybe(Allocator.get(), /*MinimalInit=*/false);
+  TSD = Registry->getTSDAndLock(&UnlockRequired);
+  EXPECT_NE(TSD, nullptr);
+  EXPECT_EQ(TSD->Cache.Canary, 0U);
+  memset(&TSD->Cache, 0x42, sizeof(TSD->Cache));
+  if (UnlockRequired)
+    TSD->unlock();
+}
+
+struct OneCache {
+  template <class Allocator>
+  using TSDRegistryT = scudo::TSDRegistrySharedT<Allocator, 1U>;
+};
+
+struct SharedCaches {
+  template <class Allocator>
+  using TSDRegistryT = scudo::TSDRegistrySharedT<Allocator, 16U>;
+};
+
+struct ExclusiveCaches {
+  template <class Allocator>
+  using TSDRegistryT = scudo::TSDRegistryExT<Allocator>;
+};
+
+TEST(ScudoTSDTest, TSDRegistryBasic) {
+  testRegistry<MockAllocator<OneCache>>();
+  testRegistry<MockAllocator<SharedCaches>>();
+  testRegistry<MockAllocator<ExclusiveCaches>>();
+}
+
+static std::mutex Mutex;
+static std::condition_variable Cv;
+static bool Ready = false;
+
+template <typename AllocatorT> static void stressCache(AllocatorT *Allocator) {
+  auto Registry = Allocator->getTSDRegistry();
+  {
+    std::unique_lock<std::mutex> Lock(Mutex);
+    while (!Ready)
+      Cv.wait(Lock);
+  }
+  Registry->initThreadMaybe(Allocator, /*MinimalInit=*/false);
+  bool UnlockRequired;
+  auto TSD = Registry->getTSDAndLock(&UnlockRequired);
+  EXPECT_NE(TSD, nullptr);
+  // For an exclusive TSD, the cache should be empty. We cannot guarantee the
+  // same for a shared TSD.
+  if (!UnlockRequired)
+    EXPECT_EQ(TSD->Cache.Canary, 0U);
+  // Transform the thread id to a uptr to use it as canary.
+  const scudo::uptr Canary = static_cast<scudo::uptr>(
+      std::hash<std::thread::id>{}(std::this_thread::get_id()));
+  TSD->Cache.Canary = Canary;
+  // Loop a few times to make sure that a concurrent thread isn't modifying it.
+  for (scudo::uptr I = 0; I < 4096U; I++)
+    EXPECT_EQ(TSD->Cache.Canary, Canary);
+  if (UnlockRequired)
+    TSD->unlock();
+}
+
+template <class AllocatorT> static void testRegistryThreaded() {
+  std::unique_ptr<AllocatorT> Allocator(new AllocatorT);
+  Allocator.get()->reset();
+  std::thread Threads[32];
+  for (scudo::uptr I = 0; I < ARRAY_SIZE(Threads); I++)
+    Threads[I] = std::thread(stressCache<AllocatorT>, Allocator.get());
+  {
+    std::unique_lock<std::mutex> Lock(Mutex);
+    Ready = true;
+    Cv.notify_all();
+  }
+  for (auto &T : Threads)
+    T.join();
+}
+
+TEST(ScudoTSDTest, TSDRegistryThreaded) {
+  testRegistryThreaded<MockAllocator<OneCache>>();
+  testRegistryThreaded<MockAllocator<SharedCaches>>();
+  testRegistryThreaded<MockAllocator<ExclusiveCaches>>();
+}

lib/scudo/standalone/tsd.h

+//===-- tsd.h ---------------------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SCUDO_TSD_H_
+#define SCUDO_TSD_H_
+
+#include "atomic_helpers.h"
+#include "common.h"
+#include "mutex.h"
+
+#include <limits.h> // for PTHREAD_DESTRUCTOR_ITERATIONS
+
+namespace scudo {
+
+template <class Allocator> struct ALIGNED(SCUDO_CACHE_LINE_SIZE) TSD {
+  typename Allocator::CacheT Cache;
+  typename Allocator::QuarantineCacheT QuarantineCache;
+  u8 DestructorIterations;
+
+  void initLinkerInitialized(Allocator *Instance) {
+    Instance->initCache(&Cache);
+    DestructorIterations = PTHREAD_DESTRUCTOR_ITERATIONS;
+  }
+  void init(Allocator *Instance) {
+    memset(this, 0, sizeof(*this));
+    initLinkerInitialized(Instance);
+  }
+
+  void commitBack(Allocator *Instance) { Instance->commitBack(this); }
+
+  INLINE bool tryLock() {
+    if (Mutex.tryLock()) {
+      atomic_store_relaxed(&Precedence, 0);
+      return true;
+    }
+    if (atomic_load_relaxed(&Precedence) == 0)
+      atomic_store_relaxed(
+          &Precedence,
+          static_cast<uptr>(getMonotonicTime() >> FIRST_32_SECOND_64(16, 0)));
+    return false;
+  }
+  INLINE void lock() {
+    atomic_store_relaxed(&Precedence, 0);
+    Mutex.lock();
+  }
+  INLINE void unlock() { Mutex.unlock(); }
+  INLINE uptr getPrecedence() { return atomic_load_relaxed(&Precedence); }
+
+private:
+  StaticSpinMutex Mutex;
+  atomic_uptr Precedence;
+};
+
+} // namespace scudo
+
+#endif // SCUDO_TSD_H_

lib/scudo/standalone/tsd_exclusive.h

+//===-- tsd_exclusive.h -----------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SCUDO_TSD_EXCLUSIVE_H_
+#define SCUDO_TSD_EXCLUSIVE_H_
+
+#include "tsd.h"
+
+#include <pthread.h>
+
+namespace scudo {
+
+enum class ThreadState : u8 {
+  NotInitialized = 0,
+  Initialized,
+  TornDown,
+};
+
+template <class Allocator> void teardownThread(void *Ptr);
+
+template <class Allocator> struct TSDRegistryExT {
+  void initLinkerInitialized() { Mutex.initLinkerInitialized(); }
+  void init() { memset(this, 0, sizeof(*this)); }
+
+  ALWAYS_INLINE void initThreadMaybe(Allocator *Instance, bool MinimalInit) {
+    if (LIKELY(State != ThreadState::NotInitialized))
+      return;
+    initThread(Instance, MinimalInit);
+  }
+
+  ALWAYS_INLINE TSD<Allocator> *getTSDAndLock(bool *UnlockRequired) {
+    if (LIKELY(State == ThreadState::Initialized)) {
+      *UnlockRequired = false;
+      return &ThreadTSD;
+    }
+    DCHECK(FallbackTSD);
+    FallbackTSD->lock();
+    *UnlockRequired = true;
+    return FallbackTSD;
+  }
+
+private:
+  NOINLINE void initOnce(Allocator *Instance) {
+    SpinMutexLock L(&Mutex);
+    if (OnceDone)
+      return;
+    CHECK_EQ(pthread_key_create(&PThreadKey, teardownThread<Allocator>), 0);
+    Instance->initLinkerInitialized();
+    FallbackTSD = reinterpret_cast<TSD<Allocator> *>(
+        map(nullptr, sizeof(TSD<Allocator>), "scudo:tsd"));
+    FallbackTSD->initLinkerInitialized(Instance);
+    OnceDone = true;
+  }
+
+  NOINLINE void initThread(Allocator *Instance, bool MinimalInit) {
+    if (UNLIKELY(!OnceDone))
+      initOnce(Instance);
+    if (MinimalInit)
+      return;
+    CHECK_EQ(
+        pthread_setspecific(PThreadKey, reinterpret_cast<void *>(Instance)), 0);
+    ThreadTSD.initLinkerInitialized(Instance);
+    State = ThreadState::Initialized;
+  }
+
+  StaticSpinMutex Mutex;
+  bool OnceDone;
+  pthread_key_t PThreadKey;
+  TSD<Allocator> *FallbackTSD;
+  static THREADLOCAL ThreadState State;
+  static THREADLOCAL TSD<Allocator> ThreadTSD;
+
+  friend void teardownThread<Allocator>(void *Ptr);
+};
+
+template <class Allocator>
+THREADLOCAL TSD<Allocator> TSDRegistryExT<Allocator>::ThreadTSD;
+template <class Allocator>
+THREADLOCAL ThreadState TSDRegistryExT<Allocator>::State;
+
+template <class Allocator> void teardownThread(void *Ptr) {
+  typedef TSDRegistryExT<Allocator> TSDRegistryT;
+  Allocator *Instance = reinterpret_cast<Allocator *>(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.
+  const u8 N = TSDRegistryT::ThreadTSD.DestructorIterations;
+  if (N > 1U) {
+    TSDRegistryT::ThreadTSD.DestructorIterations = N - 1;
+    // If pthread_setspecific fails, we will go ahead with the teardown.
+    if (LIKELY(pthread_setspecific(Instance->getTSDRegistry()->PThreadKey,
+                                   Ptr) == 0))
+      return;
+  }
+  TSDRegistryT::ThreadTSD.commitBack(Instance);
+  TSDRegistryT::State = ThreadState::TornDown;
+}
+
+} // namespace scudo
+
+#endif // SCUDO_TSD_EXCLUSIVE_H_

lib/scudo/standalone/tsd_shared.h

+//===-- tsd_shared.h --------------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SCUDO_TSD_SHARED_H_
+#define SCUDO_TSD_SHARED_H_
+
+#include "linux.h" // for getAndroidTlsPtr()
+#include "tsd.h"
+
+#include <pthread.h>
+
+namespace scudo {
+
+template <class Allocator, u32 MaxTSDCount> struct TSDRegistrySharedT {
+  void initLinkerInitialized() { Mutex.initLinkerInitialized(); }
+  void init() { memset(this, 0, sizeof(*this)); }
+
+  ALWAYS_INLINE void initThreadMaybe(Allocator *Instance,
+                                     UNUSED bool MinimalInit) {
+    if (LIKELY(getCurrentTSD()))
+      return;
+    initThread(Instance);
+  }
+
+  ALWAYS_INLINE TSD<Allocator> *getTSDAndLock(bool *UnlockRequired) {
+    TSD<Allocator> *TSD = getCurrentTSD();
+    DCHECK(TSD);
+    *UnlockRequired = true;
+    // Try to lock the currently associated context.
+    if (TSD->tryLock())
+      return TSD;
+    // If that fails, go down the slow path.
+    return getTSDAndLockSlow(TSD);
+  }
+
+private:
+  NOINLINE void initOnce(Allocator *Instance) {
+    SpinMutexLock L(&Mutex);
+    if (OnceDone)
+      return;
+    CHECK_EQ(pthread_key_create(&PThreadKey, nullptr), 0); // For non-TLS
+    Instance->initLinkerInitialized();
+    NumberOfTSDs = Min(Max(1U, getNumberOfCPUs()), MaxTSDCount);
+    TSDs = reinterpret_cast<TSD<Allocator> *>(
+        map(nullptr, sizeof(TSD<Allocator>) * NumberOfTSDs, "scudo:tsd"));
+    for (u32 I = 0; I < NumberOfTSDs; I++) {
+      TSDs[I].initLinkerInitialized(Instance);
+      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;
+    }
+    OnceDone = true;
+  }
+
+  ALWAYS_INLINE void setCurrentTSD(TSD<Allocator> *CurrentTSD) {
+#if SCUDO_ANDROID
+    *getAndroidTlsPtr() = reinterpret_cast<uptr>(CurrentTSD);
+#elif SCUDO_LINUX
+    ThreadTSD = CurrentTSD;
+#else
+    CHECK_EQ(
+        pthread_setspecific(PThreadKey, reinterpret_cast<void *>(CurrentTSD)),
+        0);
+#endif
+  }
+
+  ALWAYS_INLINE TSD<Allocator> *getCurrentTSD() {
+#if SCUDO_ANDROID
+    return reinterpret_cast<TSD<Allocator> *>(*getAndroidTlsPtr());
+#elif SCUDO_LINUX
+    return ThreadTSD;
+#else
+    return reinterpret_cast<TSD<Allocator> *>(pthread_getspecific(PThreadKey));
+#endif
+  }
+
+  NOINLINE void initThread(Allocator *Instance) {
+    if (UNLIKELY(!OnceDone))
+      initOnce(Instance);
+    // Initial context assignment is done in a plain round-robin fashion.
+    const u32 Index = atomic_fetch_add(&CurrentIndex, 1, memory_order_relaxed);
+    setCurrentTSD(&TSDs[Index % NumberOfTSDs]);
+  }
+
+  NOINLINE TSD<Allocator> *getTSDAndLockSlow(TSD<Allocator> *CurrentTSD) {
+    if (MaxTSDCount > 1U && NumberOfTSDs > 1U) {
+      // Use the Precedence of the current TSD as our random seed. Since we are
+      // in the slow path, it means that tryLock failed, and as a result it's
+      // very likely that said Precedence is non-zero.
+      u32 RandState = static_cast<u32>(CurrentTSD->getPrecedence());
+      const u32 R = getRandomU32(&RandState);
+      const u32 Inc = CoPrimes[R % NumberOfCoPrimes];
+      u32 Index = R % NumberOfTSDs;
+      uptr LowestPrecedence = UINTPTR_MAX;
+      TSD<Allocator> *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].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;
+        }
+        Index += Inc;
+        if (Index >= NumberOfTSDs)
+          Index -= NumberOfTSDs;
+      }
+      if (CandidateTSD) {
+        CandidateTSD->lock();
+        setCurrentTSD(CandidateTSD);
+        return CandidateTSD;
+      }
+    }
+    // Last resort, stick with the current one.
+    CurrentTSD->lock();
+    return CurrentTSD;
+  }
+
+  StaticSpinMutex Mutex;
+  bool OnceDone;
+  atomic_u32 CurrentIndex;
+  u32 CoPrimes[MaxTSDCount];
+  u32 NumberOfCoPrimes;
+  pthread_key_t PThreadKey;
+  TSD<Allocator> *TSDs;
+  u32 NumberOfTSDs;
+#if SCUDO_LINUX && !SCUDO_ANDROID
+  static THREADLOCAL TSD<Allocator> *ThreadTSD;
+#endif
+};
+
+#if SCUDO_LINUX && !SCUDO_ANDROID
+template <class Allocator, u32 MaxTSDCount>
+THREADLOCAL TSD<Allocator>
+    *TSDRegistrySharedT<Allocator, MaxTSDCount>::ThreadTSD;
+#endif
+
+} // namespace scudo
+
+#endif // SCUDO_TSD_SHARED_H_