[msan] Drain allocator cache when leaving thread.

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

3 files changed, 57 insertions, 1 deletions

diff --git a/lib/msan/msan.h b/lib/msan/msan.h
index 8c24de64..fe7f20ae 100644
--- a/lib/msan/msan.h
+++ b/lib/msan/msan.h
@@ -44,6 +44,7 @@ bool InitShadow(bool prot1, bool prot2, bool map_shadow, bool init_origins);
 char *GetProcSelfMaps();
 void InitializeInterceptors();
 
+void MsanAllocatorThreadFinish();
 void *MsanReallocate(StackTrace *stack, void *oldp, uptr size,
                      uptr alignment, bool zeroise);
 void MsanDeallocate(StackTrace *stack, void *ptr);
diff --git a/lib/msan/msan_allocator.cc b/lib/msan/msan_allocator.cc
index 04d6fedb..2badf712 100644
--- a/lib/msan/msan_allocator.cc
+++ b/lib/msan/msan_allocator.cc
@@ -46,6 +46,10 @@ static inline void Init() {
   allocator.Init();
 }
 
+void MsanAllocatorThreadFinish() {
+  allocator.SwallowCache(&cache);
+}
+
 static void *MsanAllocate(StackTrace *stack, uptr size,
                           uptr alignment, bool zeroise) {
   Init();
diff --git a/lib/msan/msan_interceptors.cc b/lib/msan/msan_interceptors.cc
index 46f5dd2a..c200e6c5 100644
--- a/lib/msan/msan_interceptors.cc
+++ b/lib/msan/msan_interceptors.cc
@@ -37,6 +37,8 @@ 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;
 
@@ -1040,6 +1042,39 @@ extern "C" int pthread_attr_init(void *attr);
 extern "C" int pthread_attr_destroy(void *attr);
 extern "C" int pthread_attr_setstacksize(void *attr, uptr stacksize);
 extern "C" int pthread_attr_getstack(void *attr, uptr *stack, uptr *stacksize);
+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();
+}
+
+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);
+  return callback(param);
+}
 
 INTERCEPTOR(int, pthread_create, void *th, void *attr, void *(*callback)(void*),
             void * param) {
@@ -1052,7 +1087,17 @@ INTERCEPTOR(int, pthread_create, void *th, void *attr, void *(*callback)(void*),
 
   AdjustStackSizeLinux(attr);
 
-  int res = REAL(pthread_create)(th, attr, callback, param);
+  ThreadParam p;
+  p.callback = callback;
+  p.param = param;
+  atomic_store(&p.done, 0, memory_order_relaxed);
+
+  int res = REAL(pthread_create)(th, attr, MsanThreadStartFunc, (void *)&p);
+  if (res == 0) {
+    while (atomic_load(&p.done, memory_order_acquire) != 1)
+      pthread_yield();
+  }
+
   if (attr == &myattr)
     pthread_attr_destroy(&myattr);
   if (!res) {
@@ -1387,6 +1432,12 @@ void InitializeInterceptors() {
   INTERCEPT_FUNCTION(pthread_join);
   INTERCEPT_FUNCTION(tzset);
   INTERCEPT_FUNCTION(__cxa_atexit);
+
+  if (REAL(pthread_key_create)(&g_thread_finalize_key, &thread_finalize)) {
+    Printf("MemorySanitizer: failed to create thread key\n");
+    Die();
+  }
+
   inited = 1;
 }
 }  // namespace __msan