gh-126688: Support Fork Parent and Child With Different Thread ID

Include/internal/pycore_import.h

@@ -9,6 +9,7 @@ extern "C" {
 #  error "this header requires Py_BUILD_CORE define"
 #endif
 
+#include "pycore_pythread.h"      // PyThread_ident_t
 #include "pycore_lock.h"          // PyMutex
 #include "pycore_hashtable.h"     // _Py_hashtable_t
 
@@ -21,7 +22,11 @@ extern int _PyImport_SetModuleString(const char *name, PyObject* module);
 
 extern void _PyImport_AcquireLock(PyInterpreterState *interp);
 extern void _PyImport_ReleaseLock(PyInterpreterState *interp);
-extern void _PyImport_ReInitLock(PyInterpreterState *interp);
+#ifdef HAVE_FORK
+extern void _PyImport_ReInitLock(
+    PyInterpreterState *interp,
+    PyThread_ident_t parent);
+#endif
 
 // This is used exclusively for the sys and builtins modules:
 extern int _PyImport_FixupBuiltin(

Include/internal/pycore_lock.h

@@ -151,10 +151,12 @@ _PyOnceFlag_CallOnce(_PyOnceFlag *flag, _Py_once_fn_t *fn, void *arg)
     return _PyOnceFlag_CallOnceSlow(flag, fn, arg);
 }
 
+#define _PyThread_ident_t unsigned long long
+
 // A recursive mutex. The mutex should zero-initialized.
 typedef struct {
     PyMutex mutex;
-    unsigned long long thread;  // i.e., PyThread_get_thread_ident_ex()
+    _PyThread_ident_t thread;  // i.e., PyThread_get_thread_ident_ex()
     size_t level;
 } _PyRecursiveMutex;
 
@@ -164,6 +166,20 @@ extern PyLockStatus _PyRecursiveMutex_LockTimed(_PyRecursiveMutex *m, PyTime_t t
 PyAPI_FUNC(void) _PyRecursiveMutex_Unlock(_PyRecursiveMutex *m);
 extern int _PyRecursiveMutex_TryUnlock(_PyRecursiveMutex *m);
 
+PyAPI_FUNC(_PyThread_ident_t) PyThread_get_thread_ident_ex(void);
+
+static inline void
+_PyRecursiveMutex_at_fork_reinit(_PyRecursiveMutex *m,
+                                 _PyThread_ident_t parent)
+{
+    if (m->thread == parent) {
+        m->thread = PyThread_get_thread_ident_ex();
+    }
+    else {
+        memset(m, 0, sizeof(*m));
+    }
+}
+
 // A readers-writer (RW) lock. The lock supports multiple concurrent readers or
 // a single writer. The lock is write-preferring: if a writer is waiting while
 // the lock is read-locked then, new readers will be blocked. This avoids

Include/internal/pycore_runtime.h

@@ -164,6 +164,15 @@ typedef struct pyruntimestate {
         _Py_AuditHookEntry *head;
     } audit_hooks;
 
+#ifdef HAVE_FORK
+    struct {
+        PyMutex mutex;
+        struct _os_fork_parent {
+            PyThread_ident_t tid;
+        } parent;
+    } os_fork;
+#endif
+
     struct _py_object_runtime_state object_state;
     struct _Py_float_runtime_state float_state;
     struct _Py_unicode_runtime_state unicode_state;
@@ -210,7 +219,9 @@ extern PyStatus _PyRuntimeState_Init(_PyRuntimeState *runtime);
 extern void _PyRuntimeState_Fini(_PyRuntimeState *runtime);
 
 #ifdef HAVE_FORK
-extern PyStatus _PyRuntimeState_ReInitThreads(_PyRuntimeState *runtime);
+extern PyStatus _PyRuntimeState_ReInitThreads(
+    _PyRuntimeState *runtime,
+    PyThread_ident_t parent);
 #endif
 
 /* Initialize _PyRuntimeState.

Modules/posixmodule.c

@@ -619,8 +619,12 @@ void
 PyOS_BeforeFork(void)
 {
     PyInterpreterState *interp = _PyInterpreterState_GET();
+    assert(interp->runtime == &_PyRuntime);
     run_at_forkers(interp->before_forkers, 1);
 
+    PyMutex_Lock(&_PyRuntime.os_fork.mutex);
+    _PyRuntime.os_fork.parent.tid = PyThread_get_thread_ident_ex();
+
     _PyImport_AcquireLock(interp);
     _PyEval_StopTheWorldAll(&_PyRuntime);
     HEAD_LOCK(&_PyRuntime);
@@ -629,11 +633,17 @@ PyOS_BeforeFork(void)
 void
 PyOS_AfterFork_Parent(void)
 {
+    _PyRuntime.os_fork.parent = (struct _os_fork_parent){0};
+
     HEAD_UNLOCK(&_PyRuntime);
     _PyEval_StartTheWorldAll(&_PyRuntime);
 
     PyInterpreterState *interp = _PyInterpreterState_GET();
+    assert(interp->runtime == &_PyRuntime);
     _PyImport_ReleaseLock(interp);
+
+    PyMutex_Unlock(&_PyRuntime.os_fork.mutex);
+
     run_at_forkers(interp->after_forkers_parent, 0);
 }
 
@@ -644,15 +654,25 @@ PyOS_AfterFork_Child(void)
     _PyRuntimeState *runtime = &_PyRuntime;
 
     // re-creates runtime->interpreters.mutex (HEAD_UNLOCK)
-    status = _PyRuntimeState_ReInitThreads(runtime);
+    status = _PyRuntimeState_ReInitThreads(
+                                runtime, runtime->os_fork.parent.tid);
     if (_PyStatus_EXCEPTION(status)) {
         goto fatal_error;
     }
 
     PyThreadState *tstate = _PyThreadState_GET();
     _Py_EnsureTstateNotNULL(tstate);
+    assert(tstate->interp->runtime == &_PyRuntime);
 
-    assert(tstate->thread_id == PyThread_get_thread_ident());
+    // We cannot assume that the parent and child always have
+    // the same thread ID.
+    // See https://github.com/python/cpython/issues/126688.
+    if (_PyRuntime.os_fork.parent.tid != PyThread_get_thread_ident_ex()) {
+        tstate->thread_id = PyThread_get_thread_ident();
+    }
+    else {
+        assert(tstate->thread_id == PyThread_get_thread_ident());
+    }
 #ifdef PY_HAVE_THREAD_NATIVE_ID
     tstate->native_thread_id = PyThread_get_thread_native_id();
 #endif
@@ -678,7 +698,6 @@ PyOS_AfterFork_Child(void)
     _PyEval_StartTheWorldAll(&_PyRuntime);
     _PyThreadState_DeleteList(list);
 
-    _PyImport_ReInitLock(tstate->interp);
     _PyImport_ReleaseLock(tstate->interp);
 
     _PySignal_AfterFork();
@@ -694,6 +713,9 @@ PyOS_AfterFork_Child(void)
         goto fatal_error;
     }
 
+    _PyRuntime.os_fork.parent = (struct _os_fork_parent){0};
+    PyMutex_Unlock(&_PyRuntime.os_fork.mutex);
+
     run_at_forkers(tstate->interp->after_forkers_child, 0);
     return;
 

Python/import.c

@@ -122,12 +122,15 @@ _PyImport_ReleaseLock(PyInterpreterState *interp)
     _PyRecursiveMutex_Unlock(&IMPORT_LOCK(interp));
 }
 
+#ifdef HAVE_FORK
 void
-_PyImport_ReInitLock(PyInterpreterState *interp)
+_PyImport_ReInitLock(PyInterpreterState *interp, PyThread_ident_t parent)
 {
-    // gh-126688: Thread id may change after fork() on some operating systems.
-    IMPORT_LOCK(interp).thread = PyThread_get_thread_ident_ex();
+    /* The forking thread always holds the import lock. */
+    assert(IMPORT_LOCK(interp).thread == parent);
+    _PyRecursiveMutex_at_fork_reinit(&IMPORT_LOCK(interp), parent);
 }
+#endif
 
 
 /***************/

Python/pystate.c

@@ -497,7 +497,8 @@ _PyRuntimeState_Fini(_PyRuntimeState *runtime)
 /* This function is called from PyOS_AfterFork_Child to ensure that
    newly created child processes do not share locks with the parent. */
 PyStatus
-_PyRuntimeState_ReInitThreads(_PyRuntimeState *runtime)
+_PyRuntimeState_ReInitThreads(_PyRuntimeState *runtime,
+                              PyThread_ident_t parent)
 {
     // This was initially set in _PyRuntimeState_Init().
     runtime->main_thread = PyThread_get_thread_ident();
@@ -515,6 +516,7 @@ _PyRuntimeState_ReInitThreads(_PyRuntimeState *runtime)
     for (PyInterpreterState *interp = runtime->interpreters.head;
          interp != NULL; interp = interp->next)
     {
+        _PyImport_ReInitLock(interp, parent);
         for (int i = 0; i < NUM_WEAKREF_LIST_LOCKS; i++) {
             _PyMutex_at_fork_reinit(&interp->weakref_locks[i]);
         }