[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,17 @@ 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 (intern_strings() < 0) goto fail;",
+            "if (CPyGlobalsInit() < 0) goto fail;",
+            "",
+        )
+
         if self.compiler_options.separate:
             emitter.emit_lines(
                 'capsule = PyCapsule_New(&exports, "{}.exports", NULL);'.format(
@@ -1265,9 +1288,13 @@ 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_line("intern_strings();")
+        state = self.import_state_name(module_name)
+        module_cache = emitter.static_name(module_name, None, prefix=MODULE_PREFIX)
+        emitter.emit_lines(f"static int {impl_name}(PyObject *module)", "{")
+        if not self.use_shared_lib:
+            emitter.emit_lines("if (intern_strings() < 0)", "    return -1;")
         if self.compiler_options.depends_on_librt_internal:
             emitter.emit_line("if (import_librt_internal() < 0) {")
             emitter.emit_line("return -1;")
@@ -1333,7 +1360,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 +1387,13 @@ def emit_module_exec_func(
         emitter.emit_line("return -1;")
         emitter.emit_line("}")
 
+        emitter.emit_lines(
+            declaration,
+            "{",
+            f"return CPyImport_Exec(module, {impl_name}, &{state}, &{module_cache});",
+            "}",
+        )
+
     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 +1423,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,10 +1443,11 @@ 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;")
+        emitter.emit_line("PyObject *initializing_spec = NULL;")
         emitter.emit_lines(
             f"if ({module_static}) {{",
             f"Py_INCREF({module_static});",
@@ -1453,22 +1493,35 @@ 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"initializing_spec = CPyImport_BeginInitializing({module_static});")
+        emitter.emit_line("if (initializing_spec == NULL)")
+        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("rv = CPyImport_EndInitializing(initializing_spec);")
+        emitter.emit_line("initializing_spec = NULL;")
+        emitter.emit_line("if (rv < 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("    if (initializing_spec != NULL) {")
+        emitter.emit_line("        CPyImport_EndInitializing(initializing_spec);")
+        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 +1611,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,
@@ -83,7 +85,6 @@
 from mypyc.ir.class_ir import ClassIR, NonExtClassInfo
 from mypyc.ir.func_ir import INVALID_FUNC_DEF, FuncDecl, FuncIR, FuncSignature, RuntimeArg
 from mypyc.ir.ops import (
-    NAMESPACE_MODULE,
     NAMESPACE_TYPE_VAR,
     NO_TRACEBACK_LINE_NO,
     Assign,
@@ -167,7 +168,10 @@
 from mypyc.primitives.misc_ops import (
     check_unpack_count_op,
     get_module_dict_op,
+    import_cache_get_op,
+    import_cache_set_op,
     import_op,
+    native_import_is_initialized_op,
     native_import_op,
 )
 from mypyc.primitives.registry import CFunctionDescription, function_ops
@@ -523,10 +527,39 @@ 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)
+        module_cache = self.add(
+            LoadAddress(object_pointer_rprimitive, f"{MODULE_PREFIX}{exported_name(module)}")
+        )
+        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_cache, 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 +592,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),
@@ -568,22 +603,36 @@ def gen_import(self, module: str, line: int) -> None:
         else:
             # Import using generic Python C API
             value = self.call_c(import_op, [self.load_str(module, line)], line)
-        self.add(InitStatic(value, module, namespace=NAMESPACE_MODULE))
+        self.call_c(import_cache_set_op, [module_cache, value], line)
         self.goto_and_activate(out)
 
     def check_if_module_loaded(
-        self, id: str, line: int, needs_import: BasicBlock, out: BasicBlock
+        self,
+        module_cache: Value,
+        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.
+        """Generate code that checks if a module cache has been populated.
 
         Arguments:
-            id: name of module to check if imported
+            module_cache: address of the module cache to check
             line: line number that the import occurs on
             needs_import: the BasicBlock that is run if the module has not been loaded yet
             out: the BasicBlock that is run if the module has already been loaded"""
-        first_load = self.load_module(id)
+        first_load = self.call_c(import_cache_get_op, [module_cache], line)
         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, first_load, module_cache], 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/irbuild/statement.py

@@ -58,7 +58,6 @@
     BasicBlock,
     Branch,
     Call,
-    InitStatic,
     Integer,
     LoadAddress,
     LoadErrorValue,
@@ -127,6 +126,7 @@
     check_stop_op,
     coro_op,
     get_native_attrs_op,
+    import_cache_set_op,
     import_from_many_op,
     import_many_op,
     import_op,
@@ -597,7 +597,9 @@ def transform_import_from_buckets(
                 line,
             )
     if module is not None:
-        builder.add(InitStatic(module, module_id, namespace=NAMESPACE_MODULE))
+        mod_static = LoadStatic(object_rprimitive, module_id, namespace=NAMESPACE_MODULE)
+        module_cache = builder.add(LoadAddress(object_pointer_rprimitive, mod_static))
+        builder.call_c(import_cache_set_op, [module_cache, module], line)
 
 
 def transform_import_all(builder: IRBuilder, node: ImportAll) -> None: