Protection against fork (#735)

If a thread forks, while another thread is holding an snmalloc lock, then the allocator could stop working.

This patch attempts to protect against the cases of this. There is one case that is not covered. If a fork occurs during the very first allocation. This can result in the installation of the fork handler racing with the fork, and all bets are off.

CMakeLists.txt

@@ -120,6 +120,19 @@ int main()
  SNMALLOC_LINKER_SUPPORT_NOSTDLIBXX)
 set(CMAKE_REQUIRED_LINK_OPTIONS "")
 
+# Detect if pthread_atfork works
+CHECK_CXX_SOURCE_COMPILES("
+#include <pthread.h>
+void prepare() {}
+void parent() {}
+void child() {}
+int main() {
+  pthread_atfork(prepare, parent, child);
+  return 0;
+}
+" SNMALLOC_PTHREAD_ATFORK_WORKS)
+
+
 if (NOT MSVC AND NOT (SNMALLOC_CLEANUP STREQUAL CXX11_DESTRUCTORS))
   # If the target compiler doesn't support -nostdlib++ then we must enable C at
   # the global scope for the fallbacks to work.
@@ -320,6 +333,7 @@ add_as_define(SNMALLOC_QEMU_WORKAROUND)
 add_as_define(SNMALLOC_TRACING)
 add_as_define(SNMALLOC_CI_BUILD)
 add_as_define(SNMALLOC_PLATFORM_HAS_GETENTROPY)
+add_as_define(SNMALLOC_PTHREAD_ATFORK_WORKS)
 add_as_define(SNMALLOC_HAS_LINUX_RANDOM_H)
 add_as_define(SNMALLOC_HAS_LINUX_FUTEX_H)
 if (SNMALLOC_NO_REALLOCARRAY)

src/snmalloc/ds/combininglock.h

@@ -274,26 +274,32 @@ namespace snmalloc
   template<typename F>
   inline void with(CombiningLock& lock, F&& f)
   {
-    // Test if no one is waiting
-    if (SNMALLOC_LIKELY(lock.last.load(stl::memory_order_relaxed) == nullptr))
+    // A unix fork while holding a lock can lead to deadlock. Protect against
+    // this by not allowing a fork while holding a lock.
+    PreventFork pf;
+    snmalloc::UNUSED(pf);
     {
-      // No one was waiting so low contention. Attempt to acquire the flag
-      // lock.
-      if (SNMALLOC_LIKELY(
-            lock.flag.exchange(true, stl::memory_order_acquire) == false))
+      // Test if no one is waiting
+      if (SNMALLOC_LIKELY(lock.last.load(stl::memory_order_relaxed) == nullptr))
       {
-        // We grabbed the lock.
-        // Execute the thunk.
-        f();
+        // No one was waiting so low contention. Attempt to acquire the flag
+        // lock.
+        if (SNMALLOC_LIKELY(
+              lock.flag.exchange(true, stl::memory_order_acquire) == false))
+        {
+          // We grabbed the lock.
+          // Execute the thunk.
+          f();
 
-        // Release the lock
-        lock.release();
-        return;
+          // Release the lock
+          lock.release();
+          return;
+        }
       }
-    }
 
-    // There is contention for the lock, we need to take the slow path
-    // with the queue.
-    CombiningLockNodeTempl<F> node(lock, stl::forward<F>(f));
+      // There is contention for the lock, we need to take the slow path
+      // with the queue.
+      CombiningLockNodeTempl<F> node(lock, stl::forward<F>(f));
+    }
   }
 } // namespace snmalloc

src/snmalloc/ds_aal/ds_aal.h

@@ -6,4 +6,5 @@
 #pragma once
 #include "../aal/aal.h"
 #include "flaglock.h"
+#include "prevent_fork.h"
 #include "singleton.h"

src/snmalloc/ds_aal/flaglock.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include "../aal/aal.h"
+#include "prevent_fork.h"
 #include "snmalloc/ds_core/ds_core.h"
 #include "snmalloc/stl/atomic.h"
 
