[3.13] gh-117657: Fix TSAN race involving import lock (GH-118523) (#120169)

This adds a `_PyRecursiveMutex` type based on `PyMutex` and uses that
for the import lock. This fixes some data races in the free-threaded
build and generally simplifies the import lock code.
(cherry picked from commit e21057b99967eb5323320e6d1121955e0cd2985e)

Co-authored-by: Sam Gross <colesbury@gmail.com>

7 files changed, 90 insertions, 105 deletions

diff --git a/Include/internal/pycore_import.h b/Include/internal/pycore_import.h
index bd40707..3806e0d 100644
--- a/Include/internal/pycore_import.h
+++ b/Include/internal/pycore_import.h
@@ -20,7 +20,7 @@ PyAPI_FUNC(int) _PyImport_SetModule(PyObject *name, PyObject *module);
 extern int _PyImport_SetModuleString(const char *name, PyObject* module);
 
 extern void _PyImport_AcquireLock(PyInterpreterState *interp);
-extern int _PyImport_ReleaseLock(PyInterpreterState *interp);
+extern void _PyImport_ReleaseLock(PyInterpreterState *interp);
 
 // This is used exclusively for the sys and builtins modules:
 extern int _PyImport_FixupBuiltin(
@@ -94,11 +94,7 @@ struct _import_state {
 #endif
     PyObject *import_func;
     /* The global import lock. */
-    struct {
-        PyThread_type_lock mutex;
-        unsigned long thread;
-        int level;
-    } lock;
+    _PyRecursiveMutex lock;
     /* diagnostic info in PyImport_ImportModuleLevelObject() */
     struct {
         int import_level;
@@ -123,11 +119,6 @@ struct _import_state {
 #define IMPORTS_INIT \
     { \
         DLOPENFLAGS_INIT \
-        .lock = { \
-            .mutex = NULL, \
-            .thread = PYTHREAD_INVALID_THREAD_ID, \
-            .level = 0, \
-        }, \
         .find_and_load = { \
             .header = 1, \
         }, \
@@ -180,11 +171,6 @@ extern void _PyImport_FiniCore(PyInterpreterState *interp);
 extern void _PyImport_FiniExternal(PyInterpreterState *interp);
 
 
-#ifdef HAVE_FORK
-extern PyStatus _PyImport_ReInitLock(PyInterpreterState *interp);
-#endif
-
-
 extern PyObject* _PyImport_GetBuiltinModuleNames(void);
 
 struct _module_alias {
diff --git a/Include/internal/pycore_lock.h b/Include/internal/pycore_lock.h
index a5b28e4..d5853b2 100644
--- a/Include/internal/pycore_lock.h
+++ b/Include/internal/pycore_lock.h
@@ -219,6 +219,18 @@ _PyOnceFlag_CallOnce(_PyOnceFlag *flag, _Py_once_fn_t *fn, void *arg)
     return _PyOnceFlag_CallOnceSlow(flag, fn, arg);
 }
 
+// A recursive mutex. The mutex should zero-initialized.
+typedef struct {
+    PyMutex mutex;
+    unsigned long long thread;  // i.e., PyThread_get_thread_ident_ex()
+    size_t level;
+} _PyRecursiveMutex;
+
+PyAPI_FUNC(int) _PyRecursiveMutex_IsLockedByCurrentThread(_PyRecursiveMutex *m);
+PyAPI_FUNC(void) _PyRecursiveMutex_Lock(_PyRecursiveMutex *m);
+PyAPI_FUNC(void) _PyRecursiveMutex_Unlock(_PyRecursiveMutex *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
diff --git a/Modules/_testinternalcapi/test_lock.c b/Modules/_testinternalcapi/test_lock.c
index 4900459..1544fe1 100644
--- a/Modules/_testinternalcapi/test_lock.c
+++ b/Modules/_testinternalcapi/test_lock.c
@@ -2,6 +2,7 @@
 
 #include "parts.h"
 #include "pycore_lock.h"
+#include "pycore_pythread.h"      // PyThread_get_thread_ident_ex()
 
 #include "clinic/test_lock.c.h"
 
@@ -476,6 +477,29 @@ test_lock_rwlock(PyObject *self, PyObject *obj)
     Py_RETURN_NONE;
 }
 
+static PyObject *
+test_lock_recursive(PyObject *self, PyObject *obj)
+{
+    _PyRecursiveMutex m = (_PyRecursiveMutex){0};
+    assert(!_PyRecursiveMutex_IsLockedByCurrentThread(&m));
+
+    _PyRecursiveMutex_Lock(&m);
+    assert(m.thread == PyThread_get_thread_ident_ex());
+    assert(PyMutex_IsLocked(&m.mutex));
+    assert(m.level == 0);
+
+    _PyRecursiveMutex_Lock(&m);
+    assert(m.level == 1);
+    _PyRecursiveMutex_Unlock(&m);
+
+    _PyRecursiveMutex_Unlock(&m);
+    assert(m.thread == 0);
+    assert(!PyMutex_IsLocked(&m.mutex));
+    assert(m.level == 0);
+
+    Py_RETURN_NONE;
+}
+
 static PyMethodDef test_methods[] = {
     {"test_lock_basic", test_lock_basic, METH_NOARGS},
     {"test_lock_two_threads", test_lock_two_threads, METH_NOARGS},
@@ -485,6 +509,7 @@ static PyMethodDef test_methods[] = {
     {"test_lock_benchmark", test_lock_benchmark, METH_NOARGS},
     {"test_lock_once", test_lock_once, METH_NOARGS},
     {"test_lock_rwlock", test_lock_rwlock, METH_NOARGS},
+    {"test_lock_recursive", test_lock_recursive, METH_NOARGS},
     {NULL, NULL} /* sentinel */
 };
 
diff --git a/Modules/posixmodule.c b/Modules/posixmodule.c
index 386e942..5f943d4 100644
--- a/Modules/posixmodule.c
+++ b/Modules/posixmodule.c
@@ -16,7 +16,6 @@
 #include "pycore_call.h"          // _PyObject_CallNoArgs()
 #include "pycore_ceval.h"         // _PyEval_ReInitThreads()
 #include "pycore_fileutils.h"     // _Py_closerange()
-#include "pycore_import.h"        // _PyImport_ReInitLock()
 #include "pycore_initconfig.h"    // _PyStatus_EXCEPTION()
 #include "pycore_long.h"          // _PyLong_IsNegative()
 #include "pycore_moduleobject.h"  // _PyModule_GetState()
@@ -627,10 +626,7 @@ PyOS_AfterFork_Parent(void)
     _PyEval_StartTheWorldAll(&_PyRuntime);
 
     PyInterpreterState *interp = _PyInterpreterState_GET();
-    if (_PyImport_ReleaseLock(interp) <= 0) {
-        Py_FatalError("failed releasing import lock after fork");
-    }
-
+    _PyImport_ReleaseLock(interp);
     run_at_forkers(interp->after_forkers_parent, 0);
 }
 
@@ -675,10 +671,7 @@ PyOS_AfterFork_Child(void)
     _PyEval_StartTheWorldAll(&_PyRuntime);
     _PyThreadState_DeleteList(list);
 
-    status = _PyImport_ReInitLock(tstate->interp);
-    if (_PyStatus_EXCEPTION(status)) {
-        goto fatal_error;
-    }
+    _PyImport_ReleaseLock(tstate->interp);
 
     _PySignal_AfterFork();
 
diff --git a/Python/import.c b/Python/import.c
index 6fe6df4..10ac49f 100644
--- a/Python/import.c
+++ b/Python/import.c
@@ -94,11 +94,7 @@ static struct _inittab *inittab_copy = NULL;
     (interp)->imports.import_func
 
 #define IMPORT_LOCK(interp) \
-    (interp)->imports.lock.mutex
-#define IMPORT_LOCK_THREAD(interp) \
-    (interp)->imports.lock.thread
-#define IMPORT_LOCK_LEVEL(interp) \
-    (interp)->imports.lock.level
+    (interp)->imports.lock
 
 #define FIND_AND_LOAD(interp) \
     (interp)->imports.find_and_load
@@ -115,74 +111,14 @@ static struct _inittab *inittab_copy = NULL;
 void
 _PyImport_AcquireLock(PyInterpreterState *interp)
 {
-    unsigned long me = PyThread_get_thread_ident();
-    if (me == PYTHREAD_INVALID_THREAD_ID)
-        return; /* Too bad */
-    if (IMPORT_LOCK(interp) == NULL) {
-        IMPORT_LOCK(interp) = PyThread_allocate_lock();
-        if (IMPORT_LOCK(interp) == NULL)
-            return;  /* Nothing much we can do. */
-    }
-    if (IMPORT_LOCK_THREAD(interp) == me) {
-        IMPORT_LOCK_LEVEL(interp)++;
-        return;
-    }
-    if (IMPORT_LOCK_THREAD(interp) != PYTHREAD_INVALID_THREAD_ID ||
-        !PyThread_acquire_lock(IMPORT_LOCK(interp), 0))
-    {
-        PyThreadState *tstate = PyEval_SaveThread();
-        PyThread_acquire_lock(IMPORT_LOCK(interp), WAIT_LOCK);
-        PyEval_RestoreThread(tstate);
-    }
-    assert(IMPORT_LOCK_LEVEL(interp) == 0);
-    IMPORT_LOCK_THREAD(interp) = me;
-    IMPORT_LOCK_LEVEL(interp) = 1;
+    _PyRecursiveMutex_Lock(&IMPORT_LOCK(interp));
 }
 
-int
+void
 _PyImport_ReleaseLock(PyInterpreterState *interp)
 {
-    unsigned long me = PyThread_get_thread_ident();
-    if (me == PYTHREAD_INVALID_THREAD_ID || IMPORT_LOCK(interp) == NULL)
-        return 0; /* Too bad */
-    if (IMPORT_LOCK_THREAD(interp) != me)
-        return -1;
-    IMPORT_LOCK_LEVEL(interp)--;
-    assert(IMPORT_LOCK_LEVEL(interp) >= 0);
-    if (IMPORT_LOCK_LEVEL(interp) == 0) {
-        IMPORT_LOCK_THREAD(interp) = PYTHREAD_INVALID_THREAD_ID;
-        PyThread_release_lock(IMPORT_LOCK(interp));
-    }
-    return 1;
-}
-
-#ifdef HAVE_FORK
-/* This function is called from PyOS_AfterFork_Child() to ensure that newly
-   created child processes do not share locks with the parent.
-   We now acquire the import lock around fork() calls but on some platforms
-   (Solaris 9 and earlier? see isue7242) that still left us with problems. */
-PyStatus
-_PyImport_ReInitLock(PyInterpreterState *interp)
-{
-    if (IMPORT_LOCK(interp) != NULL) {
-        if (_PyThread_at_fork_reinit(&IMPORT_LOCK(interp)) < 0) {
-            return _PyStatus_ERR("failed to create a new lock");
-        }
-    }
-
-    if (IMPORT_LOCK_LEVEL(interp) > 1) {
-        /* Forked as a side effect of import */
-        unsigned long me = PyThread_get_thread_ident();
-        PyThread_acquire_lock(IMPORT_LOCK(interp), WAIT_LOCK);
-        IMPORT_LOCK_THREAD(interp) = me;
-        IMPORT_LOCK_LEVEL(interp)--;
-    } else {
-        IMPORT_LOCK_THREAD(interp) = PYTHREAD_INVALID_THREAD_ID;
-        IMPORT_LOCK_LEVEL(interp) = 0;
-    }
-    return _PyStatus_OK();
+    _PyRecursiveMutex_Unlock(&IMPORT_LOCK(interp));
 }
-#endif
 
 
 /***************/
@@ -4111,11 +4047,6 @@ _PyImport_FiniCore(PyInterpreterState *interp)
         PyErr_FormatUnraisable("Exception ignored on clearing sys.modules");
     }
 
-    if (IMPORT_LOCK(interp) != NULL) {
-        PyThread_free_lock(IMPORT_LOCK(interp));
-        IMPORT_LOCK(interp) = NULL;
-    }
-
     _PyImport_ClearCore(interp);
 }
 
@@ -4248,8 +4179,7 @@ _imp_lock_held_impl(PyObject *module)
 /*[clinic end generated code: output=8b89384b5e1963fc input=9b088f9b217d9bdf]*/
 {
     PyInterpreterState *interp = _PyInterpreterState_GET();
-    return PyBool_FromLong(
-            IMPORT_LOCK_THREAD(interp) != PYTHREAD_INVALID_THREAD_ID);
+    return PyBool_FromLong(PyMutex_IsLocked(&IMPORT_LOCK(interp).mutex));
 }
 
 /*[clinic input]
@@ -4283,11 +4213,12 @@ _imp_release_lock_impl(PyObject *module)
 /*[clinic end generated code: output=7faab6d0be178b0a input=934fb11516dd778b]*/
 {
     PyInterpreterState *interp = _PyInterpreterState_GET();
-    if (_PyImport_ReleaseLock(interp) < 0) {
+    if (!_PyRecursiveMutex_IsLockedByCurrentThread(&IMPORT_LOCK(interp))) {
         PyErr_SetString(PyExc_RuntimeError,
                         "not holding the import lock");
         return NULL;
     }
+    _PyImport_ReleaseLock(interp);
     Py_RETURN_NONE;
 }
 
diff --git a/Python/lock.c b/Python/lock.c
index 239e56a..555f4c2 100644
--- a/Python/lock.c
+++ b/Python/lock.c
@@ -366,6 +366,48 @@ _PyOnceFlag_CallOnceSlow(_PyOnceFlag *flag, _Py_once_fn_t *fn, void *arg)
     }
 }
 
+static int
+recursive_mutex_is_owned_by(_PyRecursiveMutex *m, PyThread_ident_t tid)
+{
+    return _Py_atomic_load_ullong_relaxed(&m->thread) == tid;
+}
+
+int
+_PyRecursiveMutex_IsLockedByCurrentThread(_PyRecursiveMutex *m)
+{
+    return recursive_mutex_is_owned_by(m, PyThread_get_thread_ident_ex());
+}
+
+void
+_PyRecursiveMutex_Lock(_PyRecursiveMutex *m)
+{
+    PyThread_ident_t thread = PyThread_get_thread_ident_ex();
+    if (recursive_mutex_is_owned_by(m, thread)) {
+        m->level++;
+        return;
+    }
+    PyMutex_Lock(&m->mutex);
+    _Py_atomic_store_ullong_relaxed(&m->thread, thread);
+    assert(m->level == 0);
+}
+
+void
+_PyRecursiveMutex_Unlock(_PyRecursiveMutex *m)
+{
+    PyThread_ident_t thread = PyThread_get_thread_ident_ex();
+    if (!recursive_mutex_is_owned_by(m, thread)) {
+        Py_FatalError("unlocking a recursive mutex that is not owned by the"
+                      " current thread");
+    }
+    if (m->level > 0) {
+        m->level--;
+        return;
+    }
+    assert(m->level == 0);
+    _Py_atomic_store_ullong_relaxed(&m->thread, 0);
+    PyMutex_Unlock(&m->mutex);
+}
+
 #define _Py_WRITE_LOCKED 1
 #define _PyRWMutex_READER_SHIFT 2
 #define _Py_RWMUTEX_MAX_READERS (UINTPTR_MAX >> _PyRWMutex_READER_SHIFT)
diff --git a/Tools/tsan/suppressions_free_threading.txt b/Tools/tsan/suppressions_free_threading.txt
index 8b64d1f..cb48a30 100644
--- a/Tools/tsan/suppressions_free_threading.txt
+++ b/Tools/tsan/suppressions_free_threading.txt
@@ -26,8 +26,6 @@ race:free_threadstate
 race_top:_add_to_weak_set
 race_top:_in_weak_set
 race_top:_PyEval_EvalFrameDefault
-race_top:_PyImport_AcquireLock
-race_top:_PyImport_ReleaseLock
 race_top:_PyType_HasFeature
 race_top:assign_version_tag
 race_top:insertdict
@@ -41,9 +39,7 @@ race_top:set_discard_entry
 race_top:set_inheritable
 race_top:Py_SET_TYPE
 race_top:_PyDict_CheckConsistency
-race_top:_PyImport_AcquireLock
 race_top:_Py_dict_lookup_threadsafe
-race_top:_imp_release_lock
 race_top:_multiprocessing_SemLock_acquire_impl
 race_top:dictiter_new
 race_top:dictresize