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