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.

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,6 +214,13 @@ 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.");
@@ -221,6 +228,13 @@ namespace snmalloc
           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.");

src/test/func/protect_fork/protect_fork.cc

@@ -0,0 +1,81 @@
+#include <iostream>
+#include <snmalloc/snmalloc.h>
+
+#ifndef SNMALLOC_PTHREAD_ATFORK_WORKS
+int main()
+{
+  std::cout << "Test did not run" << std::endl;
+  return 0;
+}
+#else
+
+#  include <pthread.h>
+#  include <thread>
+
+void simulate_allocation()
+{
+  snmalloc::PreventFork pf;
+}
+
+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;
+
+  snmalloc::message<1024>("Testing PreventFork");
+
+  snmalloc::message<1024>("Adding alternative calls to pthread_atfork");
+  pthread_atfork(simulate_allocation, simulate_allocation, simulate_allocation);
+
+  snmalloc::message<1024>("Initialising PreventFork singleton");
+  {
+    // Cause initialisation of the PreventFork singleton to call pthread_atfork.
+    snmalloc::PreventFork pf;
+  }
+
+  snmalloc::message<1024>("Adding alternative calls to pthread_atfork");
+  pthread_atfork(simulate_allocation, simulate_allocation, simulate_allocation);
+
+  snmalloc::message<1024>("Creating other threads");
+  for (size_t i = 0; i < N; i++)
+  {
+    std::thread t([&block, &forking, i]() {
+      {
+        snmalloc::PreventFork pf;
+        snmalloc::message<1024>("Thread {} blocking fork", i);
+        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));
+        snmalloc::message<1024>("Thread {} releasing block", i);
+        block--;
+      }
+    });
+
+    t.detach();
+  }
+
+  snmalloc::message<1024>("Waiting for all threads to block fork");
+  while (block != N)
+    std::this_thread::yield();
+
+  snmalloc::message<1024>("Forking");
+  forking = true;
+  fork();
+
+  if (block)
+  {
+    snmalloc::message<1024>("PreventFork failed");
+    return 1;
+  }
+  snmalloc::message<1024>("PreventFork passed");
+  return 0;
+}
+
+#endif