Synchronize module cache; report error when deleting and reimporting

Author: p-sawicki

mypyc/codegen/emitmodule.py

@@ -1290,6 +1290,8 @@ def emit_module_exec_func(
         emitter.context.declarations[exec_name] = HeaderDeclaration(declaration + ";")
         impl_name = f"{exec_name}__impl"
         module_static = self.module_internal_static_name(module_name, emitter)
+        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;")
@@ -1385,9 +1387,11 @@ 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, {impl_name}, &{state});", "}"
+            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:

mypyc/irbuild/builder.py

@@ -629,7 +629,9 @@ def check_if_module_loaded(
             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)
+            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:

mypyc/lib-rt/CPy.h

@@ -42,9 +42,12 @@ int CPyImport_AcquireLock(CPyModuleLockAPI *api, PyObject *module_name,
                           PyObject **module_lock);
 int CPyImport_ReleaseLock(PyObject *module_lock);
 bool CPyImport_IsInitialized(const CPyImportState *state);
+bool CPyImport_IsInitializedForModule(const CPyImportState *state, PyObject *module,
+                                      CPyModule **module_cache);
 void CPyImport_SetInitialized(CPyImportState *state, bool initialized);
 PyObject *CPyImport_GetModuleCache(CPyModule **cache);
 void CPyImport_SetModuleCache(CPyModule **cache, PyObject *module);
+void CPyImport_ReplaceModuleCache(CPyModule **cache, PyObject *module);
 
 
 // Naming conventions:
@@ -1071,7 +1074,8 @@ PyObject *CPyImport_ImportNative(PyObject *module_name,
                                  CPyImportState *state, CPyModuleLockAPI *lock_api,
                                  PyObject *shared_lib_file, PyObject *ext_suffix,
                                  Py_ssize_t is_package);
