return existing module on Python 3.9 and up

newsfragments/3446.changed.md

@@ -1 +1 @@
-Relax multiple import check to only prevent use of subinterpreters.
+`#[pymodule]` will now return the same module object on repeated import by the same Python interpreter, on Python 3.9 and up.

pyo3-build-config/src/lib.rs

@@ -146,11 +146,6 @@ pub fn print_feature_cfgs() {
     if rustc_minor_version >= 59 {
         println!("cargo:rustc-cfg=thread_local_const_init");
     }
-
-    // Enable use of OnceLock on Rust 1.70 and greater
-    if rustc_minor_version >= 70 {
-        println!("cargo:rustc-cfg=once_lock");
-    }
 }
 
 /// Private exports used in PyO3's build.rs

pytests/tests/test_misc.py

@@ -1,5 +1,6 @@
 import importlib
 import platform
+import sys
 
 import pyo3_pytests.misc
 import pytest
@@ -10,13 +11,21 @@ def test_issue_219():
     pyo3_pytests.misc.issue_219()
 
 
+@pytest.mark.xfail(
+    platform.python_implementation() == "CPython" and sys.version_info < (3, 9),
+    reason="Cannot identify subinterpreters on Python older than 3.9",
+)
 def test_multiple_imports_same_interpreter_ok():
     spec = importlib.util.find_spec("pyo3_pytests.pyo3_pytests")
 
     module = importlib.util.module_from_spec(spec)
     assert dir(module) == dir(pyo3_pytests.pyo3_pytests)
 
 
+@pytest.mark.xfail(
+    platform.python_implementation() == "CPython" and sys.version_info < (3, 9),
+    reason="Cannot identify subinterpreters on Python older than 3.9",
+)
 @pytest.mark.skipif(
     platform.python_implementation() == "PyPy",
     reason="PyPy does not support subinterpreters",

src/impl_/pymodule.rs

@@ -1,23 +1,24 @@
 //! Implementation details of `#[pymodule]` which need to be accessible from proc-macro generated code.
 
 use std::cell::UnsafeCell;
-#[cfg(once_lock)]
-use std::sync::OnceLock;
 
-#[cfg(not(once_lock))]
-use parking_lot::Mutex;
+#[cfg(all(not(PyPy), Py_3_9))]
+use std::sync::atomic::{AtomicI64, Ordering};
 
-use crate::{exceptions::PyImportError, ffi, types::PyModule, Py, PyResult, Python};
+#[cfg(not(PyPy))]
+use crate::exceptions::PyImportError;
+use crate::{ffi, sync::GILOnceCell, types::PyModule, Py, PyResult, Python};
 
 /// `Sync` wrapper of `ffi::PyModuleDef`.
 pub struct ModuleDef {
     // wrapped in UnsafeCell so that Rust compiler treats this as interior mutability
     ffi_def: UnsafeCell<ffi::PyModuleDef>,
     initializer: ModuleInitializer,
-    #[cfg(once_lock)]
-    interpreter: OnceLock<i64>,
-    #[cfg(not(once_lock))]
-    interpreter: Mutex<Option<i64>>,
+    /// Interpreter ID where module was initialized (not applicable on PyPy).
+    #[cfg(all(not(PyPy), Py_3_9))]
+    interpreter: AtomicI64,
+    /// Initialized module object, cached to avoid reinitialization.
+    module: GILOnceCell<Py<PyModule>>,
 }
 
 /// Wrapper to enable initializer to be used in const fns.
@@ -56,10 +57,10 @@ impl ModuleDef {
         ModuleDef {
             ffi_def,
             initializer,
-            #[cfg(once_lock)]
-            interpreter: OnceLock::new(),
-            #[cfg(not(once_lock))]
-            interpreter: Mutex::new(None),
+            // -1 is never expected to be a valid interpreter ID
+            #[cfg(all(not(PyPy), Py_3_9))]
+            interpreter: AtomicI64::new(-1),
+            module: GILOnceCell::new(),
         }
     }
     /// Builds a module using user given initializer. Used for [`#[pymodule]`][crate::pymodule].
@@ -79,29 +80,47 @@ impl ModuleDef {
                 ))?;
             }
         }
-        let module = unsafe {
-            Py::<PyModule>::from_owned_ptr_or_err(py, ffi::PyModule_Create(self.ffi_def.get()))?
-        };
-        let current_interpreter =
-            unsafe { ffi::PyInterpreterState_GetID(ffi::PyInterpreterState_Get()) };
-        let initialized_interpreter = py.allow_threads(|| {
-            #[cfg(once_lock)]
-            {
-                *self.interpreter.get_or_init(|| current_interpreter)
+        #[cfg(all(not(PyPy), Py_3_9))]
+        {
+            // If APIs to do so are available, check that the interpreter has not changed.
+            let current_interpreter =
+                unsafe { ffi::PyInterpreterState_GetID(ffi::PyInterpreterState_Get()) };
+            crate::err::error_on_minusone(py, current_interpreter)?;
+            if let Err(initialized_interpreter) = self.interpreter.compare_exchange(
+                -1,
+                current_interpreter,
+                Ordering::SeqCst,
+                Ordering::SeqCst,
+            ) {
+                if initialized_interpreter != current_interpreter {
+                    return Err(PyImportError::new_err(
+                        "PyO3 modules do not yet support subinterpreters, see https://github.com/PyO3/pyo3/issues/576",
+                    ));
+                }
             }
-
-            #[cfg(not(once_lock))]
-            {
-                *self.interpreter.lock().get_or_insert(current_interpreter)
+        }
+        #[cfg(all(not(PyPy), not(Py_3_9)))]
+        {
+            // CPython before 3.9 does not have APIs to check the interpreter ID, so best that can be
+            // done to guard against subinterpreters is fail if the module is initialized twice
+            if self.module.get(py).is_some() {
+                return Err(PyImportError::new_err(
+                    "PyO3 modules compiled for CPython 3.8 or older may only be initialized once per interpreter process"
+                ));
             }
-        });
-        if current_interpreter != initialized_interpreter {
-            return Err(PyImportError::new_err(
-                "PyO3 modules do not yet support subinterpreters, see https://github.com/PyO3/pyo3/issues/576",
-            ));
         }
-        (self.initializer.0)(py, module.as_ref(py))?;
-        Ok(module)
+        self.module
+            .get_or_try_init(py, || {
+                let module = unsafe {
+                    Py::<PyModule>::from_owned_ptr_or_err(
+                        py,
+                        ffi::PyModule_Create(self.ffi_def.get()),
+                    )?
+                };
+                (self.initializer.0)(py, module.as_ref(py))?;
+                Ok(module)
+            })
+            .map(|py_module| py_module.clone_ref(py))
     }
 }