Unify C API thread VM tracking

youknowone · May 5, 2026 · Apr 21, 2026 · Apr 21, 2026 · Apr 21, 2026 · Apr 21, 2026
commit 3e65cecea1c60fa2303454117bda99a67eadbe9f
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -259,7 +259,6 @@ rustls-platform-verifier = "0.7"
 rustyline = "18"
 serde = { package = "serde_core", version = "1.0.225", default-features = false, features = ["alloc"] }
 schannel = "0.1.29"
-scoped-tls = "1"
 scopeguard = "1"
 sha-1 = "0.10.0"
 sha2 = "0.10.2"

diff --git a/crates/capi/src/pystate.rs b/crates/capi/src/pystate.rs
@@ -1,34 +1,14 @@
 use crate::pylifecycle::request_vm_from_interpreter;
-use core::cell::RefCell;
 use core::ffi::c_int;
 use core::ptr;
-use rustpython_vm::VirtualMachine;
-use rustpython_vm::vm::thread::{ThreadedVirtualMachine, VM_CURRENT, with_current_vm};
-
-thread_local! {
-    static VM: RefCell<Option<ThreadedVirtualMachine>> = const { RefCell::new(None) };
-}
-
-#[allow(dead_code)]
-pub(crate) fn with_vm<R>(f: impl FnOnce(&VirtualMachine) -> R) -> R {
-    if VM_CURRENT.is_set() {
-        // We have an active VM set, so use that.
-        with_current_vm(f)
-    } else {
-        // We do not have an active vm running in this thread. Let's use our own.
-        // This will panic if `PyGILState_Ensure` was not called beforehand.
-        VM.with(|vm_ref| {
-            let vm = vm_ref.borrow();
-            let vm = vm
-                .as_ref()
-                .expect("Thread was not attached to an interpreter");
-            vm.run(|vm| f(vm))
-        })
-    }
-}
+use rustpython_vm::vm::thread::{
+    CurrentVmAttachState, attach_current_thread, release_current_thread,
+};
 
 #[allow(non_camel_case_types)]
 type PyGILState_STATE = c_int;
