[scudo][standalone] Fork support

[cryptoad]

Summary

fork() wasn't well (or at all) supported in Scudo. This materialized
in deadlocks in children.

In order to properly support fork, we will lock the allocator pre-fork
and unlock it post-fork in parent and child. This is done via a
pthread_atfork call installing the necessary handlers.

A couple of things suck here: this function allocates - so this has to
be done post initialization as our init path is not reentrant, and it
doesn't allow for an extra pointer - so we can't pass the allocator we
are currently working with.

In order to work around this, I added a post-init template parameter
that gets executed once the allocator is initialized for the current
thread. Its job for the C wrappers is to install the atfork handlers.

I reorganized a bit the impacted area and added some tests, courtesy
of cferris@ that were deadlocking prior to this fix.

Edit: checking if this triggers presubmit.

[cryptoad]

Corrections.

[cryptoad]

Adding Quarantine & ByteMap to the things we disable/enable.

[cryptoad]

Correct several inconsistencies wrt locking orders.

[cryptoad]

Wrong version of the makefiles.

[cryptoad]

Disable GWP-ASan for now.

[eugenis]

compiler-rt/lib/scudo/standalone/wrappers_c.inc

161

why is this signalling necessary?
disable() while disabled should block on any of the allocator mutexes just fine.

[cryptoad]

compiler-rt/lib/scudo/standalone/wrappers_c.inc

161

Yes I might have overthought it.
I was thinking about a situation where one thread would do disable() then one fork() then another enable().
In my head that was going to block somewhere, but in the end I think you are right and that it's not necessary.
I will remove it.

[cryptoad]

Adding a disable()/fork()/enable() test.

[cryptoad]

compiler-rt/lib/scudo/standalone/wrappers_c.inc

161

Coming back to this.
This is actually needed for the scenario I was describing.
I added a test to illustrate it: doing enable() in a thread while another has done disable()/fork() should enable the allocator again and allow the fork() to go on.

[eugenis]

SCUDO_PREFIX(Mutex) can be acquired in the child process, and in the parent process, and is not held during fork. That's a potential deadlock.
The condvar code reduces the window for it, but does not eliminate it completely.
Why can't malloc_disable() simply directly call allocator.disable(), and the same in malloc_enable?

[cryptoad]

In D72470#1814446, @eugenis wrote:

SCUDO_PREFIX(Mutex) can be acquired in the child process, and in the parent process, and is not held during fork. That's a potential deadlock.
The condvar code reduces the window for it, but does not eliminate it completely.
Why can't malloc_disable() simply directly call allocator.disable(), and the same in malloc_enable?

You are right. Everything still works with just disable/enable. I will change this.

[cryptoad]

Applying simplification suggested by eugenis@: malloc_disable just
calls disable, same for enable.

[cryptoad]

I apologize for being pushy, but could I get a yay or nay on this to unblock Android?
Thanks!

[cryptoad]

Remove the NOINLINE for enable & disable, they were used to debug
some deadlocks.

[eugenis]

I'm concerned about the removal of gwp-asan.
Could you check if https://reviews.llvm.org/D72546 helps?

[cryptoad]

Re-enable GWP-ASan, add a comment about the potentially failing test.
We will skip it if this occurs upstream.

[eugenis]

LGTM

[cryptoad]

Adding a newline at end of file, to try and trigger presubmit.

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: unknown.

clang-format: pass.

Build artifacts: diff.json, clang-format.patch, CMakeCache.txt, console-log.txt

[cryptoad]

Rebasing.

[merge_guards_bot]

Unit tests: pass. 61794 tests passed, 0 failed and 781 were skipped.

clang-tidy: unknown.

clang-format: pass.

Build artifacts: diff.json, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[eugenis]

compiler-rt/lib/scudo/standalone/combined.h

31

A static global in a header is not a great idea - every translation unit that includes this will get its own instance of gwp-asan, and only one of those will be initialized.

This needs to be moved to the wrappers, probably.

[eugenis]

FWIW I'm ok with this landing in the previous version, disabling GWP-ASan, because its integration in scudo is so broken right now.

[hctim]

In D72470#1818424, @eugenis wrote:

FWIW I'm ok with this landing in the previous version, disabling GWP-ASan, because its integration in scudo is so broken right now.

+1 - we can fix GWP-ASan-in-scudo-standalone later.

[eugenis]

compiler-rt/lib/scudo/standalone/combined.h

31

I'm not sure I completely understand this comment about constexpr initialization.
But if the problem is that gwp-asan is not zero initialized, then it seems really easy to fix. There are only two non-zero fields:

• AdjustedSampleRate can be stored in a pre-decremented form. It is not used on the hot path.
• GuardedPagePool - I don't see why it can not start as 0

Then we make GuardedAlloc a member of Allocator, and we are done. What am I missing?

[cryptoad]

Disable GWP-ASan, FIXME added.

[merge_guards_bot]

Unit tests: pass. 61822 tests passed, 0 failed and 781 were skipped.

clang-tidy: unknown.

clang-format: pass.

Build artifacts: diff.json, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

compiler-rt/lib/scudo/standalone/CMakeLists.txt

 add_compiler_rt_component(scudo_standalone)
