[scudo][standalone] Implement TSD registry disabling

Summary:
In order to implement `malloc_{enable|disable}` we were just disabling
(or really locking) the Primary and the Secondary. That meant that
allocations could still be serviced from the TSD as long as the cache
wouldn't have to be filled from the Primary.

This wasn't working out for Android tests, so this change implements
registry disabling (eg: locking) so that `getTSDAndLock` doesn't
return a TSD if the allocator is disabled. This also means that the
Primary doesn't have to be disabled in this situation.

For the Shared Registry, we loop through all the TSDs and lock them.
For the Exclusive Registry, we add a `Disabled` boolean to the Registry
that forces `getTSDAndLock` to use the Fallback TSD instead of the
thread local one. Disabling the Registry is then done by locking the
Fallback TSD and setting the boolean in question (I don't think this
needed an atomic variable but I might be wrong).

I clang-formatted the whole thing as usual hence the couple of extra
whiteline changes in this CL.

Reviewers: cferris, pcc, hctim, morehouse, eugenis

Subscribers: jfb, #sanitizers, llvm-commits

Tags: #sanitizers, #llvm

Differential Revision: https://reviews.llvm.org/D71719

1 files changed, 15 insertions, 1 deletions

diff --git a/compiler-rt/lib/scudo/standalone/tsd_exclusive.h b/compiler-rt/lib/scudo/standalone/tsd_exclusive.h
index 971ae4857fc..89b001a739c 100644
--- a/compiler-rt/lib/scudo/standalone/tsd_exclusive.h
+++ b/compiler-rt/lib/scudo/standalone/tsd_exclusive.h
@@ -48,7 +48,8 @@ template <class Allocator> struct TSDRegistryExT {
   }
 
   ALWAYS_INLINE TSD<Allocator> *getTSDAndLock(bool *UnlockRequired) {
-    if (LIKELY(State == ThreadState::Initialized)) {
+    if (LIKELY(State == ThreadState::Initialized &&
+               !atomic_load(&Disabled, memory_order_acquire))) {
       *UnlockRequired = false;
       return &ThreadTSD;
     }
@@ -58,6 +59,18 @@ template <class Allocator> struct TSDRegistryExT {
     return FallbackTSD;
   }
 
+  // To disable the exclusive TSD registry, we effectively lock the fallback TSD
+  // and force all threads to attempt to use it instead of their local one.
+  void disable() {
+    FallbackTSD->lock();
+    atomic_store(&Disabled, 1U, memory_order_release);
+  }
+
+  void enable() {
+    atomic_store(&Disabled, 0U, memory_order_release);
+    FallbackTSD->unlock();
+  }
+
 private:
   void initOnceMaybe(Allocator *Instance) {
     ScopedLock L(Mutex);
@@ -81,6 +94,7 @@ private:
 
   pthread_key_t PThreadKey;
   bool Initialized;
+  atomic_u8 Disabled;
   TSD<Allocator> *FallbackTSD;
   HybridMutex Mutex;
   static THREADLOCAL ThreadState State;