diff --git a/crates/stdlib/src/contextvars.rs b/crates/stdlib/src/contextvars.rs index 0c1adf9d7..06c0e7585 100644 --- a/crates/stdlib/src/contextvars.rs +++ b/crates/stdlib/src/contextvars.rs @@ -1,14 +1,6 @@ pub(crate) use _contextvars::PyContext; pub(crate) use _contextvars::module_def; -use crate::vm::PyRef; -use core::cell::RefCell; - -thread_local! { - // TODO: Vec doesn't seem to match copy behavior - static CONTEXTS: RefCell>> = RefCell::default(); -} - #[pymodule] mod _contextvars { use crate::vm::{ @@ -110,10 +102,9 @@ mod _contextvars { ))); } - super::CONTEXTS.with_borrow_mut(|ctxs| { - zelf.inner.idx.store(ctxs.len(), Ordering::Relaxed); - ctxs.push(zelf.to_owned()); - }); + let mut ctxs = vm.context_stack.borrow_mut(); + zelf.inner.idx.store(ctxs.len(), Ordering::Relaxed); + ctxs.push(zelf.to_owned().into()); Ok(()) } @@ -126,14 +117,11 @@ mod _contextvars { ))); } - super::CONTEXTS.with_borrow_mut(|ctxs| { - ctxs.pop_if(|ctx| ctx.get_id() == zelf.get_id()) - .map(drop) - .ok_or_else(|| { - vm.new_runtime_error( - "cannot exit context: thread state references a different context object" - ) - }) + let mut ctxs = vm.context_stack.borrow_mut(); + ctxs.pop_if(|ctx| ctx.is(zelf)).map(drop).ok_or_else(|| { + vm.new_runtime_error( + "cannot exit context: thread state references a different context object", + ) })?; zelf.inner.entered.store(false, Ordering::Release); @@ -141,18 +129,17 @@ mod _contextvars { } fn current(vm: &VirtualMachine) -> PyRef { - super::CONTEXTS.with_borrow_mut(|ctxs| { - if let Some(ctx) = ctxs.last() { - ctx.clone() - } else { - let ctx = Self::empty(vm); - ctx.inner.idx.store(0, Ordering::Relaxed); - ctx.inner.entered.store(true, Ordering::Release); - let ctx = ctx.into_ref(&vm.ctx); - ctxs.push(ctx); - ctxs[0].clone() - } - }) + let mut ctxs = vm.context_stack.borrow_mut(); + if let Some(ctx) = ctxs.last() { + ctx.clone().downcast::().unwrap() + } else { + let ctx = Self::empty(vm); + ctx.inner.idx.store(0, Ordering::Relaxed); + ctx.inner.entered.store(true, Ordering::Release); + let ctx = ctx.into_ref(&vm.ctx); + ctxs.push(ctx.clone().into()); + ctx + } } fn contains(&self, needle: &Py) -> bool { @@ -445,10 +432,12 @@ mod _contextvars { ) -> PyResult> { // The replaced cache entry comes back out so that dropping it, which // can run a __del__ that calls back in, happens with no lock held. - let (found, replaced) = super::CONTEXTS.with_borrow(|ctxs| { - let Some(ctx) = ctxs.last() else { + let (found, replaced) = (|| { + let ctxs = vm.context_stack.borrow(); + let Some(ctx_obj) = ctxs.last() else { return (None, None); }; + let ctx = ctx_obj.downcast_ref::().unwrap(); let mut cached = zelf.cached.lock(); if let Some(cached) = &*cached && zelf.cached_id.load(Ordering::SeqCst) == ctx.get_id() @@ -467,7 +456,7 @@ mod _contextvars { }); (Some(obj), replaced) - }); + })(); drop(replaced); let value = if let Some(value) = found { diff --git a/crates/vm/src/vm/mod.rs b/crates/vm/src/vm/mod.rs index 1382eaac4..dbdd5ba7b 100644 --- a/crates/vm/src/vm/mod.rs +++ b/crates/vm/src/vm/mod.rs @@ -123,6 +123,8 @@ pub struct VirtualMachine { pub asyncio_running_loop: RefCell>, /// Current running asyncio task for this thread pub asyncio_running_task: RefCell>, + /// Active Context stack for this thread and interpreter (PEP 567 / contextvars) + pub context_stack: RefCell>, pub(crate) callable_cache: CallableCache, /// Side channel for TailCall: the bytecode loop stores the new frame /// pointer here before returning `ExecutionResult::TailCall`. @@ -1273,6 +1275,7 @@ impl VirtualMachine { async_gen_finalizer: RefCell::new(None), asyncio_running_loop: RefCell::new(None), asyncio_running_task: RefCell::new(None), + context_stack: RefCell::default(), callable_cache: CallableCache::default(), pending_tailcall_frame: Cell::new(None), pending_tailcall_owner: core::cell::UnsafeCell::new(None), diff --git a/crates/vm/src/vm/thread.rs b/crates/vm/src/vm/thread.rs index 7e6b1a5b6..e1bc40c44 100644 --- a/crates/vm/src/vm/thread.rs +++ b/crates/vm/src/vm/thread.rs @@ -1485,6 +1485,7 @@ impl VirtualMachine { async_gen_finalizer: RefCell::new(None), asyncio_running_loop: RefCell::new(None), asyncio_running_task: RefCell::new(None), + context_stack: RefCell::default(), callable_cache: self.callable_cache.clone(), pending_tailcall_frame: Cell::new(None), pending_tailcall_owner: core::cell::UnsafeCell::new(None), diff --git a/extra_tests/snippets/stdlib_contextvars.py b/extra_tests/snippets/stdlib_contextvars.py index e0dc0d2c5..915f732d1 100644 --- a/extra_tests/snippets/stdlib_contextvars.py +++ b/extra_tests/snippets/stdlib_contextvars.py @@ -54,3 +54,33 @@ class ReentrantEq: ctx1.run(nested.set, ReentrantEq()) ctx2.run(nested.set, object()) assert ctx1 == ctx2 + +try: + import gc + + import _interpreters +except ImportError: + _interpreters = None + +if _interpreters is not None: + child = _interpreters.create() + try: + assert ( + _interpreters.run_string( + child, + """ +import contextvars +contextvars.ContextVar('child_marker').set(42) +""", + ) + is None + ) + finally: + _interpreters.destroy(child) + + gc.collect() + assert [ + (var.name, value) + for var, value in copy_context().items() + if var.name == "child_marker" + ] == []