Switch to newer thread::LocalKey convenience methods #6123
Conversation
WalkthroughIntroduces scoped TLS and scopeguard dependencies. Refactors multiple thread-local RefCell access patterns to with_borrow/with_borrow_mut helpers. Converts VM current-thread storage to scoped TLS, revising VM entry/access control flow. Adjusts sys/thread context accessors and cleanup logic. Minor iterator-based update in Levenshtein buffer initialization. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant Caller
participant VM as VirtualMachine
participant TLS as VM_CURRENT (scoped TLS)
participant Stack as VM_STACK
rect rgb(240,248,255)
note over Caller,VM: enter_vm(vm, f)
Caller->>VM: enter_vm(vm, f)
VM->>Stack: push(vm)
VM->>TLS: set(vm) {scoped}
activate TLS
TLS-->>VM: run f()
VM->>VM: scopeguard defer Stack.pop()
deactivate TLS
VM-->>Caller: result
end
rect rgb(255,250,240)
note over Caller,TLS: with_current_vm(f)
Caller->>TLS: with(f) (requires is_set)
TLS-->>Caller: f(&VM)
end
sequenceDiagram
autonumber
participant Py as Python code
participant CV as ContextVar
participant Ctxs as CONTEXTS (TLS)
participant Ctx as PyContext
rect rgb(245,255,250)
note over Py,Ctxs: current(vm)
Py->>Ctxs: with_borrow_mut
alt last exists
Ctxs-->>Py: &last
else create
Ctxs->>Ctx: new default
Ctxs->>Ctxs: push(Ctx)
Ctxs-->>Py: &first
end
end
rect rgb(250,245,255)
note over Py,CV: CV.get(default)
Py->>Ctxs: with_borrow
Ctxs-->>Py: last context ref
Py->>CV: check cache or lookup vars
CV-->>Py: value or default
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
Poem
✨ Finishing Touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
Status, Documentation and Community
|
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
stdlib/src/contextvars.rs (1)
377-399: Fix unsound read of AtomicCell via as_ptr; use swap(None)/store to avoid UB and racesDereferencing as_ptr without synchronization can observe freed/stale memory if another thread swaps the cell concurrently; it also dereferences a raw pointer (UB risk). Use atomic swap to take ownership, validate, and restore or replace.
- let found = super::CONTEXTS.with_borrow(|ctxs| { - let ctx = ctxs.last()?; - let cached_ptr = zelf.cached.as_ptr(); - debug_assert!(!cached_ptr.is_null()); - if let Some(cached) = unsafe { &*cached_ptr } { - if zelf.cached_id.load(Ordering::SeqCst) == ctx.get_id() - && cached.idx + 1 == ctxs.len() - { - return Some(cached.object.clone()); - } - } - let vars = ctx.borrow_vars(); - let obj = vars.get(zelf)?; - zelf.cached_id.store(ctx.get_id(), Ordering::SeqCst); - - // TODO: ensure cached is not changed - let _removed = zelf.cached.swap(Some(ContextVarCache { - object: obj.clone(), - idx: ctxs.len() - 1, - })); - - Some(obj.clone()) - }); + let found = super::CONTEXTS.with_borrow(|ctxs| { + let ctx = ctxs.last()?; + // Take a snapshot of the cache without aliasing; restore if still valid. + if let Some(cached) = zelf.cached.swap(None) { + if zelf.cached_id.load(Ordering::SeqCst) == ctx.get_id() + && cached.idx + 1 == ctxs.len() + { + // Put it back and return fast-path. + let obj = cached.object.clone(); + zelf.cached.store(Some(cached)); + return Some(obj); + } + // stale cache is dropped + } + // Miss: read from current context and repopulate cache. + let obj = { + let vars = ctx.borrow_vars(); + vars.get(zelf)?.clone() + }; + zelf.cached_id.store(ctx.get_id(), Ordering::SeqCst); + zelf.cached.store(Some(ContextVarCache { + object: obj.clone(), + idx: ctxs.len() - 1, + })); + Some(obj) + });This keeps the lock-free happy path and eliminates raw-pointer aliasing. It also prevents dropping a valid cache entry on the fast path.
🧹 Nitpick comments (5)
common/src/str.rs (1)
536-539: Borrow-helper + iterator init: cleaner and safeSwitching to
with_borrow_mutremoves needless RefCell ceremony, and the iterator-based fill respectsa_endbounds checked above. Nice touch.If you want to shave a tiny bit of work, you can avoid repeated multiplies by using an accumulator:
- for (i, x) in buffer.iter_mut().take(a_end).enumerate() { - *x = (i + 1) * MOVE_COST; - } + let mut v = MOVE_COST; + for x in buffer.iter_mut().take(a_end) { + *x = v; + v += MOVE_COST; + }vm/src/stdlib/sys.rs (1)
829-835: Depth storage type mismatch is fine; consider unifying types to drop castsConverting i32→u32 on set and u32→i32 on get is safe given the earlier non-negative check and the i32 parameter bound. Consider storing
i32inCOROUTINE_ORIGIN_TRACKING_DEPTHto remove the casts and make intent explicit.- crate::vm::thread::COROUTINE_ORIGIN_TRACKING_DEPTH.set(depth as u32); + crate::vm::thread::COROUTINE_ORIGIN_TRACKING_DEPTH.set(depth);and
- crate::vm::thread::COROUTINE_ORIGIN_TRACKING_DEPTH.get() as i32 + crate::vm::thread::COROUTINE_ORIGIN_TRACKING_DEPTH.get()Also applies to: 839-840
vm/src/vm/thread.rs (1)
20-23: Clarify panic message; behavior change is acceptablePanic-on-unset is fine, but “null” is misleading now. Suggest clearer wording.
- if !VM_CURRENT.is_set() { - panic!("call with_current_vm() but VM_CURRENT is null"); - } + if !VM_CURRENT.is_set() { + panic!("with_current_vm() called but VM_CURRENT is not set"); + }wasm/lib/src/vm_class.rs (1)
167-169: Avoid unnecessary mutable borrow in with_uncheckedYou only clone the Rc; shared borrow suffices.
- let stored_vm = STORED_VMS.with_borrow_mut(|vms| vms.get_mut(&self.id).unwrap().clone()); + let stored_vm = STORED_VMS.with_borrow(|vms| vms.get(&self.id).unwrap().clone());stdlib/src/contextvars.rs (1)
141-153: Nit: avoid extra clone and index by returning ctx directlySlightly simpler and avoids one clone of the PyRef.
- super::CONTEXTS.with_borrow_mut(|ctxs| { + super::CONTEXTS.with_borrow_mut(|ctxs| { if let Some(ctx) = ctxs.last() { ctx.clone() } else { let ctx = Self::empty(vm); ctx.inner.idx.set(0); ctx.inner.entered.set(true); let ctx = ctx.into_ref(&vm.ctx); ctxs.push(ctx); - ctxs[0].clone() + ctxs.last().unwrap().clone() } })
📜 Review details
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
⛔ Files ignored due to path filters (1)
Cargo.lockis excluded by!**/*.lock
📒 Files selected for processing (8)
Cargo.toml(1 hunks)common/src/str.rs(1 hunks)stdlib/src/contextvars.rs(3 hunks)vm/Cargo.toml(1 hunks)vm/src/stdlib/sys.rs(3 hunks)vm/src/stdlib/thread.rs(2 hunks)vm/src/vm/thread.rs(2 hunks)wasm/lib/src/vm_class.rs(8 hunks)
🧰 Additional context used
📓 Path-based instructions (1)
**/*.rs
📄 CodeRabbit inference engine (.github/copilot-instructions.md)
**/*.rs: Follow the default rustfmt code style (cargo fmtto format)
Always run clippy to lint code (cargo clippy) before completing tasks. Fix any warnings or lints that are introduced by your changes
Follow Rust best practices for error handling and memory management
Use the macro system (pyclass,pymodule,pyfunction, etc.) when implementing Python functionality in Rust
Files:
stdlib/src/contextvars.rsvm/src/stdlib/sys.rscommon/src/str.rsvm/src/stdlib/thread.rswasm/lib/src/vm_class.rsvm/src/vm/thread.rs
🧬 Code graph analysis (1)
vm/src/stdlib/thread.rs (1)
wasm/lib/src/vm_class.rs (1)
new(42-75)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (11)
- GitHub Check: Run rust tests (windows-latest)
- GitHub Check: Run snippets and cpython tests (ubuntu-latest)
- GitHub Check: Run snippets and cpython tests on wasm-wasi
- GitHub Check: Run rust tests (ubuntu-latest)
- GitHub Check: Run rust tests (macos-latest)
- GitHub Check: Check Rust code with rustfmt and clippy
- GitHub Check: Run snippets and cpython tests (windows-latest)
- GitHub Check: Run snippets and cpython tests (macos-latest)
- GitHub Check: Run tests under miri
- GitHub Check: Ensure compilation on various targets
- GitHub Check: Check the WASM package and demo
🔇 Additional comments (20)
vm/Cargo.toml (1)
68-69: Confirm non-native target builds
Workspace deps for scoped-tls and scopeguard look correct; please manually verify thatrustpython-vmbuilds cleanly forwasm32-unknown-unknownand your standard desktop targets.Cargo.toml (1)
210-211: Add to [workspace.dependencies]: LGTMPinning to major "1" for scoped-tls and scopeguard is appropriate and consistent with the vm crate workspace usage.
vm/src/stdlib/sys.rs (2)
910-918: Borrow-helper read: concise and correct
with_borrow(Clone::clone).to_pyobject(vm)cleanly maps Option → Python object, avoiding temporary borrows. No concerns.
889-895: Approve hook setter semantics preserved
OptionalArg<Option> paths correctly distinguish “unset” vs explicit None..set(finalizer)/.set(firstiter)will clear vs set hooks as intended, andget_asyncgen_hooksmaps an underlyingNoneto PythonNonevia theOption<T>::to_pyobjectimpl.vm/src/stdlib/thread.rs (3)
337-341: Cleanup via SENTINELS.take() looks goodDrains TLS cleanly and avoids nested borrows; unlock logic preserved.
356-356: TLS declaration OKConst init and type remain unchanged; style nit-only change.
361-361: Nice ergonomic switch to with_borrow_mutRemoves manual RefCell dance and reduces borrow scope.
vm/src/vm/thread.rs (3)
17-18: Scoped TLS for VM_CURRENT is a solid safety winEliminates raw pointer handling and simplifies access.
29-31: defer! pop + set() is tidy and exception-safeStack restoration is guaranteed; no objections.
54-55: Consistent VM binding via scoped TLSKeeps VM_CURRENT coherent with VM_STACK; reads cleanly.
wasm/lib/src/vm_class.rs (9)
61-65: Init hooks via with_borrow are fineBorrow scope is minimal; matches prior behavior.
81-82: add_init_func switch to with_borrow_mut is correctDirect push without manual borrow boilerplate.
100-105: stored_vm_from_wasm read path looks goodClones Rc under shared borrow; safe and concise.
108-109: weak_vm read path LGTMDowngrade under shared borrow is appropriate.
117-123: VMStore::init mutation path is cleanAvoids explicit borrow_mut management; idempotent insert preserved.
128-129: _get simplified nicelythen_some keeps things compact.
139-147: destroy path OKRemoval + explicit drop is clear; commented drop handlers retained.
151-152: ids uses with_borrow appropriatelyCollects keys without extra cloning.
187-188: valid read path LGTMShared borrow + contains_key is perfect here.
stdlib/src/contextvars.rs (1)
110-113: LGTM: switch to with_borrow_mut keeps semantics and drops an extra RefCell danceIndexing before push still matches the cached idx logic used later. No functional change.
| 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)) | ||
| })?; |
There was a problem hiding this comment.
💡 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.
There was a problem hiding this comment.
There was a problem hiding this comment.
.
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.
| 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)) | ||
| })?; |
There was a problem hiding this comment.
This branch has been sitting around in my local checkout for a while, and I was finally reminded of it today.
THREAD_LOCAL.with(|x: &RefCell<_>| x.borrow_mut().foo())->THREAD_LOCAL.with_borrow_mut(|x: &mut _| x.foo()), since Rust 1.73.Summary by CodeRabbit
Release Notes
Dependencies
scoped-tlsandscopeguardPerformance and Optimization
Technical Improvements
Maintenance
Note: These changes primarily focus on internal infrastructure improvements.