+const PYGILSTATE_LOCKED: PyGILState_STATE = 0;
+const PYGILSTATE_UNLOCKED: PyGILState_STATE = 1;
 
 #[repr(C)]
 pub struct PyThreadState {
@@ -37,24 +17,24 @@ pub struct PyThreadState {
 
 /// Make sure this thread has a running vm attached. This only creates a new vm if we don't already
 /// have one. So this will only create a new vm when we are in a new thread created outside RustPython.
-pub(crate) fn ensure_thread_has_vm_attached() {
-    if !VM_CURRENT.is_set() {
-        VM.with(|vm| {
-            vm.borrow_mut()
-                .get_or_insert_with(request_vm_from_interpreter);
-        });
-    }
+pub(crate) fn ensure_thread_has_vm_attached() -> CurrentVmAttachState {
+    attach_current_thread(request_vm_from_interpreter)
 }
 
 #[unsafe(no_mangle)]
 pub extern "C" fn PyGILState_Ensure() -> PyGILState_STATE {
-    ensure_thread_has_vm_attached();
-
-    0
+    match ensure_thread_has_vm_attached() {
+        CurrentVmAttachState::AlreadyAttached => PYGILSTATE_LOCKED,
+        CurrentVmAttachState::Attached => PYGILSTATE_UNLOCKED,
+    }
 }
 
 #[unsafe(no_mangle)]
-pub extern "C" fn PyGILState_Release(_state: PyGILState_STATE) {}
+pub extern "C" fn PyGILState_Release(state: PyGILState_STATE) {
+    if state == PYGILSTATE_UNLOCKED {
+        release_current_thread(CurrentVmAttachState::Attached);
+    }
+}
 
 #[unsafe(no_mangle)]
 pub extern "C" fn PyEval_SaveThread() -> *mut PyThreadState {
@@ -64,25 +44,25 @@ pub extern "C" fn PyEval_SaveThread() -> *mut PyThreadState {
 #[cfg(test)]
 mod tests {
     use crate::get_main_interpreter;
-    use crate::pystate::{VM, with_vm};
+    use crate::pystate::{PyGILState_Ensure, PyGILState_Release};
     use pyo3::prelude::*;
-    use rustpython_vm::vm::thread::VM_CURRENT;
+    use rustpython_vm::vm::thread::{current_vm_is_set, with_current_vm};
 
     #[test]
     fn test_new_thread() {
         Python::attach(|_py| {
-            with_vm(|_vm| {
+            with_current_vm(|_vm| {
                 assert!(
-                    VM_CURRENT.is_set(),
+                    current_vm_is_set(),
                     "This thread did not have a vm attached"
                 )
             });
 
             std::thread::spawn(move || {
                 Python::attach(|_py| {
-                    with_vm(|_vm| {
+                    with_current_vm(|_vm| {
                         assert!(
-                            VM_CURRENT.is_set(),
+                            current_vm_is_set(),
                             "This thread did not have a vm attached"
                         )
                     });
@@ -97,9 +77,6 @@ mod tests {
     fn test_current_vm_main_thread() {
         Python::initialize();
 
-        // Detach vm from thread, because initialize attaches one automatically.
-        VM.with(|vm| vm.borrow_mut().take());
-
         // let RustPython create a vm for this thread.
         let vm = get_main_interpreter()
             .as_ref()
@@ -108,17 +85,27 @@ mod tests {
 
         // Attach the vm using RustPython
         vm.run(|_vm| {
-            assert!(VM_CURRENT.is_set(), "This thread should have a vm attached");
+            assert!(current_vm_is_set(), "This thread should have a vm attached");
 
             Python::attach(|_py| {
-                with_vm(|_vm| {
-                    assert!(VM_CURRENT.is_set());
-                    assert!(
-                        VM.with(|vm| vm.borrow().is_none()),
-                        "We should not create a new vm when there is already a vm active"
-                    )
+                with_current_vm(|_vm| {
+                    assert!(current_vm_is_set());
                 })
             })
         });
     }
+
+    #[test]
+    fn test_gilstate_release_detaches_external_thread() {
+        Python::initialize();
+
+        std::thread::spawn(|| {
+            let state = PyGILState_Ensure();
+            assert!(current_vm_is_set());
+            PyGILState_Release(state);
+            assert!(!current_vm_is_set());
+        })
+        .join()
+        .unwrap();
+    }
 }
diff --git a/crates/vm/Cargo.toml b/crates/vm/Cargo.toml
@@ -65,7 +65,6 @@ num-traits = { workspace = true }
 num_enum = { workspace = true }
 parking_lot = { workspace = true }
 paste = { workspace = true }
-scoped-tls = { workspace = true }
 scopeguard = { workspace = true }
 serde = { workspace = true, optional = true }
 static_assertions = { workspace = true }

diff --git a/crates/vm/src/vm/thread.rs b/crates/vm/src/vm/thread.rs
@@ -50,6 +50,20 @@ pub type CurrentFrameSlot = Arc<ThreadSlot>;
 thread_local! {
     pub(super) static VM_STACK: RefCell<Vec<NonNull<VirtualMachine>>> = Vec::with_capacity(1).into();
 
+    /// Thread state created through the GILState-style C API.
+    ///
+    /// This is separate from the current VM stack: it only means "attached now",
+    /// while this owns the per-thread VM that may be detached and re-attached.
+    /// Despite the historical CPython "GILState" name, this does not model a
+    /// GIL; it stores the VM used by that compatibility API.
+    ///
+    /// The Box keeps the VM address stable while VM_STACK holds a raw pointer to it.
+    /// This matters when release_current_thread() moves the owner out of TLS and
+    /// drops it while the VM is still current, so object destructors can still find
+    /// their VM.
+    #[cfg(feature = "threading")]
+    static GILSTATE_VM: RefCell<Option<Box<ThreadedVirtualMachine>>> = const { RefCell::new(None) };
+
     pub(crate) static COROUTINE_ORIGIN_TRACKING_DEPTH: Cell<u32> = const { Cell::new(0) };
 
     /// Current thread's slot for sys._current_frames() and sys._current_exceptions()
@@ -66,42 +80,114 @@ thread_local! {
 
 }
 
-scoped_tls::scoped_thread_local!(pub static VM_CURRENT: VirtualMachine);
+#[must_use]
+pub fn current_vm_is_set() -> bool {
+    VM_STACK.with(|vms| !vms.borrow().is_empty())
+}
 
 pub fn with_current_vm<R>(f: impl FnOnce(&VirtualMachine) -> R) -> R {
-    if !VM_CURRENT.is_set() {
-        panic!("call with_current_vm() but VM_CURRENT is null");
-    }
-    VM_CURRENT.with(f)
+    VM_STACK.with(|vms| {
+        let vm = vms
+            .borrow()
+            .last()
+            .copied()
+            .expect("call with_current_vm() but no current VM is attached");
+        // SAFETY: entries in VM_STACK either borrow a VM for the dynamic
+        // scope of a set_current_vm()/enter_vm() call or point at GILSTATE_VM.
+        f(unsafe { vm.as_ref() })
+    })
 }
 
-pub fn enter_vm<R>(vm: &VirtualMachine, f: impl FnOnce() -> R) -> R {
+fn set_current_vm<R>(vm: &VirtualMachine, f: impl FnOnce() -> R) -> R {
     VM_STACK.with(|vms| {
-        // Outermost enter_vm: transition DETACHED → ATTACHED
-        #[cfg(all(unix, feature = "threading"))]
-        let was_outermost = vms.borrow().is_empty();
-
         vms.borrow_mut().push(vm.into());
+        scopeguard::defer! {
+            vms.borrow_mut().pop();
+        }
+        f()
+    })
+}
 
-        // Initialize thread slot for this thread if not already done
-        #[cfg(feature = "threading")]
-        init_thread_slot_if_needed(vm);
+pub fn enter_vm<R>(vm: &VirtualMachine, f: impl FnOnce() -> R) -> R {
+    // Outermost enter_vm: transition DETACHED → ATTACHED
+    #[cfg(all(unix, feature = "threading"))]
+    let was_outermost = !current_vm_is_set();
 
+    // Initialize thread slot for this thread if not already done
+    #[cfg(feature = "threading")]
+    init_thread_slot_if_needed(vm);
+
+    #[cfg(all(unix, feature = "threading"))]
+    if was_outermost {
+        attach_thread(vm);
+    }
+
+    scopeguard::defer! {
+        // Outermost exit: transition ATTACHED → DETACHED
         #[cfg(all(unix, feature = "threading"))]
         if was_outermost {
-            attach_thread(vm);
+            detach_thread();
         }
+    }
+    set_current_vm(vm, f)
+}
 
-        scopeguard::defer! {
-            // Outermost exit: transition ATTACHED → DETACHED
-            #[cfg(all(unix, feature = "threading"))]
-            if vms.borrow().len() == 1 {
-                detach_thread();
-            }
-            vms.borrow_mut().pop();
-        }
-        VM_CURRENT.set(vm, f)
-    })
+#[cfg(feature = "threading")]
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub enum CurrentVmAttachState {
+    AlreadyAttached,
+    Attached,
+}
+
+/// Attach the current native thread to a RustPython VM until
+/// `release_current_thread()` is called.
+#[cfg(feature = "threading")]
+pub fn attach_current_thread(
+    make_vm: impl FnOnce() -> ThreadedVirtualMachine,
+) -> CurrentVmAttachState {
+    if current_vm_is_set() {
+        return CurrentVmAttachState::AlreadyAttached;
+    }
+
+    GILSTATE_VM.with(|gilstate_vm| {
+        let mut gilstate_vm = gilstate_vm.borrow_mut();
+        let threaded_vm = gilstate_vm.get_or_insert_with(|| Box::new(make_vm()));
+        let vm = &threaded_vm.vm;
+
+        vm.c_stack_soft_limit
+            .set(VirtualMachine::calculate_c_stack_soft_limit());
+
+        init_thread_slot_if_needed(vm);
+
+        #[cfg(unix)]
+        attach_thread(vm);
+
+        VM_STACK.with(|vms| {
+            debug_assert!(vms.borrow().is_empty());
+            vms.borrow_mut().push(vm.into());
+        });
+    });
+
+    CurrentVmAttachState::Attached
+}
+
+#[cfg(feature = "threading")]
+pub fn release_current_thread(state: CurrentVmAttachState) {
+    if state == CurrentVmAttachState::AlreadyAttached {
+        return;
+    }
+
+    let gilstate_vm = GILSTATE_VM.with(|gilstate_vm| gilstate_vm.borrow_mut().take());
+    drop(gilstate_vm);
+
+    VM_STACK.with(|vms| {
+        vms.borrow_mut()
+            .pop()
+            .expect("release_current_thread() called without an attached VM");
+    });
+
+    #[cfg(unix)]
+    detach_thread();
 }
 
 /// Initialize thread slot for current thread if not already initialized.
@@ -509,17 +595,20 @@ where
         obj.fast_isinstance(vm.ctx.types.object_type)
     };
     VM_STACK.with(|vms| {
-        let interp = match vms.borrow().iter().copied().exactly_one() {
-            Ok(x) => {
-                debug_assert!(vm_owns_obj(x));
-                x
+        let interp = {
+            let vms = vms.borrow();
+            match vms.iter().copied().exactly_one() {
+                Ok(x) => {
+                    debug_assert!(vm_owns_obj(x));
+                    x
+                }
+                Err(mut others) => others.find(|x| vm_owns_obj(*x))?,
             }
-            Err(mut others) => others.find(|x| vm_owns_obj(*x))?,
         };
         // SAFETY: all references in VM_STACK should be valid, and should not be changed or moved
         // at least until this function returns and the stack unwinds to an enter_vm() call
         let vm = unsafe { interp.as_ref() };
-        Some(VM_CURRENT.set(vm, || f(vm)))
+        Some(set_current_vm(vm, || f(vm)))
     })
 }