[scudo] Add Android support

Summary:
This change adds Android support to the allocator (but doesn't yet enable it in
the cmake config), and should be the last fragment of the rewritten change
D31947.

Android has more memory constraints than other platforms, so the idea of a
unique context per thread would not have worked. The alternative chosen is to
allocate a set of contexts based on the number of cores on the machine, and
share those contexts within the threads. Contexts can be dynamically reassigned
to threads to prevent contention, based on a scheme suggested by @dvyuokv in
the initial review.

Additionally, given that Android doesn't support ELF TLS (only emutls for now),
we use the TSan TLS slot to make things faster: Scudo is mutually exclusive
with other sanitizers so this shouldn't cause any problem.

An additional change made here, is replacing `thread_local` by `THREADLOCAL`
and using the initial-exec thread model in the non-Android version to prevent
extraneous weak definition and checks on the relevant variables.

Reviewers: kcc, dvyukov, alekseyshl

Reviewed By: alekseyshl

Subscribers: srhines, mgorny, llvm-commits

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

git-svn-id: https://llvm.org/svn/llvm-project/compiler-rt/trunk@302300 91177308-0d34-0410-b5e6-96231b3b80d8

10 files changed, 261 insertions, 20 deletions

diff --git a/lib/scudo/CMakeLists.txt b/lib/scudo/CMakeLists.txt
index 3a8f4ae4f..14c199fa8 100644
--- a/lib/scudo/CMakeLists.txt
+++ b/lib/scudo/CMakeLists.txt
@@ -14,6 +14,7 @@ set(SCUDO_SOURCES
   scudo_interceptors.cpp
   scudo_new_delete.cpp
   scudo_termination.cpp
+  scudo_tls_android.cpp
   scudo_tls_linux.cpp
   scudo_utils.cpp)
 
