Skip to content

Commit 3e65cec

Browse files
youknowoneyouknowone
committed
Unify C API thread VM tracking
1 parent 4e88888 commit 3e65cec

5 files changed

Lines changed: 158 additions & 91 deletions

File tree

‎Cargo.lock‎

Lines changed: 0 additions & 7 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎Cargo.toml‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -259,7 +259,6 @@ rustls-platform-verifier = "0.7"
259259
rustyline = "18"
260260
serde = { package = "serde_core", version = "1.0.225", default-features = false, features = ["alloc"] }
261261
schannel = "0.1.29"
262-
scoped-tls = "1"
263262
scopeguard = "1"
264263
sha-1 = "0.10.0"
265264
sha2 = "0.10.2"

‎crates/capi/src/pystate.rs‎

Lines changed: 39 additions & 52 deletions
Original file line numberDiff line numberDiff line change
@@ -1,34 +1,14 @@
11
use crate::pylifecycle::request_vm_from_interpreter;
2-
use core::cell::RefCell;
32
use core::ffi::c_int;
43
use core::ptr;
5-
use rustpython_vm::VirtualMachine;
6-
use rustpython_vm::vm::thread::{ThreadedVirtualMachine, VM_CURRENT, with_current_vm};
7-
8-
thread_local! {
9-
static VM: RefCell<Option<ThreadedVirtualMachine>> = const { RefCell::new(None) };
10-
}
11-
12-
#[allow(dead_code)]
13-
pub(crate) fn with_vm<R>(f: impl FnOnce(&VirtualMachine) -> R) -> R {
14-
if VM_CURRENT.is_set() {
15-
// We have an active VM set, so use that.
16-
with_current_vm(f)
17-
} else {
18-
// We do not have an active vm running in this thread. Let's use our own.
19-
// This will panic if `PyGILState_Ensure` was not called beforehand.
20-
VM.with(|vm_ref| {
21-
let vm = vm_ref.borrow();
22-
let vm = vm
23-
.as_ref()
24-
.expect("Thread was not attached to an interpreter");
25-
vm.run(|vm| f(vm))
26-
})
27-
}
28-
}
4+
use rustpython_vm::vm::thread::{
5+
CurrentVmAttachState, attach_current_thread, release_current_thread,
6+
};
297

308
#[allow(non_camel_case_types)]
319
type PyGILState_STATE = c_int;
10+
const PYGILSTATE_LOCKED: PyGILState_STATE = 0;
11+
const PYGILSTATE_UNLOCKED: PyGILState_STATE = 1;
3212

