Protection against fork

CMakeLists.txt

@@ -114,6 +114,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.
@@ -300,6 +313,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

@@ -270,26 +270,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,125 @@
+#pragma once
+
+#include <snmalloc/aal/aal.h>
+#include <snmalloc/stl/atomic.h>
+#include <stddef.h>
+
+#ifdef SNMALLOC_PTHREAD_ATFORK_WORKS
+#  include <pthread.h>
+
+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};
+
+    // 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_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();
+      }
+    }
+
+    // Unsets the flag that allows threads to enter PreventFork regions
+    // and for another thread to request a fork.
+    static void postfork()
+    {
+      threads_preventing_fork = 0;
+    }
+
+    // This is called by the Singleton class to ensure that it is
+    // only called once. Argument ignored, and only used for Singleton
+    // class design.
+    static void init(size_t*) noexcept
+    {
+      pthread_atfork(prefork, postfork, postfork);
+    }
+
+    // This function requires the Singleton class,
+    // which also depends on PreventFork. We define this
+    // implementation in singleton.h
+    static void ensure_init();
+
+  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 current = threads_preventing_fork.load();
+
+          if (
+            (current % 2 == 0) &&
+            threads_preventing_fork.compare_exchange_weak(current, current + 2))
+          {
+            break;
+          }
+          Aal::pause();
+        };
+      }
+    }
+
+    ~PreventFork()
+    {
+      if (--depth_of_prevention == 0)
+      {
+        threads_preventing_fork -= 2;
+      }
+    }
+  };
+} // namespace snmalloc
+#else
+namespace snmalloc
+{
+  class PreventFork
+  {
+  public:
+    static void init() {}
+  };
+} // namespace snmalloc
+#endif

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))
         {
@@ -55,4 +61,10 @@ namespace snmalloc
       return obj;
     }
   };
+
+  inline void PreventFork::ensure_init()
+  {
+    static Singleton<size_t, &PreventFork::init> singleton;
+    singleton.get();
+  }
 } // namespace snmalloc

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/test/func/protect_fork/protect_fork.cc

@@ -0,0 +1,60 @@
+#include <snmalloc/snmalloc.h>
+#include <iostream>
+
+#ifndef SNMALLOC_PTHREAD_ATFORK_WORKS
+int main()
+{
+  std::cout << "Test did not run" << std::endl;
+  return 0;
+}
+#else
+
+
+#include <pthread.h>
+#include <thread>
+
+int main()
+{
+  // Counter for the number of threads that are blocking the fork
+  std::atomic<size_t> block = false;
+  // Specifies that the forking thread has observed that all the blocking
+  // threads are in place.
+  std::atomic<bool> forking = false;
+
+  size_t N = 3;
+
+  for (size_t i = 0; i < N; i ++)
+  {
+    std::thread t([&block, &forking]() {
+      {
+        snmalloc::PreventFork pf;
+        block++;
+        std::this_thread::sleep_for(std::chrono::milliseconds(100));
+        while (!forking)
+          std::this_thread::yield();
+        std::this_thread::sleep_for(std::chrono::milliseconds(100));
+
+        block--;
+      }
+    });
+
+    t.detach();
+  }
+
+  while(block != N)
+    std::this_thread::yield();
+
+  forking = true;
+
+  fork();
+
+  if (block)
+  {
+    snmalloc::message<1024>("PreventFork failed");
+    return 1;
+  }
+  snmalloc::message<1024>("PreventFork passed");
+  return 0;
+}
+
+#endif