@@ -108,6 +109,10 @@ namespace snmalloc
   private:
     FlagWord& lock;
 
+    // A unix fork while holding a lock can lead to deadlock. Protect against
+    // this by not allowing a fork while holding a lock.
+    PreventFork pf{};
+
   public:
     FlagLock(FlagWord& lock) : lock(lock)
     {

src/snmalloc/ds_aal/prevent_fork.h

@@ -0,0 +1,161 @@
+#pragma once
+
+#include <snmalloc/aal/aal.h>
+#include <snmalloc/stl/atomic.h>
+#include <stddef.h>
+
+#ifdef SNMALLOC_PTHREAD_ATFORK_WORKS
+#  include <pthread.h>
+#endif
+
+namespace snmalloc
+{
+  // This is a simple implementation of a class that can be
+  // used to prevent a process from forking. Holding a lock
+  // in the allocator while forking can lead to deadlocks.
+  // This causes the fork to wait out any other threads inside
+  // the allocators locks.
+  //
+  // The use is
+  // ```
+  //  {
+  //     PreventFork pf;
+  //     // Code that should not be running during a fork.
+  //  }
+  // ```
+  class PreventFork
+  {
+    // Global atomic counter of the number of threads currently preventing the
+    // system from forking. The bottom bit is used to signal that a thread is
+    // wanting to fork.
+    static inline stl::Atomic<size_t> threads_preventing_fork{0};
+
+    // The depth of the current thread's prevention of forking.
+    // This is used to enable reentrant prevention of forking.
+    static inline thread_local size_t depth_of_prevention{0};
+
+    // There could be multiple copies of the atfork handler installed.
+    // Only perform work for the first prefork and final postfork.
+    static inline thread_local size_t depth_of_handlers{0};
+
+    // This function ensures that the fork handler has been installed at least
+    // once. It might be installed more than once, this is safe. As subsequent
+    // calls would be ignored.
+    static void ensure_init()
+    {
+#ifdef SNMALLOC_PTHREAD_ATFORK_WORKS
+      static stl::Atomic<bool> initialised{false};
+
+      if (initialised.load(stl::memory_order_acquire))
+        return;
+
+      pthread_atfork(prefork, postfork_parent, postfork_child);
+      initialised.store(true, stl::memory_order_release);
+#endif
+    };
+
+  public:
+    PreventFork()
+    {
+      if (depth_of_prevention++ == 0)
+      {
+        // Ensure that the system is initialised before we start.
+        // Don't do this on nested Prevent calls.
+        ensure_init();
+        while (true)
+        {
+          auto prev = threads_preventing_fork.fetch_add(2);
+          if (prev % 2 == 0)
+            break;
+
+          threads_preventing_fork.fetch_sub(2);
+
+          while ((threads_preventing_fork.load() % 2) == 1)
+          {
+            Aal::pause();
+          }
+        };
+      }
+    }
+
+    ~PreventFork()
+    {
+      if (--depth_of_prevention == 0)
+      {
+        threads_preventing_fork -= 2;
+      }
+    }
+
+    // The function that notifies new threads not to enter PreventFork regions
+    // It waits until all threads are no longer in a PreventFork region before
+    // returning.
+    static void prefork()
+    {
+      if (depth_of_handlers++ != 0)
+        return;
+
+      if (depth_of_prevention != 0)
+        error("Fork attempted while in PreventFork region.");
+
+      while (true)
+      {
+        auto current = threads_preventing_fork.load();
+        if (
+          (current % 2 == 0) &&
+          (threads_preventing_fork.compare_exchange_weak(current, current + 1)))
+        {
+          break;
+        }
+        Aal::pause();
+      };
+
+      while (threads_preventing_fork.load() != 1)
+      {
+        Aal::pause();
+      }
+
+      // Finally set the flag that allows this thread to enter PreventFork
+      // regions This is safe as the only other calls here are to other prefork
+      // handlers.
+      depth_of_prevention++;
+    }
+
+    // Unsets the flag that allows threads to enter PreventFork regions
+    // and for another thread to request a fork.
+    static void postfork_child()
+    {
+      // Count out the number of handlers that have been called, and
+      // only perform on the last.
+      if (--depth_of_handlers != 0)
+        return;
+
+      // This thread is no longer preventing a fork, so decrement the counter.
+      depth_of_prevention--;
+
+      // Allow other threads to allocate
+      // There could have been threads spinning in the prefork handler having
+      // optimistically increasing thread_preventing_fork by 2, but now the
+      // threads do not exist due to the fork. So restart the counter in the
+      // child.
+      threads_preventing_fork = 0;
+    }
+
+    // Unsets the flag that allows threads to enter PreventFork regions
+    // and for another thread to request a fork.
+    static void postfork_parent()
+    {
+      // Count out the number of handlers that have been called, and
+      // only perform on the last.
+      if (--depth_of_handlers != 0)
+        return;
+
+      // This thread is no longer preventing a fork, so decrement the counter.
+      depth_of_prevention--;
+
+      // Allow other threads to allocate
+      // Just remove the bit, and let the potential other threads in prefork
+      // remove their counts.
+      threads_preventing_fork--;
+    }
+  };
+} // namespace snmalloc

src/snmalloc/ds_aal/singleton.h

@@ -1,5 +1,6 @@
 #pragma once
 
+#include "prevent_fork.h"
 #include "snmalloc/stl/atomic.h"
 
 namespace snmalloc
@@ -36,6 +37,11 @@ namespace snmalloc
       auto state = initialised.load(stl::memory_order_acquire);
       if (SNMALLOC_UNLIKELY(state == State::Uninitialised))
       {
+        // A unix fork while initialising a singleton can lead to deadlock.
+        // Protect against this by not allowing a fork while attempting
+        // initialisation.
+        PreventFork pf;
+        snmalloc::UNUSED(pf);
         if (initialised.compare_exchange_strong(
               state, State::Initialising, stl::memory_order_relaxed))
         {

src/snmalloc/mem/freelist_queue.h

@@ -110,6 +110,14 @@ namespace snmalloc
       invariant();
       freelist::Object::atomic_store_null(last, Key, Key_tweak);
 
+      // The following non-linearisable effect is normally benign,
+      // but could lead to a remote list become completely detached
+      // during a fork in a multi-threaded process. This would lead
+      // to a memory leak, which is probably the least of your problems
+      // if you forked in during a deallocation.
+      PreventFork pf;
+      snmalloc::UNUSED(pf);
+
       // Exchange needs to be acq_rel.
       // *  It needs to be a release, so nullptr in next is visible.
       // *  Needs to be acquire, so linking into the list does not race with

src/snmalloc/stl/gnu/atomic.h

@@ -214,13 +214,27 @@ namespace snmalloc
           addressof(val), 1, order(MemoryOrder::SEQ_CST));
       }
 
+      SNMALLOC_FAST_PATH T operator--()
+      {
+        static_assert(stl::is_integral_v<T>, "T must be an integral type.");
+        return __atomic_sub_fetch(
+          addressof(val), 1, order(MemoryOrder::SEQ_CST));
+      }
+
       SNMALLOC_FAST_PATH const T operator++(int)
       {
         static_assert(stl::is_integral_v<T>, "T must be an integral type.");
         return __atomic_fetch_add(
           addressof(val), 1, order(MemoryOrder::SEQ_CST));
       }
 
+      SNMALLOC_FAST_PATH const T operator--(int)
+      {
+        static_assert(stl::is_integral_v<T>, "T must be an integral type.");
+        return __atomic_fetch_sub(
+          addressof(val), 1, order(MemoryOrder::SEQ_CST));
+      }
+
       SNMALLOC_FAST_PATH T operator-=(T decrement)
       {
         static_assert(stl::is_integral_v<T>, "T must be an integral type.");