diff --git a/crates/capi/src/lib.rs b/crates/capi/src/lib.rs index 374fd3301d6..2aff75fe15b 100644 --- a/crates/capi/src/lib.rs +++ b/crates/capi/src/lib.rs @@ -1,7 +1,8 @@ #![allow(clippy::missing_safety_doc)] use crate::pyerrors::init_exception_statics; -use crate::pylifecycle::MAIN_INTERP; +use crate::pylifecycle::{MAIN_INTERP, MAIN_INTERP_PTR}; +use core::sync::atomic::Ordering; pub use rustpython_vm::PyObject; use rustpython_vm::{Context, Interpreter}; use std::sync::MutexGuard; @@ -61,4 +62,8 @@ pub fn init_main_interpreter(interpreter: Interpreter) { // Safety: Interpreter was not initialized before, so we can safely assume the statics are not used unsafe { init_exception_statics(&Context::genesis().exceptions) }; *interp = Some(interpreter); + MAIN_INTERP_PTR.store( + interp.as_ref().unwrap() as *const _ as *mut _, + Ordering::Release, + ); } diff --git a/crates/capi/src/pylifecycle.rs b/crates/capi/src/pylifecycle.rs index 534c97956fe..7255a0e34a5 100644 --- a/crates/capi/src/pylifecycle.rs +++ b/crates/capi/src/pylifecycle.rs @@ -3,6 +3,7 @@ use crate::pyerrors::init_exception_statics; use crate::pystate::ensure_thread_has_vm_attached; use alloc::ffi::CString; use core::ffi::{c_char, c_int, c_ulong}; +use core::sync::atomic::{AtomicPtr, Ordering}; use rustpython_vm::common::rc::PyRc; use rustpython_vm::stdlib::sys; use rustpython_vm::version::{MAJOR, MICRO, MINOR, RUSTPYTHON_BUILD_INFO, VERSION_HEX}; @@ -11,6 +12,7 @@ use rustpython_vm::{Context, Interpreter}; use std::sync::{LazyLock, Mutex}; pub(crate) static MAIN_INTERP: Mutex> = Mutex::new(None); +pub(crate) static MAIN_INTERP_PTR: AtomicPtr = AtomicPtr::new(core::ptr::null_mut()); /// Request a thread local vm from the main interpreter pub(crate) fn request_vm_from_interpreter() -> ThreadedVirtualMachine { @@ -25,7 +27,7 @@ pub static Py_Version: c_ulong = VERSION_HEX as c_ulong; #[unsafe(no_mangle)] pub extern "C" fn Py_IsInitialized() -> c_int { - get_main_interpreter().is_some() as c_int + !MAIN_INTERP_PTR.load(Ordering::Acquire).is_null() as c_int } #[unsafe(no_mangle)] @@ -52,6 +54,10 @@ pub extern "C" fn Py_InitializeEx(_initsigs: c_int) { }) .build() .into(); + MAIN_INTERP_PTR.store( + interp.as_ref().unwrap() as *const _ as *mut _, + Ordering::Release, + ); drop(interp); ensure_thread_has_vm_attached(); } diff --git a/crates/capi/src/pystate.rs b/crates/capi/src/pystate.rs index 173c26088cf..865a116443b 100644 --- a/crates/capi/src/pystate.rs +++ b/crates/capi/src/pystate.rs @@ -1,7 +1,7 @@ -use crate::get_main_interpreter; -use crate::pylifecycle::request_vm_from_interpreter; +use crate::pylifecycle::{MAIN_INTERP_PTR, request_vm_from_interpreter}; use crate::util::FfiResult; use core::ffi::c_int; +use core::sync::atomic::Ordering; use rustpython_vm::vm::thread::{ CurrentVmAttachState, SavedThreadState, attach_current_thread, release_current_thread, restore_current_thread, save_current_thread, with_current_vm, @@ -67,11 +67,12 @@ pub unsafe extern "C" fn PyEval_RestoreThread(state: *mut PyThreadState) { #[unsafe(no_mangle)] pub extern "C" fn PyInterpreterState_Get() -> *mut PyInterpreterState { - get_main_interpreter() - .as_ref() - .map(|interp| interp as *const PyInterpreterState) - .expect("PyInterpreterState_Get called but no main interpreter was found") - .cast_mut() + let ptr = MAIN_INTERP_PTR.load(Ordering::Acquire); + assert!( + !ptr.is_null(), + "PyInterpreterState_Get() called but the interpreter is not initialized" + ); + ptr } #[unsafe(no_mangle)]