From 02b681ce76f62ef9d45e6ac509b018bdb63f6aca Mon Sep 17 00:00:00 2001
From: Eric Snow <ericsnowcurrently@gmail.com>
Date: Thu, 6 Apr 2023 11:32:16 -0600
Subject: [PATCH 1/2] Add a granular lock for _PyRuntime.imports.extensions.

---
 Include/internal/pycore_import.h | 2 ++
 Python/pystate.c                 | 5 ++++-
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/Include/internal/pycore_import.h b/Include/internal/pycore_import.h
index 7a78a91aa617e6..0a9f24efbdb908 100644
--- a/Include/internal/pycore_import.h
+++ b/Include/internal/pycore_import.h
@@ -19,6 +19,8 @@ struct _import_runtime_state {
            used exclusively for when the extensions dict is access/modified
            from an arbitrary thread. */
         PyThreadState main_tstate;
+        /* A lock to guard the dict. */
+        PyThread_type_lock mutex;
         /* A dict mapping (filename, name) to PyModuleDef for modules.
            Only legacy (single-phase init) extension modules are added
            and only if they support multiple initialization (m_size >- 0)
diff --git a/Python/pystate.c b/Python/pystate.c
index 37cef972e0feb4..6be1fa8f72442d 100644
--- a/Python/pystate.c
+++ b/Python/pystate.c
@@ -354,7 +354,7 @@ _Py_COMP_DIAG_IGNORE_DEPR_DECLS
 static const _PyRuntimeState initial = _PyRuntimeState_INIT(_PyRuntime);
 _Py_COMP_DIAG_POP
 
-#define NUMLOCKS 4
+#define NUMLOCKS 5
 
 static int
 alloc_for_runtime(PyThread_type_lock locks[NUMLOCKS])
@@ -408,6 +408,7 @@ init_runtime(_PyRuntimeState *runtime,
         &runtime->xidregistry.mutex,
         &runtime->getargs.mutex,
         &runtime->unicode_state.ids.lock,
+        &runtime->imports.extensions.mutex,
     };
     for (int i = 0; i < NUMLOCKS; i++) {
         assert(locks[i] != NULL);
@@ -492,6 +493,7 @@ _PyRuntimeState_Fini(_PyRuntimeState *runtime)
         &runtime->xidregistry.mutex,
         &runtime->getargs.mutex,
         &runtime->unicode_state.ids.lock,
+        &runtime->imports.extensions.mutex,
     };
     for (int i = 0; i < NUMLOCKS; i++) {
         FREE_LOCK(*lockptrs[i]);
@@ -520,6 +522,7 @@ _PyRuntimeState_ReInitThreads(_PyRuntimeState *runtime)
         &runtime->xidregistry.mutex,
         &runtime->getargs.mutex,
         &runtime->unicode_state.ids.lock,
+        &runtime->imports.extensions.mutex,
     };
     int reinit_err = 0;
     for (int i = 0; i < NUMLOCKS; i++) {

From 17553790058bf33c8d6185d66311cc6597458141 Mon Sep 17 00:00:00 2001
From: Eric Snow <ericsnowcurrently@gmail.com>
Date: Thu, 6 Apr 2023 11:32:43 -0600
Subject: [PATCH 2/2] Use the lock.

---
 Python/import.c | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/Python/import.c b/Python/import.c
index 1db5b9333bbba1..a3ab140b9f6450 100644
--- a/Python/import.c
+++ b/Python/import.c
@@ -413,8 +413,11 @@ remove_module(PyThreadState *tstate, PyObject *name)
 Py_ssize_t
 _PyImport_GetNextModuleIndex(void)
 {
+    PyThread_acquire_lock(EXTENSIONS.mutex, WAIT_LOCK);
     LAST_MODULE_INDEX++;
-    return LAST_MODULE_INDEX;
+    Py_ssize_t index = LAST_MODULE_INDEX;
+    PyThread_release_lock(EXTENSIONS.mutex);
+    return index;
 }
 
 static const char *
@@ -703,6 +706,7 @@ _PyImport_ClearModulesByIndex(PyInterpreterState *interp)
 const char *
 _PyImport_ResolveNameWithPackageContext(const char *name)
 {
+    PyThread_acquire_lock(EXTENSIONS.mutex, WAIT_LOCK);
     if (PKGCONTEXT != NULL) {
         const char *p = strrchr(PKGCONTEXT, '.');
         if (p != NULL && strcmp(name, p+1) == 0) {
@@ -710,14 +714,17 @@ _PyImport_ResolveNameWithPackageContext(const char *name)
             PKGCONTEXT = NULL;
         }
     }
+    PyThread_release_lock(EXTENSIONS.mutex);
     return name;
 }
 
 const char *
 _PyImport_SwapPackageContext(const char *newcontext)
 {
+    PyThread_acquire_lock(EXTENSIONS.mutex, WAIT_LOCK);
     const char *oldcontext = PKGCONTEXT;
     PKGCONTEXT = newcontext;
+    PyThread_release_lock(EXTENSIONS.mutex);
     return oldcontext;
 }
 
@@ -865,13 +872,13 @@ gets even messier.
 static inline void
 extensions_lock_acquire(void)
 {
-    // XXX For now the GIL is sufficient.
+    PyThread_acquire_lock(_PyRuntime.imports.extensions.mutex, WAIT_LOCK);
 }
 
 static inline void
 extensions_lock_release(void)
 {
-    // XXX For now the GIL is sufficient.
+    PyThread_release_lock(_PyRuntime.imports.extensions.mutex);
 }
 
 /* Magic for extension modules (built-in as well as dynamically