Proper fix of pthread_atfork deadlock in Musl

PartitionRoot::EnableThreadCacheIfSupported()
  ::partition_alloc::internal::ScopedGuard guard{lock_};
  ThreadCache::Create(this);
    ThreadCache::ThreadCache()
      PlatformThread::CurrentId()
        InitAtFork::InitAtFork()
          pthread_atfork()
            malloc()
              ShimMalloc()
                PartitionAllocFunctionsInternal::Malloc()
                  PartitionRoot::AllocInternal()
                    PartitionRoot::AllocInternalNoHooks()
                      PartitionRoot::RawAlloc()
                        ::partition_alloc::internal::ScopedGuard guard{internal::PartitionRootLock(this)};

src/base/allocator/partition_allocator/src/partition_alloc/partition_alloc_base/threading/platform_thread_posix.cc

@@ -28,6 +28,10 @@
 #include <zircon/process.h>
 #endif
 
+#if defined(__MUSL__)
+#include "partition_alloc/shim/allocator_shim.h"
+#endif
+
 namespace partition_alloc::internal::base {
 
 #if PA_BUILDFLAG(IS_LINUX) || PA_BUILDFLAG(IS_CHROMEOS)
@@ -59,8 +63,21 @@ thread_local bool g_is_main_thread = true;
 class InitAtFork {
  public:
   InitAtFork() {
-#if !defined(__MUSL__)
+#if defined(__MUSL__)
+    allocator_shim::AllocatorDispatch d =
+        *allocator_shim::GetAllocatorDispatchChainHeadForTesting();
+    d.alloc_function = +[](size_t size, void*) -> void* {
+      // The size of the scratch fits struct atfork_funcs in Musl pthread_atfork.c.
+      static char scratch[5 * sizeof(void*)];
+      return size != sizeof(scratch) ? nullptr : scratch;
+    };
+    allocator_shim::InsertAllocatorDispatch(&d);
+#endif
+
     pthread_atfork(nullptr, nullptr, internal::InvalidateTidCache);
+
+#if defined(__MUSL__)
+    allocator_shim::RemoveAllocatorDispatchForTesting(&d);
 #endif
   }
 };

src/base/allocator/partition_allocator/src/partition_alloc/partition_root.cc

@@ -51,6 +51,10 @@
 #endif  // PA_CONFIG(ENABLE_SHADOW_METADATA)
 #endif  // PA_BUILDFLAG(IS_LINUX) || PA_BUILDFLAG(IS_CHROMEOS)
 
+#if defined(__MUSL__)
+#include "partition_alloc/shim/allocator_shim.h"
+#endif
+
 namespace partition_alloc::internal {
 
 #if PA_BUILDFLAG(RECORD_ALLOC_INFO)
@@ -297,11 +301,7 @@ void PartitionAllocMallocInitOnce() {
     return;
   }
 
-#if defined(__MUSL__)
-  static_cast<void>(BeforeForkInParent);
-  static_cast<void>(AfterForkInParent);
-  static_cast<void>(AfterForkInChild);
-#elif PA_BUILDFLAG(IS_LINUX) || PA_BUILDFLAG(IS_CHROMEOS)
+#if PA_BUILDFLAG(IS_LINUX) || PA_BUILDFLAG(IS_CHROMEOS)
   // When fork() is called, only the current thread continues to execute in the
   // child process. If the lock is held, but *not* by this thread when fork() is
   // called, we have a deadlock.
@@ -323,9 +323,25 @@ void PartitionAllocMallocInitOnce() {
   // However, no perfect solution really exists to make threads + fork()
   // cooperate, but deadlocks are real (and fork() is used in DEATH_TEST()s),
   // and other malloc() implementations use the same techniques.
+
+#if defined(__MUSL__)
+    allocator_shim::AllocatorDispatch d =
+        *allocator_shim::GetAllocatorDispatchChainHeadForTesting();
+    d.alloc_function = +[](size_t size, void*) -> void* {
+      // The size of the scratch fits struct atfork_funcs in Musl pthread_atfork.c.
+      static char scratch[5 * sizeof(void*)];
+      return size != sizeof(scratch) ? nullptr : scratch;
+    };
+    allocator_shim::InsertAllocatorDispatch(&d);
+#endif
+
   int err =
       pthread_atfork(BeforeForkInParent, AfterForkInParent, AfterForkInChild);
   PA_CHECK(err == 0);
+
+#if defined(__MUSL__)
+    allocator_shim::RemoveAllocatorDispatchForTesting(&d);
+#endif
 #endif  // PA_BUILDFLAG(IS_LINUX) || PA_BUILDFLAG(IS_CHROMEOS)
 }