@@ -2,12 +2,30 @@ use crate::pylifecycle::request_vm_from_interpreter;
22use core:: cell:: RefCell ;
33use core:: ffi:: c_int;
44use 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
78thread_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) ]
1230type 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) ]
2749pub 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) ]
4264mod 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 ( )
0 commit comments