diff --git a/lib/scudo/scudo_allocator.cpp b/lib/scudo/scudo_allocator.cpp
index 2ccdcd903..5420fc964 100644
--- a/lib/scudo/scudo_allocator.cpp
+++ b/lib/scudo/scudo_allocator.cpp
@@ -368,11 +368,12 @@ struct ScudoAllocator {
     void *Ptr;
     uptr Salt;
     uptr AllocationAlignment = FromPrimary ? MinAlignment : Alignment;
-    ScudoThreadContext *ThreadContext = getThreadContext();
+    ScudoThreadContext *ThreadContext = getThreadContextAndLock();
     if (LIKELY(ThreadContext)) {
       Salt = getPrng(ThreadContext)->getNext();
       Ptr = BackendAllocator.Allocate(getAllocatorCache(ThreadContext),
                                       NeededSize, AllocationAlignment);
+      ThreadContext->unlock();
     } else {
       SpinMutexLock l(&FallbackMutex);
       Salt = FallbackPrng.getNext();
@@ -434,9 +435,10 @@ struct ScudoAllocator {
     if (BypassQuarantine) {
       Chunk->eraseHeader();
       void *Ptr = Chunk->getAllocBeg(Header);
-      ScudoThreadContext *ThreadContext = getThreadContext();
+      ScudoThreadContext *ThreadContext = getThreadContextAndLock();
       if (LIKELY(ThreadContext)) {
         getBackendAllocator().Deallocate(getAllocatorCache(ThreadContext), Ptr);
+        ThreadContext->unlock();
       } else {
         SpinMutexLock Lock(&FallbackMutex);
         getBackendAllocator().Deallocate(&FallbackAllocatorCache, Ptr);
@@ -445,12 +447,13 @@ struct ScudoAllocator {
       UnpackedHeader NewHeader = *Header;
       NewHeader.State = ChunkQuarantine;
       Chunk->compareExchangeHeader(&NewHeader, Header);
-      ScudoThreadContext *ThreadContext = getThreadContext();
+      ScudoThreadContext *ThreadContext = getThreadContextAndLock();
       if (LIKELY(ThreadContext)) {
         AllocatorQuarantine.Put(getQuarantineCache(ThreadContext),
                                 QuarantineCallback(
                                     getAllocatorCache(ThreadContext)),
                                 Chunk, Size);
+        ThreadContext->unlock();
       } else {
         SpinMutexLock l(&FallbackMutex);
         AllocatorQuarantine.Put(&FallbackQuarantineCache,
diff --git a/lib/scudo/scudo_allocator.h b/lib/scudo/scudo_allocator.h
index 2cac2de71..f159deffb 100644
--- a/lib/scudo/scudo_allocator.h
+++ b/lib/scudo/scudo_allocator.h
@@ -72,7 +72,13 @@ const uptr AlignedChunkHeaderSize =
 
 #if SANITIZER_CAN_USE_ALLOCATOR64
 const uptr AllocatorSpace = ~0ULL;
-const uptr AllocatorSize = 0x40000000000ULL;  // 4TB.
+# if defined(__aarch64__) && SANITIZER_ANDROID
+const uptr AllocatorSize = 0x4000000000ULL;  // 256G.
+# elif defined(__aarch64__)
+const uptr AllocatorSize = 0x10000000000ULL;  // 1T.
+# else
+const uptr AllocatorSize = 0x40000000000ULL;  // 4T.
+# endif
 typedef DefaultSizeClassMap SizeClassMap;
 struct AP {
   static const uptr kSpaceBeg = AllocatorSpace;
diff --git a/lib/scudo/scudo_tls.h b/lib/scudo/scudo_tls.h
index 0d7d1bffd..f6039bebe 100644
--- a/lib/scudo/scudo_tls.h
+++ b/lib/scudo/scudo_tls.h
@@ -19,10 +19,16 @@
 #include "scudo_allocator.h"
 #include "scudo_utils.h"
 
+#include "sanitizer_common/sanitizer_linux.h"
+#include "sanitizer_common/sanitizer_platform.h"
+
 namespace __scudo {
 
-struct ALIGNED(64) ScudoThreadContext {
- public:
+// Platform specific base thread context definitions.
+#include "scudo_tls_context_android.inc"
+#include "scudo_tls_context_linux.inc"
+
+struct ALIGNED(64) ScudoThreadContext : public ScudoThreadContextPlatform {
   AllocatorCache Cache;
   Xorshift128Plus Prng;
   uptr QuarantineCachePlaceHolder[4];
@@ -32,8 +38,9 @@ struct ALIGNED(64) ScudoThreadContext {
 
 void initThread();
 
-// Fastpath functions are defined in the following platform specific headers.
-#include "scudo_tls_linux.h"
+// Platform specific dastpath functions definitions.
+#include "scudo_tls_android.inc"
+#include "scudo_tls_linux.inc"
 
 }  // namespace __scudo
 
diff --git a/lib/scudo/scudo_tls_android.cpp b/lib/scudo/scudo_tls_android.cpp
new file mode 100644
index 000000000..0e3602b2f
--- /dev/null
+++ b/lib/scudo/scudo_tls_android.cpp
@@ -0,0 +1,95 @@
+//===-- scudo_tls_android.cpp -----------------------------------*- 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 implementation for Android.
+///
+//===----------------------------------------------------------------------===//
+
+#include "sanitizer_common/sanitizer_platform.h"
+
+#if SANITIZER_LINUX && SANITIZER_ANDROID
+
+#include "scudo_tls.h"
+
+#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;
+
+// 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 (int i = 0; i < NumberOfContexts; i++)
+    ThreadContexts[i].init();
+}
+
+void initThread() {
+  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);
+}
+
+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;
+    }
+  }
+  // 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;
+    }
+  }
+  if (LIKELY(Precedence != UINT64_MAX)) {
+    ThreadContext->lock();
+    *get_android_tls_ptr() = reinterpret_cast<uptr>(ThreadContext);
+    return ThreadContext;
+  }
+  // Last resort (can this happen?), stick with the current one.
+  ThreadContext =
+      reinterpret_cast<ScudoThreadContext *>(*get_android_tls_ptr());
+  ThreadContext->lock();
+  return ThreadContext;
+}
+
+}  // namespace __scudo
+
+#endif  // SANITIZER_LINUX && SANITIZER_ANDROID
diff --git a/lib/scudo/scudo_tls_android.inc b/lib/scudo/scudo_tls_android.inc
new file mode 100644
index 000000000..8ecad7a30
--- /dev/null
+++ b/lib/scudo/scudo_tls_android.inc
@@ -0,0 +1,44 @@
+//===-- 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() {
+  if (LIKELY(*get_android_tls_ptr()))
+    return;
+  initThread();
+}
+
+ScudoThreadContext *getThreadContextAndLockSlow();
+
+ALWAYS_INLINE ScudoThreadContext *getThreadContextAndLock() {
+  ScudoThreadContext *ThreadContext =
+      reinterpret_cast<ScudoThreadContext *>(*get_android_tls_ptr());
+  CHECK(ThreadContext);
+  // Try to lock the currently associated context.
+  if (ThreadContext->tryLock())
+    return ThreadContext;
+  // If it failed, go the slow path.
+  return getThreadContextAndLockSlow();
+}
+
+#endif  // SANITIZER_LINUX && SANITIZER_ANDROID
+
+#endif  // SCUDO_TLS_ANDROID_H_
diff --git a/lib/scudo/scudo_tls_context_android.inc b/lib/scudo/scudo_tls_context_android.inc
new file mode 100644
index 000000000..f1951319d
--- /dev/null
+++ b/lib/scudo/scudo_tls_context_android.inc
@@ -0,0 +1,54 @@
+//===-- 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 ScudoThreadContextPlatform {
+  INLINE bool tryLock() {
+    if (Mutex.TryLock()) {
+      atomic_store_relaxed(&SlowLockPrecedence, 0);
+      return true;
+    }
+    if (atomic_load_relaxed(&SlowLockPrecedence) == 0)
+      atomic_store_relaxed(&SlowLockPrecedence, NanoTime());
+    return false;
+  }
+
+  INLINE void lock() {
+    Mutex.Lock();
+    atomic_store_relaxed(&SlowLockPrecedence, 0);
+  }
+
+  INLINE void unlock() {
+    Mutex.Unlock();
+  }
+
+  INLINE u64 getSlowLockPrecedence() {
+    return atomic_load_relaxed(&SlowLockPrecedence);
+  }
+
+ private:
+  StaticSpinMutex Mutex;
+  atomic_uint64_t SlowLockPrecedence;
+};
+
+#endif  // SANITIZER_LINUX && SANITIZER_ANDROID
+
+#endif  // SCUDO_TLS_CONTEXT_ANDROID_INC_
diff --git a/lib/scudo/scudo_tls_context_linux.inc b/lib/scudo/scudo_tls_context_linux.inc
new file mode 100644
index 000000000..8d292bdbc
--- /dev/null
+++ b/lib/scudo/scudo_tls_context_linux.inc
@@ -0,0 +1,29 @@
+//===-- 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 ScudoThreadContextPlatform {
+  ALWAYS_INLINE void unlock() {}
+};
+
+#endif  // SANITIZER_LINUX && !SANITIZER_ANDROID
+
+#endif  // SCUDO_TLS_CONTEXT_LINUX_INC_
diff --git a/lib/scudo/scudo_tls_linux.cpp b/lib/scudo/scudo_tls_linux.cpp
index 3453367f8..5a9cc998b 100644
--- a/lib/scudo/scudo_tls_linux.cpp
+++ b/lib/scudo/scudo_tls_linux.cpp
@@ -14,7 +14,7 @@
 
 #include "sanitizer_common/sanitizer_platform.h"
 
-#if SANITIZER_LINUX
+#if SANITIZER_LINUX && !SANITIZER_ANDROID
 
 #include "scudo_tls.h"
 
@@ -26,8 +26,10 @@ namespace __scudo {
 static pthread_once_t GlobalInitialized = PTHREAD_ONCE_INIT;
 static pthread_key_t PThreadKey;
 
-thread_local ThreadState ScudoThreadState = ThreadNotInitialized;
-thread_local ScudoThreadContext ThreadLocalContext;
+__attribute__((tls_model("initial-exec")))
+THREADLOCAL ThreadState ScudoThreadState = ThreadNotInitialized;
+__attribute__((tls_model("initial-exec")))
+THREADLOCAL ScudoThreadContext ThreadLocalContext;
 
 static void teardownThread(void *Ptr) {
   uptr Iteration = reinterpret_cast<uptr>(Ptr);
@@ -59,4 +61,4 @@ void initThread() {
 
 }  // namespace __scudo
 
-#endif  // SANITIZER_LINUX
+#endif  // SANITIZER_LINUX && !SANITIZER_ANDROID
diff --git a/lib/scudo/scudo_tls_linux.h b/lib/scudo/scudo_tls_linux.inc
index 0994f2d7b..242ee3329 100644
--- a/lib/scudo/scudo_tls_linux.h
+++ b/lib/scudo/scudo_tls_linux.inc
@@ -1,4 +1,4 @@
-//===-- scudo_tls_linux.h ---------------------------------------*- C++ -*-===//
+//===-- scudo_tls_linux.inc -------------------------------------*- C++ -*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
@@ -19,17 +19,17 @@
 # error "This file must be included inside scudo_tls.h."
 #endif  // SCUDO_TLS_H_
 
-#include "sanitizer_common/sanitizer_platform.h"
-
-#if SANITIZER_LINUX
+#if SANITIZER_LINUX && !SANITIZER_ANDROID
 
 enum ThreadState : u8 {
   ThreadNotInitialized = 0,
   ThreadInitialized,
   ThreadTornDown,
 };
-extern thread_local ThreadState ScudoThreadState;
-extern thread_local ScudoThreadContext ThreadLocalContext;
+__attribute__((tls_model("initial-exec")))
+extern THREADLOCAL ThreadState ScudoThreadState;
+__attribute__((tls_model("initial-exec")))
+extern THREADLOCAL ScudoThreadContext ThreadLocalContext;
 
 ALWAYS_INLINE void initThreadMaybe() {
   if (LIKELY(ScudoThreadState != ThreadNotInitialized))
@@ -37,12 +37,12 @@ ALWAYS_INLINE void initThreadMaybe() {
   initThread();
 }
 
-ALWAYS_INLINE ScudoThreadContext *getThreadContext() {
+ALWAYS_INLINE ScudoThreadContext *getThreadContextAndLock() {
   if (UNLIKELY(ScudoThreadState == ThreadTornDown))
     return nullptr;
   return &ThreadLocalContext;
 }
 
-#endif  // SANITIZER_LINUX
+#endif  // SANITIZER_LINUX && !SANITIZER_ANDROID
 
 #endif  // SCUDO_TLS_LINUX_H_