[sanitizer] Avoid possible deadlock in child process after fork

[m.ostapenko]

Summary

This patch addresses https://github.com/google/sanitizers/issues/774
When we fork a multi-threaded process it's possible to deadlock if some thread acquired StackDepot or allocator internal lock just before fork. In this case the lock will never be released in child process causing deadlock on following memory alloc/dealloc routine.

[kcc]

is a test possible?

[m.ostapenko]

In D33325#758753, @kcc wrote:

is a test possible?

I'm trying to cook. Need to have concurrent fork with malloc/free and the StackDepot lock to be acquired when fork is called. But the window is too small :(. Can we widen it somehow?

[eugenis]

Look at test/msan/fork.cc, it should be possible to replace copy_uninit_thread2 with an allocation loop.

[m.ostapenko]

In D33325#758766, @eugenis wrote:

Look at test/msan/fork.cc, it should be possible to replace copy_uninit_thread2 with an allocation loop.

Thank you, will try tomorrow.

[m.ostapenko]

Updating. When adapting test/msan/fork.cc for ASan I've noticed that deadlock can also happen if the whole allocator is locked when we calling fork. Corresponding backtrace looks like this:

#0  atomic_exchange<__sanitizer::atomic_uint32_t> () at /home/max/src/llvm/projects/compiler-rt/lib/sanitizer_common/sanitizer_atomic_clang.h:68
#1  Lock () at /home/max/src/llvm/projects/compiler-rt/lib/sanitizer_common/sanitizer_linux.cc:541
#2  0x000000000041e9f9 in GenericScopedLock () at /home/max/src/llvm/projects/compiler-rt/lib/asan/../sanitizer_common/sanitizer_mutex.h:187
#3  GetFromAllocator () at /home/max/src/llvm/projects/compiler-rt/lib/asan/../sanitizer_common/sanitizer_allocator_primary64.h:134
#4  0x000000000041e99b in Refill () at /home/max/src/llvm/projects/compiler-rt/lib/asan/../sanitizer_common/sanitizer_allocator_local_cache.h:108
#5  0x000000000041e550 in Allocate () at /home/max/src/llvm/projects/compiler-rt/lib/asan/../sanitizer_common/sanitizer_allocator_local_cache.h:50
#6  0x000000000041e443 in Allocate () at /home/max/src/llvm/projects/compiler-rt/lib/asan/../sanitizer_common/sanitizer_allocator_combined.h:68
#7  0x000000000041b6cb in Allocate () at /home/max/src/llvm/projects/compiler-rt/lib/asan/asan_allocator.cc:407
#8  0x00000000004b9148 in __interceptor_malloc () at /home/max/src/llvm/projects/compiler-rt/lib/asan/asan_malloc_linux.cc:67
#9  0x00000000004e878d in child () at fork.cc:47
#10 0x00000000004e8968 in test () at fork.cc:67
#11 0x00000000004e8a38 in main () at fork.cc:78

Thus we need to lock allocator as well.

[eugenis]

Good catch. This is a problem for MSan as well, is not it? Do you mind fixing it in both tools?

For bonus points, TSan. I see it locks something
void ForkBefore(ThreadState *thr, uptr pc) {

ctx->thread_registry->Lock();
ctx->report_mtx.Lock();

}

But again, not the allocator.

[m.ostapenko]

In D33325#759703, @eugenis wrote:

Good catch. This is a problem for MSan as well, is not it? Do you mind fixing it in both tools?

Sure, we can fix it for all sanitizers if the problem exists, I'll check this more carefully. FWIW I've run the attached testcase with MSan and TSan, but it didn't hang (as opposed to ASan). Do they use allocator caches (that can lock on refilling)?

For bonus points, TSan. I see it locks something
void ForkBefore(ThreadState *thr, uptr pc) {

ctx->thread_registry->Lock();
ctx->report_mtx.Lock();

}

But again, not the allocator.

[eugenis]

I think they do.

[m.ostapenko]

Updating. Lock allocator on fork for {A, M, T}San and moving testcase to sanitizer_common. I've also managed to reproduce the deadlock with MSan, but still not with TSan.

[eugenis]

lib/msan/msan_interceptors.cc

1231

just a nit: ASan takes these locks in a different order, please be consistent.

test/sanitizer_common/TestCases/Linux/allocator_fork_no_hang.cc

33

A volatile or -O0 flag would not hurt to make sure this is not optimized out.
I'm not sure how this can cause cache refill without quarantine. I think it would just keep reallocating the same chunk. Could it be the reason the problem does not reproduce on TSan?

42

What are these hash table cells?

[m.ostapenko]

test/sanitizer_common/TestCases/Linux/allocator_fork_no_hang.cc

42

Oh, this comment reflects to StackDepotLockAll(). I'll remove it.

[m.ostapenko]

Updating. Reproduced the deadlock with TSan and LSan on adjusted testcase. Adding Dmitry since the patch now affects TSan code as well.

