Add helper function.

Add _PyRecursiveMutex_at_fork_reinit() helper function.

Author: nascheme

Include/internal/pycore_lock.h

@@ -173,6 +173,9 @@ PyAPI_FUNC(void) _PyRecursiveMutex_Lock(_PyRecursiveMutex *m);
 extern PyLockStatus _PyRecursiveMutex_LockTimed(_PyRecursiveMutex *m, PyTime_t timeout, _PyLockFlags flags);
 PyAPI_FUNC(void) _PyRecursiveMutex_Unlock(_PyRecursiveMutex *m);
 extern int _PyRecursiveMutex_TryUnlock(_PyRecursiveMutex *m);
+#ifdef HAVE_FORK
+extern void _PyRecursiveMutex_at_fork_reinit(_PyRecursiveMutex *m, int owned);
+#endif
 
 // 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

Python/_warnings.c

@@ -282,24 +282,14 @@ void
 _PyWarnings_BeforeFork(PyInterpreterState *interp)
 {
     WarningsState *st = warnings_get_state(interp);
-    // Record ownership before fork(): the thread ID can change in the child.
     st->lock_held_at_fork = _PyRecursiveMutex_IsLockedByCurrentThread(&st->lock);
 }
 
 void
 _PyWarnings_AfterFork(PyInterpreterState *interp)
 {
     WarningsState *st = warnings_get_state(interp);
-    if (st->lock_held_at_fork) {
-        // The surviving thread will still release the lock as its stack
-        // unwinds. Preserve the recursion depth, but discard dead waiters.
-        st->lock.mutex = (PyMutex){._bits = _Py_LOCKED};
-        st->lock.thread = PyThread_get_thread_ident_ex();
-    }
-    else {
-        // The owner (if any) no longer exists in the child.
-        st->lock = (_PyRecursiveMutex){0};
-    }
+    _PyRecursiveMutex_at_fork_reinit(&st->lock, st->lock_held_at_fork);
     st->lock_held_at_fork = false;
 }
 #endif

Python/lock.c

@@ -453,6 +453,25 @@ _PyRecursiveMutex_TryUnlock(_PyRecursiveMutex *m)
     return 0;
 }
 
+#ifdef HAVE_FORK
+// Reinitialize the mutex in the child process after fork(). The caller must
+// record before fork() whether the forking thread owns the mutex, since the
+// thread ident can change in the child (gh-126688). If it does, the child
+// keeps the mutex with the same recursion level; otherwise the owner and any
+// waiters no longer exist in the child and the mutex is reset.
+void
+_PyRecursiveMutex_at_fork_reinit(_PyRecursiveMutex *m, int owned)
+{
+    if (owned) {
+        m->mutex = (PyMutex){._bits = _Py_LOCKED};
+        _Py_atomic_store_ullong_relaxed(&m->thread, PyThread_get_thread_ident_ex());
+    }
+    else {
+        memset(m, 0, sizeof(*m));
+    }
+}
+#endif
+
 #define _Py_WRITE_LOCKED 1
 #define _PyRWMutex_READER_SHIFT 2
 #define _Py_RWMUTEX_MAX_READERS (UINTPTR_MAX >> _PyRWMutex_READER_SHIFT)