fix(contextvars): isolate context stack per interpreter and thread on VirtualMachine (#8838)

* fix(contextvars): isolate context stack per interpreter and thread on VirtualMachine

* style: format and sort imports in extra_tests/snippets/stdlib_contextvars.py

Assisted-by: Antigravity:Gemini-3.8-Flash

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<Vec<PyRef<PyContext>>> = 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<Self> {
-            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::<Self>().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<ContextVar>) -> bool {
@@ -445,10 +432,12 @@ mod _contextvars {
         ) -> PyResult<Option<PyObjectRef>> {
             // 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::<PyContext>().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 {

crates/vm/src/vm/mod.rs

@@ -123,6 +123,8 @@ pub struct VirtualMachine {
     pub asyncio_running_loop: RefCell<Option<PyObjectRef>>,
     /// Current running asyncio task for this thread
     pub asyncio_running_task: RefCell<Option<PyObjectRef>>,
+    /// Active Context stack for this thread and interpreter (PEP 567 / contextvars)
+    pub context_stack: RefCell<Vec<PyObjectRef>>,
     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),

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),

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"
+    ] == []