gh-81057: Move the Extension Modules Cache to _PyRuntimeState (gh-99355)

We also move the closely related max_module_number and add comments documenting the group of struct members.

https://github.com/python/cpython/issues/81057

Include/internal/pycore_import.h

@@ -5,6 +5,23 @@
 extern "C" {
 #endif
 
+
+struct _import_runtime_state {
+    /* The most recent value assigned to a PyModuleDef.m_base.m_index.
+       This is incremented each time PyModuleDef_Init() is called,
+       which is just about every time an extension module is imported.
+       See PyInterpreterState.modules_by_index for more info. */
+    Py_ssize_t last_module_index;
+    /* 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)
+       or are imported in the main interpreter.
+       This is initialized lazily in _PyImport_FixupExtensionObject().
+       Modules are added there and looked up in _imp.find_extension(). */
+    PyObject *extensions;
+};
+
+
 #ifdef HAVE_FORK
 extern PyStatus _PyImport_ReInitLock(void);
 #endif

Include/internal/pycore_interp.h

@@ -123,6 +123,25 @@ struct _is {
 
     // sys.modules dictionary
     PyObject *modules;
+    /* This is the list of module objects for all legacy (single-phase init)
+       extension modules ever loaded in this process (i.e. imported
+       in this interpreter or in any other).  Py_None stands in for
+       modules that haven't actually been imported in this interpreter.
+
+       A module's index (PyModuleDef.m_base.m_index) is used to look up
+       the corresponding module object for this interpreter, if any.
+       (See PyState_FindModule().)  When any extension module
+       is initialized during import, its moduledef gets initialized by
+       PyModuleDef_Init(), and the first time that happens for each
+       PyModuleDef, its index gets set to the current value of
+       a global counter (see _PyRuntimeState.imports.last_module_index).
+       The entry for that index in this interpreter remains unset until
+       the module is actually imported here.  (Py_None is used as
+       a placeholder.)  Note that multi-phase init modules always get
+       an index for which there will never be a module set.
+
+       This is initialized lazily in _PyState_AddModule(), which is also
+       where modules get added. */
     PyObject *modules_by_index;
     // Dictionary of the sys module
     PyObject *sysdict;

Include/internal/pycore_runtime.h

@@ -11,6 +11,7 @@ extern "C" {
 #include "pycore_atomic.h"          /* _Py_atomic_address */
 #include "pycore_gil.h"             // struct _gil_runtime_state
 #include "pycore_global_objects.h"  // struct _Py_global_objects
+#include "pycore_import.h"          // struct _import_runtime_state
 #include "pycore_interp.h"          // PyInterpreterState
 #include "pycore_unicodeobject.h"   // struct _Py_unicode_runtime_ids
 
@@ -115,6 +116,7 @@ typedef struct pyruntimestate {
     void (*exitfuncs[NEXITFUNCS])(void);
     int nexitfuncs;
 
+    struct _import_runtime_state imports;
     struct _ceval_runtime_state ceval;
     struct _gilstate_runtime_state gilstate;
     struct _getargs_runtime_state getargs;

Include/moduleobject.h

@@ -43,8 +43,22 @@ PyAPI_DATA(PyTypeObject) PyModuleDef_Type;
 
 typedef struct PyModuleDef_Base {
   PyObject_HEAD
+  /* The function used to re-initialize the module.
+     This is only set for legacy (single-phase init) extension modules
+     and only used for those that support multiple initializations
+     (m_size >= 0).
+     It is set by _PyImport_LoadDynamicModuleWithSpec()
+     and _imp.create_builtin(). */
   PyObject* (*m_init)(void);
+  /* The module's index into its interpreter's modules_by_index cache.
+     This is set for all extension modules but only used for legacy ones.
+     (See PyInterpreterState.modules_by_index for more info.)
+     It is set by PyModuleDef_Init(). */
   Py_ssize_t m_index;
+  /* A copy of the module's __dict__ after the first time it was loaded.
+     This is only set/used for legacy modules that do not support
+     multiple initializations.
+     It is set by _PyImport_FixupExtensionObject(). */
   PyObject* m_copy;
 } PyModuleDef_Base;
 

Objects/moduleobject.c

@@ -9,7 +9,6 @@
 #include "pycore_moduleobject.h"  // _PyModule_GetDef()
 #include "structmember.h"         // PyMemberDef
 
-static Py_ssize_t max_module_number;
 
 static PyMemberDef module_members[] = {
     {"__dict__", T_OBJECT, offsetof(PyModuleObject, md_dict), READONLY},
@@ -43,10 +42,10 @@ PyModuleDef_Init(PyModuleDef* def)
 {
     assert(PyModuleDef_Type.tp_flags & Py_TPFLAGS_READY);
     if (def->m_base.m_index == 0) {
-        max_module_number++;
+        _PyRuntime.imports.last_module_index++;
         Py_SET_REFCNT(def, 1);
         Py_SET_TYPE(def, &PyModuleDef_Type);
-        def->m_base.m_index = max_module_number;
+        def->m_base.m_index = _PyRuntime.imports.last_module_index;
     }
     return (PyObject*)def;
 }

Python/import.c

@@ -27,9 +27,6 @@ extern "C" {
 /* Forward references */
 static PyObject *import_add_module(PyThreadState *tstate, PyObject *name);
 
-/* See _PyImport_FixupExtensionObject() below */
-static PyObject *extensions = NULL;
-
 /* This table is defined in config.c: */
 extern struct _inittab _PyImport_Inittab[];
 
@@ -221,10 +218,12 @@ _imp_release_lock_impl(PyObject *module)
     Py_RETURN_NONE;
 }
 
+static inline void _extensions_cache_clear(void);
+
 void
 _PyImport_Fini(void)
 {
-    Py_CLEAR(extensions);
+    _extensions_cache_clear();
     if (import_lock != NULL) {
         PyThread_free_lock(import_lock);
         import_lock = NULL;
@@ -398,6 +397,51 @@ PyImport_GetMagicTag(void)
    dictionary, to avoid loading shared libraries twice.
 */
 
+static PyModuleDef *
+_extensions_cache_get(PyObject *filename, PyObject *name)
+{
+    PyObject *extensions = _PyRuntime.imports.extensions;
+    if (extensions == NULL) {
+        return NULL;
+    }
+    PyObject *key = PyTuple_Pack(2, filename, name);
+    if (key == NULL) {
+        return NULL;
+    }
+    PyModuleDef *def = (PyModuleDef *)PyDict_GetItemWithError(extensions, key);
+    Py_DECREF(key);
+    return def;
+}
+
+static int
+_extensions_cache_set(PyObject *filename, PyObject *name, PyModuleDef *def)
+{
+    PyObject *extensions = _PyRuntime.imports.extensions;
+    if (extensions == NULL) {
+        extensions = PyDict_New();
+        if (extensions == NULL) {
+            return -1;
+        }
+        _PyRuntime.imports.extensions = extensions;
+    }
+    PyObject *key = PyTuple_Pack(2, filename, name);
+    if (key == NULL) {
+        return -1;
+    }
+    int res = PyDict_SetItem(extensions, key, (PyObject *)def);
+    Py_DECREF(key);
+    if (res < 0) {
+        return -1;
+    }
+    return 0;
+}
+
+static void
+_extensions_cache_clear(void)
+{
+    Py_CLEAR(_PyRuntime.imports.extensions);
+}
+
 int
 _PyImport_FixupExtensionObject(PyObject *mod, PyObject *name,
                                PyObject *filename, PyObject *modules)
@@ -442,20 +486,7 @@ _PyImport_FixupExtensionObject(PyObject *mod, PyObject *name,
             }
         }
 
-        if (extensions == NULL) {
+        if (_extensions_cache_set(filename, name, def) < 0) {
-            extensions = PyDict_New();
-            if (extensions == NULL) {
-                return -1;
-            }
-        }
-
-        PyObject *key = PyTuple_Pack(2, filename, name);
-        if (key == NULL) {
-            return -1;
-        }
-        int res = PyDict_SetItem(extensions, key, (PyObject *)def);
-        Py_DECREF(key);
-        if (res < 0) {
             return -1;
         }
     }
@@ -480,16 +511,7 @@ static PyObject *
 import_find_extension(PyThreadState *tstate, PyObject *name,
                       PyObject *filename)
 {
-    if (extensions == NULL) {
+    PyModuleDef *def = _extensions_cache_get(filename, name);
-        return NULL;
-    }
-
-    PyObject *key = PyTuple_Pack(2, filename, name);
-    if (key == NULL) {
-        return NULL;
-    }
-    PyModuleDef* def = (PyModuleDef *)PyDict_GetItemWithError(extensions, key);
-    Py_DECREF(key);
     if (def == NULL) {
         return NULL;
     }

Tools/c-analyzer/cpython/globals-to-fix.tsv

@@ -317,7 +317,6 @@ Python/hamt.c	-	_empty_hamt	-
 
 # state
 Objects/typeobject.c	resolve_slotdups	pname	-
-Python/import.c	-	extensions	-
 
 
 ##################################
@@ -449,7 +448,6 @@ Python/getargs.c	-	static_arg_parsers	-
 Objects/dictobject.c	-	_pydict_global_version	-
 Objects/dictobject.c	-	next_dict_keys_version	-
 Objects/funcobject.c	-	next_func_version	-
-Objects/moduleobject.c	-	max_module_number	-
 Objects/object.c	-	_Py_RefTotal	-
 Python/perf_trampoline.c	-	perf_status	-
 Python/perf_trampoline.c	-	extra_code_index	-