3313
#[repr(C)]
3414
pub struct PyThreadState {
@@ -37,24 +17,24 @@ pub struct PyThreadState {
3717

3818
/// Make sure this thread has a running vm attached. This only creates a new vm if we don't already
3919
/// have one. So this will only create a new vm when we are in a new thread created outside RustPython.
40-
pub(crate) fn ensure_thread_has_vm_attached() {
41-
if !VM_CURRENT.is_set() {
42-
VM.with(|vm| {
43-
vm.borrow_mut()
44-
.get_or_insert_with(request_vm_from_interpreter);
45-
});
46-
}
20+
pub(crate) fn ensure_thread_has_vm_attached() -> CurrentVmAttachState {
21+
attach_current_thread(request_vm_from_interpreter)
4722
}
4823

4924
#[unsafe(no_mangle)]
5025
pub extern "C" fn PyGILState_Ensure() -> PyGILState_STATE {
51-
ensure_thread_has_vm_attached();
52-
53-
0
26+
match ensure_thread_has_vm_attached() {
27+
CurrentVmAttachState::AlreadyAttached => PYGILSTATE_LOCKED,
28+
CurrentVmAttachState::Attached => PYGILSTATE_UNLOCKED,
29+
}
5430
}
5531

5632
#[unsafe(no_mangle)]
57-
pub extern "C" fn PyGILState_Release(_state: PyGILState_STATE) {}
33+
pub extern "C" fn PyGILState_Release(state: PyGILState_STATE) {
34+
if state == PYGILSTATE_UNLOCKED {
35+
release_current_thread(CurrentVmAttachState::Attached);
36+
}
37+
}
5838

5939
#[unsafe(no_mangle)]
6040
pub extern "C" fn PyEval_SaveThread() -> *mut PyThreadState {
@@ -64,25 +44,25 @@ pub extern "C" fn PyEval_SaveThread() -> *mut PyThreadState {
6444
#[cfg(test)]
6545
mod tests {
6646
use crate::get_main_interpreter;
67-
use crate::pystate::{VM, with_vm};
47+
use crate::pystate::{PyGILState_Ensure, PyGILState_Release};
6848
use pyo3::prelude::*;
69-
use rustpython_vm::vm::thread::VM_CURRENT;
49+
use rustpython_vm::vm::thread::{current_vm_is_set, with_current_vm};
7050

7151
#[test]
7252
fn test_new_thread() {
7353
Python::attach(|_py| {
74-
with_vm(|_vm| {
54+
with_current_vm(|_vm| {
7555
assert!(
76-
VM_CURRENT.is_set(),
56+
current_vm_is_set(),
7757
"This thread did not have a vm attached"
7858
)
7959
});
8060

8161
std::thread::spawn(move || {
8262
Python::attach(|_py| {
83-
with_vm(|_vm| {
63+
with_current_vm(|_vm| {
8464
assert!(
85-
VM_CURRENT.is_set(),
65+
current_vm_is_set(),
8666
"This thread did not have a vm attached"
8767
)
8868
});
@@ -97,9 +77,6 @@ mod tests {
9777
fn test_current_vm_main_thread() {
9878
Python::initialize();
9979

100-
// Detach vm from thread, because initialize attaches one automatically.
101-
VM.with(|vm| vm.borrow_mut().take());
102-
10380
// let RustPython create a vm for this thread.
10481
let vm = get_main_interpreter()
10582
.as_ref()
@@ -108,17 +85,27 @@ mod tests {
10885

10986
// Attach the vm using RustPython
11087
vm.run(|_vm| {
111-
assert!(VM_CURRENT.is_set(), "This thread should have a vm attached");
88+
assert!(current_vm_is_set(), "This thread should have a vm attached");
11289

11390
Python::attach(|_py| {
114-
with_vm(|_vm| {
115-
assert!(VM_CURRENT.is_set());
116-
assert!(
117-
VM.with(|vm| vm.borrow().is_none()),
118-
"We should not create a new vm when there is already a vm active"
119-
)
91+
with_current_vm(|_vm| {
92+
assert!(current_vm_is_set());
12093
})
12194
})
12295
});
12396
}
97+
98+
#[test]
99+
fn test_gilstate_release_detaches_external_thread() {
100+
Python::initialize();
101+
102+
std::thread::spawn(|| {
103+
let state = PyGILState_Ensure();
104+
assert!(current_vm_is_set());
105+
PyGILState_Release(state);
106+
assert!(!current_vm_is_set());
107+
})
108+
.join()
109+
.unwrap();
110+
}
124111
}

‎crates/vm/Cargo.toml‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,6 @@ num-traits = { workspace = true }
6565
num_enum = { workspace = true }
6666
parking_lot = { workspace = true }
6767
paste = { workspace = true }
68-
scoped-tls = { workspace = true }
6968
scopeguard = { workspace = true }
7069
serde = { workspace = true, optional = true }
7170
static_assertions = { workspace = true }

‎crates/vm/src/vm/thread.rs‎

Lines changed: 119 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,20 @@ pub type CurrentFrameSlot = Arc<ThreadSlot>;
5050
thread_local! {
5151
pub(super) static VM_STACK: RefCell<Vec<NonNull<VirtualMachine>>> = Vec::with_capacity(1).into();
5252

53+
/// Thread state created through the GILState-style C API.
54+
///
55+
/// This is separate from the current VM stack: it only means "attached now",
56+
/// while this owns the per-thread VM that may be detached and re-attached.
57+
/// Despite the historical CPython "GILState" name, this does not model a
58+
/// GIL; it stores the VM used by that compatibility API.
59+
///
60+
/// The Box keeps the VM address stable while VM_STACK holds a raw pointer to it.
61+
/// This matters when release_current_thread() moves the owner out of TLS and
62+
/// drops it while the VM is still current, so object destructors can still find
63+
/// their VM.
64+
#[cfg(feature = "threading")]
65+
static GILSTATE_VM: RefCell<Option<Box<ThreadedVirtualMachine>>> = const { RefCell::new(None) };
66+
5367
pub(crate) static COROUTINE_ORIGIN_TRACKING_DEPTH: Cell<u32> = const { Cell::new(0) };
5468

5569
/// Current thread's slot for sys._current_frames() and sys._current_exceptions()
@@ -66,42 +80,114 @@ thread_local! {
6680

6781
}
6882

69-
scoped_tls::scoped_thread_local!(pub static VM_CURRENT: VirtualMachine);
83+
#[must_use]
84+
pub fn current_vm_is_set() -> bool {
85+
VM_STACK.with(|vms| !vms.borrow().is_empty())
86+
}
7087

7188
pub fn with_current_vm<R>(f: impl FnOnce(&VirtualMachine) -> R) -> R {
72-
if !VM_CURRENT.is_set() {
73-
panic!("call with_current_vm() but VM_CURRENT is null");
74-
}
75-
VM_CURRENT.with(f)
89+
VM_STACK.with(|vms| {
90+
let vm = vms
91+
.borrow()
92+
.last()
93+
.copied()
94+
.expect("call with_current_vm() but no current VM is attached");
95+
// SAFETY: entries in VM_STACK either borrow a VM for the dynamic
96+
// scope of a set_current_vm()/enter_vm() call or point at GILSTATE_VM.
97+
f(unsafe { vm.as_ref() })
98+
})
7699
}
77100

78-
pub fn enter_vm<R>(vm: &VirtualMachine, f: impl FnOnce() -> R) -> R {
101+
fn set_current_vm<R>(vm: &VirtualMachine, f: impl FnOnce() -> R) -> R {
79102
VM_STACK.with(|vms| {
80-
// Outermost enter_vm: transition DETACHED → ATTACHED
81-
#[cfg(all(unix, feature = "threading"))]
82-
let was_outermost = vms.borrow().is_empty();
83-
84103
vms.borrow_mut().push(vm.into());
104+
scopeguard::defer! {
105+
vms.borrow_mut().pop();
106+
}
107+
f()
108+
})
109+
}
85110

86-
// Initialize thread slot for this thread if not already done
87-
#[cfg(feature = "threading")]
88-
init_thread_slot_if_needed(vm);
111+
pub fn enter_vm<R>(vm: &VirtualMachine, f: impl FnOnce() -> R) -> R {
112+
// Outermost enter_vm: transition DETACHED → ATTACHED
113+
#[cfg(all(unix, feature = "threading"))]
114+
let was_outermost = !current_vm_is_set();
89115

116+
// Initialize thread slot for this thread if not already done
117+
#[cfg(feature = "threading")]
118+
init_thread_slot_if_needed(vm);
119+
120+
#[cfg(all(unix, feature = "threading"))]
121+
if was_outermost {
122+
attach_thread(vm);
123+
}
124+
125+
scopeguard::defer! {
126+
// Outermost exit: transition ATTACHED → DETACHED
90127
#[cfg(all(unix, feature = "threading"))]
91128
if was_outermost {
92-
attach_thread(vm);
129+
detach_thread();
93130
}
131+
}
132+
set_current_vm(vm, f)
133+
}
94134

95-
scopeguard::defer! {
96-
// Outermost exit: transition ATTACHED → DETACHED
97-
#[cfg(all(unix, feature = "threading"))]
98-
if vms.borrow().len() == 1 {
99-
detach_thread();
100-
}
101-
vms.borrow_mut().pop();
102-
}
103-
VM_CURRENT.set(vm, f)
104-
})
135+
#[cfg(feature = "threading")]
136+
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
137+
pub enum CurrentVmAttachState {
138+
AlreadyAttached,
139+
Attached,
140+
}
141+
142+
/// Attach the current native thread to a RustPython VM until
143+
/// `release_current_thread()` is called.
144+
#[cfg(feature = "threading")]
145+
pub fn attach_current_thread(
146+
make_vm: impl FnOnce() -> ThreadedVirtualMachine,
147+
) -> CurrentVmAttachState {
148+
if current_vm_is_set() {
149+
return CurrentVmAttachState::AlreadyAttached;
150+
}
151+
152+
GILSTATE_VM.with(|gilstate_vm| {
153+
let mut gilstate_vm = gilstate_vm.borrow_mut();
154+
let threaded_vm = gilstate_vm.get_or_insert_with(|| Box::new(make_vm()));
155+
let vm = &threaded_vm.vm;
156+
157+
vm.c_stack_soft_limit
158+
.set(VirtualMachine::calculate_c_stack_soft_limit());
159+
160+
init_thread_slot_if_needed(vm);
161+
162+
#[cfg(unix)]
163+
attach_thread(vm);
164+
165+
VM_STACK.with(|vms| {
166+
debug_assert!(vms.borrow().is_empty());
167+
vms.borrow_mut().push(vm.into());
168+
});
169+
});
170+
171+
CurrentVmAttachState::Attached
172+
}
173+
174+
#[cfg(feature = "threading")]
175+
pub fn release_current_thread(state: CurrentVmAttachState) {
176+
if state == CurrentVmAttachState::AlreadyAttached {
177+
return;
178+
}
179+
180+
let gilstate_vm = GILSTATE_VM.with(|gilstate_vm| gilstate_vm.borrow_mut().take());
181+
drop(gilstate_vm);
182+
183+
VM_STACK.with(|vms| {
184+
vms.borrow_mut()
185+
.pop()
186+
.expect("release_current_thread() called without an attached VM");
187+
});
188+
189+
#[cfg(unix)]
190+
detach_thread();
105191
}
106192

107193
/// Initialize thread slot for current thread if not already initialized.
@@ -509,17 +595,20 @@ where
509595
obj.fast_isinstance(vm.ctx.types.object_type)
510596
};
511597
VM_STACK.with(|vms| {
512-
let interp = match vms.borrow().iter().copied().exactly_one() {
513-
Ok(x) => {
514-
debug_assert!(vm_owns_obj(x));
515-
x
598+
let interp = {
599+
let vms = vms.borrow();
600+
match vms.iter().copied().exactly_one() {
601+
Ok(x) => {
602+
debug_assert!(vm_owns_obj(x));
603+
x
604+
}
605+
Err(mut others) => others.find(|x| vm_owns_obj(*x))?,
516606
}
517-
Err(mut others) => others.find(|x| vm_owns_obj(*x))?,
518607
};
519608
// SAFETY: all references in VM_STACK should be valid, and should not be changed or moved
520609
// at least until this function returns and the stack unwinds to an enter_vm() call
521610
let vm = unsafe { interp.as_ref() };
522-
Some(VM_CURRENT.set(vm, || f(vm)))
611+
Some(set_current_vm(vm, || f(vm)))
523612
})
524613
}
525614

0 commit comments

Comments
 (0)