Skip to content

Commit cca4fe6

Browse files
coolreader18coolreader18
authored
Switch to newer thread::LocalKey convenience methods (#6123)
1 parent d17dcd8 commit cca4fe6

File tree

9 files changed

+65
-94
lines changed

9 files changed

+65
-94
lines changed

‎Cargo.lock

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

‎Cargo.toml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -207,6 +207,8 @@ rustix = { version = "1.0", features = ["event"] }
207207
rustyline = "17.0.0"
208208
serde = { version = "1.0.133", default-features = false }
209209
schannel = "0.1.27"
210+
scoped-tls = "1"
211+
scopeguard = "1"
210212
static_assertions = "1.1"
211213
strum = "0.27"
212214
strum_macros = "0.27"

‎common/src/str.rs

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -533,10 +533,9 @@ pub mod levenshtein {
533533
return max_cost + 1;
534534
}
535535

536-
BUFFER.with(|buffer| {
537-
let mut buffer = buffer.borrow_mut();
538-
for i in 0..a_end {
539-
buffer[i] = (i + 1) * MOVE_COST;
536+
BUFFER.with_borrow_mut(|buffer| {
537+
for (i, x) in buffer.iter_mut().take(a_end).enumerate() {
538+
*x = (i + 1) * MOVE_COST;
540539
}
541540

542541
let mut result = 0usize;

‎stdlib/src/contextvars.rs

Lines changed: 9 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -107,8 +107,7 @@ mod _contextvars {
107107
return Err(vm.new_runtime_error(msg));
108108
}
109109

110-
super::CONTEXTS.with(|ctxs| {
111-
let mut ctxs = ctxs.borrow_mut();
110+
super::CONTEXTS.with_borrow_mut(|ctxs| {
112111
zelf.inner.idx.set(ctxs.len());
113112
ctxs.push(zelf.to_owned());
114113
});
@@ -126,27 +125,20 @@ mod _contextvars {
126125
return Err(vm.new_runtime_error(msg));
127126
}
128127

129-
super::CONTEXTS.with(|ctxs| {
130-
let mut ctxs = ctxs.borrow_mut();
131-
// TODO: use Vec::pop_if once stabilized
132-
if ctxs.last().is_some_and(|ctx| ctx.get_id() == zelf.get_id()) {
133-
let _ = ctxs.pop();
134-
Ok(())
135-
} else {
136-
let msg =
137-
"cannot exit context: thread state references a different context object"
138-
.to_owned();
139-
Err(vm.new_runtime_error(msg))
140-
}
128+
super::CONTEXTS.with_borrow_mut(|ctxs| {
129+
let err_msg =
130+
"cannot exit context: thread state references a different context object";
131+
ctxs.pop_if(|ctx| ctx.get_id() == zelf.get_id())
132+
.map(drop)
133+
.ok_or_else(|| vm.new_runtime_error(err_msg))
141134
})?;
142135
zelf.inner.entered.set(false);
143136

144137
Ok(())
145138
}
146139

147140
fn current(vm: &VirtualMachine) -> PyRef<Self> {
148-
super::CONTEXTS.with(|ctxs| {
149-
let mut ctxs = ctxs.borrow_mut();
141+
super::CONTEXTS.with_borrow_mut(|ctxs| {
150142
if let Some(ctx) = ctxs.last() {
151143
ctx.clone()
152144
} else {
@@ -382,8 +374,7 @@ mod _contextvars {
382374
default: OptionalArg<PyObjectRef>,
383375
vm: &VirtualMachine,
384376
) -> PyResult<Option<PyObjectRef>> {
385-
let found = super::CONTEXTS.with(|ctxs| {
386-
let ctxs = ctxs.borrow();
377+
let found = super::CONTEXTS.with_borrow(|ctxs| {
387378
let ctx = ctxs.last()?;
388379
let cached_ptr = zelf.cached.as_ptr();
389380
debug_assert!(!cached_ptr.is_null());

‎vm/Cargo.toml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,8 @@ num_enum = { workspace = true }
6565
once_cell = { workspace = true }
6666
parking_lot = { workspace = true }
6767
paste = { workspace = true }
68+
scoped-tls = { workspace = true }
69+
scopeguard = { workspace = true }
6870
serde = { workspace = true, optional = true }
6971
static_assertions = { workspace = true }
7072
strum = { workspace = true }

‎vm/src/stdlib/sys.rs

Lines changed: 8 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -830,13 +830,13 @@ mod sys {
830830
if depth < 0 {
831831
return Err(vm.new_value_error("depth must be >= 0"));
832832
}
833-
crate::vm::thread::COROUTINE_ORIGIN_TRACKING_DEPTH.with(|cell| cell.set(depth as _));
833+
crate::vm::thread::COROUTINE_ORIGIN_TRACKING_DEPTH.set(depth as u32);
834834
Ok(())
835835
}
836836

837837
#[pyfunction]
838838
fn get_coroutine_origin_tracking_depth() -> i32 {
839-
crate::vm::thread::COROUTINE_ORIGIN_TRACKING_DEPTH.with(|cell| cell.get()) as _
839+
crate::vm::thread::COROUTINE_ORIGIN_TRACKING_DEPTH.get() as i32
840840
}
841841

842842
#[pyfunction]
@@ -887,14 +887,10 @@ mod sys {
887887
}
888888

889889
if let Some(finalizer) = args.finalizer.into_option() {
890-
crate::vm::thread::ASYNC_GEN_FINALIZER.with(|cell| {
891-
cell.replace(finalizer);
892-
});
890+
crate::vm::thread::ASYNC_GEN_FINALIZER.set(finalizer);
893891
}
894892
if let Some(firstiter) = args.firstiter.into_option() {
895-
crate::vm::thread::ASYNC_GEN_FIRSTITER.with(|cell| {
896-
cell.replace(firstiter);
897-
});
893+
crate::vm::thread::ASYNC_GEN_FIRSTITER.set(firstiter);
898894
}
899895

900896
Ok(())
@@ -914,9 +910,11 @@ mod sys {
914910
fn get_asyncgen_hooks(vm: &VirtualMachine) -> PyAsyncgenHooks {
915911
PyAsyncgenHooks {
916912
firstiter: crate::vm::thread::ASYNC_GEN_FIRSTITER
917-
.with(|cell| cell.borrow().clone().to_pyobject(vm)),
913+
.with_borrow(Clone::clone)
914+
.to_pyobject(vm),
918915
finalizer: crate::vm::thread::ASYNC_GEN_FINALIZER
919-
.with(|cell| cell.borrow().clone().to_pyobject(vm)),
916+
.with_borrow(Clone::clone)
917+
.to_pyobject(vm),
920918
}
921919
}
922920

‎vm/src/stdlib/thread.rs

Lines changed: 6 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -334,13 +334,11 @@ pub(crate) mod _thread {
334334
);
335335
}
336336
}
337-
SENTINELS.with(|sentinels| {
338-
for lock in sentinels.replace(Default::default()) {
339-
if lock.mu.is_locked() {
340-
unsafe { lock.mu.unlock() };
341-
}
337+
for lock in SENTINELS.take() {
338+
if lock.mu.is_locked() {
339+
unsafe { lock.mu.unlock() };
342340
}
343-
});
341+
}
344342
vm.state.thread_count.fetch_sub(1);
345343
}
346344

@@ -355,14 +353,12 @@ pub(crate) mod _thread {
355353
Err(vm.new_exception_empty(vm.ctx.exceptions.system_exit.to_owned()))
356354
}
357355

358-
thread_local! {
359-
static SENTINELS: RefCell<Vec<PyRef<Lock>>> = const { RefCell::new(Vec::new()) };
360-
}
356+
thread_local!(static SENTINELS: RefCell<Vec<PyRef<Lock>>> = const { RefCell::new(Vec::new()) });
361357

362358
#[pyfunction]
363359
fn _set_sentinel(vm: &VirtualMachine) -> PyRef<Lock> {
364360
let lock = Lock { mu: RawMutex::INIT }.into_ref(&vm.ctx);
365-
SENTINELS.with(|sentinels| sentinels.borrow_mut().push(lock.clone()));
361+
SENTINELS.with_borrow_mut(|sentinels| sentinels.push(lock.clone()));
366362
lock
367363
}
368364

‎vm/src/vm/thread.rs

Lines changed: 9 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -8,30 +8,26 @@ use std::{
88

99
thread_local! {
1010
pub(super) static VM_STACK: RefCell<Vec<NonNull<VirtualMachine>>> = Vec::with_capacity(1).into();
11-
static VM_CURRENT: RefCell<*const VirtualMachine> = std::ptr::null::<VirtualMachine>().into();
1211

1312
pub(crate) static COROUTINE_ORIGIN_TRACKING_DEPTH: Cell<u32> = const { Cell::new(0) };
1413
pub(crate) static ASYNC_GEN_FINALIZER: RefCell<Option<PyObjectRef>> = const { RefCell::new(None) };
1514
pub(crate) static ASYNC_GEN_FIRSTITER: RefCell<Option<PyObjectRef>> = const { RefCell::new(None) };
1615
}
1716

17+
scoped_tls::scoped_thread_local!(static VM_CURRENT: VirtualMachine);
18+
1819
pub fn with_current_vm<R>(f: impl FnOnce(&VirtualMachine) -> R) -> R {
19-
VM_CURRENT.with(|x| unsafe {
20-
f(x.clone()
21-
.into_inner()
22-
.as_ref()
23-
.expect("call with_current_vm() but VM_CURRENT is null"))
24-
})
20+
if !VM_CURRENT.is_set() {
21+
panic!("call with_current_vm() but VM_CURRENT is null");
22+
}
23+
VM_CURRENT.with(f)
2524
}
2625

2726
pub fn enter_vm<R>(vm: &VirtualMachine, f: impl FnOnce() -> R) -> R {
2827
VM_STACK.with(|vms| {
2928
vms.borrow_mut().push(vm.into());
30-
let prev = VM_CURRENT.with(|current| current.replace(vm));
31-
let ret = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f));
32-
vms.borrow_mut().pop();
33-
VM_CURRENT.with(|current| current.replace(prev));
34-
ret.unwrap_or_else(|e| std::panic::resume_unwind(e))
29+
scopeguard::defer! { vms.borrow_mut().pop(); }
30+
VM_CURRENT.set(vm, f)
3531
})
3632
}
3733

@@ -55,10 +51,7 @@ where
5551
// SAFETY: all references in VM_STACK should be valid, and should not be changed or moved
5652
// at least until this function returns and the stack unwinds to an enter_vm() call
5753
let vm = unsafe { interp.as_ref() };
58-
let prev = VM_CURRENT.with(|current| current.replace(vm));
59-
let ret = f(vm);
60-
VM_CURRENT.with(|current| current.replace(prev));
61-
Some(ret)
54+
Some(VM_CURRENT.set(vm, || f(vm)))
6255
})
6356
}
6457

‎wasm/lib/src/vm_class.rs

Lines changed: 18 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -58,8 +58,8 @@ impl StoredVirtualMachine {
5858
setup_browser_module(vm);
5959
}
6060

61-
VM_INIT_FUNCS.with(|cell| {
62-
for f in cell.borrow().iter() {
61+
VM_INIT_FUNCS.with_borrow(|funcs| {
62+
for f in funcs {
6363
f(vm)
6464
}
6565
});
@@ -78,7 +78,7 @@ impl StoredVirtualMachine {
7878
/// Add a hook to add builtins or frozen modules to the RustPython VirtualMachine while it's
7979
/// initializing.
8080
pub fn add_init_func(f: fn(&mut VirtualMachine)) {
81-
VM_INIT_FUNCS.with(|cell| cell.borrow_mut().push(f))
81+
VM_INIT_FUNCS.with_borrow_mut(|funcs| funcs.push(f))
8282
}
8383

8484
// It's fine that it's thread local, since WASM doesn't even have threads yet. thread_local!
@@ -97,17 +97,15 @@ pub fn get_vm_id(vm: &VirtualMachine) -> &str {
9797
.expect("VirtualMachine inside of WASM crate should have wasm_id set")
9898
}
9999
pub(crate) fn stored_vm_from_wasm(wasm_vm: &WASMVirtualMachine) -> Rc<StoredVirtualMachine> {
100-
STORED_VMS.with(|cell| {
101-
cell.borrow()
102-
.get(&wasm_vm.id)
100+
STORED_VMS.with_borrow(|vms| {
101+
vms.get(&wasm_vm.id)
103102
.expect("VirtualMachine is not valid")
104103
.clone()
105104
})
106105
}
107106
pub(crate) fn weak_vm(vm: &VirtualMachine) -> Weak<StoredVirtualMachine> {
108107
let id = get_vm_id(vm);
109-
STORED_VMS
110-
.with(|cell| Rc::downgrade(cell.borrow().get(id).expect("VirtualMachine is not valid")))
108+
STORED_VMS.with_borrow(|vms| Rc::downgrade(vms.get(id).expect("VirtualMachine is not valid")))
111109
}
112110

113111
#[wasm_bindgen(js_name = vmStore)]
@@ -116,8 +114,7 @@ pub struct VMStore;
116114
#[wasm_bindgen(js_class = vmStore)]
117115
impl VMStore {
118116
pub fn init(id: String, inject_browser_module: Option<bool>) -> WASMVirtualMachine {
119-
STORED_VMS.with(|cell| {
120-
let mut vms = cell.borrow_mut();
117+
STORED_VMS.with_borrow_mut(|vms| {
121118
if !vms.contains_key(&id) {
122119
let stored_vm =
123120
StoredVirtualMachine::new(id.clone(), inject_browser_module.unwrap_or(true));
@@ -128,14 +125,7 @@ impl VMStore {
128125
}
129126

130127
pub(crate) fn _get(id: String) -> Option<WASMVirtualMachine> {
131-
STORED_VMS.with(|cell| {
132-
let vms = cell.borrow();
133-
if vms.contains_key(&id) {
134-
Some(WASMVirtualMachine { id })
135-
} else {
136-
None
137-
}
138-
})
128+
STORED_VMS.with_borrow(|vms| vms.contains_key(&id).then_some(WASMVirtualMachine { id }))
139129
}
140130

141131
pub fn get(id: String) -> JsValue {
@@ -146,24 +136,19 @@ impl VMStore {
146136
}
147137

148138
pub fn destroy(id: String) {
149-
STORED_VMS.with(|cell| {
150-
use std::collections::hash_map::Entry;
151-
match cell.borrow_mut().entry(id) {
152-
Entry::Occupied(o) => {
153-
let (_k, stored_vm) = o.remove_entry();
154-
// for f in stored_vm.drop_handlers.iter() {
155-
// f();
156-
// }
157-
// deallocate the VM
158-
drop(stored_vm);
159-
}
160-
Entry::Vacant(_v) => {}
139+
STORED_VMS.with_borrow_mut(|vms| {
140+
if let Some(stored_vm) = vms.remove(&id) {
141+
// for f in stored_vm.drop_handlers.iter() {
142+
// f();
143+
// }
144+
// deallocate the VM
145+
drop(stored_vm);
161146
}
162147
});
163148
}
164149

165150
pub fn ids() -> Vec<JsValue> {
166-
STORED_VMS.with(|cell| cell.borrow().keys().map(|k| k.into()).collect())
151+
STORED_VMS.with_borrow(|vms| vms.keys().map(|k| k.into()).collect())
167152
}
168153
}
169154

@@ -179,10 +164,7 @@ impl WASMVirtualMachine {
179164
where
180165
F: FnOnce(&StoredVirtualMachine) -> R,
181166
{
182-
let stored_vm = STORED_VMS.with(|cell| {
183-
let mut vms = cell.borrow_mut();
184-
vms.get_mut(&self.id).unwrap().clone()
185-
});
167+
let stored_vm = STORED_VMS.with_borrow_mut(|vms| vms.get_mut(&self.id).unwrap().clone());
186168
f(&stored_vm)
187169
}
188170

@@ -202,7 +184,7 @@ impl WASMVirtualMachine {
202184
}
203185

204186
pub fn valid(&self) -> bool {
205-
STORED_VMS.with(|cell| cell.borrow().contains_key(&self.id))
187+
STORED_VMS.with_borrow(|vms| vms.contains_key(&self.id))
206188
}
207189

208190
pub(crate) fn push_held_rc(

0 commit comments

Comments
 (0)