[dvyukov]

I don't completely understand this. allocator_fork_no_hang.cc is broken and deserves deadlocking. Why are we trying to fix this program only under sanitizers rather than detecting and reporting the bug? Sanitizers are meant to be less forgiving than production environment. Why are we trying to conceal the bug?

[m.ostapenko]

In D33325#761941, @dvyukov wrote:

I don't completely understand this. allocator_fork_no_hang.cc is broken and deserves deadlocking. Why are we trying to fix this program only under sanitizers rather than detecting and reporting the bug? Sanitizers are meant to be less forgiving than production environment. Why are we trying to conceal the bug?

Yes, the testcase deserves deadlocking (malloc is not async signal safe). And the actual bug is reported here: https://bugzilla.gnome.org/show_bug.cgi?id=738620.
The motivation for "fixing" this in sanitizers is that most of modern allocators (Glibc, tcmalloc, jemalloc) are actually fork safe, so why not to do the same in sanitizers. And despite the fact that calling malloc after multi-threaded fork is not allowed, many people still would expect their code to work (since it works with Glibc/tcmalloc/jemalloc etc).

[dvyukov]

In D33325#761980, @m.ostapenko wrote:

In D33325#761941, @dvyukov wrote:

I don't completely understand this. allocator_fork_no_hang.cc is broken and deserves deadlocking. Why are we trying to fix this program only under sanitizers rather than detecting and reporting the bug? Sanitizers are meant to be less forgiving than production environment. Why are we trying to conceal the bug?

Yes, the testcase deserves deadlocking (malloc is not async signal safe). And the actual bug is reported here: https://bugzilla.gnome.org/show_bug.cgi?id=738620.
The motivation for "fixing" this in sanitizers is that most of modern allocators (Glibc, tcmalloc, jemalloc) are actually fork safe, so why not to do the same in sanitizers. And despite the fact that calling malloc after multi-threaded fork is not allowed, many people still would expect their code to work (since it works with Glibc/tcmalloc/jemalloc etc).

okay

[dvyukov]

lib/tsan/rtl/tsan_dense_alloc.h

94

This is no more than Lock, please call it Lock. ForceLock produces a false impression that this is something more than just Lock.

98

Same here.

lib/tsan/rtl/tsan_rtl.cc

459

Also StackDepotLockAll.

lib/tsan/rtl/tsan_sync.h

128

Same here.

130

Same here.

[m.ostapenko]

Addressing Dmitry's comments.

[dvyukov]

lib/tsan/rtl/tsan_rtl.cc

460

Now looking at this more closely and paging in some context.
This change is not necessary for tsan. To make it work we need to lock whole lot of mutexes. Memory allocator is less of a problem, when plain memory accesses can cause deadlock. We would need to lock trace mutexes of all threads, all sync objects, deadlock detector, internal allocator, annotations mutex, atexit mutex, global proc mutex and probably some more. Moreover we need to lock then in the proper order. For example, this change will cause deadlocks in otherwise working programs because allocator mutex needs to be locked after report mutex (see CanLockTab in tsan_mutex.cc for some ordering constraints).
Because of that tsan uses a different mechanism to systematically prevent deadlocks (grep for after_multithreaded_fork). So if the test works under tsan, let's leave tsan alone.

[m.ostapenko]

Removing TSan part. The attached testcase still hangs with TSan, but since locking allocator is problematic let's leave TSan alone.

[m.ostapenko]

Rebased. Gentle reminder.

[eugenis]

LGTM

[eugenis]

This change is causing problems on Linux. The immediate issue is Tcl, which creates new threads from a pthread_atfork() child handler. The handler is called from libc's fork(), and pthread_create tries to acquire the StackDepot lock before it is released in the fork interceptor.

Tcl is broken, of course:

If a fork() call in a multi-threaded process leads to a child fork handler calling any function that is not async-signal-safe, the behavior is undefined.

But in general, even malloc() called from pthread_atfork() will deadlock with this change.
It looks like the right way to do this is call pthread_atfork() as early as possible to setup lock/unlock handlers for the allocator and the stack depot.

I'll revert this change in the meantime.

[eugenis]

Reverted in r304735

[m.ostapenko]

Oh, sorry.

In D33325#773286, @eugenis wrote:

This change is causing problems on Linux. The immediate issue is Tcl, which creates new threads from a pthread_atfork() child handler. The handler is called from libc's fork(), and pthread_create tries to acquire the StackDepot lock before it is released in the fork interceptor.

Won't this deadlock with MSan as well? AFAIK it also locks StackDepot in fork interceptor...

Tcl is broken, of course:

If a fork() call in a multi-threaded process leads to a child fork handler calling any function that is not async-signal-safe, the behavior is undefined.

But in general, even malloc() called from pthread_atfork() will deadlock with this change.
It looks like the right way to do this is call pthread_atfork() as early as possible to setup lock/unlock handlers for the allocator and the stack depot.

