From ccc03ce0fc1f6ab6f8c2fdb613572a87e97ec7e1 Mon Sep 17 00:00:00 2001 From: Matthew Parkinson Date: Tue, 4 Mar 2025 13:44:59 +0000 Subject: [PATCH] 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 | 14 ++ src/snmalloc/ds/combininglock.h | 38 +++-- src/snmalloc/ds_aal/ds_aal.h | 1 + src/snmalloc/ds_aal/flaglock.h | 5 + src/snmalloc/ds_aal/prevent_fork.h | 161 +++++++++++++++++++++ src/snmalloc/ds_aal/singleton.h | 6 + src/snmalloc/mem/freelist_queue.h | 8 + src/snmalloc/stl/gnu/atomic.h | 14 ++ src/test/func/protect_fork/protect_fork.cc | 81 +++++++++++ 9 files changed, 312 insertions(+), 16 deletions(-) create mode 100644 src/snmalloc/ds_aal/prevent_fork.h create mode 100644 src/test/func/protect_fork/protect_fork.cc diff --git a/CMakeLists.txt b/CMakeLists.txt index da7d02f..8506e17 100644 --- a/CMakeLists.txt +++ b/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 +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) diff --git a/src/snmalloc/ds/combininglock.h b/src/snmalloc/ds/combininglock.h index a86d0cc..e979db2 100644 --- a/src/snmalloc/ds/combininglock.h +++ b/src/snmalloc/ds/combininglock.h @@ -274,26 +274,32 @@ namespace snmalloc template 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 node(lock, stl::forward(f)); + // There is contention for the lock, we need to take the slow path + // with the queue. + CombiningLockNodeTempl node(lock, stl::forward(f)); + } } } // namespace snmalloc diff --git a/src/snmalloc/ds_aal/ds_aal.h b/src/snmalloc/ds_aal/ds_aal.h index bf16de4..21eeb8d 100644 --- a/src/snmalloc/ds_aal/ds_aal.h +++ b/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" \ No newline at end of file diff --git a/src/snmalloc/ds_aal/flaglock.h b/src/snmalloc/ds_aal/flaglock.h index e8dee95..a939f8e 100644 --- a/src/snmalloc/ds_aal/flaglock.h +++ b/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) { diff --git a/src/snmalloc/ds_aal/prevent_fork.h b/src/snmalloc/ds_aal/prevent_fork.h new file mode 100644 index 0000000..2ef0c24 --- /dev/null +++ b/src/snmalloc/ds_aal/prevent_fork.h @@ -0,0 +1,161 @@ +#pragma once + +#include +#include +#include + +#ifdef SNMALLOC_PTHREAD_ATFORK_WORKS +# include +#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 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 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 \ No newline at end of file diff --git a/src/snmalloc/ds_aal/singleton.h b/src/snmalloc/ds_aal/singleton.h index b2efbfa..aebbcd5 100644 --- a/src/snmalloc/ds_aal/singleton.h +++ b/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)) { diff --git a/src/snmalloc/mem/freelist_queue.h b/src/snmalloc/mem/freelist_queue.h index 53fb698..a032508 100644 --- a/src/snmalloc/mem/freelist_queue.h +++ b/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 diff --git a/src/snmalloc/stl/gnu/atomic.h b/src/snmalloc/stl/gnu/atomic.h index 7e0e342..7a19397 100644 --- a/src/snmalloc/stl/gnu/atomic.h +++ b/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 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 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 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 must be an integral type."); diff --git a/src/test/func/protect_fork/protect_fork.cc b/src/test/func/protect_fork/protect_fork.cc new file mode 100644 index 0000000..28a8f4c --- /dev/null +++ b/src/test/func/protect_fork/protect_fork.cc @@ -0,0 +1,81 @@ +#include +#include + +#ifndef SNMALLOC_PTHREAD_ATFORK_WORKS +int main() +{ + std::cout << "Test did not run" << std::endl; + return 0; +} +#else + +# include +# include + +void simulate_allocation() +{ + snmalloc::PreventFork pf; +} + +int main() +{ + // Counter for the number of threads that are blocking the fork + std::atomic block = false; + // Specifies that the forking thread has observed that all the blocking + // threads are in place. + std::atomic 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 \ No newline at end of file