From a574d37f0ec00eda245dbf32795afec4f79a04e8 Mon Sep 17 00:00:00 2001 From: Matthew Parkinson Date: Fri, 24 Jan 2025 16:33:39 +0000 Subject: [PATCH 1/8] Protection against fork 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 | 125 +++++++++++++++++++++ src/snmalloc/ds_aal/singleton.h | 12 ++ src/snmalloc/mem/freelist_queue.h | 8 ++ src/test/func/protect_fork/protect_fork.cc | 60 ++++++++++ 8 files changed, 247 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 834bc827f..550e4999e 100644 --- a/CMakeLists.txt +++ b/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 +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) diff --git a/src/snmalloc/ds/combininglock.h b/src/snmalloc/ds/combininglock.h index 190c90deb..3df9f6a80 100644 --- a/src/snmalloc/ds/combininglock.h +++ b/src/snmalloc/ds/combininglock.h @@ -270,26 +270,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 bf16de4b3..21eeb8dd6 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 e8dee95bd..a939f8e49 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 000000000..92d95c836 --- /dev/null +++ b/src/snmalloc/ds_aal/prevent_fork.h @@ -0,0 +1,125 @@ +#pragma once + +#include +#include +#include + +#ifdef SNMALLOC_PTHREAD_ATFORK_WORKS +# include + +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}; + + // 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 \ No newline at end of file diff --git a/src/snmalloc/ds_aal/singleton.h b/src/snmalloc/ds_aal/singleton.h index b2efbfa68..2f946d0fe 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)) { @@ -55,4 +61,10 @@ namespace snmalloc return obj; } }; + + inline void PreventFork::ensure_init() + { + static Singleton singleton; + singleton.get(); + } } // namespace snmalloc diff --git a/src/snmalloc/mem/freelist_queue.h b/src/snmalloc/mem/freelist_queue.h index 0688a75a5..914d9f704 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/test/func/protect_fork/protect_fork.cc b/src/test/func/protect_fork/protect_fork.cc new file mode 100644 index 000000000..058424172 --- /dev/null +++ b/src/test/func/protect_fork/protect_fork.cc @@ -0,0 +1,60 @@ +#include +#include + +#ifndef SNMALLOC_PTHREAD_ATFORK_WORKS +int main() +{ + std::cout << "Test did not run" << std::endl; + return 0; +} +#else + + +#include +#include + +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; + + 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 \ No newline at end of file From 04924837fa75842437f3c9402ac8d032d32c91f3 Mon Sep 17 00:00:00 2001 From: Matthew Parkinson Date: Wed, 29 Jan 2025 11:22:30 +0000 Subject: [PATCH 2/8] Extend test to simulate allocation. --- src/test/func/protect_fork/protect_fork.cc | 20 +++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/src/test/func/protect_fork/protect_fork.cc b/src/test/func/protect_fork/protect_fork.cc index 058424172..773d58f3f 100644 --- a/src/test/func/protect_fork/protect_fork.cc +++ b/src/test/func/protect_fork/protect_fork.cc @@ -1,5 +1,5 @@ -#include #include +#include #ifndef SNMALLOC_PTHREAD_ATFORK_WORKS int main() @@ -9,9 +9,13 @@ int main() } #else +# include +# include -#include -#include +void simulate_allocation() +{ + snmalloc::PreventFork pf; +} int main() { @@ -23,7 +27,13 @@ int main() size_t N = 3; - for (size_t i = 0; i < N; i ++) + pthread_atfork(simulate_allocation, simulate_allocation, simulate_allocation); + { + snmalloc::PreventFork pf; + } + pthread_atfork(simulate_allocation, simulate_allocation, simulate_allocation); + + for (size_t i = 0; i < N; i++) { std::thread t([&block, &forking]() { { @@ -41,7 +51,7 @@ int main() t.detach(); } - while(block != N) + while (block != N) std::this_thread::yield(); forking = true; From 15227e3d5e8cf19e9d705c760eebad0d78c99201 Mon Sep 17 00:00:00 2001 From: Matthew Parkinson Date: Wed, 29 Jan 2025 11:19:28 +0000 Subject: [PATCH 3/8] Fix allocation during in other pre/post fork handlers This allows for the thread that is executing the prefork and postfork handlers to continue to allocate right up to the point of fork. --- src/snmalloc/ds_aal/prevent_fork.h | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/snmalloc/ds_aal/prevent_fork.h b/src/snmalloc/ds_aal/prevent_fork.h index 92d95c836..e1ac6ccad 100644 --- a/src/snmalloc/ds_aal/prevent_fork.h +++ b/src/snmalloc/ds_aal/prevent_fork.h @@ -59,12 +59,20 @@ namespace snmalloc { 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() { + // This thread is no longer preventing a fork, so decrement the counter. + depth_of_prevention--; + + // Allow other threads to allocate threads_preventing_fork = 0; } From 791462c5860279ef0a128337d90c4cbaaebdaa07 Mon Sep 17 00:00:00 2001 From: Matthew Parkinson Date: Wed, 29 Jan 2025 11:32:13 +0000 Subject: [PATCH 4/8] Fix build for Windows. --- src/snmalloc/ds_aal/prevent_fork.h | 1 + 1 file changed, 1 insertion(+) diff --git a/src/snmalloc/ds_aal/prevent_fork.h b/src/snmalloc/ds_aal/prevent_fork.h index e1ac6ccad..2d11601f1 100644 --- a/src/snmalloc/ds_aal/prevent_fork.h +++ b/src/snmalloc/ds_aal/prevent_fork.h @@ -128,6 +128,7 @@ namespace snmalloc { public: static void init() {} + static void ensure_init(); }; } // namespace snmalloc #endif \ No newline at end of file From f42f66546720aad3cb1fa50d804268067ba42713 Mon Sep 17 00:00:00 2001 From: Matthew Parkinson Date: Wed, 29 Jan 2025 11:44:13 +0000 Subject: [PATCH 5/8] Add comment. --- src/test/func/protect_fork/protect_fork.cc | 1 + 1 file changed, 1 insertion(+) diff --git a/src/test/func/protect_fork/protect_fork.cc b/src/test/func/protect_fork/protect_fork.cc index 773d58f3f..565870a56 100644 --- a/src/test/func/protect_fork/protect_fork.cc +++ b/src/test/func/protect_fork/protect_fork.cc @@ -29,6 +29,7 @@ int main() pthread_atfork(simulate_allocation, simulate_allocation, simulate_allocation); { + // Cause initialisation of the PreventFork singleton to call pthread_atfork. snmalloc::PreventFork pf; } pthread_atfork(simulate_allocation, simulate_allocation, simulate_allocation); From 8ed0ee43889d36f0badfea37e68d7fdc8471809e Mon Sep 17 00:00:00 2001 From: Matthew Parkinson Date: Wed, 29 Jan 2025 11:46:02 +0000 Subject: [PATCH 6/8] Fix build for Windows. --- src/snmalloc/ds_aal/prevent_fork.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/snmalloc/ds_aal/prevent_fork.h b/src/snmalloc/ds_aal/prevent_fork.h index 2d11601f1..9feea331f 100644 --- a/src/snmalloc/ds_aal/prevent_fork.h +++ b/src/snmalloc/ds_aal/prevent_fork.h @@ -127,7 +127,7 @@ namespace snmalloc class PreventFork { public: - static void init() {} + static void init(size_t*) noexcept {} static void ensure_init(); }; } // namespace snmalloc From 846f349a19910a105844b79e3558f68c60fee413 Mon Sep 17 00:00:00 2001 From: Matthew Parkinson Date: Wed, 29 Jan 2025 11:47:49 +0000 Subject: [PATCH 7/8] Clangformat --- src/snmalloc/ds_aal/prevent_fork.h | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/snmalloc/ds_aal/prevent_fork.h b/src/snmalloc/ds_aal/prevent_fork.h index 9feea331f..2e08da0f5 100644 --- a/src/snmalloc/ds_aal/prevent_fork.h +++ b/src/snmalloc/ds_aal/prevent_fork.h @@ -60,8 +60,9 @@ namespace snmalloc 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. + // 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++; } @@ -128,6 +129,7 @@ namespace snmalloc { public: static void init(size_t*) noexcept {} + static void ensure_init(); }; } // namespace snmalloc From 92287a8f345a545a9716d2dd105153d2885f32a1 Mon Sep 17 00:00:00 2001 From: Matthew Parkinson Date: Wed, 29 Jan 2025 13:21:55 +0000 Subject: [PATCH 8/8] Add logging to help with CI deadlock. --- src/test/func/protect_fork/protect_fork.cc | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/src/test/func/protect_fork/protect_fork.cc b/src/test/func/protect_fork/protect_fork.cc index 565870a56..28a8f4c7a 100644 --- a/src/test/func/protect_fork/protect_fork.cc +++ b/src/test/func/protect_fork/protect_fork.cc @@ -27,24 +27,33 @@ int main() 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]() { + 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--; } }); @@ -52,11 +61,12 @@ int main() 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)