[mypyc] Synchronize native-to-native imports

mypy/nativeparse.py

@@ -19,7 +19,6 @@
 from __future__ import annotations
 
 import os
-import time
 from typing import Final, cast
 
 import ast_serialize
@@ -280,13 +279,6 @@ def parse_to_binary_ast(
     source: str | bytes | None = None,
     skip_function_bodies: bool = False,
 ) -> tuple[bytes, list[ParseError], TypeIgnores, bytes, bool, bool, str, list[tuple[int, str]]]:
-    # This is a horrible hack to work around a mypyc bug where imported
-    # module may be not ready in a thread sometimes.
-    t0 = time.time()
-    while ast_serialize is None:
-        time.sleep(0.0001)  # type: ignore[unreachable]
-        if time.time() - t0 > 10.0:
-            raise ImportError("Cannot import ast_serialize")
     ast_bytes, errors, ignores, import_bytes, ast_data = ast_serialize.parse(
         filename,
         source,

mypyc/codegen/emitmodule.py

@@ -47,7 +47,9 @@
 from mypyc.codegen.literals import Literals
 from mypyc.common import (
     EXT_SUFFIX,
+    IMPORT_STATE_PREFIX,
     IS_FREE_THREADED,
+    MODULE_LOCK_API_PREFIX,
     MODULE_PREFIX,
     PREFIX,
     RUNTIME_C_FILES,
@@ -656,6 +658,13 @@ def __init__(
         # probably want to enable it always, but we'll wait until it's stable.
         self.multi_phase_init = IS_FREE_THREADED
 
+    def import_state_name(self, module_name: str) -> str:
+        return f"{IMPORT_STATE_PREFIX}{exported_name(module_name)}"
+
+    def module_lock_api_name(self) -> str:
+        assert self.group_name is not None
+        return f"{MODULE_LOCK_API_PREFIX}{exported_name(self.group_name)}"
+
     @property
     def group_suffix(self) -> str:
         return "_" + exported_name(self.group_name) if self.group_name else ""
@@ -692,6 +701,9 @@ def generate_c_for_modules(self) -> list[tuple[str, str]]:
         base_emitter.emit_line(f'#include "__native_internal{self.short_group_suffix}.h"')
         emitter = base_emitter
 
+        if self.use_shared_lib:
+            self.declare_module_lock_api()
+
         self.generate_literal_tables()
 
         for module_name, module in self.modules.items():
@@ -947,6 +959,16 @@ def generate_shared_lib_init(self, emitter: Emitter) -> None:
             "",
         )
 
+        lock_api = self.module_lock_api_name()
+        emitter.emit_lines(
+            f"if ({lock_api} == NULL) {{",
+            f"{lock_api} = CPyModuleLockAPI_Alloc();",
+            f"if ({lock_api} == NULL) goto fail;",
+            "}",
+            "if (CPyGlobalsInit() < 0) goto fail;",
+            "",
+        )
+
         if self.compiler_options.separate:
             emitter.emit_lines(
                 'capsule = PyCapsule_New(&exports, "{}.exports", NULL);'.format(
@@ -1265,8 +1287,9 @@ def emit_module_exec_func(
         exec_name = f"CPyExec_{exported_name(module_name)}"
         declaration = f"int {exec_name}(PyObject *module)"
         emitter.context.declarations[exec_name] = HeaderDeclaration(declaration + ";")
+        impl_name = f"{exec_name}__impl"
         module_static = self.module_internal_static_name(module_name, emitter)
-        emitter.emit_lines(declaration, "{")
+        emitter.emit_lines(f"static int {impl_name}(PyObject *module)", "{")
         emitter.emit_line("intern_strings();")
         if self.compiler_options.depends_on_librt_internal:
             emitter.emit_line("if (import_librt_internal() < 0) {")
@@ -1333,7 +1356,10 @@ def emit_module_exec_func(
                 name_prefix = cl.name_prefix(emitter.names)
                 emitter.emit_line(f"CPyDef_{name_prefix}_trait_vtable_setup();")
 
-        emitter.emit_lines("if (CPyGlobalsInit() < 0)", "    goto fail;")
+        if not self.use_shared_lib:
+            # With shared lib we initialize globals in its init function in case
+            # modules are executed concurrently.
+            emitter.emit_lines("if (CPyGlobalsInit() < 0)", "    goto fail;")
 
         self.generate_top_level_call(module, emitter)
 
@@ -1357,6 +1383,14 @@ def emit_module_exec_func(
         emitter.emit_line("return -1;")
         emitter.emit_line("}")
 
+        state = self.import_state_name(module_name)
+        emitter.emit_lines(
+            declaration,
+            "{",
+            f'return CPyImport_Exec(module, "{module_name}", {impl_name}, &{state});',
+            "}",
+        )
+
     def emit_init_only_func(self, emitter: Emitter, module_name: str, module_prefix: str) -> None:
         """Emit CPyInitOnly_* which creates the module object without executing the body.
 
@@ -1386,10 +1420,12 @@ def emit_module_init_func(
     ) -> None:
         if not self.use_shared_lib:
             declaration = f"PyMODINIT_FUNC PyInit_{module_name}(void)"
+            impl_declaration = declaration
         else:
             n = f"CPyInit_{exported_name(module_name)}"
             declaration = f"PyObject *{n}(void)"
             emitter.context.declarations[n] = HeaderDeclaration(declaration + ";")
+            impl_declaration = declaration
 
         if self.multi_phase_init:
             emitter.emit_lines(declaration, "{")
@@ -1404,7 +1440,7 @@ def emit_module_init_func(
             self.emit_init_only_func(emitter, module_name, module_prefix)
 
         # Emit CPyInit_* / PyInit_* which creates the module and executes the body.
-        emitter.emit_lines(declaration, "{")
+        emitter.emit_lines(impl_declaration, "{")
         module_static = self.module_internal_static_name(module_name, emitter)
 
         emitter.emit_line("PyObject* modname = NULL;")
@@ -1453,22 +1489,32 @@ def emit_module_init_func(
         emitter.emit_line("Py_DECREF(shared_lib_file);")
         emitter.emit_line("if (rv < 0) goto fail;")
 
-        # Register in sys.modules early so that circular imports via
-        # CPyImport_ImportNative can detect that this module is already
-        # being initialized and avoid re-executing the module body.
+        # Mark the module as initializing before publishing it so that CPython's
+        # import fast path waits on the module lock. Publishing early also lets
+        # CPyImport_ImportNative detect circular imports.
+        emitter.emit_line(f"if (CPyImport_SetInitializing({module_static}, 1) < 0)")
+        emitter.emit_line("    goto fail;")
         emitter.emit_line(
             f"if (PyObject_SetItem(PyImport_GetModuleDict(), modname, {module_static}) < 0)"
         )
         emitter.emit_line("    goto fail;")
         emitter.emit_line("Py_CLEAR(modname);")
         emitter.emit_lines(f"if ({exec_func}({module_static}) != 0)", "    goto fail;")
+        emitter.emit_line(f"if (CPyImport_SetInitializing({module_static}, 0) < 0)")
+        emitter.emit_line("    goto fail;")
         emitter.emit_line(f"return {module_static};")
         emitter.emit_lines("fail:")
         # Clean up on failure: remove from sys.modules and clear the static
         # so that a subsequent import attempt will retry initialization.
         emitter.emit_line("{")
         emitter.emit_line("    PyObject *exc_type, *exc_val, *exc_tb;")
         emitter.emit_line("    PyErr_Fetch(&exc_type, &exc_val, &exc_tb);")
+        emitter.emit_line(f"    if ({module_static} != NULL) {{")
+        emitter.emit_line(f"        CPyImport_SetInitializing({module_static}, 0);")
+        emitter.emit_line("        PyErr_Clear();")
+        emitter.emit_line("    }")
+        state = self.import_state_name(module_name)
+        emitter.emit_line(f"    CPyImport_SetInitialized(&{state}, 0);")
         emitter.emit_line("    if (modname == NULL) {")
         emitter.emit_line(f'        modname = PyUnicode_FromString("{module_name}");')
         emitter.emit_line("        if (modname == NULL) CPyError_OutOfMemory();")
@@ -1558,10 +1604,22 @@ def declare_module(self, module_name: str, emitter: Emitter) -> None:
         if module_name in self.modules:
             internal_static_name = self.module_internal_static_name(module_name, emitter)
             self.declare_global("CPyModule *", internal_static_name, initializer="NULL")
+            state_name = self.import_state_name(module_name)
+            if state_name not in self.context.declarations:
+                self.context.declarations[state_name] = HeaderDeclaration(
+                    f"CPyImportState {state_name};", defn=[f"CPyImportState {state_name} = {{0}};"]
+                )
         static_name = emitter.static_name(module_name, None, prefix=MODULE_PREFIX)
         self.declare_global("CPyModule *", static_name)
         self.simple_inits.append((static_name, "Py_None"))
 
+    def declare_module_lock_api(self) -> None:
+        name = self.module_lock_api_name()
+        if name not in self.context.declarations:
+            self.context.declarations[name] = HeaderDeclaration(
+                f"CPyModuleLockAPI *{name};", defn=[f"CPyModuleLockAPI *{name} = NULL;"]
+            )
+
     def declare_imports(self, imps: Iterable[str], emitter: Emitter) -> None:
         for imp in imps:
             self.declare_module(imp, emitter)

mypyc/common.py

@@ -14,6 +14,8 @@
 STATIC_PREFIX: Final = "CPyStatic_"  # Static variables (for literals etc.)
 TYPE_PREFIX: Final = "CPyType_"  # Type object struct
 MODULE_PREFIX: Final = "CPyModule_"  # Cached modules
+IMPORT_STATE_PREFIX: Final = "CPyImportState_"  # Native module initialization state
+MODULE_LOCK_API_PREFIX: Final = "CPyModuleLockAPI_"  # CPython module-lock API cache
 TYPE_VAR_PREFIX: Final = "CPyTypeVar_"  # Type variables when using new-style Python 3.12 syntax
 ATTR_PREFIX: Final = "_"  # Attributes
 FAST_PREFIX: Final = "__mypyc_fast_"  # Optimized methods in non-extension classes
@@ -92,6 +94,7 @@
     "tuple_ops.c",
     "exc_ops.c",
     "misc_ops.c",
+    "locks.c",
     "generic_ops.c",
     "pythonsupport.c",
     "function_wrapper.c",

mypyc/irbuild/builder.py

@@ -70,9 +70,11 @@
     BITMAP_BITS,
     EXT_SUFFIX,
     GENERATOR_ATTRIBUTE_PREFIX,
+    IMPORT_STATE_PREFIX,
     IS_FREE_THREADED,
     KEEP_ALIVE_SHORT_LIVED,
     KEEP_ALIVE_WHOLE_EXPRESSION,
+    MODULE_LOCK_API_PREFIX,
     MODULE_PREFIX,
     SELF_NAME,
     TEMP_ATTR_NAME,
@@ -168,6 +170,7 @@
     check_unpack_count_op,
     get_module_dict_op,
     import_op,
+    native_import_is_initialized_op,
     native_import_op,
 )
 from mypyc.primitives.registry import CFunctionDescription, function_ops
@@ -523,10 +526,36 @@ def gen_import(self, module: str, line: int) -> None:
         self.imports[module] = None
 
         needs_import, out = BasicBlock(), BasicBlock()
-        self.check_if_module_loaded(module, line, needs_import, out)
+        is_native_module = self.is_native_module(module)
+        is_same_group_native = is_native_module and self.is_same_group_module(module)
+        import_state: Value | None = None
+        module_lock_api: Value | None = None
+        if is_same_group_native:
+            import_state = self.add(
+                LoadAddress(c_pointer_rprimitive, f"{IMPORT_STATE_PREFIX}{exported_name(module)}")
+            )
+            group_name = self.mapper.group_map.get(self.module_name)
+            if group_name is not None:
+                module_lock_api = self.add(
+                    LoadGlobal(
+                        c_pointer_rprimitive,
+                        f"{MODULE_LOCK_API_PREFIX}{exported_name(group_name)}",
+                    )
+                )
+            else:
+                module_lock_api = Integer(0, c_pointer_rprimitive)
+        if is_native_module and not is_same_group_native:
+            # A sys.modules entry may still be executing in another thread. For
+            # native modules in another compilation group, use CPython's import
+            # path so that its per-module lock waits for initialization.
+            self.goto(needs_import)
+        else:
+            self.check_if_module_loaded(module, line, needs_import, out, import_state)
 
         self.activate_block(needs_import)
-        if self.is_native_module(module) and self.is_same_group_module(module):
+        if is_same_group_native:
+            assert import_state is not None
+            assert module_lock_api is not None
             # Use custom import machinery for native-to-native imports in the same group
             init_only_func = self.add(
                 LoadGlobal(c_pointer_rprimitive, f"CPyInitOnly_{exported_name(module)}")
@@ -559,6 +588,8 @@ def gen_import(self, module: str, line: int) -> None:
                     init_only_func,
                     exec_func,
                     module_static,
+                    import_state,
+                    module_lock_api,
                     shared_lib_file,
                     ext_suffix,
                     Integer(1 if is_pkg else 0, c_pyssize_t_rprimitive),
@@ -572,7 +603,12 @@ def gen_import(self, module: str, line: int) -> None:
         self.goto_and_activate(out)
 
     def check_if_module_loaded(
-        self, id: str, line: int, needs_import: BasicBlock, out: BasicBlock
+        self,
+        id: str,
+        line: int,
+        needs_import: BasicBlock,
+        out: BasicBlock,
+        import_state: Value | None = None,
     ) -> None:
         """Generate code that checks if the module `id` has been loaded yet.
 
@@ -583,7 +619,14 @@ def check_if_module_loaded(
             out: the BasicBlock that is run if the module has already been loaded"""
         first_load = self.load_module(id)
         comparison = self.translate_is_op(first_load, self.none_object(line), "is not", line)
-        self.add_bool_branch(comparison, out, needs_import)
+        if import_state is None:
+            self.add_bool_branch(comparison, out, needs_import)
+        else:
+            check_initialized = BasicBlock()
+            self.add_bool_branch(comparison, check_initialized, needs_import)
+            self.activate_block(check_initialized)
+            initialized = self.call_c(native_import_is_initialized_op, [import_state], line)
+            self.add_bool_branch(initialized, out, needs_import)
 
     def get_module(self, module: str, line: int) -> Value:
         # Python 3.7 has a nice 'PyImport_GetModule' function that we can't use :(

mypyc/lib-rt/CPy.h

@@ -22,6 +22,29 @@ extern "C" {
 #define CPYTHON_LARGE_INT_ERRMSG "Python int too large to convert to C ssize_t"
 
 
+// Native module import synchronization
+
+typedef struct CPyModuleLockAPI CPyModuleLockAPI;
+
+typedef struct {
+    int32_t initialized;
+} CPyImportState;
+
+enum {
+    CPY_LOCK_ERROR = -1,
+    CPY_LOCK_ACQUIRED = 0,
+    CPY_LOCK_DEADLOCK = 1,
+};
+
+CPyModuleLockAPI *CPyModuleLockAPI_Alloc(void);
+void CPyModuleLockAPI_Free(CPyModuleLockAPI *api);
+int CPyImport_AcquireLock(CPyModuleLockAPI *api, PyObject *module_name,
+                          PyObject **module_lock);
+int CPyImport_ReleaseLock(PyObject *module_lock);
+bool CPyImport_IsInitialized(const CPyImportState *state);
+void CPyImport_SetInitialized(CPyImportState *state, bool initialized);
+
+
 // Naming conventions:
 //
 // Tagged: tagged int
@@ -1043,8 +1066,12 @@ PyObject *CPyImport_ImportNative(PyObject *module_name,
                                  PyObject *(*init_only_fn)(void),
                                  int (*exec_fn)(PyObject *),
                                  CPyModule **module_static,
+                                 CPyImportState *state, CPyModuleLockAPI *lock_api,
                                  PyObject *shared_lib_file, PyObject *ext_suffix,
                                  Py_ssize_t is_package);
+int CPyImport_Exec(PyObject *module, const char *module_name,
+                   int (*exec_fn)(PyObject *), CPyImportState *state);
+int CPyImport_SetInitializing(PyObject *module, bool initializing);
 int CPyImport_SetDunderAttrs(PyObject *module, PyObject *module_name, PyObject *shared_lib_file,
                              PyObject *ext_suffix, Py_ssize_t is_package);