Skip to content

Commit e06327f

Browse files
Use VM_CURRENT when already active
1 parent 3eb587b commit e06327f

3 files changed

Lines changed: 46 additions & 19 deletions

File tree

‎crates/capi/src/pylifecycle.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
use crate::get_main_interpreter;
2-
use crate::pystate::attach_vm_to_thread;
2+
use crate::pystate::ensure_thread_has_vm_attached;
33
use core::ffi::c_int;
44
use rustpython_vm::Interpreter;
55
use rustpython_vm::vm::thread::ThreadedVirtualMachine;
@@ -31,7 +31,7 @@ pub extern "C" fn Py_InitializeEx(_initsigs: c_int) {
3131
if interp.is_none() {
3232
*interp = Interpreter::with_init(Default::default(), |_vm| {}).into();
3333
drop(interp);
34-
attach_vm_to_thread();
34+
ensure_thread_has_vm_attached();
3535
}
3636
}
3737

‎crates/capi/src/pystate.rs‎

Lines changed: 43 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,30 @@ use crate::pylifecycle::request_vm_from_interpreter;
22
use core::cell::RefCell;
33
use core::ffi::c_int;
44
use core::ptr;
5-
use rustpython_vm::vm::thread::ThreadedVirtualMachine;
5+
use rustpython_vm::VirtualMachine;
6+
use rustpython_vm::vm::thread::{ThreadedVirtualMachine, VM_CURRENT, with_current_vm};
67

78
thread_local! {
89
static VM: RefCell<Option<ThreadedVirtualMachine>> = const { RefCell::new(None) };
910
}
1011

12+
pub(crate) fn with_vm<R>(f: impl FnOnce(&VirtualMachine) -> R) -> R {
13+
if VM_CURRENT.is_set() {
14+
// We have an active VM set, so use that.
15+
with_current_vm(f)
16+
} else {
17+
// We do not have an active vm running in this thread. Let's use our own.
18+
// This will panic if `PyGILState_Ensure` was not called beforehand.
19+
VM.with(|vm_ref| {
20+
let vm = vm_ref.borrow();
21+
let vm = vm
22+
.as_ref()
23+
.expect("Thread was not attached to an interpreter");
24+
vm.run(|vm| f(vm))
25+
})
26+
}
27+
}
28+
1129
#[allow(non_camel_case_types)]
1230
type PyGILState_STATE = c_int;
1331

@@ -16,16 +34,20 @@ pub struct PyThreadState {
1634
_interp: *mut core::ffi::c_void,
1735
}
1836

19-
pub(crate) fn attach_vm_to_thread() {
20-
VM.with(|vm| {
21-
vm.borrow_mut()
22-
.get_or_insert_with(request_vm_from_interpreter);
23-
});
37+
/// Make sure this thread has a running vm attached. This only creates a new vm if we don't already
38+
/// have one. So this will only create a new vm when we are in a new thread created outside RustPython.
39+
pub(crate) fn ensure_thread_has_vm_attached() {
40+
if !VM_CURRENT.is_set() {
41+
VM.with(|vm| {
42+
vm.borrow_mut()
43+
.get_or_insert_with(request_vm_from_interpreter);
44+
});
45+
}
2446
}
2547

2648
#[unsafe(no_mangle)]
2749
pub extern "C" fn PyGILState_Ensure() -> PyGILState_STATE {
28-
attach_vm_to_thread();
50+
ensure_thread_has_vm_attached();
2951

3052
0
3153
}
@@ -40,23 +62,28 @@ pub extern "C" fn PyEval_SaveThread() -> *mut PyThreadState {
4062

4163
#[cfg(test)]
4264
mod tests {
43-
use crate::pystate::VM;
65+
use crate::pystate::{VM, with_vm};
4466
use pyo3::prelude::*;
67+
use rustpython_vm::vm::thread::VM_CURRENT;
4568

4669
#[test]
4770
fn test_new_thread() {
4871
Python::attach(|_py| {
49-
assert!(
50-
VM.with(|vm| vm.borrow().is_some()),
51-
"This thread did not have a vm attached"
52-
);
72+
with_vm(|vm| {
73+
assert!(
74+
VM_CURRENT.is_set(),
75+
"This thread did not have a vm attached"
76+
)
77+
});
5378

5479
std::thread::spawn(move || {
5580
Python::attach(|_py| {
56-
assert!(
57-
VM.with(|vm| vm.borrow().is_some()),
58-
"This thread did not have a vm attached"
59-
);
81+
with_vm(|vm| {
82+
assert!(
83+
VM_CURRENT.is_set(),
84+
"This thread did not have a vm attached"
85+
)
86+
});
6087
});
6188
})
6289
.join()

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,7 @@ thread_local! {
6666

6767
}
6868

69-
scoped_tls::scoped_thread_local!(static VM_CURRENT: VirtualMachine);
69+
scoped_tls::scoped_thread_local!(pub static VM_CURRENT: VirtualMachine);
7070

7171
pub fn with_current_vm<R>(f: impl FnOnce(&VirtualMachine) -> R) -> R {
7272
if !VM_CURRENT.is_set() {

0 commit comments

Comments
 (0)