Perhaps we can do it from asan_init? From man:
"The order of calls to pthread_atfork() is significant. The parent and child fork handlers shall be called in the order in which they were established by calls to pthread_atfork(). The prepare fork handlers shall be called in the opposite order."

So in this case our unlock routines will be called before others. This still can deadlock when preloading asan to noinstrumented binary though.

I'll revert this change in the meantime.

[eugenis]

In D33325#773322, @m.ostapenko wrote:

Oh, sorry.

In D33325#773286, @eugenis wrote:

This change is causing problems on Linux. The immediate issue is Tcl, which creates new threads from a pthread_atfork() child handler. The handler is called from libc's fork(), and pthread_create tries to acquire the StackDepot lock before it is released in the fork interceptor.

Won't this deadlock with MSan as well? AFAIK it also locks StackDepot in fork interceptor...

Maybe it does not use StackDepot in pthread_create?

Tcl is broken, of course:

If a fork() call in a multi-threaded process leads to a child fork handler calling any function that is not async-signal-safe, the behavior is undefined.

But in general, even malloc() called from pthread_atfork() will deadlock with this change.
It looks like the right way to do this is call pthread_atfork() as early as possible to setup lock/unlock handlers for the allocator and the stack depot.

Perhaps we can do it from asan_init? From man:
"The order of calls to pthread_atfork() is significant. The parent and child fork handlers shall be called in the order in which they were established by calls to pthread_atfork(). The prepare fork handlers shall be called in the opposite order."

So in this case our unlock routines will be called before others. This still can deadlock when preloading asan to noinstrumented binary though.

Sounds like that's the best we can do. If we are lucky, sanitizer initialization will happen early enough due to ENSURE_ASAN_INIT in one of the interceptors.

I'll revert this change in the meantime.

[m.ostapenko]

In D33325#773332, @eugenis wrote:

In D33325#773322, @m.ostapenko wrote:

Oh, sorry.

In D33325#773286, @eugenis wrote:

This change is causing problems on Linux. The immediate issue is Tcl, which creates new threads from a pthread_atfork() child handler. The handler is called from libc's fork(), and pthread_create tries to acquire the StackDepot lock before it is released in the fork interceptor.

Won't this deadlock with MSan as well? AFAIK it also locks StackDepot in fork interceptor...

Maybe it does not use StackDepot in pthread_create?

Tcl is broken, of course:

If a fork() call in a multi-threaded process leads to a child fork handler calling any function that is not async-signal-safe, the behavior is undefined.

But in general, even malloc() called from pthread_atfork() will deadlock with this change.
It looks like the right way to do this is call pthread_atfork() as early as possible to setup lock/unlock handlers for the allocator and the stack depot.

Perhaps we can do it from asan_init? From man:
"The order of calls to pthread_atfork() is significant. The parent and child fork handlers shall be called in the order in which they were established by calls to pthread_atfork(). The prepare fork handlers shall be called in the opposite order."

So in this case our unlock routines will be called before others. This still can deadlock when preloading asan to noinstrumented binary though.

Sounds like that's the best we can do. If we are lucky, sanitizer initialization will happen early enough due to ENSURE_ASAN_INIT in one of the interceptors.

Hm, perhaps we can just intercept pthread_atfork and force ENSURE_ASAN_INITED? And install allocator hooks in AsanInitFromRtl?

I'll revert this change in the meantime.

[eugenis]

Sounds good.

lib/asan/asan_allocator.h

 
 uptr asan_mz_size(const void *ptr);
 void asan_mz_force_lock();
 void asan_mz_force_unlock();
 
 void PrintInternalAllocatorStats();
 void AsanSoftRssLimitExceededCallback(bool exceeded);
 
+AsanAllocator &get_allocator();
+
 }  // namespace __asan
 #endif  // ASAN_ALLOCATOR_H

lib/asan/asan_allocator.cc

static u32 RZSize2Log(u32 rz_size) {
   CHECK_GE(rz_size, 16);
   CHECK_LE(rz_size, 2048);
   CHECK(IsPowerOfTwo(rz_size));
   u32 res = Log2(rz_size) - 4;
   CHECK_EQ(rz_size, RZLog2Size(res));
   return res;
 }
 
-static AsanAllocator &get_allocator();
-
 // The memory chunk allocated from the underlying allocator looks like this:
 // L L L L L L H H U U U U U U R R
 //   L -- left redzone words (0 or more bytes)
 //   H -- ChunkHeader (16 bytes), which is also a part of the left redzone.
 //   U -- user memory.
 //   R -- right redzone (0 or more bytes)
 // ChunkBase consists of ChunkHeader and other bytes that overlap with user
 // memory.
#endif
   void ForceUnlock() {
     fallback_mutex.Unlock();
     allocator.ForceUnlock();
   }
 };
 
 static Allocator instance(LINKER_INITIALIZED);
 
