Ensure fork safety

Author: pitrou

Include/internal/pycore_pythread.h

@@ -137,6 +137,12 @@ PyAPI_FUNC(int) PyThread_join_thread(Py_uintptr_t);
  */
 PyAPI_FUNC(int) PyThread_detach_thread(Py_uintptr_t);
 
+/*
+ * Obtain the new thread ident and handle in a forked child process.
+ */
+PyAPI_FUNC(void) PyThread_update_thread_after_fork(unsigned long long* ident,
+                                                   Py_uintptr_t* handle);
+
 #ifdef __cplusplus
 }
 #endif

Lib/threading.py

@@ -590,7 +590,7 @@ def __repr__(self):
         return f"<{cls.__module__}.{cls.__qualname__} at {id(self):#x}: {status}>"
 
     def _at_fork_reinit(self):
-        # Private method called by Thread._reset_internal_locks()
+        # Private method called by Thread._after_fork()
         self._cond._at_fork_reinit()
 
     def is_set(self):
@@ -936,11 +936,15 @@ class is implemented.
         # For debugging and _after_fork()
         _dangling.add(self)
 
-    def _reset_internal_locks(self, is_alive):
-        # private!  Called by _after_fork() to reset our internal locks as
-        # they may be in an invalid state leading to a deadlock or crash.
+    def _after_fork(self, new_ident=None):
+        # Private!  Called by threading._after_fork().
         self._started._at_fork_reinit()
-        if is_alive:
+        if new_ident is not None:
+            # This thread is alive.
+            self._ident = new_ident
+            if self._handle is not None:
+                self._handle.after_fork_alive()
+                assert self._handle.ident == new_ident
             # bpo-42350: If the fork happens when the thread is already stopped
             # (ex: after threading._shutdown() has been called), _tstate_lock
             # is None. Do nothing in this case.
@@ -950,11 +954,14 @@ def _reset_internal_locks(self, is_alive):
             if self._join_lock is not None:
                 self._join_lock._at_fork_reinit()
         else:
-            # The thread isn't alive after fork: it doesn't have a tstate
+            # This thread isn't alive after fork: it doesn't have a tstate
             # anymore.
             self._is_stopped = True
             self._tstate_lock = None
             self._join_lock = None
+            if self._handle is not None:
+                self._handle.after_fork_dead()
+                self._handle = None
 
     def __repr__(self):
         assert self._initialized, "Thread.__init__() was not called"
@@ -1707,15 +1714,13 @@ def _after_fork():
             # Any lock/condition variable may be currently locked or in an
             # invalid state, so we reinitialize them.
             if thread is current:
-                # There is only one active thread. We reset the ident to
-                # its new value since it can have changed.
-                thread._reset_internal_locks(True)
+                # This is the one and only active thread.
                 ident = get_ident()
-                thread._ident = ident
+                thread._after_fork(new_ident=ident)
                 new_active[ident] = thread
             else:
                 # All the others are already stopped.
-                thread._reset_internal_locks(False)
+                thread._after_fork()
                 thread._stop()
 
         _limbo.clear()

Modules/_threadmodule.c

@@ -90,6 +90,22 @@ ThreadHandle_get_ident(ThreadHandleObject *self, void *ignored)
     return PyLong_FromUnsignedLongLong(self->ident);
 }
 
+
+static PyObject *
+ThreadHandle_after_fork_alive(ThreadHandleObject *self, void* ignored)
+{
+    PyThread_update_thread_after_fork(&self->ident, &self->handle);
+    Py_RETURN_NONE;
+}
+
+static PyObject *
+ThreadHandle_after_fork_dead(ThreadHandleObject *self, void* ignored)
+{
+    // Disallow calls to detach() and join() as they could crash.
+    self->joinable = 0;
+    Py_RETURN_NONE;
+}
+
 static PyObject *
 ThreadHandle_detach(ThreadHandleObject *self, void* ignored)
 {
@@ -141,6 +157,8 @@ static PyGetSetDef ThreadHandle_getsetlist[] = {
 
 static PyMethodDef ThreadHandle_methods[] =
 {
+    {"after_fork_alive", (PyCFunction)ThreadHandle_after_fork_alive, METH_NOARGS},
+    {"after_fork_dead", (PyCFunction)ThreadHandle_after_fork_dead, METH_NOARGS},
     {"detach", (PyCFunction)ThreadHandle_detach, METH_NOARGS},
     {"join", (PyCFunction)ThreadHandle_join, METH_NOARGS},
     {0, 0}

Python/thread_nt.h

@@ -240,6 +240,10 @@ PyThread_detach_thread(Py_uintptr_t handle) {
     return (CloseHandle(hThread) == 0);
 }
 
+void
+PyThread_update_thread_after_fork(unsigned long long* ident, Py_uintptr_t* handle) {
+}
+
 /*
  * Return the thread Id instead of a handle. The Id is said to uniquely identify the
  * thread in the system

Python/thread_pthread.h

@@ -337,6 +337,16 @@ PyThread_detach_thread(Py_uintptr_t th) {
     return pthread_detach((pthread_t) th);
 }
 
+void
+PyThread_update_thread_after_fork(unsigned long long* ident, Py_uintptr_t* handle) {
+    // The thread id might have been updated in the forked child
+    pthread_t th = pthread_self();
+    *ident = (unsigned long long) th;
+    *handle = (Py_uintptr_t) th;
+    assert(th == (pthread_t) *ident);
+    assert(th == (pthread_t) *handle);
+}
+
 /* XXX This implementation is considered (to quote Tim Peters) "inherently
    hosed" because:
      - It does not guarantee the promise that a non-zero integer is returned.