Create TLS vars before installing fork handlers

The fork handlers now depend (at least on macOS) on calling
`pthread_getspecific` and `pthread_setspecific` for our recursion guard,
we need to ensure that the TLS key has been created before we install
our fork handlers.

Signed-off-by: Matt Wozniski <[email protected]>

Author: godlygeek

src/memray/_memray.pyx

@@ -55,6 +55,7 @@ from _memray.source cimport FileSource
 from _memray.source cimport SocketSource
 from _memray.tracking_api cimport Tracker as NativeTracker
 from _memray.tracking_api cimport install_trace_function
+from _memray.tracking_api cimport set_up_pthread_fork_handlers
 from cpython cimport PyErr_CheckSignals
 from libc.math cimport ceil
 from libc.stdint cimport uint64_t
@@ -76,16 +77,12 @@ from ._metadata import Metadata
 from ._stats import Stats
 from ._thread_name_interceptor import ThreadNameInterceptor
 
-
-cdef extern from "pthread.h" nogil:
-    int pthread_atfork(void (*prepare)(), void (*parent)(), void (*child)())
-
 # NOTE: We can't reinitialize tracking in a child process until the interpreter
 #       has reinitialized its locks, so we do it in a Python fork handler.
 #       But, Python fork handlers aren't guaranteed to run for every fork, and
 #       the child process must never inherit an active tracker, so we must use
 #       a pthread fork handler to disable tracking before forking.
-pthread_atfork(&NativeTracker.prepareFork, &NativeTracker.parentFork, NULL)
+set_up_pthread_fork_handlers()
 os.register_at_fork(after_in_child=NativeTracker.childFork)
 
 

src/memray/_memray/tracking_api.cpp

@@ -559,7 +559,6 @@ Tracker::Tracker(
 {
     static std::once_flag once;
     call_once(once, [] {
-        RecursionGuard::initialize();
         // We use the pthread TLS API for this vector because we must be able
         // to re-create it while TLS destructors are running (a destructor can
         // call malloc, hitting our malloc hook). POSIX guarantees multiple
@@ -1235,6 +1234,18 @@ Tracker::handleGreenletSwitch(PyObject* from, PyObject* to)
     PythonStackTracker::get().handleGreenletSwitch(from, to);
 }
 
+void
+set_up_pthread_fork_handlers()
+{
+    static std::once_flag once;
+    call_once(once, [] {
+        // On macOS the recursion guard uses pthread TLS keys, so we
+        // need to create those keys before we install our handlers.
+        RecursionGuard::initialize();
+        pthread_atfork(&Tracker::prepareFork, &Tracker::parentFork, NULL);
+    });
+}
+
 void
 install_trace_function()
 {

src/memray/_memray/tracking_api.h

@@ -134,6 +134,12 @@ PyTraceTrampoline(PyObject* obj, PyFrameObject* frame, int what, PyObject* arg);
 void
 install_trace_function();
 
+/**
+ * Install our pthread fork handlers.
+ */
+void
+set_up_pthread_fork_handlers();
+
 class NativeTrace
 {
   public:

src/memray/_memray/tracking_api.pxd

@@ -6,6 +6,7 @@ from libcpp.string cimport string
 
 
 cdef extern from "tracking_api.h" namespace "memray::tracking_api":
+    void set_up_pthread_fork_handlers() except+
     void install_trace_function() except*
 
     cdef cppclass Tracker:
@@ -36,11 +37,5 @@ cdef extern from "tracking_api.h" namespace "memray::tracking_api":
         @staticmethod
         void registerThreadNameById(uint64_t, const char*) except+
 
-        @staticmethod
-        void prepareFork() noexcept nogil
-
-        @staticmethod
-        void parentFork() noexcept nogil
-
         @staticmethod
         void childFork() noexcept nogil