-static AsanAllocator &get_allocator() {
+AsanAllocator &get_allocator() {
   return instance.allocator;
 }
 
 bool AsanChunkView::IsValid() const {
   return chunk_ && chunk_->chunk_state != CHUNK_AVAILABLE;
 }
 bool AsanChunkView::IsAllocated() const {
   return chunk_ && chunk_->chunk_state == CHUNK_ALLOCATED;

lib/asan/asan_interceptors.cc

 #include "asan_internal.h"
 #include "asan_mapping.h"
 #include "asan_poisoning.h"
 #include "asan_report.h"
 #include "asan_stack.h"
 #include "asan_stats.h"
 #include "asan_suppressions.h"
 #include "lsan/lsan_common.h"
+#include "sanitizer_common/sanitizer_stackdepot.h"
 #include "sanitizer_common/sanitizer_libc.h"
 
 #if SANITIZER_POSIX
 #include "sanitizer_common/sanitizer_posix.h"
 #endif
 
 #if defined(__i386) && SANITIZER_LINUX
 #define ASAN_PTHREAD_CREATE_VERSION "GLIBC_2.1"
#endif
   ENSURE_ASAN_INITED();
   int res = REAL(__cxa_atexit)(func, arg, dso_handle);
   REAL(__cxa_atexit)(AtCxaAtexit, nullptr, nullptr);
   return res;
 }
 #endif  // ASAN_INTERCEPT___CXA_ATEXIT
 
 #if ASAN_INTERCEPT_FORK
+static void BeforeFork() {
+  get_allocator().ForceLock();
+  StackDepotLockAll();
+}
+
+static void AfterFork() {
+  StackDepotUnlockAll();
+  get_allocator().ForceUnlock();
+}
+
 INTERCEPTOR(int, fork, void) {
   ENSURE_ASAN_INITED();
+  BeforeFork();
   if (common_flags()->coverage) CovBeforeFork();
   int pid = REAL(fork)();
   if (common_flags()->coverage) CovAfterFork(pid);
+  AfterFork();
   return pid;
 }
 #endif  // ASAN_INTERCEPT_FORK
 
 // ---------------------- InitializeAsanInterceptors ---------------- {{{1
 namespace __asan {
 void InitializeAsanInterceptors() {
   static bool was_called_once;

lib/lsan/lsan_interceptors.cc

 #include "sanitizer_common/sanitizer_atomic.h"
 #include "sanitizer_common/sanitizer_common.h"
 #include "sanitizer_common/sanitizer_flags.h"
 #include "sanitizer_common/sanitizer_internal_defs.h"
 #include "sanitizer_common/sanitizer_linux.h"
 #include "sanitizer_common/sanitizer_platform_interceptors.h"
 #include "sanitizer_common/sanitizer_platform_limits_posix.h"
 #include "sanitizer_common/sanitizer_posix.h"
+#include "sanitizer_common/sanitizer_stackdepot.h"
 #include "sanitizer_common/sanitizer_tls_get_addr.h"
 #include "lsan.h"
 #include "lsan_allocator.h"
 #include "lsan_common.h"
 #include "lsan_thread.h"
 
 #include <stddef.h>
 

 
 INTERCEPTOR(void*, valloc, uptr size) {
   ENSURE_LSAN_INITED;
   GET_STACK_TRACE_MALLOC;
   return lsan_valloc(size, stack);
 }
 #endif
 
+static void BeforeFork() {
+  LockAllocator();
+  StackDepotLockAll();
+}
+
+static void AfterFork() {
+  StackDepotUnlockAll();
+  UnlockAllocator();
+}
+
+INTERCEPTOR(int, fork, void) {
+  ENSURE_LSAN_INITED;
+  BeforeFork();
+  int pid = REAL(fork)();
+  AfterFork();
+  return pid;
+}
+
 #if SANITIZER_INTERCEPT_MEMALIGN
 INTERCEPTOR(void*, memalign, uptr alignment, uptr size) {
   ENSURE_LSAN_INITED;
   GET_STACK_TRACE_MALLOC;
   return lsan_memalign(alignment, size, stack);
 }
 #define LSAN_MAYBE_INTERCEPT_MEMALIGN INTERCEPT_FUNCTION(memalign)
 
void InitializeInterceptors() {
   INTERCEPT_FUNCTION(posix_memalign);
   INTERCEPT_FUNCTION(valloc);
   LSAN_MAYBE_INTERCEPT_PVALLOC;
   LSAN_MAYBE_INTERCEPT_MALLOC_USABLE_SIZE;
   LSAN_MAYBE_INTERCEPT_MALLINFO;
   LSAN_MAYBE_INTERCEPT_MALLOPT;
   INTERCEPT_FUNCTION(pthread_create);
   INTERCEPT_FUNCTION(pthread_join);
+  INTERCEPT_FUNCTION(fork);
 
   if (pthread_key_create(&g_thread_finalize_key, &thread_finalize)) {
     Report("LeakSanitizer: failed to create thread key.\n");
     Die();
   }
 }
 
 } // namespace __lsan

lib/msan/msan_allocator.h

 // This file is a part of MemorySanitizer.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MSAN_ALLOCATOR_H
 #define MSAN_ALLOCATOR_H
 
 #include "sanitizer_common/sanitizer_common.h"
+#include "sanitizer_common/sanitizer_allocator.h"
+#include "sanitizer_common/sanitizer_allocator_interface.h"
 
 namespace __msan {
 
+struct Metadata {
+  uptr requested_size;
+};
+
+struct MsanMapUnmapCallback {
+  void OnMap(uptr p, uptr size) const {}
+  void OnUnmap(uptr p, uptr size) const {
+    __msan_unpoison((void *)p, size);
+
+    // We are about to unmap a chunk of user memory.
+    // Mark the corresponding shadow memory as not needed.
+    uptr shadow_p = MEM_TO_SHADOW(p);
+    ReleaseMemoryPagesToOS(shadow_p, shadow_p + size);
+    if (__msan_get_track_origins()) {
+      uptr origin_p = MEM_TO_ORIGIN(p);
+      ReleaseMemoryPagesToOS(origin_p, origin_p + size);
+    }
+  }
+};
+
+#if defined(__mips64)
+  static const uptr kMaxAllowedMallocSize = 2UL << 30;
+  static const uptr kRegionSizeLog = 20;
+  static const uptr kNumRegions = SANITIZER_MMAP_RANGE_SIZE >> kRegionSizeLog;
+  typedef TwoLevelByteMap<(kNumRegions >> 12), 1 << 12> ByteMap;
+
+  struct AP32 {
+    static const uptr kSpaceBeg = 0;
+    static const u64 kSpaceSize = SANITIZER_MMAP_RANGE_SIZE;
+    static const uptr kMetadataSize = sizeof(Metadata);
+    typedef __sanitizer::CompactSizeClassMap SizeClassMap;
+    static const uptr kRegionSizeLog = __msan::kRegionSizeLog;
+    typedef __msan::ByteMap ByteMap;
+    typedef MsanMapUnmapCallback MapUnmapCallback;
+    static const uptr kFlags = 0;
+  };
+  typedef SizeClassAllocator32<AP32> PrimaryAllocator;
+#elif defined(__x86_64__)
+#if SANITIZER_LINUX && !defined(MSAN_LINUX_X86_64_OLD_MAPPING)
+  static const uptr kAllocatorSpace = 0x700000000000ULL;
+#else
+  static const uptr kAllocatorSpace = 0x600000000000ULL;
+#endif
+  static const uptr kMaxAllowedMallocSize = 8UL << 30;
+
+  struct AP64 {  // Allocator64 parameters. Deliberately using a short name.
+    static const uptr kSpaceBeg = kAllocatorSpace;
+    static const uptr kSpaceSize = 0x40000000000; // 4T.
+    static const uptr kMetadataSize = sizeof(Metadata);
+    typedef DefaultSizeClassMap SizeClassMap;
+    typedef MsanMapUnmapCallback MapUnmapCallback;
+    static const uptr kFlags = 0;
+  };
+
+  typedef SizeClassAllocator64<AP64> PrimaryAllocator;
+
+#elif defined(__powerpc64__)
+  static const uptr kMaxAllowedMallocSize = 2UL << 30;  // 2G
+
+  struct AP64 {  // Allocator64 parameters. Deliberately using a short name.
+    static const uptr kSpaceBeg = 0x300000000000;
+    static const uptr kSpaceSize = 0x020000000000; // 2T.
+    static const uptr kMetadataSize = sizeof(Metadata);
+    typedef DefaultSizeClassMap SizeClassMap;
+    typedef MsanMapUnmapCallback MapUnmapCallback;
+    static const uptr kFlags = 0;
+  };
+
+  typedef SizeClassAllocator64<AP64> PrimaryAllocator;
+#elif defined(__aarch64__)
+  static const uptr kMaxAllowedMallocSize = 2UL << 30;  // 2G
+  static const uptr kRegionSizeLog = 20;
+  static const uptr kNumRegions = SANITIZER_MMAP_RANGE_SIZE >> kRegionSizeLog;
+  typedef TwoLevelByteMap<(kNumRegions >> 12), 1 << 12> ByteMap;
+
+  struct AP32 {
+    static const uptr kSpaceBeg = 0;
+    static const u64 kSpaceSize = SANITIZER_MMAP_RANGE_SIZE;
+    static const uptr kMetadataSize = sizeof(Metadata);
+    typedef __sanitizer::CompactSizeClassMap SizeClassMap;
+    static const uptr kRegionSizeLog = __msan::kRegionSizeLog;
+    typedef __msan::ByteMap ByteMap;
+    typedef MsanMapUnmapCallback MapUnmapCallback;
+    static const uptr kFlags = 0;
+  };
+  typedef SizeClassAllocator32<AP32> PrimaryAllocator;
+#endif
+typedef SizeClassAllocatorLocalCache<PrimaryAllocator> AllocatorCache;
+typedef LargeMmapAllocator<MsanMapUnmapCallback> SecondaryAllocator;
+typedef CombinedAllocator<PrimaryAllocator, AllocatorCache,
+                          SecondaryAllocator> Allocator;
+
+
+Allocator &get_allocator();
+
 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

lib/msan/msan_allocator.cc

 //===-- msan_allocator.cc --------------------------- ---------------------===//
 //
 //                     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.
 //
 // MemorySanitizer allocator.
 //===----------------------------------------------------------------------===//
 
-#include "sanitizer_common/sanitizer_allocator.h"
-#include "sanitizer_common/sanitizer_allocator_interface.h"
 #include "msan.h"
 #include "msan_allocator.h"
 #include "msan_origin.h"
 #include "msan_thread.h"
 #include "msan_poisoning.h"
 
 namespace __msan {
 
-struct Metadata {
-  uptr requested_size;
-};
-
-struct MsanMapUnmapCallback {
-  void OnMap(uptr p, uptr size) const {}
-  void OnUnmap(uptr p, uptr size) const {
-    __msan_unpoison((void *)p, size);
-
-    // We are about to unmap a chunk of user memory.
-    // Mark the corresponding shadow memory as not needed.
-    uptr shadow_p = MEM_TO_SHADOW(p);
-    ReleaseMemoryPagesToOS(shadow_p, shadow_p + size);
-    if (__msan_get_track_origins()) {
-      uptr origin_p = MEM_TO_ORIGIN(p);
-      ReleaseMemoryPagesToOS(origin_p, origin_p + size);
-    }
-  }
-};
-
-#if defined(__mips64)
-  static const uptr kMaxAllowedMallocSize = 2UL << 30;
-  static const uptr kRegionSizeLog = 20;
-  static const uptr kNumRegions = SANITIZER_MMAP_RANGE_SIZE >> kRegionSizeLog;
-  typedef TwoLevelByteMap<(kNumRegions >> 12), 1 << 12> ByteMap;
-
-  struct AP32 {
-    static const uptr kSpaceBeg = 0;
-    static const u64 kSpaceSize = SANITIZER_MMAP_RANGE_SIZE;
-    static const uptr kMetadataSize = sizeof(Metadata);
-    typedef __sanitizer::CompactSizeClassMap SizeClassMap;
-    static const uptr kRegionSizeLog = __msan::kRegionSizeLog;
-    typedef __msan::ByteMap ByteMap;
-    typedef MsanMapUnmapCallback MapUnmapCallback;
-    static const uptr kFlags = 0;
-  };
-  typedef SizeClassAllocator32<AP32> PrimaryAllocator;
-#elif defined(__x86_64__)
-#if SANITIZER_LINUX && !defined(MSAN_LINUX_X86_64_OLD_MAPPING)
-  static const uptr kAllocatorSpace = 0x700000000000ULL;
-#else
-  static const uptr kAllocatorSpace = 0x600000000000ULL;
-#endif
-  static const uptr kMaxAllowedMallocSize = 8UL << 30;
-
-  struct AP64 {  // Allocator64 parameters. Deliberately using a short name.
-    static const uptr kSpaceBeg = kAllocatorSpace;
-    static const uptr kSpaceSize = 0x40000000000; // 4T.
-    static const uptr kMetadataSize = sizeof(Metadata);
-    typedef DefaultSizeClassMap SizeClassMap;
-    typedef MsanMapUnmapCallback MapUnmapCallback;
-    static const uptr kFlags = 0;
-  };
-
-  typedef SizeClassAllocator64<AP64> PrimaryAllocator;
-
-#elif defined(__powerpc64__)
-  static const uptr kMaxAllowedMallocSize = 2UL << 30;  // 2G
-
-  struct AP64 {  // Allocator64 parameters. Deliberately using a short name.
-    static const uptr kSpaceBeg = 0x300000000000;
-    static const uptr kSpaceSize = 0x020000000000; // 2T.
-    static const uptr kMetadataSize = sizeof(Metadata);
-    typedef DefaultSizeClassMap SizeClassMap;
-    typedef MsanMapUnmapCallback MapUnmapCallback;
-    static const uptr kFlags = 0;
-  };
-
-  typedef SizeClassAllocator64<AP64> PrimaryAllocator;
-#elif defined(__aarch64__)
-  static const uptr kMaxAllowedMallocSize = 2UL << 30;  // 2G
-  static const uptr kRegionSizeLog = 20;
-  static const uptr kNumRegions = SANITIZER_MMAP_RANGE_SIZE >> kRegionSizeLog;
-  typedef TwoLevelByteMap<(kNumRegions >> 12), 1 << 12> ByteMap;
-
-  struct AP32 {
-    static const uptr kSpaceBeg = 0;
-    static const u64 kSpaceSize = SANITIZER_MMAP_RANGE_SIZE;
-    static const uptr kMetadataSize = sizeof(Metadata);
-    typedef __sanitizer::CompactSizeClassMap SizeClassMap;
-    static const uptr kRegionSizeLog = __msan::kRegionSizeLog;
-    typedef __msan::ByteMap ByteMap;
-    typedef MsanMapUnmapCallback MapUnmapCallback;
-    static const uptr kFlags = 0;
-  };
-  typedef SizeClassAllocator32<AP32> PrimaryAllocator;
-#endif
-typedef SizeClassAllocatorLocalCache<PrimaryAllocator> AllocatorCache;
-typedef LargeMmapAllocator<MsanMapUnmapCallback> SecondaryAllocator;
-typedef CombinedAllocator<PrimaryAllocator, AllocatorCache,
-                          SecondaryAllocator> Allocator;
-
 static Allocator allocator;
 static AllocatorCache fallback_allocator_cache;
 static SpinMutex fallback_mutex;
 
+Allocator &get_allocator() { return allocator; }
+
 void MsanAllocatorInit() {
   allocator.Init(
       common_flags()->allocator_may_return_null,
       common_flags()->allocator_release_to_os_interval_ms);
 }
 
 AllocatorCache *GetAllocatorCache(MsanThreadLocalMallocStorage *ms) {
   CHECK(ms);

lib/msan/msan_interceptors.cc

if (p != (void *)-1) {
     if (!res) {
       __msan_unpoison(p, ds.shm_segsz);
     }
   }
   return p;
 }
 
 static void BeforeFork() {
+  get_allocator().ForceLock();
   StackDepotLockAll();
   ChainedOriginDepotLockAll();
 }
 
 static void AfterFork() {
   ChainedOriginDepotUnlockAll();
   StackDepotUnlockAll();
+  get_allocator().ForceUnlock();
 }
 
 INTERCEPTOR(int, fork, void) {
   ENSURE_MSAN_INITED();
   BeforeFork();
   int pid = REAL(fork)();
   AfterFork();
   return pid;

lib/tsan/rtl/tsan_dense_alloc.h

void FlushCache(Cache *c) {
     }
   }
 
   void InitCache(Cache *c) {
     c->pos = 0;
     internal_memset(c->cache, 0, sizeof(c->cache));
   }
 
+  void Lock() {
+    mtx_.Lock();
+  }
+
+  void Unlock() {
+    mtx_.Unlock();
+  }
+
  private:
   T *map_[kL1Size];
   SpinMutex mtx_;
   IndexT freelist_;
   uptr fillpos_;
 
   void Refill(Cache *c) {
     SpinMutexLock lock(&mtx_);

lib/tsan/rtl/tsan_rtl.cc

 // Main file (entry points) for the TSan run-time.
 //===----------------------------------------------------------------------===//
 
 #include "sanitizer_common/sanitizer_atomic.h"
 #include "sanitizer_common/sanitizer_common.h"
 #include "sanitizer_common/sanitizer_libc.h"
 #include "sanitizer_common/sanitizer_stackdepot.h"
 #include "sanitizer_common/sanitizer_placement_new.h"
+#include "sanitizer_common/sanitizer_stackdepot.h"
 #include "sanitizer_common/sanitizer_symbolizer.h"
 #include "tsan_defs.h"
 #include "tsan_platform.h"
 #include "tsan_rtl.h"
 #include "tsan_mman.h"
 #include "tsan_suppressions.h"
 #include "tsan_symbolize.h"
 #include "ubsan/ubsan_init.h"
#if TSAN_COLLECT_STATS
   StatOutput(ctx->stat);
 #endif
 
   return failed ? common_flags()->exitcode : 0;
 }
 
 #if !SANITIZER_GO
 void ForkBefore(ThreadState *thr, uptr pc) {
+  allocator()->ForceLock();
+  ctx->metamap.Lock();
+  StackDepotLockAll();
   ctx->thread_registry->Lock();
   ctx->report_mtx.Lock();
 }
 
 void ForkParentAfter(ThreadState *thr, uptr pc) {
   ctx->report_mtx.Unlock();
   ctx->thread_registry->Unlock();
+  StackDepotUnlockAll();
+  ctx->metamap.Unlock();
+  allocator()->ForceUnlock();
 }
 
 void ForkChildAfter(ThreadState *thr, uptr pc) {
   ctx->report_mtx.Unlock();
   ctx->thread_registry->Unlock();
+  StackDepotUnlockAll();
+  ctx->metamap.Unlock();
+  allocator()->ForceUnlock();
 
   uptr nthread = 0;
   ctx->thread_registry->GetNumberOfThreads(0, 0, &nthread /* alive threads */);
   VPrintf(1, "ThreadSanitizer: forked new process with pid %d,"
       " parent had %d threads\n", (int)internal_getpid(), (int)nthread);
   if (nthread == 1) {
     StartBackgroundThread();
   } else {

lib/tsan/rtl/tsan_sync.h

public:
   SyncVar* GetOrCreateAndLock(ThreadState *thr, uptr pc,
                               uptr addr, bool write_lock);
   SyncVar* GetIfExistsAndLock(uptr addr, bool write_lock);
 
   void MoveMemory(uptr src, uptr dst, uptr sz);
 
   void OnProcIdle(Processor *proc);
 
+  void Lock();
+
+  void Unlock();
+
  private:
   static const u32 kFlagMask  = 3u << 30;
   static const u32 kFlagBlock = 1u << 30;
   static const u32 kFlagSync  = 2u << 30;
   typedef DenseSlabAlloc<MBlock, 1<<16, 1<<12> BlockAlloc;
   typedef DenseSlabAlloc<SyncVar, 1<<16, 1<<10> SyncAlloc;
   BlockAlloc block_alloc_;
   SyncAlloc sync_alloc_;

lib/tsan/rtl/tsan_sync.cc

void MetaMap::MoveMemory(uptr src, uptr dst, uptr sz) {
   }
 }
 
 void MetaMap::OnProcIdle(Processor *proc) {
   block_alloc_.FlushCache(&proc->block_cache);
   sync_alloc_.FlushCache(&proc->sync_cache);
 }
 
+void MetaMap::Lock() {
+  block_alloc_.Lock();
+}
+
+void MetaMap::Unlock() {
+  block_alloc_.Unlock();
+}
+
 }  // namespace __tsan

test/sanitizer_common/TestCases/Linux/allocator_fork_no_hang.cc

+// https://github.com/google/sanitizers/issues/774
+// Test that sanitizer allocator is fork-safe.
+// Run a number of threads that perform memory allocation/deallocation, then fork
+// and verify that malloc/free do not deadlock in the child process.
+
+// RUN: %clangxx -std=c++11 -O0 %s -o %t
+// RUN: ASAN_OPTIONS=detect_leaks=0 %run %t 2>&1 | FileCheck %s
+
+// Fun fact: if test output is redirected to a file (as opposed to
+// being piped directly to FileCheck), we may lose some "done"s due to
+// a kernel bug:
+// https://lkml.org/lkml/2014/2/17/324
+
+// Flaky on PPC64.
+// UNSUPPORTED: powerpc64-target-arch
+// UNSUPPORTED: powerpc64le-target-arch
+
+#include <pthread.h>
+#include <unistd.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <sys/types.h>
+#include <sys/wait.h>
+#include <sys/time.h>
+#include <signal.h>
+#include <errno.h>
+
+int done;
+
+void *worker(void *arg) {
+  while (true) {
+    void *p = malloc(4);
+    if (__atomic_load_n(&done, __ATOMIC_RELAXED))
+      return 0;
+  }
+  return 0;
+}
+
+// Run through malloc/free in the child process.
+// This can deadlock on allocator cache refilling.
+void child() {
+  for (int i = 0; i < 10000; ++i) {
+    void *p = malloc(4);
+  }
+  write(2, "done\n", 5);
+}
+
+void test() {
+  const int kThreads = 10;
+  pthread_t t[kThreads];
+  for (int i = 0; i < kThreads; ++i)
+    pthread_create(&t[i], NULL, worker, (void*)(long)i);
+  usleep(100000);
+  pid_t pid = fork();
+  if (pid) {
+    // parent
+    __atomic_store_n(&done, 1, __ATOMIC_RELAXED);
+    pid_t p;
+    while ((p = wait(NULL)) == -1) {  }
+  } else {
+    // child
+    child();
+  }
+}
+
+int main() {
+  const int kChildren = 30;
+  for (int i = 0; i < kChildren; ++i) {
+    pid_t pid = fork();
+    if (pid) {
+      // parent
+    } else {
+      test();
+      exit(0);
+    }
+  }
+
+  for (int i = 0; i < kChildren; ++i) {
+    pid_t p;
+    while ((p = wait(NULL)) == -1) {  }
+  }
+
+  return 0;
+}
+
+// Expect 30 (== kChildren) "done" messages.
+// CHECK: done
+// CHECK: done
+// CHECK: done
+// CHECK: done
+// CHECK: done
+// CHECK: done
+// CHECK: done
+// CHECK: done
+// CHECK: done
+// CHECK: done
+// CHECK: done
+// CHECK: done
+// CHECK: done
+// CHECK: done
+// CHECK: done
+// CHECK: done
+// CHECK: done
+// CHECK: done
+// CHECK: done
+// CHECK: done
+// CHECK: done
+// CHECK: done
+// CHECK: done
+// CHECK: done
+// CHECK: done
+// CHECK: done
+// CHECK: done
+// CHECK: done
+// CHECK: done
+// CHECK: done