-int CPyImport_Exec(PyObject *module, int (*exec_fn)(PyObject *), CPyImportState *state);
+int CPyImport_Exec(PyObject *module, int (*exec_fn)(PyObject *), CPyImportState *state,
+                   CPyModule **module_cache);
 PyObject *CPyImport_BeginInitializing(PyObject *module);
 int CPyImport_EndInitializing(PyObject *spec);
 int CPyImport_SetDunderAttrs(PyObject *module, PyObject *module_name, PyObject *shared_lib_file,

mypyc/lib-rt/locks.c

@@ -89,6 +89,14 @@ bool CPyImport_IsInitialized(const CPyImportState *state) {
 #endif
 }
 
+bool CPyImport_IsInitializedForModule(const CPyImportState *state, PyObject *module,
+                                      CPyModule **module_cache) {
+    // Read initialized before re-reading the cache. If a retry completed after
+    // the caller's first cache load, this load rejects its stale pointer.
+    return CPyImport_IsInitialized(state)
+        && CPyImport_GetModuleCache(module_cache) == module;
+}
+
 void CPyImport_SetInitialized(CPyImportState *state, bool initialized) {
 #ifdef _WIN32
     InterlockedExchange((volatile LONG *)&state->initialized, initialized);
@@ -97,6 +105,19 @@ void CPyImport_SetInitialized(CPyImportState *state, bool initialized) {
 #endif
 }
 
+static void CPyImport_DecRefOld(PyObject *previous) {
+    if (previous == NULL) {
+        return;
+    }
+#ifdef Py_GIL_DISABLED
+    // Atomic loads return borrowed references, so defer releasing the old
+    // reference until concurrent readers have passed a quiescent point.
+    CPy_DecRefAttrOld(previous);
+#else
+    Py_DECREF(previous);
+#endif
+}
+
 PyObject *CPyImport_GetModuleCache(CPyModule **cache) {
 #ifdef _WIN32
     return InterlockedCompareExchangePointer((PVOID volatile *)cache, NULL, NULL);
@@ -121,3 +142,17 @@ void CPyImport_SetModuleCache(CPyModule **cache, PyObject *module) {
     }
 #endif
 }
+
+void CPyImport_ReplaceModuleCache(CPyModule **cache, PyObject *module) {
+    Py_INCREF(module);
+    CPyModule *previous;
+#ifdef _WIN32
+    previous = InterlockedExchangePointer((PVOID volatile *)cache, module);
+#else
+    previous = __atomic_exchange_n(cache, (CPyModule *)module, __ATOMIC_ACQ_REL);
+#endif
+    if (previous == NULL || previous == (CPyModule *)Py_None) {
+        return;
+    }
+    CPyImport_DecRefOld((PyObject *)previous);
+}

mypyc/lib-rt/misc_ops.c

@@ -1585,13 +1585,28 @@ static int CPyImport_ReleaseLockPreservingException(PyObject *module_lock) {
 }
 
 // Execute a module once; caller holds the module lock.
-int CPyImport_Exec(PyObject *module, int (*exec_fn)(PyObject *), CPyImportState *state) {
+int CPyImport_Exec(PyObject *module, int (*exec_fn)(PyObject *), CPyImportState *state,
+                   CPyModule **module_cache) {
     if (CPyImport_IsInitialized(state)) {
-        return 0;
+        const char *module_name = PyModule_GetName(module);
+        if (module_name != NULL) {
+            PyErr_Format(PyExc_ImportError,
+                         "native module '%s' does not support reinitialization",
+                         module_name);
+        }
+        return -1;
     }
 
     int result = exec_fn(module);
     if (result == 0) {
+        // Keep the cache lazy. A normal shim import should not populate it, since
+        // the first compiled native import must still validate sys.modules. If a
+        // compiled import already populated the cache (including with a partial
+        // module during a circular import), refresh it to this instance.
+        PyObject *cached_module = CPyImport_GetModuleCache(module_cache);
+        if (cached_module != NULL && cached_module != Py_None) {
+            CPyImport_ReplaceModuleCache(module_cache, module);
+        }
         CPyImport_SetInitialized(state, true);
     }
     return result;
@@ -1737,8 +1752,8 @@ PyObject *CPyImport_ImportNative(PyObject *module_name,
     PyObject_DelItem(module_dict, module_name);
     PyErr_Clear();
     PyErr_Restore(exc_type, exc_val, exc_tb);
-    Py_CLEAR(*module_static);
     CPyImport_SetInitialized(state, false);
+    Py_CLEAR(*module_static);
     CPyImport_ReleaseLockPreservingException(module_lock);
     return NULL;
 }

mypyc/primitives/misc_ops.py

@@ -160,9 +160,9 @@
 )
 
 native_import_is_initialized_op = custom_op(
-    arg_types=[c_pointer_rprimitive],
+    arg_types=[c_pointer_rprimitive, object_rprimitive, object_pointer_rprimitive],
     return_type=bit_rprimitive,
-    c_function_name="CPyImport_IsInitialized",
+    c_function_name="CPyImport_IsInitializedForModule",
     error_kind=ERR_NEVER,
 )
 

mypyc/test-data/run-multimodule.test

@@ -2415,6 +2415,93 @@ assert errors == ["first initialization failed"]
 assert failure_state.attempts == 2
 assert other_importer.load() == 42
 
+[case testFailedConcurrentCircularNativeImportRefreshesCache]
+# separate: [(["native.py", "other_a.py", "other_b.py"], "testgroup")]
+pass
+
+[file import_sync.py]
+from threading import Event
+
+a_started = Event()
+b_started = Event()
+a_importing_b = Event()
+b_waiting_for_driver = Event()
+b_may_import_a = Event()
+
+[file failure_state.py]
+attempts = 0
+
+[file other_a.py]
+import failure_state
+
+failure_state.attempts += 1
+if failure_state.attempts == 1:
+    import import_sync
+
+    import_sync.a_started.set()
+    assert import_sync.b_started.wait(timeout=5)
+    import_sync.a_importing_b.set()
+    import other_b
+    raise RuntimeError("first initialization failed")
+
+value = 42
+
+[file other_b.py]
+import import_sync
+
+import_sync.b_started.set()
+assert import_sync.a_started.wait(timeout=5)
+assert import_sync.a_importing_b.wait(timeout=5)
+import_sync.b_waiting_for_driver.set()
+assert import_sync.b_may_import_a.wait(timeout=5)
+import other_a
+
+saw_partial_a = not hasattr(other_a, "value")
+
+def get_a_value() -> int:
+    return other_a.value
+
+[file driver.py]
+from concurrent.futures import ThreadPoolExecutor
+import importlib
+from time import monotonic, sleep
+
+from testutil import assertRaises
+
+import failure_state
+import import_sync
+import native  # Preload the compilation group's shared library.
+
+with ThreadPoolExecutor(max_workers=2) as executor:
+    future_a = executor.submit(importlib.import_module, "other_a")
+    assert import_sync.a_started.wait(timeout=5)
+    future_b = executor.submit(importlib.import_module, "other_b")
+    assert import_sync.a_importing_b.wait(timeout=5)
+    assert import_sync.b_waiting_for_driver.wait(timeout=5)
+
+    # Wait until other_a is blocked on other_b's module lock. Letting other_b
+    # import other_a then deterministically closes the lock cycle and returns
+    # other_a's partial module.
+    bootstrap = importlib.import_module("importlib._bootstrap")
+    lock = bootstrap._get_module_lock("other_b")
+    deadline = monotonic() + 5
+    while not lock.waiters and monotonic() < deadline:
+        sleep(0.001)
+    has_waiter = bool(lock.waiters)
+    import_sync.b_may_import_a.set()
+    assert has_waiter
+
+    with assertRaises(RuntimeError, "first initialization failed"):
+        future_a.result(timeout=10)
+    other_b = future_b.result(timeout=10)
+
+assert failure_state.attempts == 1
+assert other_b.saw_partial_a
+other_a = importlib.import_module("other_a")
+assert failure_state.attempts == 2
+assert other_a.value == 42
+assert other_b.get_a_value() == 42
+
 [case testConcurrentCircularNativeImports]
 # separate: [(["other_a.py", "other_b.py"], "testgroup")]
 pass
@@ -2457,6 +2544,33 @@ with ThreadPoolExecutor(max_workers=2) as executor:
 assert other_a.other_value == "b"
 assert other_b.other_value == "a"
 
+[case testNativeModuleReimportBehavior]
+value = 42
+
+def get_value() -> int:
+    return value
+
+[file driver.py]
+import importlib
+import sys
+
+from testutil import assertRaises
+
+import native
+
+first = native
+assert first.get_value() == 42
+del sys.modules["native"]
+if hasattr(sys, "_is_gil_enabled") and not sys._is_gil_enabled():
+    with assertRaises(ImportError, "native module 'native' does not support reinitialization"):
+        importlib.import_module("native")
+    assert "native" not in sys.modules
+else:
+    # Legacy single-phase initialization returns the original module object.
+    second = importlib.import_module("native")
+    assert second is first
+assert first.get_value() == 42
+
 [case testTopLevelThreadImportsNativeModuleFromSameGroup]
 # separate: [(["native.py", "other_target.py"], "testgroup")]
 from threading import Thread