Skip to content

Switch to newer thread::LocalKey convenience methods #6123

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
youknowone merged 1 commit into from
Sep 3, 2025
Merged

Switch to newer thread::LocalKey convenience methods #6123

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 2 additions & 0 deletions Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -207,6 +207,8 @@ rustix = { version = "1.0", features = ["event"] }
rustyline = "17.0.0"
serde = { version = "1.0.133", default-features = false }
schannel = "0.1.27"
scoped-tls = "1"
scopeguard = "1"
static_assertions = "1.1"
strum = "0.27"
strum_macros = "0.27"
Expand Down
7 changes: 3 additions & 4 deletions common/src/str.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -533,10 +533,9 @@ pub mod levenshtein {
return max_cost + 1;
}

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

let mut result = 0usize;
Expand Down
27 changes: 9 additions & 18 deletions stdlib/src/contextvars.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -107,8 +107,7 @@ mod _contextvars {
return Err(vm.new_runtime_error(msg));
}

super::CONTEXTS.with(|ctxs| {
let mut ctxs = ctxs.borrow_mut();
super::CONTEXTS.with_borrow_mut(|ctxs| {
zelf.inner.idx.set(ctxs.len());
ctxs.push(zelf.to_owned());
});
Expand All @@ -126,27 +125,20 @@ mod _contextvars {
return Err(vm.new_runtime_error(msg));
}

super::CONTEXTS.with(|ctxs| {
let mut ctxs = ctxs.borrow_mut();
// TODO: use Vec::pop_if once stabilized
if ctxs.last().is_some_and(|ctx| ctx.get_id() == zelf.get_id()) {
let _ = ctxs.pop();
Ok(())
} else {
let msg =
"cannot exit context: thread state references a different context object"
.to_owned();
Err(vm.new_runtime_error(msg))
}
super::CONTEXTS.with_borrow_mut(|ctxs| {
let err_msg =
"cannot exit context: thread state references a different context object";
ctxs.pop_if(|ctx| ctx.get_id() == zelf.get_id())
.map(drop)
.ok_or_else(|| vm.new_runtime_error(err_msg))
})?;
Comment on lines +128 to 134
Copy link
Contributor

@coderabbitai coderabbitai bot Sep 3, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Verification agent

🧩 Analysis chain

Don’t rely on non-standard Vec::pop_if unless an explicit extension trait is in scope

pop_if is not in std. If you don’t have a local extension trait implementing it for Vec, this won’t compile. Even if you do, the intent (LIFO-only pop when last matches) can be expressed clearly with std-only code.

Apply this std-only fallback:

-            super::CONTEXTS.with_borrow_mut(|ctxs| {
-                let err_msg =
-                    "cannot exit context: thread state references a different context object";
-                ctxs.pop_if(|ctx| ctx.get_id() == zelf.get_id())
-                    .map(drop)
-                    .ok_or_else(|| vm.new_runtime_error(err_msg))
-            })?;
+            super::CONTEXTS.with_borrow_mut(|ctxs| {
+                let err_msg =
+                    "cannot exit context: thread state references a different context object";
+                if ctxs.last().map_or(false, |ctx| ctx.get_id() == zelf.get_id()) {
+                    ctxs.pop();
+                    Ok(())
+                } else {
+                    Err(vm.new_runtime_error(err_msg))
+                }
+            })?;

To confirm whether pop_if exists in the repo, run:

🏁 Script executed:

#!/bin/bash
# Look for a Vec extension providing `pop_if` and all callsites.
rg -nP -C2 '\bpop_if\s*\('
rg -nP -C3 'trait\s+\w+\s*{[^}]*pop_if\s*\('

Length of output: 512

Replace non-existent Vec::pop_if with standard code
Vec::pop_if isn’t defined in std or this repo, so the call will fail to compile. Replace:

ctxs.pop_if(|ctx| ctx.get_id() == zelf.get_id())
    .map(drop)
    .ok_or_else(|| vm.new_runtime_error(err_msg))

with:

-if ctxs.pop_if(|ctx| ctx.get_id() == zelf.get_id())
-    .map(drop)
-    .ok_or_else(|| vm.new_runtime_error(err_msg))
+if ctxs.last().map_or(false, |ctx| ctx.get_id() == zelf.get_id()) {
+    ctxs.pop();
+    Ok(())
+} else {
+    Err(vm.new_runtime_error(err_msg))
+}
🤖 Prompt for AI Agents 
In stdlib/src/contextvars.rs around lines 128–134, Vec::pop_if does not exist so
replace that call with a search for the matching element using
ctxs.iter().rposition(|ctx| ctx.get_id() == zelf.get_id()) and then remove the
found index (ctxs.remove(pos)) to drop it; if rposition returns None, return the
vm.new_runtime_error(err_msg) as before. Ensure you call map/drop on the removed
element (or simply drop it) and keep the overall ? error propagation semantics.

Copy link
Member

@youknowone youknowone Sep 3, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

https://doc.rust-lang.org/std/vec/struct.Vec.html#method.pop_if

Copy link
Contributor

@coderabbitai coderabbitai bot Sep 3, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

.

The current implementation using pop_if in the diff is correct and will compile successfully. My suggestion to replace it with standard library code was based on incorrect information. The code change from the traditional pattern to using with_borrow_mut combined with pop_if is a valid improvement that takes advantage of newer Rust convenience methods, which aligns perfectly with the PR's objective.

Thank you for the correction and providing the documentation link!

✏️ Learnings added 
Learnt from: youknowone
PR: RustPython/RustPython#6123
File: stdlib/src/contextvars.rs:128-134
Timestamp: 2025-09-03T13:15:23.384Z
Learning: Vec::pop_if was stabilized in Rust 1.86.0 (released April 3, 2025) and is available in standard library for conditionally removing the last element from a vector based on a predicate.

Note: Learnings are effective only in the context of similar code segments. To apply general coding guidelines, please configure review instructions. You can manage existing learnings in the app.

zelf.inner.entered.set(false);

Ok(())
}

fn current(vm: &VirtualMachine) -> PyRef<Self> {
super::CONTEXTS.with(|ctxs| {
let mut ctxs = ctxs.borrow_mut();
super::CONTEXTS.with_borrow_mut(|ctxs| {
if let Some(ctx) = ctxs.last() {
ctx.clone()
} else {
Expand Down Expand Up @@ -382,8 +374,7 @@ mod _contextvars {
default: OptionalArg<PyObjectRef>,
vm: &VirtualMachine,
) -> PyResult<Option<PyObjectRef>> {
let found = super::CONTEXTS.with(|ctxs| {
let ctxs = ctxs.borrow();
let found = super::CONTEXTS.with_borrow(|ctxs| {
let ctx = ctxs.last()?;
let cached_ptr = zelf.cached.as_ptr();
debug_assert!(!cached_ptr.is_null());
Expand Down
2 changes: 2 additions & 0 deletions vm/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,8 @@ num_enum = { workspace = true }
once_cell = { 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 }
strum = { workspace = true }
Expand Down
18 changes: 8 additions & 10 deletions vm/src/stdlib/sys.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -830,13 +830,13 @@ mod sys {
if depth < 0 {
return Err(vm.new_value_error("depth must be >= 0"));
}
crate::vm::thread::COROUTINE_ORIGIN_TRACKING_DEPTH.with(|cell| cell.set(depth as _));
crate::vm::thread::COROUTINE_ORIGIN_TRACKING_DEPTH.set(depth as u32);
Ok(())
}

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

#[pyfunction]
Expand Down Expand Up @@ -887,14 +887,10 @@ mod sys {
}

if let Some(finalizer) = args.finalizer.into_option() {
crate::vm::thread::ASYNC_GEN_FINALIZER.with(|cell| {
cell.replace(finalizer);
});
crate::vm::thread::ASYNC_GEN_FINALIZER.set(finalizer);
}
if let Some(firstiter) = args.firstiter.into_option() {
crate::vm::thread::ASYNC_GEN_FIRSTITER.with(|cell| {
cell.replace(firstiter);
});
crate::vm::thread::ASYNC_GEN_FIRSTITER.set(firstiter);
}

Ok(())
Expand All @@ -914,9 +910,11 @@ mod sys {
fn get_asyncgen_hooks(vm: &VirtualMachine) -> PyAsyncgenHooks {
PyAsyncgenHooks {
firstiter: crate::vm::thread::ASYNC_GEN_FIRSTITER
.with(|cell| cell.borrow().clone().to_pyobject(vm)),
.with_borrow(Clone::clone)
.to_pyobject(vm),
finalizer: crate::vm::thread::ASYNC_GEN_FINALIZER
.with(|cell| cell.borrow().clone().to_pyobject(vm)),
.with_borrow(Clone::clone)
.to_pyobject(vm),
}
}

Expand Down
16 changes: 6 additions & 10 deletions vm/src/stdlib/thread.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -334,13 +334,11 @@ pub(crate) mod _thread {
);
}
}
SENTINELS.with(|sentinels| {
for lock in sentinels.replace(Default::default()) {
if lock.mu.is_locked() {
unsafe { lock.mu.unlock() };
}
for lock in SENTINELS.take() {
if lock.mu.is_locked() {
unsafe { lock.mu.unlock() };
}
});
}
vm.state.thread_count.fetch_sub(1);
}

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

thread_local! {
static SENTINELS: RefCell<Vec<PyRef<Lock>>> = const { RefCell::new(Vec::new()) };
}
thread_local!(static SENTINELS: RefCell<Vec<PyRef<Lock>>> = const { RefCell::new(Vec::new()) });

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

Expand Down
25 changes: 9 additions & 16 deletions vm/src/vm/thread.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,30 +8,26 @@ use std::{

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

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

scoped_tls::scoped_thread_local!(static VM_CURRENT: VirtualMachine);

pub fn with_current_vm<R>(f: impl FnOnce(&VirtualMachine) -> R) -> R {
VM_CURRENT.with(|x| unsafe {
f(x.clone()
.into_inner()
.as_ref()
.expect("call with_current_vm() but VM_CURRENT is null"))
})
if !VM_CURRENT.is_set() {
panic!("call with_current_vm() but VM_CURRENT is null");
}
VM_CURRENT.with(f)
}

pub fn enter_vm<R>(vm: &VirtualMachine, f: impl FnOnce() -> R) -> R {
VM_STACK.with(|vms| {
vms.borrow_mut().push(vm.into());
let prev = VM_CURRENT.with(|current| current.replace(vm));
let ret = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f));
vms.borrow_mut().pop();
VM_CURRENT.with(|current| current.replace(prev));
ret.unwrap_or_else(|e| std::panic::resume_unwind(e))
scopeguard::defer! { vms.borrow_mut().pop(); }
VM_CURRENT.set(vm, f)
})
}

Expand All @@ -55,10 +51,7 @@ where
// 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() };
let prev = VM_CURRENT.with(|current| current.replace(vm));
let ret = f(vm);
VM_CURRENT.with(|current| current.replace(prev));
Some(ret)
Some(VM_CURRENT.set(vm, || f(vm)))
})
}

Expand Down
54 changes: 18 additions & 36 deletions wasm/lib/src/vm_class.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,8 +58,8 @@ impl StoredVirtualMachine {
setup_browser_module(vm);
}

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

// It's fine that it's thread local, since WASM doesn't even have threads yet. thread_local!
Expand All @@ -97,17 +97,15 @@ pub fn get_vm_id(vm: &VirtualMachine) -> &str {
.expect("VirtualMachine inside of WASM crate should have wasm_id set")
}
pub(crate) fn stored_vm_from_wasm(wasm_vm: &WASMVirtualMachine) -> Rc<StoredVirtualMachine> {
STORED_VMS.with(|cell| {
cell.borrow()
.get(&wasm_vm.id)
STORED_VMS.with_borrow(|vms| {
vms.get(&wasm_vm.id)
.expect("VirtualMachine is not valid")
.clone()
})
}
pub(crate) fn weak_vm(vm: &VirtualMachine) -> Weak<StoredVirtualMachine> {
let id = get_vm_id(vm);
STORED_VMS
.with(|cell| Rc::downgrade(cell.borrow().get(id).expect("VirtualMachine is not valid")))
STORED_VMS.with_borrow(|vms| Rc::downgrade(vms.get(id).expect("VirtualMachine is not valid")))
}

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

pub(crate) fn _get(id: String) -> Option<WASMVirtualMachine> {
STORED_VMS.with(|cell| {
let vms = cell.borrow();
if vms.contains_key(&id) {
Some(WASMVirtualMachine { id })
} else {
None
}
})
STORED_VMS.with_borrow(|vms| vms.contains_key(&id).then_some(WASMVirtualMachine { id }))
}

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

pub fn destroy(id: String) {
STORED_VMS.with(|cell| {
use std::collections::hash_map::Entry;
match cell.borrow_mut().entry(id) {
Entry::Occupied(o) => {
let (_k, stored_vm) = o.remove_entry();
// for f in stored_vm.drop_handlers.iter() {
// f();
// }
// deallocate the VM
drop(stored_vm);
}
Entry::Vacant(_v) => {}
STORED_VMS.with_borrow_mut(|vms| {
if let Some(stored_vm) = vms.remove(&id) {
// for f in stored_vm.drop_handlers.iter() {
// f();
// }
// deallocate the VM
drop(stored_vm);
}
});
}

pub fn ids() -> Vec<JsValue> {
STORED_VMS.with(|cell| cell.borrow().keys().map(|k| k.into()).collect())
STORED_VMS.with_borrow(|vms| vms.keys().map(|k| k.into()).collect())
}
}

Expand All @@ -179,10 +164,7 @@ impl WASMVirtualMachine {
where
F: FnOnce(&StoredVirtualMachine) -> R,
{
let stored_vm = STORED_VMS.with(|cell| {
let mut vms = cell.borrow_mut();
vms.get_mut(&self.id).unwrap().clone()
});
let stored_vm = STORED_VMS.with_borrow_mut(|vms| vms.get_mut(&self.id).unwrap().clone());
f(&stored_vm)
}

Expand All @@ -202,7 +184,7 @@ impl WASMVirtualMachine {
}

pub fn valid(&self) -> bool {
STORED_VMS.with(|cell| cell.borrow().contains_key(&self.id))
STORED_VMS.with_borrow(|vms| vms.contains_key(&self.id))
}

pub(crate) fn push_held_rc(
Expand Down
Loading