-if (COMPILER_RT_HAS_GWP_ASAN)
+if (FALSE AND COMPILER_RT_HAS_GWP_ASAN)
   add_dependencies(scudo_standalone gwp_asan)
 endif()
 
 include_directories(../..)
 
 set(SCUDO_CFLAGS)
 
 list(APPEND SCUDO_CFLAGS
set(SCUDO_SOURCES_C_WRAPPERS
   )
 
 set(SCUDO_SOURCES_CXX_WRAPPERS
   wrappers_cpp.cpp
   )
 
 set(SCUDO_OBJECT_LIBS)
 
-if (COMPILER_RT_HAS_GWP_ASAN)
+if (FALSE AND COMPILER_RT_HAS_GWP_ASAN)
   list(APPEND SCUDO_OBJECT_LIBS RTGwpAsan)
   list(APPEND SCUDO_CFLAGS -DGWP_ASAN_HOOKS)
 endif()
 
 if(COMPILER_RT_HAS_SCUDO_STANDALONE)
   add_compiler_rt_object_libraries(RTScudoStandalone
     ARCHS ${SCUDO_STANDALONE_SUPPORTED_ARCH}
     SOURCES ${SCUDO_SOURCES}

compiler-rt/lib/scudo/standalone/bytemap.h

void set(uptr Index, u8 Value) {
     DCHECK_EQ(0U, Map[Index]);
     Map[Index] = Value;
   }
   u8 operator[](uptr Index) {
     DCHECK_LT(Index, Size);
     return Map[Index];
   }
 
+  void disable() {}
+  void enable() {}
+
 private:
   u8 *Map;
 };
 
 template <uptr Level1Size, uptr Level2Size> class TwoLevelByteMap {
 public:
   void initLinkerInitialized() {
     Level1Map = reinterpret_cast<atomic_uptr *>(
public:
   u8 operator[](uptr Index) const {
     DCHECK_LT(Index, Level1Size * Level2Size);
     u8 *Level2Map = get(Index / Level2Size);
     if (!Level2Map)
       return 0;
     return Level2Map[Index % Level2Size];
   }
 
+  NOINLINE void disable() { Mutex.lock(); }
+  NOINLINE void enable() { Mutex.unlock(); }
+
 private:
   u8 *get(uptr Index) const {
     DCHECK_LT(Index, Level1Size);
     return reinterpret_cast<u8 *>(
         atomic_load(&Level1Map[Index], memory_order_acquire));
   }
 
   u8 *getOrCreate(uptr Index) {

compiler-rt/lib/scudo/standalone/combined.h

 
 #ifdef GWP_ASAN_HOOKS
 #include "gwp_asan/guarded_pool_allocator.h"
 // GWP-ASan is declared here in order to avoid indirect call overhead. It's also
 // instantiated outside of the Allocator class, as the allocator is only
 // zero-initialised. GWP-ASan requires constant initialisation, and the Scudo
 // allocator doesn't have a constexpr constructor (see discussion here:
 // https://reviews.llvm.org/D69265#inline-624315).
 static gwp_asan::GuardedPoolAllocator GuardedAlloc;
 #endif // GWP_ASAN_HOOKS
 
+extern "C" inline void EmptyCallback() {}
+
 namespace scudo {
 
-template <class Params> class Allocator {
+template <class Params, void (*PostInitCallback)(void) = EmptyCallback>
+class Allocator {
 public:
   using PrimaryT = typename Params::Primary;
   using CacheT = typename PrimaryT::CacheT;
-  typedef Allocator<Params> ThisT;
+  typedef Allocator<Params, PostInitCallback> ThisT;
   typedef typename Params::template TSDRegistryT<ThisT> TSDRegistryT;
 
+  void callPostInitCallback() {
+    static pthread_once_t OnceControl = PTHREAD_ONCE_INIT;
+    pthread_once(&OnceControl, PostInitCallback);
+  }
+
   struct QuarantineCallback {
     explicit QuarantineCallback(ThisT &Instance, CacheT &LocalCache)
         : Allocator(Instance), Cache(LocalCache) {}
 
     // Chunk recycling function, returns a quarantined chunk to the backend,
     // first making sure it hasn't been tampered with.
     void recycle(void *Ptr) {
       Chunk::UnpackedHeader Header;
if (NewPtr) {
       quarantineOrDeallocateChunk(OldPtr, &OldHeader, OldSize);
     }
     return NewPtr;
   }
 
   // TODO(kostyak): disable() is currently best-effort. There are some small
   //                windows of time when an allocation could still succeed after
   //                this function finishes. We will revisit that later.
-  void disable() {
+  NOINLINE void disable() {
     initThreadMaybe();
     TSDRegistry.disable();
+    Stats.disable();
+    Quarantine.disable();
+    Primary.disable();
     Secondary.disable();
   }
 
-  void enable() {
+  NOINLINE void enable() {
     initThreadMaybe();
     Secondary.enable();
+    Primary.enable();
+    Quarantine.enable();
+    Stats.enable();
     TSDRegistry.enable();
   }
 
   // The function returns the amount of bytes required to store the statistics,
   // which might be larger than the amount of bytes provided. Note that the
   // statistics buffer is not necessarily constant between calls to this
   // function. This can be called with a null buffer or zero size for buffer
   // sizing purposes.

compiler-rt/lib/scudo/standalone/primary32.h

void pushBatch(uptr ClassId, TransferBatch *B) {
     SizeClassInfo *Sci = getSizeClassInfo(ClassId);
     ScopedLock L(Sci->Mutex);
     Sci->FreeList.push_front(B);
     Sci->Stats.PushedBlocks += B->getCount();
     if (Sci->CanRelease)
       releaseToOSMaybe(Sci, ClassId);
   }
 
-  void disable() {
-    for (uptr I = 0; I < NumClasses; I++)
-      getSizeClassInfo(I)->Mutex.lock();
+  NOINLINE void disable() {
+    // The BatchClassId must be locked last since other classes can use it.
+    for (sptr I = static_cast<sptr>(NumClasses) - 1; I >= 0; I--) {
+      if (static_cast<uptr>(I) == SizeClassMap::BatchClassId)
+        continue;
+      getSizeClassInfo(static_cast<uptr>(I))->Mutex.lock();
+    }
+    getSizeClassInfo(SizeClassMap::BatchClassId)->Mutex.lock();
+    RegionsStashMutex.lock();
+    PossibleRegions.disable();
   }
 
-  void enable() {
-    for (sptr I = static_cast<sptr>(NumClasses) - 1; I >= 0; I--)
-      getSizeClassInfo(static_cast<uptr>(I))->Mutex.unlock();
+  NOINLINE void enable() {
+    PossibleRegions.enable();
+    RegionsStashMutex.unlock();
+    getSizeClassInfo(SizeClassMap::BatchClassId)->Mutex.unlock();
+    for (uptr I = 0; I < NumClasses; I++) {
+      if (I == SizeClassMap::BatchClassId)
+        continue;
+      getSizeClassInfo(I)->Mutex.unlock();
+    }
   }
 
   template <typename F> void iterateOverBlocks(F Callback) {
     for (uptr I = MinRegionIndex; I <= MaxRegionIndex; I++)
       if (PossibleRegions[I]) {
         const uptr BlockSize = getSizeByClassId(PossibleRegions[I]);
         const uptr From = I * RegionSize;
         const uptr To = From + (RegionSize / BlockSize) * BlockSize;

compiler-rt/lib/scudo/standalone/primary64.h

void pushBatch(uptr ClassId, TransferBatch *B) {
     RegionInfo *Region = getRegionInfo(ClassId);
     ScopedLock L(Region->Mutex);
     Region->FreeList.push_front(B);
     Region->Stats.PushedBlocks += B->getCount();
     if (Region->CanRelease)
       releaseToOSMaybe(Region, ClassId);
   }
 
-  void disable() {
-    for (uptr I = 0; I < NumClasses; I++)
-      getRegionInfo(I)->Mutex.lock();
+  NOINLINE void disable() {
+    // The BatchClassId must be locked last since other classes can use it.
+    for (sptr I = static_cast<sptr>(NumClasses) - 1; I >= 0; I--) {
+      if (static_cast<uptr>(I) == SizeClassMap::BatchClassId)
+        continue;
+      getRegionInfo(static_cast<uptr>(I))->Mutex.lock();
+    }
+    getRegionInfo(SizeClassMap::BatchClassId)->Mutex.lock();
   }
 
-  void enable() {
-    for (sptr I = static_cast<sptr>(NumClasses) - 1; I >= 0; I--)
-      getRegionInfo(static_cast<uptr>(I))->Mutex.unlock();
+  NOINLINE void enable() {
+    getRegionInfo(SizeClassMap::BatchClassId)->Mutex.unlock();
+    for (uptr I = 0; I < NumClasses; I++) {
+      if (I == SizeClassMap::BatchClassId)
+        continue;
+      getRegionInfo(I)->Mutex.unlock();
+    }
   }
 
   template <typename F> void iterateOverBlocks(F Callback) const {
     for (uptr I = 0; I < NumClasses; I++) {
       if (I == SizeClassMap::BatchClassId)
         continue;
       const RegionInfo *Region = getRegionInfo(I);
       const uptr BlockSize = getSizeByClassId(I);

compiler-rt/lib/scudo/standalone/quarantine.h

if (C->getSize() > getCacheSize())
       drain(C, Cb);
   }
 
   void NOINLINE drain(CacheT *C, Callback Cb) {
     {
       ScopedLock L(CacheMutex);
       Cache.transfer(C);
     }
-    if (Cache.getSize() > getMaxSize() && RecyleMutex.tryLock())
+    if (Cache.getSize() > getMaxSize() && RecycleMutex.tryLock())
       recycle(atomic_load_relaxed(&MinSize), Cb);
   }
 
   void NOINLINE drainAndRecycle(CacheT *C, Callback Cb) {
     {
       ScopedLock L(CacheMutex);
       Cache.transfer(C);
     }
-    RecyleMutex.lock();
+    RecycleMutex.lock();
     recycle(0, Cb);
   }
 
   void getStats(ScopedString *Str) const {
     // It assumes that the world is stopped, just as the allocator's printStats.
     Cache.getStats(Str);
     Str->append("Quarantine limits: global: %zuK; thread local: %zuK\n",
                 getMaxSize() >> 10, getCacheSize() >> 10);
   }
 
+  NOINLINE void disable() {
+    // RecycleMutex must be locked 1st since we grab CacheMutex within recycle.
+    RecycleMutex.lock();
+    CacheMutex.lock();
+  }
+
+  NOINLINE void enable() {
+    CacheMutex.unlock();
+    RecycleMutex.unlock();
+  }
+
 private:
   // Read-only data.
   alignas(SCUDO_CACHE_LINE_SIZE) HybridMutex CacheMutex;
   CacheT Cache;
-  alignas(SCUDO_CACHE_LINE_SIZE) HybridMutex RecyleMutex;
+  alignas(SCUDO_CACHE_LINE_SIZE) HybridMutex RecycleMutex;
   atomic_uptr MinSize;
   atomic_uptr MaxSize;
   alignas(SCUDO_CACHE_LINE_SIZE) atomic_uptr MaxCacheSize;
 
   void NOINLINE recycle(uptr MinSize, Callback Cb) {
     CacheT Tmp;
     Tmp.init();
     {
Tmp.init();
               CacheSize * OverheadThresholdPercents) {
         Cache.mergeBatches(&Tmp);
       }
       // Extract enough chunks from the quarantine to get below the max
       // quarantine size and leave some leeway for the newly quarantined chunks.
       while (Cache.getSize() > MinSize)
         Tmp.enqueueBatch(Cache.dequeueBatch());
     }
-    RecyleMutex.unlock();
+    RecycleMutex.unlock();
     doRecycle(&Tmp, Cb);
   }
 
   void NOINLINE doRecycle(CacheT *C, Callback Cb) {
     while (QuarantineBatch *B = C->dequeueBatch()) {
       const u32 Seed = static_cast<u32>(
           (reinterpret_cast<uptr>(B) ^ reinterpret_cast<uptr>(C)) >> 4);
       B->shuffle(Seed);

compiler-rt/lib/scudo/standalone/secondary.h

public:
   }
 
   static uptr getBlockSize(void *Ptr) {
     return getBlockEnd(Ptr) - reinterpret_cast<uptr>(Ptr);
   }
 
   void getStats(ScopedString *Str) const;
 
-  void disable() { Mutex.lock(); }
+  NOINLINE void disable() { Mutex.lock(); }
 
-  void enable() { Mutex.unlock(); }
+  NOINLINE void enable() { Mutex.unlock(); }
 
   template <typename F> void iterateOverBlocks(F Callback) const {
     for (const auto &H : InUseBlocks)
       Callback(reinterpret_cast<uptr>(&H) + LargeBlock::getHeaderSize());
   }
 
   static uptr getMaxFreeListSize(void) { return MaxFreeListSize; }
 

compiler-rt/lib/scudo/standalone/stats.h

for (const auto &Stats : StatsList) {
       for (uptr I = 0; I < StatCount; I++)
         S[I] += Stats.get(static_cast<StatType>(I));
     }
     // All stats must be non-negative.
     for (uptr I = 0; I < StatCount; I++)
       S[I] = static_cast<sptr>(S[I]) >= 0 ? S[I] : 0;
   }
 
+  NOINLINE void disable() { Mutex.lock(); }
+  NOINLINE void enable() { Mutex.unlock(); }
+
 private:
   mutable HybridMutex Mutex;
   DoublyLinkedList<LocalStats> StatsList;
 };
 
 } // namespace scudo
 
 #endif // SCUDO_STATS_H_

compiler-rt/lib/scudo/standalone/tests/CMakeLists.txt

set(SCUDO_UNITTEST_CFLAGS
   # Extra flags for the C++ tests
   # TODO(kostyak): find a way to make -fsized-deallocation work
   -Wno-mismatched-new-delete)
 
 if(ANDROID)
   list(APPEND SCUDO_UNITTEST_CFLAGS -fno-emulated-tls)
 endif()
 
-if (COMPILER_RT_HAS_GWP_ASAN)
+if (FALSE AND COMPILER_RT_HAS_GWP_ASAN)
   list(APPEND SCUDO_UNITTEST_CFLAGS -DGWP_ASAN_HOOKS)
 endif()
 
 set(SCUDO_TEST_ARCH ${SCUDO_STANDALONE_SUPPORTED_ARCH})
 
 # gtests requires c++
 set(LINK_FLAGS ${COMPILER_RT_UNITTEST_LINK_FLAGS})
 foreach(lib ${SANITIZER_TEST_CXX_LIBRARIES})
   list(APPEND LINK_FLAGS -l${lib})
 endforeach()
 list(APPEND LINK_FLAGS -pthread)
 # Linking against libatomic is required with some compilers
 list(APPEND LINK_FLAGS -latomic)
 
 set(SCUDO_TEST_HEADERS)
 foreach (header ${SCUDO_HEADERS})
   list(APPEND SCUDO_TEST_HEADERS ${CMAKE_CURRENT_SOURCE_DIR}/../${header})
 endforeach()
 
 macro(add_scudo_unittest testname)
   cmake_parse_arguments(TEST "" "" "SOURCES;ADDITIONAL_RTOBJECTS" ${ARGN})
-  if (COMPILER_RT_HAS_GWP_ASAN)
+  if (FALSE AND COMPILER_RT_HAS_GWP_ASAN)
     list(APPEND TEST_ADDITIONAL_RTOBJECTS RTGwpAsan)
   endif()
 
   if(COMPILER_RT_HAS_SCUDO_STANDALONE)
     foreach(arch ${SCUDO_TEST_ARCH})
       # Additional runtime objects get added along RTScudoStandalone
       set(SCUDO_TEST_RTOBJECTS $<TARGET_OBJECTS:RTScudoStandalone.${arch}>)
       foreach(rtobject ${TEST_ADDITIONAL_RTOBJECTS})

compiler-rt/lib/scudo/standalone/tests/tsd_test.cpp

void initLinkerInitialized() {
     Initialized = true;
   }
   void reset() { memset(this, 0, sizeof(*this)); }
 
   void unmapTestOnly() { TSDRegistry.unmapTestOnly(); }
   void initCache(CacheT *Cache) { memset(Cache, 0, sizeof(*Cache)); }
   void commitBack(scudo::TSD<MockAllocator> *TSD) {}
   TSDRegistryT *getTSDRegistry() { return &TSDRegistry; }
+  void callPostInitCallback() {}
 
   bool isInitialized() { return Initialized; }
 
 private:
   bool Initialized;
   TSDRegistryT TSDRegistry;
 };
 

compiler-rt/lib/scudo/standalone/tests/wrappers_c_test.cpp

EXPECT_DEATH(
         malloc_disable();
         alarm(1);
         P = malloc(Size);
         malloc_enable();
       },
       "");
 }
 
+// Fuchsia doesn't have fork or malloc_info.
 #if !SCUDO_FUCHSIA
+
 TEST(ScudoWrappersCTest, MallocInfo) {
   char Buffer[64];
   FILE *F = fmemopen(Buffer, sizeof(Buffer), "w+");
   EXPECT_NE(F, nullptr);
   errno = 0;
   EXPECT_EQ(malloc_info(0, F), 0);
   EXPECT_EQ(errno, 0);
   fclose(F);
   EXPECT_EQ(strncmp(Buffer, "<malloc version=\"scudo-", 23), 0);
 }
-#endif
+
+TEST(ScudoWrappersCTest, MallocFork) {
+  void *P;
+  pid_t Pid = fork();
+  EXPECT_GE(Pid, 0);
+  if (Pid == 0) {
+    P = malloc(Size);
+    EXPECT_NE(P, nullptr);
+    memset(P, 0x42, Size);
+    free(P);
+    _exit(0);
+  }
+  waitpid(Pid, nullptr, 0);
+  P = malloc(Size);
+  EXPECT_NE(P, nullptr);
+  memset(P, 0x42, Size);
+  free(P);
+
+  // fork should stall if the allocator has been disabled.
+  EXPECT_DEATH(
+      {
+        malloc_disable();
+        alarm(1);
+        Pid = fork();
+        EXPECT_GE(Pid, 0);
+      },
+      "");
+}
+
+static pthread_mutex_t Mutex;
+static pthread_cond_t Conditional = PTHREAD_COND_INITIALIZER;
+
+static void *enableMalloc(void *Unused) {
+  // Initialize the allocator for this thread.
+  void *P = malloc(Size);
+  EXPECT_NE(P, nullptr);
+  memset(P, 0x42, Size);
+  free(P);
+
+  // Signal the main thread we are ready.
+  pthread_mutex_lock(&Mutex);
+  pthread_cond_signal(&Conditional);
+  pthread_mutex_unlock(&Mutex);
+
+  // Wait for the malloc_disable & fork, then enable the allocator again.
+  sleep(1);
+  malloc_enable();
+
+  return nullptr;
+}
+
+TEST(ScudoWrappersCTest, MallocDisableForkEnable) {
+  pthread_t ThreadId;
+  EXPECT_EQ(pthread_create(&ThreadId, nullptr, &enableMalloc, nullptr), 0);
+
+  // Wait for the thread to be warmed up.
+  pthread_mutex_lock(&Mutex);
+  pthread_cond_wait(&Conditional, &Mutex);
+  pthread_mutex_unlock(&Mutex);
+
+  // Disable the allocator and fork. fork should succeed after malloc_enable.
+  malloc_disable();
+  pid_t Pid = fork();
+  EXPECT_GE(Pid, 0);
+  if (Pid == 0) {
+    void *P = malloc(Size);
+    EXPECT_NE(P, nullptr);
+    memset(P, 0x42, Size);
+    free(P);
+    _exit(0);
+  }
+  waitpid(Pid, nullptr, 0);
+  EXPECT_EQ(pthread_join(ThreadId, 0), 0);
+}
+
+#endif // SCUDO_FUCHSIA

compiler-rt/lib/scudo/standalone/tests/wrappers_cpp_test.cpp

 //===-- wrappers_cpp_test.cpp -----------------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "tests/scudo_unit_test.h"
 
+#include <atomic>
 #include <condition_variable>
 #include <mutex>
 #include <thread>
 #include <vector>
 
 void operator delete(void *, size_t) noexcept;
 void operator delete[](void *, size_t) noexcept;
 
TEST(ScudoWrappersCppTest, ThreadedNew) {
   {
     std::unique_lock<std::mutex> Lock(Mutex);
     Ready = true;
     Cv.notify_all();
   }
   for (auto &T : Threads)
     T.join();
 }
+
+#if !SCUDO_FUCHSIA
+TEST(ScudoWrappersCppTest, AllocAfterFork) {
+  std::atomic_bool Stop;
+
+  // Create threads that simply allocate and free different sizes.
+  std::vector<std::thread *> Threads;
+  for (size_t N = 0; N < 5; N++) {
+    std::thread *T = new std::thread([&Stop] {
+      while (!Stop) {
+        for (size_t SizeLog = 3; SizeLog <= 21; SizeLog++) {
+          char *P = new char[1UL << SizeLog];
+          EXPECT_NE(P, nullptr);
+          // Make sure this value is not optimized away.
+          asm volatile("" : : "r,m"(P) : "memory");
+          delete[] P;
+        }
+      }
+    });
+    Threads.push_back(T);
+  }
+
+  // Create a thread to fork and allocate.
+  for (size_t N = 0; N < 100; N++) {
+    pid_t Pid;
+    if ((Pid = fork()) == 0) {
+      for (size_t SizeLog = 3; SizeLog <= 21; SizeLog++) {
+        char *P = new char[1UL << SizeLog];
+        EXPECT_NE(P, nullptr);
+        // Make sure this value is not optimized away.
+        asm volatile("" : : "r,m"(P) : "memory");
+        // Make sure we can touch all of the allocation.
+        memset(P, 0x32, 1U << SizeLog);
+        // EXPECT_LE(1U << SizeLog, malloc_usable_size(ptr));
+        delete[] P;
+      }
+      _exit(10);
+    }
+    EXPECT_NE(-1, Pid);
+    int Status;
+    EXPECT_EQ(Pid, waitpid(Pid, &Status, 0));
+    EXPECT_FALSE(WIFSIGNALED(Status));
+    EXPECT_EQ(10, WEXITSTATUS(Status));
+  }
+
+  printf("Waiting for threads to complete\n");
+  Stop = true;
+  for (auto Thread : Threads)
+    Thread->join();
+  Threads.clear();
+}
+#endif
 No newline at end of file

compiler-rt/lib/scudo/standalone/tsd.h

 //===-- tsd.h ---------------------------------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef SCUDO_TSD_H_
 #define SCUDO_TSD_H_
 
 #include "atomic_helpers.h"
 #include "common.h"
 #include "mutex.h"
 
 #include <limits.h> // for PTHREAD_DESTRUCTOR_ITERATIONS
+#include <pthread.h>
 
 // With some build setups, this might still not be defined.
 #ifndef PTHREAD_DESTRUCTOR_ITERATIONS
 #define PTHREAD_DESTRUCTOR_ITERATIONS 4
 #endif
 
 namespace scudo {
 

compiler-rt/lib/scudo/standalone/tsd_exclusive.h

 //===-- tsd_exclusive.h -----------------------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef SCUDO_TSD_EXCLUSIVE_H_
 #define SCUDO_TSD_EXCLUSIVE_H_
 
 #include "tsd.h"
 
-#include <pthread.h>
-
 namespace scudo {
 
 enum class ThreadState : u8 {
   NotInitialized = 0,
   Initialized,
   TornDown,
 };
 
ALWAYS_INLINE TSD<Allocator> *getTSDAndLock(bool *UnlockRequired) {
     DCHECK(FallbackTSD);
     FallbackTSD->lock();
     *UnlockRequired = true;
     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() {
+  NOINLINE void disable() {
+    Mutex.lock();
     FallbackTSD->lock();
     atomic_store(&Disabled, 1U, memory_order_release);
   }
 
-  void enable() {
+  NOINLINE void enable() {
     atomic_store(&Disabled, 0U, memory_order_release);
     FallbackTSD->unlock();
+    Mutex.unlock();
   }
 
 private:
   void initOnceMaybe(Allocator *Instance) {
     ScopedLock L(Mutex);
     if (LIKELY(Initialized))
       return;
     initLinkerInitialized(Instance); // Sets Initialized.
   }
 
   // Using minimal initialization allows for global initialization while keeping
   // the thread specific structure untouched. The fallback structure will be
   // used instead.
   NOINLINE void initThread(Allocator *Instance, bool MinimalInit) {
     initOnceMaybe(Instance);
     if (UNLIKELY(MinimalInit))
       return;
     CHECK_EQ(
         pthread_setspecific(PThreadKey, reinterpret_cast<void *>(Instance)), 0);
     ThreadTSD.initLinkerInitialized(Instance);
     State = ThreadState::Initialized;
+    Instance->callPostInitCallback();
   }
 
   pthread_key_t PThreadKey;
   bool Initialized;
   atomic_u8 Disabled;
   TSD<Allocator> *FallbackTSD;
   HybridMutex Mutex;
   static THREADLOCAL ThreadState State;

compiler-rt/lib/scudo/standalone/tsd_shared.h

 //===-- tsd_shared.h --------------------------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef SCUDO_TSD_SHARED_H_
 #define SCUDO_TSD_SHARED_H_
 
 #include "linux.h" // for getAndroidTlsPtr()
 #include "tsd.h"
 
-#include <pthread.h>
-
 namespace scudo {
 
 template <class Allocator, u32 MaxTSDCount> struct TSDRegistrySharedT {
   void initLinkerInitialized(Allocator *Instance) {
     Instance->initLinkerInitialized();
     CHECK_EQ(pthread_key_create(&PThreadKey, nullptr), 0); // For non-TLS
     NumberOfTSDs = Min(Max(1U, getNumberOfCPUs()), MaxTSDCount);
     TSDs = reinterpret_cast<TSD<Allocator> *>(
ALWAYS_INLINE TSD<Allocator> *getTSDAndLock(bool *UnlockRequired) {
     *UnlockRequired = true;
     // Try to lock the currently associated context.
     if (TSD->tryLock())
       return TSD;
     // If that fails, go down the slow path.
     return getTSDAndLockSlow(TSD);
   }
 
-  void disable() {
+  NOINLINE void disable() {
+    Mutex.lock();
     for (u32 I = 0; I < NumberOfTSDs; I++)
       TSDs[I].lock();
   }
 
-  void enable() {
+  NOINLINE void enable() {
     for (u32 I = 0; I < NumberOfTSDs; I++)
       TSDs[I].unlock();
+    Mutex.unlock();
   }
 
 private:
   ALWAYS_INLINE void setCurrentTSD(TSD<Allocator> *CurrentTSD) {
 #if _BIONIC
     *getAndroidTlsPtr() = reinterpret_cast<uptr>(CurrentTSD);
 #elif SCUDO_LINUX
     ThreadTSD = CurrentTSD;
void initOnceMaybe(Allocator *Instance) {
     initLinkerInitialized(Instance); // Sets Initialized.
   }
 
   NOINLINE void initThread(Allocator *Instance) {
     initOnceMaybe(Instance);
     // Initial context assignment is done in a plain round-robin fashion.
     const u32 Index = atomic_fetch_add(&CurrentIndex, 1U, memory_order_relaxed);
     setCurrentTSD(&TSDs[Index % NumberOfTSDs]);
+    Instance->callPostInitCallback();
   }
 
   NOINLINE TSD<Allocator> *getTSDAndLockSlow(TSD<Allocator> *CurrentTSD) {
     if (MaxTSDCount > 1U && NumberOfTSDs > 1U) {
       // Use the Precedence of the current TSD as our random seed. Since we are
       // in the slow path, it means that tryLock failed, and as a result it's
       // very likely that said Precedence is non-zero.
       const u32 R = static_cast<u32>(CurrentTSD->getPrecedence());

compiler-rt/lib/scudo/standalone/wrappers_c.cpp

 
 #include "allocator_config.h"
 #include "wrappers_c.h"
 #include "wrappers_c_checks.h"
 
 #include <stdint.h>
 #include <stdio.h>
 
-static scudo::Allocator<scudo::Config> Allocator;
+#define SCUDO_PREFIX(name) name
+#define SCUDO_ALLOCATOR Allocator
+
+extern "C" void SCUDO_PREFIX(malloc_postinit)();
+static scudo::Allocator<scudo::Config, SCUDO_PREFIX(malloc_postinit)>
+    SCUDO_ALLOCATOR;
 // Pointer to the static allocator so that the C++ wrappers can access it.
 // Technically we could have a completely separated heap for C & C++ but in
 // reality the amount of cross pollination between the two is staggering.
-scudo::Allocator<scudo::Config> *AllocatorPtr = &Allocator;
-
-extern "C" {
+scudo::Allocator<scudo::Config, SCUDO_PREFIX(malloc_postinit)> *
+    CONCATENATE(SCUDO_ALLOCATOR, Ptr) = &SCUDO_ALLOCATOR;
 
-#define SCUDO_PREFIX(name) name
-#define SCUDO_ALLOCATOR Allocator
 #include "wrappers_c.inc"
+
 #undef SCUDO_ALLOCATOR
 #undef SCUDO_PREFIX
 
-INTERFACE void __scudo_print_stats(void) { Allocator.printStats(); }
-
-} // extern "C"
+extern "C" INTERFACE void __scudo_print_stats(void) { Allocator.printStats(); }
 
 #endif // !SCUDO_ANDROID || !_BIONIC

compiler-rt/lib/scudo/standalone/wrappers_c.inc

 
 // malloc-type functions have to be aligned to std::max_align_t. This is
 // distinct from (1U << SCUDO_MIN_ALIGNMENT_LOG), since C++ new-type functions
 // do not have to abide by the same requirement.
 #ifndef SCUDO_MALLOC_ALIGNMENT
 #define SCUDO_MALLOC_ALIGNMENT FIRST_32_SECOND_64(8U, 16U)
 #endif
 
+extern "C" {
+
 INTERFACE WEAK void *SCUDO_PREFIX(calloc)(size_t nmemb, size_t size) {
   scudo::uptr Product;
   if (UNLIKELY(scudo::checkForCallocOverflow(size, nmemb, &Product))) {
     if (SCUDO_ALLOCATOR.canReturnNull()) {
       errno = ENOMEM;
       return nullptr;
     }
     scudo::reportCallocOverflow(nmemb, size);

 
 INTERFACE WEAK int SCUDO_PREFIX(malloc_iterate)(
     uintptr_t base, size_t size,
     void (*callback)(uintptr_t base, size_t size, void *arg), void *arg) {
   SCUDO_ALLOCATOR.iterateOverChunks(base, size, callback, arg);
   return 0;
 }
 
+static pthread_mutex_t SCUDO_PREFIX(Mutex) = PTHREAD_MUTEX_INITIALIZER;
+static pthread_cond_t SCUDO_PREFIX(Conditional) = PTHREAD_COND_INITIALIZER;
+static bool SCUDO_PREFIX(Disabled) = false;
+
+INTERFACE WEAK void SCUDO_PREFIX(malloc_enable)() {
+  pthread_mutex_lock(&SCUDO_PREFIX(Mutex));
+  SCUDO_ALLOCATOR.enable();
+  SCUDO_PREFIX(Disabled) = false;
+  pthread_cond_broadcast(&SCUDO_PREFIX(Conditional));
+  pthread_mutex_unlock(&SCUDO_PREFIX(Mutex));
+}
+
 INTERFACE WEAK void SCUDO_PREFIX(malloc_disable)() {
+  pthread_mutex_lock(&SCUDO_PREFIX(Mutex));
+  while (SCUDO_PREFIX(Disabled))
+    pthread_cond_wait(&SCUDO_PREFIX(Conditional), &SCUDO_PREFIX(Mutex));
   SCUDO_ALLOCATOR.disable();
+  SCUDO_PREFIX(Disabled) = true;
+  pthread_mutex_unlock(&SCUDO_PREFIX(Mutex));
 }
 
-INTERFACE WEAK void SCUDO_PREFIX(malloc_enable)() { SCUDO_ALLOCATOR.enable(); }
+void SCUDO_PREFIX(malloc_postinit)() {
+  pthread_atfork(SCUDO_PREFIX(malloc_disable), SCUDO_PREFIX(malloc_enable),
+                 SCUDO_PREFIX(malloc_enable));
+}
 
 INTERFACE WEAK int SCUDO_PREFIX(mallopt)(int param, UNUSED int value) {
   if (param == M_DECAY_TIME) {
     // TODO(kostyak): set release_to_os_interval_ms accordingly.
     return 1;
   } else if (param == M_PURGE) {
     SCUDO_ALLOCATOR.releaseToOS();
     return 1;
return scudo::setErrnoOnNull(
       SCUDO_ALLOCATOR.allocate(size, scudo::Chunk::Origin::Malloc, alignment));
 }
 
 INTERFACE WEAK int SCUDO_PREFIX(malloc_info)(UNUSED int options, FILE *stream) {
   fputs("<malloc version=\"scudo-1\">", stream);
   fputs("</malloc>", stream);
   return 0;
 }
+
+} // extern "C"

compiler-rt/lib/scudo/standalone/wrappers_c_bionic.cpp

 
 #include "allocator_config.h"
 #include "wrappers_c.h"
 #include "wrappers_c_checks.h"
 
 #include <stdint.h>
 #include <stdio.h>
 
-static scudo::Allocator<scudo::AndroidConfig> Allocator;
-static scudo::Allocator<scudo::AndroidSvelteConfig> SvelteAllocator;
-
-extern "C" {
-
 // Regular MallocDispatch definitions.
 #define SCUDO_PREFIX(name) CONCATENATE(scudo_, name)
 #define SCUDO_ALLOCATOR Allocator
+
+extern "C" void SCUDO_PREFIX(malloc_postinit)();
+static scudo::Allocator<scudo::AndroidConfig, SCUDO_PREFIX(malloc_postinit)>
+    SCUDO_ALLOCATOR;
+// Pointer to the static allocator so that the C++ wrappers can access it.
+// Technically we could have a completely separated heap for C & C++ but in
+// reality the amount of cross pollination between the two is staggering.
+scudo::Allocator<scudo::AndroidConfig, SCUDO_PREFIX(malloc_postinit)> *
+    CONCATENATE(SCUDO_ALLOCATOR, Ptr) = &SCUDO_ALLOCATOR;
+
 #include "wrappers_c.inc"
+
 #undef SCUDO_ALLOCATOR
 #undef SCUDO_PREFIX
 
 // Svelte MallocDispatch definitions.
 #define SCUDO_PREFIX(name) CONCATENATE(scudo_svelte_, name)
 #define SCUDO_ALLOCATOR SvelteAllocator
+
+extern "C" void SCUDO_PREFIX(malloc_postinit)();
+static scudo::Allocator<scudo::AndroidSvelteConfig,
+                        SCUDO_PREFIX(malloc_postinit)>
+    SCUDO_ALLOCATOR;
+// Pointer to the static allocator so that the C++ wrappers can access it.
+// Technically we could have a completely separated heap for C & C++ but in
+// reality the amount of cross pollination between the two is staggering.
+scudo::Allocator<scudo::AndroidSvelteConfig, SCUDO_PREFIX(malloc_postinit)> *
+    CONCATENATE(SCUDO_ALLOCATOR, Ptr) = &SCUDO_ALLOCATOR;
+
 #include "wrappers_c.inc"
+
 #undef SCUDO_ALLOCATOR
 #undef SCUDO_PREFIX
 
 // The following is the only function that will end up initializing both
 // allocators, which will result in a slight increase in memory footprint.
 INTERFACE void __scudo_print_stats(void) {
   Allocator.printStats();
   SvelteAllocator.printStats();
 }
 
-} // extern "C"
-
 #endif // SCUDO_ANDROID && _BIONIC

compiler-rt/lib/scudo/standalone/wrappers_cpp.cpp

 
 // Skip this compilation unit if compiled as part of Bionic.
 #if !SCUDO_ANDROID || !_BIONIC
 
 #include "allocator_config.h"
 
 #include <stdint.h>
 
-extern scudo::Allocator<scudo::Config> *AllocatorPtr;
+extern "C" INTERFACE void malloc_postinit();
+extern scudo::Allocator<scudo::Config, malloc_postinit> *AllocatorPtr;
 
 namespace std {
 struct nothrow_t {};
 enum class align_val_t : size_t {};
 } // namespace std
 
 INTERFACE WEAK void *operator new(size_t size) {
   return AllocatorPtr->allocate(size, scudo::Chunk::Origin::New);