[msan] Introduce MsanThread. Move thread-local allocator cache out of TLS.

This reduces .tbss from 109K down to almost nothing.

llvm-svn: 205618

compiler-rt/lib/msan/CMakeLists.txt

@@ -8,6 +8,7 @@ set(MSAN_RTL_SOURCES
   msan_linux.cc
   msan_new_delete.cc
   msan_report.cc
+  msan_thread.cc
   )
 
 set(MSAN_RTL_CFLAGS ${SANITIZER_COMMON_CFLAGS})

compiler-rt/lib/msan/msan.cc

@@ -13,6 +13,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "msan.h"
+#include "msan_thread.h"
 #include "sanitizer_common/sanitizer_atomic.h"
 #include "sanitizer_common/sanitizer_common.h"
 #include "sanitizer_common/sanitizer_flags.h"
@@ -59,8 +60,6 @@ THREADLOCAL u64 __msan_va_arg_overflow_size_tls;
 SANITIZER_INTERFACE_ATTRIBUTE
 THREADLOCAL u32 __msan_origin_tls;
 
-THREADLOCAL MsanStackBounds msan_stack_bounds;
-
 static THREADLOCAL int is_in_symbolizer;
 static THREADLOCAL int is_in_loader;
 
@@ -154,14 +153,14 @@ static void InitializeFlags(Flags *f, const char *options) {
 
 void GetStackTrace(StackTrace *stack, uptr max_s, uptr pc, uptr bp,
                    bool request_fast_unwind) {
-  if (!StackTrace::WillUseFastUnwind(request_fast_unwind)) {
+  MsanThread *t = GetCurrentThread();
+  if (!t || !StackTrace::WillUseFastUnwind(request_fast_unwind)) {
     // Block reports from our interceptors during _Unwind_Backtrace.
     SymbolizerScope sym_scope;
     return stack->Unwind(max_s, pc, bp, 0, 0, 0, request_fast_unwind);
   }
-  uptr stack_bottom = msan_stack_bounds.stack_addr;
-  uptr stack_top = stack_bottom + msan_stack_bounds.stack_size;
-  stack->Unwind(max_s, pc, bp, 0, stack_top, stack_bottom, request_fast_unwind);
+  stack->Unwind(max_s, pc, bp, 0, t->stack_top(), t->stack_bottom(),
+                request_fast_unwind);
 }
 
 void PrintWarning(uptr pc, uptr bp) {
@@ -315,10 +314,12 @@ void __msan_init() {
   Symbolizer::Init(common_flags()->external_symbolizer_path);
   Symbolizer::Get()->AddHooks(EnterSymbolizer, ExitSymbolizer);
 
-  GetThreadStackAndTls(/* main */ true, &msan_stack_bounds.stack_addr,
-                       &msan_stack_bounds.stack_size,
-                       &msan_stack_bounds.tls_addr,
-                       &msan_stack_bounds.tls_size);
+  MsanTSDInit(MsanTSDDtor);
+
+  MsanThread *main_thread = MsanThread::Create(0, 0);
+  SetCurrentThread(main_thread);
+  main_thread->ThreadStart();
+
   VPrintf(1, "MemorySanitizer init done\n");
 
   msan_init_is_running = 0;

compiler-rt/lib/msan/msan.h

@@ -128,6 +128,11 @@ class ScopedThreadLocalStateBackup {
 
 extern void (*death_callback)(void);
 
+void MsanTSDInit(void (*destructor)(void *tsd));
+void *MsanTSDGet();
+void MsanTSDSet(void *tsd);
+void MsanTSDDtor(void *tsd);
+
 }  // namespace __msan
 
 #define MSAN_MALLOC_HOOK(ptr, size) \
@@ -135,11 +140,4 @@ extern void (*death_callback)(void);
 #define MSAN_FREE_HOOK(ptr) \
   if (&__msan_free_hook) __msan_free_hook(ptr)
 
-struct MsanStackBounds {
-  uptr stack_addr, stack_size;
-  uptr tls_addr, tls_size;
-};
-
-extern THREADLOCAL MsanStackBounds msan_stack_bounds;
-
 #endif  // MSAN_H

compiler-rt/lib/msan/msan_allocator.cc

@@ -15,6 +15,8 @@
 #include "sanitizer_common/sanitizer_allocator.h"
 #include "sanitizer_common/sanitizer_stackdepot.h"
 #include "msan.h"
+#include "msan_allocator.h"
+#include "msan_thread.h"
 
 namespace __msan {
 
@@ -48,8 +50,9 @@ typedef LargeMmapAllocator<MsanMapUnmapCallback> SecondaryAllocator;
 typedef CombinedAllocator<PrimaryAllocator, AllocatorCache,
                           SecondaryAllocator> Allocator;
 
-static THREADLOCAL AllocatorCache cache;
 static Allocator allocator;
+static AllocatorCache fallback_allocator_cache;
+static SpinMutex fallback_mutex;
 
 static int inited = 0;
 
@@ -60,35 +63,51 @@ static inline void Init() {
   allocator.Init();
 }
 
-void MsanAllocatorThreadFinish() {
-  allocator.SwallowCache(&cache);
+AllocatorCache *GetAllocatorCache(MsanThreadLocalMallocStorage *ms) {
+  CHECK(ms);
+  CHECK_LE(sizeof(AllocatorCache), sizeof(ms->allocator_cache));
+  return reinterpret_cast<AllocatorCache *>(ms->allocator_cache);
 }
 
-static void *MsanAllocate(StackTrace *stack, uptr size,
-                          uptr alignment, bool zeroise) {
+void MsanThreadLocalMallocStorage::CommitBack() {
+  allocator.SwallowCache(GetAllocatorCache(this));
+}
+
+static void *MsanAllocate(StackTrace *stack, uptr size, uptr alignment,
+                          bool zeroise) {
   Init();
   if (size > kMaxAllowedMallocSize) {
     Report("WARNING: MemorySanitizer failed to allocate %p bytes\n",
            (void *)size);
     return AllocatorReturnNull();
   }
-  void *res = allocator.Allocate(&cache, size, alignment, false);
-  Metadata *meta = reinterpret_cast<Metadata*>(allocator.GetMetaData(res));
+  MsanThread *t = GetCurrentThread();
+  void *allocated;
+  if (t) {
+    AllocatorCache *cache = GetAllocatorCache(&t->malloc_storage());
+    allocated = allocator.Allocate(cache, size, alignment, false);
+  } else {
+    SpinMutexLock l(&fallback_mutex);
+    AllocatorCache *cache = &fallback_allocator_cache;
+    allocated = allocator.Allocate(cache, size, alignment, false);
+  }
+  Metadata *meta =
+      reinterpret_cast<Metadata *>(allocator.GetMetaData(allocated));
   meta->requested_size = size;
   if (zeroise) {
-    __msan_clear_and_unpoison(res, size);
+    __msan_clear_and_unpoison(allocated, size);
   } else if (flags()->poison_in_malloc) {
-    __msan_poison(res, size);
+    __msan_poison(allocated, size);
     if (__msan_get_track_origins()) {
       u32 stack_id = StackDepotPut(stack->trace, stack->size);
       CHECK(stack_id);
       CHECK_EQ((stack_id >> 31),
                0);  // Higher bit is occupied by stack origins.
-      __msan_set_origin(res, size, stack_id);
+      __msan_set_origin(allocated, size, stack_id);
     }
   }
-  MSAN_MALLOC_HOOK(res, size);
-  return res;
+  MSAN_MALLOC_HOOK(allocated, size);
+  return allocated;
 }
 
 void MsanDeallocate(StackTrace *stack, void *p) {
@@ -110,7 +129,15 @@ void MsanDeallocate(StackTrace *stack, void *p) {
       __msan_set_origin(p, size, stack_id);
     }
   }
-  allocator.Deallocate(&cache, p);
+  MsanThread *t = GetCurrentThread();
+  if (t) {
+    AllocatorCache *cache = GetAllocatorCache(&t->malloc_storage());
+    allocator.Deallocate(cache, p);
+  } else {
+    SpinMutexLock l(&fallback_mutex);
+    AllocatorCache *cache = &fallback_allocator_cache;
+    allocator.Deallocate(cache, p);
+  }
 }
 
 void *MsanReallocate(StackTrace *stack, void *old_p, uptr new_size,

compiler-rt/lib/msan/msan_allocator.h

@@ -0,0 +1,33 @@
+//===-- msan_allocator.h ----------------------------------------*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// This file is a part of MemorySanitizer.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MSAN_ALLOCATOR_H
+#define MSAN_ALLOCATOR_H
+
+#include "sanitizer_common/sanitizer_common.h"
+
+namespace __msan {
+
+struct MsanThreadLocalMallocStorage {
+  uptr quarantine_cache[16];
+  // Allocator cache contains atomic_uint64_t which must be 8-byte aligned.
+  ALIGNED(8) uptr allocator_cache[96 * (512 * 8 + 16)];  // Opaque.
+  void CommitBack();
+
+ private:
+  // These objects are allocated via mmap() and are zero-initialized.
+  MsanThreadLocalMallocStorage() {}
+};
+
+} // namespace __msan
+#endif // MSAN_ALLOCATOR_H

compiler-rt/lib/msan/msan_interceptors.cc

@@ -16,6 +16,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "msan.h"
+#include "msan_thread.h"
 #include "sanitizer_common/sanitizer_platform_limits_posix.h"
 #include "sanitizer_common/sanitizer_allocator.h"
 #include "sanitizer_common/sanitizer_allocator_internal.h"
@@ -37,8 +38,6 @@ using __sanitizer::atomic_load;
 using __sanitizer::atomic_store;
 using __sanitizer::atomic_uintptr_t;
 
-static unsigned g_thread_finalize_key;
-
 // True if this is a nested interceptor.
 static THREADLOCAL int in_interceptor_scope;
 
@@ -1038,48 +1037,11 @@ INTERCEPTOR(int, signal, int signo, uptr cb) {
 
 extern "C" int pthread_attr_init(void *attr);
 extern "C" int pthread_attr_destroy(void *attr);
-extern "C" int pthread_setspecific(unsigned key, const void *v);
-extern "C" int pthread_yield();
-
-static void thread_finalize(void *v) {
-  uptr iter = (uptr)v;
-  if (iter > 1) {
-    if (pthread_setspecific(g_thread_finalize_key, (void*)(iter - 1))) {
-      Printf("MemorySanitizer: failed to set thread key\n");
-      Die();
-    }
-    return;
-  }
-  MsanAllocatorThreadFinish();
-  __msan_unpoison((void *)msan_stack_bounds.stack_addr,
-                  msan_stack_bounds.stack_size);
-  if (msan_stack_bounds.tls_size)
-    __msan_unpoison((void *)msan_stack_bounds.tls_addr,
-                    msan_stack_bounds.tls_size);
-}
-
-struct ThreadParam {
-  void* (*callback)(void *arg);
-  void *param;
-  atomic_uintptr_t done;
-};
 
 static void *MsanThreadStartFunc(void *arg) {
-  ThreadParam *p = (ThreadParam *)arg;
-  void* (*callback)(void *arg) = p->callback;
-  void *param = p->param;
-  if (pthread_setspecific(g_thread_finalize_key,
-          (void *)kPthreadDestructorIterations)) {
-    Printf("MemorySanitizer: failed to set thread key\n");
-    Die();
-  }
-  atomic_store(&p->done, 1, memory_order_release);
-
-  GetThreadStackAndTls(/* main */ false, &msan_stack_bounds.stack_addr,
-                       &msan_stack_bounds.stack_size,
-                       &msan_stack_bounds.tls_addr,
-                       &msan_stack_bounds.tls_size);
-  return IndirectExternCall(callback)(param);
+  MsanThread *t = (MsanThread *)arg;
+  SetCurrentThread(t);
+  return t->ThreadStart();
 }
 
 INTERCEPTOR(int, pthread_create, void *th, void *attr, void *(*callback)(void*),
@@ -1093,16 +1055,9 @@ INTERCEPTOR(int, pthread_create, void *th, void *attr, void *(*callback)(void*),
 
   AdjustStackSize(attr);
 
-  ThreadParam p;
-  p.callback = callback;
-  p.param = param;
-  atomic_store(&p.done, 0, memory_order_relaxed);
+  MsanThread *t = MsanThread::Create(callback, param);
 
-  int res = REAL(pthread_create)(th, attr, MsanThreadStartFunc, (void *)&p);
-  if (res == 0) {
-    while (atomic_load(&p.done, memory_order_acquire) != 1)
-      pthread_yield();
-  }
+  int res = REAL(pthread_create)(th, attr, MsanThreadStartFunc, t);
 
   if (attr == &myattr)
     pthread_attr_destroy(&myattr);
@@ -1114,6 +1069,7 @@ INTERCEPTOR(int, pthread_create, void *th, void *attr, void *(*callback)(void*),
 
 INTERCEPTOR(int, pthread_key_create, __sanitizer_pthread_key_t *key,
             void (*dtor)(void *value)) {
+  if (msan_init_is_running) return REAL(pthread_key_create)(key, dtor);
   ENSURE_MSAN_INITED();
   int res = REAL(pthread_key_create)(key, dtor);
   if (!res && key)
@@ -1368,6 +1324,8 @@ void __msan_clear_and_unpoison(void *a, uptr size) {
 }
 
 void *__msan_memcpy(void *dest, const void *src, SIZE_T n) {
+  if (!msan_inited) return internal_memcpy(dest, src, n);
+  if (msan_init_is_running) return REAL(memcpy)(dest, src, n);
   ENSURE_MSAN_INITED();
   GET_STORE_STACK_TRACE;
   void *res = fast_memcpy(dest, src, n);
@@ -1376,6 +1334,8 @@ void *__msan_memcpy(void *dest, const void *src, SIZE_T n) {
 }
 
 void *__msan_memset(void *s, int c, SIZE_T n) {
+  if (!msan_inited) return internal_memset(s, c, n);
+  if (msan_init_is_running) return REAL(memset)(s, c, n);
   ENSURE_MSAN_INITED();
   void *res = fast_memset(s, c, n);
   __msan_unpoison(s, n);
@@ -1383,6 +1343,8 @@ void *__msan_memset(void *s, int c, SIZE_T n) {
 }
 
 void *__msan_memmove(void *dest, const void *src, SIZE_T n) {
+  if (!msan_inited) return internal_memmove(dest, src, n);
+  if (msan_init_is_running) return REAL(memmove)(dest, src, n);
   ENSURE_MSAN_INITED();
   GET_STORE_STACK_TRACE;
   void *res = REAL(memmove)(dest, src, n);
@@ -1603,11 +1565,6 @@ void InitializeInterceptors() {
   INTERCEPT_FUNCTION(__cxa_atexit);
   INTERCEPT_FUNCTION(shmat);
 
-  if (REAL(pthread_key_create)(&g_thread_finalize_key, &thread_finalize)) {
-    Printf("MemorySanitizer: failed to create thread key\n");
-    Die();
-  }
-
   inited = 1;
 }
 }  // namespace __msan

compiler-rt/lib/msan/msan_linux.cc

@@ -16,9 +16,11 @@
 #if SANITIZER_LINUX
 
 #include "msan.h"
+#include "msan_thread.h"
 
 #include <elf.h>
 #include <link.h>
+#include <pthread.h>
 #include <stdio.h>
 #include <stdlib.h>
 #include <signal.h>
@@ -97,6 +99,36 @@ void InstallAtExitHandler() {
   atexit(MsanAtExit);
 }
 
+// ---------------------- TSD ---------------- {{{1
+
+static pthread_key_t tsd_key;
+static bool tsd_key_inited = false;
+void MsanTSDInit(void (*destructor)(void *tsd)) {
+  CHECK(!tsd_key_inited);
+  tsd_key_inited = true;
+  CHECK_EQ(0, pthread_key_create(&tsd_key, destructor));
+}
+
+void *MsanTSDGet() {
+  CHECK(tsd_key_inited);
+  return pthread_getspecific(tsd_key);
+}
+
+void MsanTSDSet(void *tsd) {
+  CHECK(tsd_key_inited);
+  pthread_setspecific(tsd_key, tsd);
+}
+
+void MsanTSDDtor(void *tsd) {
+  MsanThread *t = (MsanThread*)tsd;
+  if (t->destructor_iterations_ > 1) {
+    t->destructor_iterations_--;
+    CHECK_EQ(0, pthread_setspecific(tsd_key, tsd));
+    return;
+  }
+  MsanThread::TSDDtor(tsd);
+}
+
 }  // namespace __msan
 
 #endif  // __linux__

compiler-rt/lib/msan/msan_thread.cc

@@ -0,0 +1,86 @@
+
+#include "msan.h"
+#include "msan_thread.h"
+#include "msan_interface_internal.h"
+
+namespace __msan {
+
+MsanThread *MsanThread::Create(thread_callback_t start_routine,
+                               void *arg) {
+  uptr PageSize = GetPageSizeCached();
+  uptr size = RoundUpTo(sizeof(MsanThread), PageSize);
+  MsanThread *thread = (MsanThread*)MmapOrDie(size, __func__);
+  thread->start_routine_ = start_routine;
+  thread->arg_ = arg;
+  thread->destructor_iterations_ = kPthreadDestructorIterations;
+
+  return thread;
+}
+
+void MsanThread::SetThreadStackAndTls() {
+  uptr tls_size = 0;
+  uptr stack_size = 0;
+  GetThreadStackAndTls(IsMainThread(), &stack_bottom_, &stack_size,
+                       &tls_begin_, &tls_size);
+  stack_top_ = stack_bottom_ + stack_size;
+  tls_end_ = tls_begin_ + tls_size;
+
+  int local;
+  CHECK(AddrIsInStack((uptr)&local));
+}
+
+void MsanThread::ClearShadowForThreadStackAndTLS() {
+  __msan_unpoison((void *)stack_bottom_, stack_top_ - stack_bottom_);
+  if (tls_begin_ != tls_end_)
+    __msan_unpoison((void *)tls_begin_, tls_end_ - tls_begin_);
+}
+
+void MsanThread::Init() {
+  SetThreadStackAndTls();
+  CHECK(MEM_IS_APP(stack_bottom_));
+  CHECK(MEM_IS_APP(stack_top_ - 1));
+  ClearShadowForThreadStackAndTLS();
+}
+
+void MsanThread::TSDDtor(void *tsd) {
+  MsanThread *t = (MsanThread*)tsd;
+  t->Destroy();
+}
+
+void MsanThread::Destroy() {
+  malloc_storage().CommitBack();
+  // We also clear the shadow on thread destruction because
+  // some code may still be executing in later TSD destructors
+  // and we don't want it to have any poisoned stack.
+  ClearShadowForThreadStackAndTLS();
+  uptr size = RoundUpTo(sizeof(MsanThread), GetPageSizeCached());
+  UnmapOrDie(this, size);
+}
+
+thread_return_t MsanThread::ThreadStart() {
+  Init();
+
+  if (!start_routine_) {
+    // start_routine_ == 0 if we're on the main thread or on one of the
+    // OS X libdispatch worker threads. But nobody is supposed to call
+    // ThreadStart() for the worker threads.
+    return 0;
+  }
+
+  thread_return_t res = IndirectExternCall(start_routine_)(arg_);
+
+  return res;
+}
+
+MsanThread *GetCurrentThread() {
+  return reinterpret_cast<MsanThread *>(MsanTSDGet());
+}
+
+void SetCurrentThread(MsanThread *t) {
+  // Make sure we do not reset the current MsanThread.
+  CHECK_EQ(0, MsanTSDGet());
+  MsanTSDSet(t);
+  CHECK_EQ(t, MsanTSDGet());
+}
+
+} // namespace __msan

compiler-rt/lib/msan/msan_thread.h

@@ -0,0 +1,65 @@
+//===-- msan_thread.h -------------------------------------------*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// This file is a part of MemorySanitizer.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MSAN_THREAD_H
+#define MSAN_THREAD_H
+
+#include "msan_allocator.h"
+#include "sanitizer_common/sanitizer_common.h"
+
+namespace __msan {
+
+class MsanThread {
+ public:
+  static MsanThread *Create(thread_callback_t start_routine, void *arg);
+  static void TSDDtor(void *tsd);
+  void Destroy();
+
+  void Init();  // Should be called from the thread itself.
+  thread_return_t ThreadStart();
+
+  uptr stack_top() { return stack_top_; }
+  uptr stack_bottom() { return stack_bottom_; }
+  uptr tls_begin() { return tls_begin_; }
+  uptr tls_end() { return tls_end_; }
+  bool IsMainThread() { return start_routine_ == 0; }
+
+  bool AddrIsInStack(uptr addr) {
+    return addr >= stack_bottom_ && addr < stack_top_;
+  }
+
+  MsanThreadLocalMallocStorage &malloc_storage() { return malloc_storage_; }
+
+  int destructor_iterations_;
+
+ private:
+  // NOTE: There is no MsanThread constructor. It is allocated
+  // via mmap() and *must* be valid in zero-initialized state.
+  void SetThreadStackAndTls();
+  void ClearShadowForThreadStackAndTLS();
+  thread_callback_t start_routine_;
+  void *arg_;
+  uptr stack_top_;
+  uptr stack_bottom_;
+  uptr tls_begin_;
+  uptr tls_end_;
+
+  MsanThreadLocalMallocStorage malloc_storage_;
+};
+
+MsanThread *GetCurrentThread();
+void SetCurrentThread(MsanThread *t);
+
+} // namespace __msan
+
+#endif // MSAN_THREAD_H

compiler-rt/lib/msan/tests/msan_test.cc

@@ -2818,22 +2818,22 @@ TEST(MemorySanitizer, SmallStackThread) {
   ASSERT_EQ(0, res);
 }
 
-TEST(MemorySanitizer, PreAllocatedStackThread) {
+TEST(MemorySanitizer, SmallPreAllocatedStackThread) {
   pthread_attr_t attr;
   pthread_t t;
   int res;
   res = pthread_attr_init(&attr);
   ASSERT_EQ(0, res);
   void *stack;
-  const size_t kStackSize = 64 * 1024;
+  const size_t kStackSize = 16 * 1024;
   res = posix_memalign(&stack, 4096, kStackSize);
   ASSERT_EQ(0, res);
   res = pthread_attr_setstack(&attr, stack, kStackSize);
   ASSERT_EQ(0, res);
-  // A small self-allocated stack can not be extended by the tool.
-  // In this case pthread_create is expected to fail.
   res = pthread_create(&t, &attr, SmallStackThread_threadfn, NULL);
-  EXPECT_NE(0, res);
+  EXPECT_EQ(0, res);
+  res = pthread_join(t, NULL);
+  ASSERT_EQ(0, res);
   res = pthread_attr_destroy(&attr);
   ASSERT_EQ(0, res);
 }