fixes #5900 -- address race condition with initialization and site.py loading

newsfragments/5903.fixed.md

@@ -0,0 +1 @@
+fixed a race condition where `Python::attach` or `try_attach` could return before `site.py` had finished running

src/internal/state.rs

@@ -99,6 +99,11 @@ impl AttachGuard {
             return Err(AttachError::NotInitialized);
         }
 
+        // Py_IsInitialized() can return 1 while Py_InitializeEx is still
+        // running (e.g. importing site.py). Block until any in-progress PyO3
+        // initialization has fully completed.
+        crate::interpreter_lifecycle::wait_for_initialization();
+
         // Calling `PyGILState_Ensure` while finalizing may crash CPython in unpredictable
         // ways, we'll make a best effort attempt here to avoid that. (There's a time of
         // check to time-of-use issue, but it's better than nothing.)

src/interpreter_lifecycle.rs

@@ -80,6 +80,28 @@ where
     result
 }
 
+/// If PyO3 is currently running `Py_InitializeEx` inside the `Once` guard,
+/// block until it completes. Needed because `Py_InitializeEx` sets the
+/// `initialized` flag in the interpreter to true before it finishes all its
+/// steps (in particular, before it imports `site.py`).
+///
+/// This must only be called after `Py_IsInitialized()` has returned true.
+///
+/// If the `Once` was never started (e.g. the interpreter was initialized
+/// externally, not through PyO3), `call_once` runs the empty closure and
+/// returns — this is fine because `initialize()` checks
+/// `Py_IsInitialized()` inside its closure and skips `Py_InitializeEx` if
+/// the interpreter is already running. If the `Once` is currently in
+/// progress (another thread is inside `initialize()`), `call_once` blocks
+/// until it completes.
+pub(crate) fn wait_for_initialization() {
+    // TODO: use START.wait_force() on MSRV 1.86
+    // TODO: may not be needed on Python 3.15 (https://github.com/python/cpython/pull/146303)
+    START.call_once(|| {
+        assert_ne!(unsafe { crate::ffi::Py_IsInitialized() }, 0);
+    });
+}
+
 pub(crate) fn ensure_initialized() {
     // Maybe auto-initialize the interpreter:
     //  - If auto-initialize feature set and supported, try to initialize the interpreter.

tests/test_init_race.rs

@@ -0,0 +1,46 @@
+#![cfg(not(any(PyPy, GraalPy)))]
+#![cfg(not(target_arch = "wasm32"))]
+
+use pyo3::types::PyAnyMethods;
+
+// Regression test for `try_attach` fast-path returning before site.py had
+// finished initializing.
+#[test]
+fn test_concurrent_init_site_race() {
+    let tmpdir = std::env::temp_dir().join("pyo3_test_init_race");
+    std::fs::create_dir_all(&tmpdir).unwrap();
+
+    std::fs::write(
+        tmpdir.join("sitecustomize.py"),
+        "import sys, time\n\
+         sys._pyo3_site_done = False\n\
+         time.sleep(2)\n\
+         sys._pyo3_site_done = True\n",
+    )
+    .unwrap();
+
+    std::env::set_var("PYTHONPATH", &tmpdir);
+
+    std::thread::scope(|s| {
+        s.spawn(|| {
+            pyo3::Python::initialize();
+        });
+
+        s.spawn(|| loop {
+            let result = pyo3::Python::try_attach(|py| {
+                let done = py
+                    .import("sys")
+                    .unwrap()
+                    .getattr("_pyo3_site_done")
+                    .unwrap()
+                    .extract::<bool>()
+                    .unwrap();
+                assert!(done);
+            });
+            if result.is_some() {
+                break;
+            }
+            std::hint::spin_loop();
+        });
+    });
+}