diff --git a/libc/malloc_debug/DebugData.cpp b/libc/malloc_debug/DebugData.cpp index 92b866d03..0447566b0 100644 --- a/libc/malloc_debug/DebugData.cpp +++ b/libc/malloc_debug/DebugData.cpp @@ -82,3 +82,21 @@ bool DebugData::Initialize() { } return true; } + +void DebugData::PrepareFork() { + if (track != nullptr) { + track->PrepareFork(); + } +} + +void DebugData::PostForkParent() { + if (track != nullptr) { + track->PostForkParent(); + } +} + +void DebugData::PostForkChild() { + if (track != nullptr) { + track->PostForkChild(); + } +} diff --git a/libc/malloc_debug/DebugData.h b/libc/malloc_debug/DebugData.h index 40978dbe3..4600b33bc 100644 --- a/libc/malloc_debug/DebugData.h +++ b/libc/malloc_debug/DebugData.h @@ -82,6 +82,10 @@ class DebugData { bool need_header() { return need_header_; } size_t extra_bytes() { return extra_bytes_; } + void PrepareFork(); + void PostForkParent(); + void PostForkChild(); + std::unique_ptr backtrace; std::unique_ptr track; std::unique_ptr front_guard; diff --git a/libc/malloc_debug/TrackData.h b/libc/malloc_debug/TrackData.h index 45e68927e..dcf0ede1a 100644 --- a/libc/malloc_debug/TrackData.h +++ b/libc/malloc_debug/TrackData.h @@ -58,6 +58,10 @@ class TrackData { void DisplayLeaks(DebugData& debug); + void PrepareFork() { pthread_mutex_lock(&mutex_); } + void PostForkParent() { pthread_mutex_unlock(&mutex_); } + void PostForkChild() { pthread_mutex_init(&mutex_, NULL); } + private: pthread_mutex_t mutex_ = PTHREAD_MUTEX_INITIALIZER; std::unordered_set headers_; diff --git a/libc/malloc_debug/malloc_debug.cpp b/libc/malloc_debug/malloc_debug.cpp index 0c0907d00..dcc6048f7 100644 --- a/libc/malloc_debug/malloc_debug.cpp +++ b/libc/malloc_debug/malloc_debug.cpp @@ -85,6 +85,28 @@ void* debug_valloc(size_t size); __END_DECLS // ------------------------------------------------------------------------ +static void InitAtfork() { + static pthread_once_t atfork_init = PTHREAD_ONCE_INIT; + pthread_once(&atfork_init, [](){ + pthread_atfork( + [](){ + if (g_debug != nullptr) { + g_debug->PrepareFork(); + } + }, + [](){ + if (g_debug != nullptr) { + g_debug->PostForkParent(); + } + }, + [](){ + if (g_debug != nullptr) { + g_debug->PostForkChild(); + } + } + ); + }); +} static void LogTagError(const Header* header, const void* pointer, const char* name) { ScopedDisableDebugCalls disable; @@ -156,6 +178,9 @@ bool debug_initialize(const MallocDispatch* malloc_dispatch, int* malloc_zygote_ if (malloc_zygote_child == nullptr) { return false; } + + InitAtfork(); + g_malloc_zygote_child = malloc_zygote_child; g_dispatch = malloc_dispatch;