Fix the itertools.tee leak and race, the contextvars and generator races, _asyncio's InvalidStateError lookup, and _imp's frozen data (#8518)

* _imp: read get_frozen_object's data as a whole marshal value

The data argument went through deserialize_code(), which reads a code body
without the type byte in front of it, so nothing marshal.dumps() produces was
accepted. Read it with marshal.loads() and require a code object back.

Assisted-by: Claude

* _imp: bind find_frozen's arguments with FromArgs

The arity check was done by hand on a FuncArgs. Take the arguments through a
FromArgs struct instead, which makes withdata keyword-only, and fill in the
data it asks for: the frozen encoding is not marshal, so the code is
re-serialized into what get_frozen_object() reads back.

Assisted-by: Claude

* itertools: give tee's shared buffer a Python type

The buffer was a PyRc<PyItertoolsTeeData>, which is not a Python object, so the
collector could not walk into it and any cycle running through a tee was
uncollectable. Make it the _tee_dataobject type with a traverse, held by
PyRef, and split the rest along the same lines: tee() is a function, _tee is
the iterator type it builds, and _tee takes a single iterable rather than
returning a tuple from __new__. _tee is weak-referenceable, tee() rejects a
negative n with ValueError and reserves its tuple fallibly.

test_itertools.test_tee passes now; its expectedFailure marker is removed.

Assisted-by: Claude

* contextvars: hold the shared context state under locks

The variable map was a RefCell, the enter flag, the context index, the token's
used flag and the variable hash were Cells, and each carried an unsafe impl
Sync. A Context or ContextVar shared between threads overlapped their borrows
and panicked. The map and the per-variable cache are now PyMutex, the flags and
the index are atomics, entering a context is a compare_exchange, and the three
unsafe impl Sync are gone.

The cache also stopped being read through AtomicCell::as_ptr, which raced a
concurrent store on a value holding a PyObjectRef.

Values displaced from the map or the cache are dropped after the lock is
released: __del__ can come straight back into the same context, and the locks
are not reentrant.

Assisted-by: Claude

* generator: read the frame state under the running claim

send(), send_none(), throw() and close() read `closed` and `frame.lasti()`
before `running` was compare_exchanged, so the frame they went on to resume
could be one another thread had already advanced. A resume that decided from
`lasti() == 0` that the generator had not started pushes no value onto the
value stack, and the code after the yield pops one, which underflows the stack.

The compare_exchange now hands back a guard, taken before those reads and
released after maybe_close(), so the generator is retired while it is still
claimed.

Assisted-by: Claude

* itertools: claim a tee's position and its buffer as one step

`_tee::next` read `index` and moved it on afterwards, and
`_tee_dataobject::get_item` released `running` before the value it fetched
from the source was cached.

Two callers on one `_tee` then read the same index, hand out the same value
twice and advance past a value that was never cached, which leaves `index`
past `values.len()` and indexes the buffer out of bounds. Two callers at the
same index on separate tees each fetch a value from the source, and one of the
two is dropped without reaching a caller.

Both claims now cover the read and the update.

Assisted-by: Claude

* _asyncio: require InvalidStateError to be a type

new_invalid_state_error() called whatever `asyncio.exceptions.InvalidStateError`
names and unwrapped the downcast of the result, so a future asked for a result
it does not have panicked once that attribute was rebound to something that is
not an exception:

    asyncio.InvalidStateError = lambda *args: 42
    _asyncio.Future(loop=object()).result()

The type is looked up the way get_cancelled_error_type() looks its own up, and
raised with new_exception_msg; a lookup that does not produce an exception type
falls back to RuntimeError, as the other arms already did.

Assisted-by: Claude

Lib/test/test_itertools.py

@@ -1270,7 +1270,6 @@ class TestBasicOps(unittest.TestCase):
         self.assertRaises(TypeError, next, dropwhile(10, [(4,5)]))
         self.assertRaises(ValueError, next, dropwhile(errfunc, [(4,5)]))
 
-    @unittest.expectedFailure  # TODO: RUSTPYTHON
     def test_tee(self):
         n = 200
 

crates/stdlib/src/_asyncio.rs

@@ -2745,16 +2745,20 @@ pub(crate) mod _asyncio {
         }
     }
 
+    fn get_invalid_state_error_type(vm: &VirtualMachine) -> PyResult<PyTypeRef> {
+        let module = vm.import("asyncio.exceptions", 0)?;
+        let exc_type = vm
+            .get_attribute_opt(module, vm.ctx.intern_str("InvalidStateError"))?
+            .ok_or_else(|| vm.new_attribute_error("InvalidStateError not found"))?;
+        exc_type
+            .downcast()
+            .map_err(|_| vm.new_type_error("InvalidStateError is not a type"))
+    }
+
     fn new_invalid_state_error(vm: &VirtualMachine, msg: &str) -> PyBaseExceptionRef {
-        match vm.import("asyncio.exceptions", 0) {
-            Ok(module) => {
-                match vm.get_attribute_opt(module, vm.ctx.intern_str("InvalidStateError")) {
-                    Ok(Some(exc_type)) => match exc_type.call((msg,), vm) {
-                        Ok(exc) => exc.downcast().unwrap(),
-                        Err(_) => vm.new_runtime_error(msg.to_string()),
-                    },
-                    _ => vm.new_runtime_error(msg.to_string()),
-                }
+        match get_invalid_state_error_type(vm) {
+            Ok(invalid_state_error) => {
+                vm.new_exception_msg(invalid_state_error, msg.to_string().into())
             }
             Err(_) => vm.new_runtime_error(msg.to_string()),
         }

crates/stdlib/src/contextvars.rs

@@ -15,16 +15,16 @@ mod _contextvars {
         AsObject, Py, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, atomic_func,
         builtins::{PyGenericAlias, PyList, PyStrRef, PyType, PyTypeRef},
         class::StaticType,
-        common::{hash::PyHash, lock::LazyLock, wtf8::Wtf8Buf},
+        common::{
+            hash::PyHash,
+            lock::{LazyLock, PyMutex},
+            wtf8::Wtf8Buf,
+        },
         function::{ArgCallable, FuncArgs, OptionalArg},
         protocol::{PyMappingMethods, PySequenceMethods},
         types::{AsMapping, AsSequence, Constructor, Hashable, Iterable, Representable},
     };
-    use core::{
-        cell::{Cell, RefCell, UnsafeCell},
-        sync::atomic::Ordering,
-    };
-    use crossbeam_utils::atomic::AtomicCell;
+    use core::sync::atomic::{AtomicBool, AtomicI64, AtomicUsize, Ordering};
     use indexmap::IndexMap;
 
     // TODO: Real hamt implementation
@@ -33,7 +33,7 @@ mod _contextvars {
     #[pyclass(no_attr, name = "Hamt", module = "contextvars")]
     #[derive(Debug, PyPayload)]
     pub(crate) struct HamtObject {
-        hamt: RefCell<Hamt>,
+        hamt: PyMutex<Hamt>,
     }
 
     #[pyclass]
@@ -42,23 +42,19 @@ mod _contextvars {
     impl Default for HamtObject {
         fn default() -> Self {
             Self {
-                hamt: RefCell::new(Hamt::default()),
+                hamt: PyMutex::new(Hamt::default()),
             }
         }
     }
 
-    unsafe impl Sync for HamtObject {}
-
     #[derive(Debug)]
     struct ContextInner {
-        idx: Cell<usize>,
+        idx: AtomicUsize,
         vars: PyRef<HamtObject>,
         // PyObject *ctx_weakreflist;
-        entered: Cell<bool>,
+        entered: AtomicBool,
     }
 
-    unsafe impl Sync for ContextInner {}
-
     #[pyattr]
     #[pyclass(name = "Context")]
     #[derive(Debug, PyPayload)]
@@ -71,23 +67,30 @@ mod _contextvars {
         fn empty(vm: &VirtualMachine) -> Self {
             Self {
                 inner: ContextInner {
-                    idx: Cell::new(usize::MAX),
+                    idx: AtomicUsize::new(usize::MAX),
                     vars: HamtObject::default().into_ref(&vm.ctx),
-                    entered: Cell::new(false),
+                    entered: AtomicBool::new(false),
                 },
             }
         }
 
-        fn borrow_vars(&self) -> impl core::ops::Deref<Target = Hamt> + '_ {
-            self.inner.vars.hamt.borrow()
+        fn borrow_vars(&self) -> impl core::ops::DerefMut<Target = Hamt> + '_ {
+            self.inner.vars.hamt.lock()
         }
 
         fn borrow_vars_mut(&self) -> impl core::ops::DerefMut<Target = Hamt> + '_ {
-            self.inner.vars.hamt.borrow_mut()
+            self.inner.vars.hamt.lock()
         }
 
         fn enter(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<()> {
-            if zelf.inner.entered.get() {
+            // A context is entered by one thread at a time, so the check and the
+            // claim have to be a single step.
+            if zelf
+                .inner
+                .entered
+                .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
+                .is_err()
+            {
                 return Err(vm.new_runtime_error(format!(
                     "cannot enter context: {} is already entered",
                     zelf.as_object().repr(vm)?
@@ -95,16 +98,15 @@ mod _contextvars {
             }
 
             super::CONTEXTS.with_borrow_mut(|ctxs| {
-                zelf.inner.idx.set(ctxs.len());
+                zelf.inner.idx.store(ctxs.len(), Ordering::Relaxed);
                 ctxs.push(zelf.to_owned());
             });
-            zelf.inner.entered.set(true);
 
             Ok(())
         }
 
         fn exit(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<()> {
-            if !zelf.inner.entered.get() {
+            if !zelf.inner.entered.load(Ordering::Acquire) {
                 return Err(vm.new_runtime_error(format!(
                     "cannot exit context: {} is not entered",
                     zelf.as_object().repr(vm)?
@@ -120,7 +122,7 @@ mod _contextvars {
                         )
                     })
             })?;
-            zelf.inner.entered.set(false);
+            zelf.inner.entered.store(false, Ordering::Release);
 
             Ok(())
         }
@@ -131,8 +133,8 @@ mod _contextvars {
                     ctx.clone()
                 } else {
                     let ctx = Self::empty(vm);
-                    ctx.inner.idx.set(0);
-                    ctx.inner.entered.set(true);
+                    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()
@@ -170,13 +172,13 @@ mod _contextvars {
         fn copy(&self, vm: &VirtualMachine) -> Self {
             // Deep copy the vars - clone the underlying Hamt data, not just the PyRef
             let vars_copy = HamtObject {
-                hamt: RefCell::new(self.inner.vars.hamt.borrow().clone()),
+                hamt: PyMutex::new(self.inner.vars.hamt.lock().clone()),
             };
             Self {
                 inner: ContextInner {
-                    idx: Cell::new(usize::MAX),
+                    idx: AtomicUsize::new(usize::MAX),
                     vars: vars_copy.into_ref(&vm.ctx),
-                    entered: Cell::new(false),
+                    entered: AtomicBool::new(false),
                 },
             }
         }
@@ -186,11 +188,8 @@ mod _contextvars {
             var: PyRef<ContextVar>,
             vm: &VirtualMachine,
         ) -> PyResult<PyObjectRef> {
-            let vars = self.borrow_vars();
-            let item = vars
-                .get(&*var)
-                .ok_or_else(|| vm.new_key_error(var.into()))?;
-            Ok(item.to_owned())
+            let item = self.borrow_vars().get(&*var).map(|item| item.to_owned());
+            item.ok_or_else(|| vm.new_key_error(var.into()))
         }
 
         fn __len__(&self) -> usize {
@@ -290,11 +289,11 @@ mod _contextvars {
         name: String,
         default: Option<PyObjectRef>,
         #[pytraverse(skip)]
-        cached: AtomicCell<Option<ContextVarCache>>,
+        cached: PyMutex<Option<ContextVarCache>>,
         #[pytraverse(skip)]
-        cached_id: core::sync::atomic::AtomicUsize, // cached_tsid in CPython
+        cached_id: AtomicUsize, // cached_tsid in CPython
         #[pytraverse(skip)]
-        hash: UnsafeCell<PyHash>,
+        hash: AtomicI64,
     }
 
     impl core::fmt::Debug for ContextVar {
@@ -303,8 +302,6 @@ mod _contextvars {
         }
     }
 
-    unsafe impl Sync for ContextVar {}
-
     impl PartialEq for ContextVar {
         fn eq(&self, other: &Self) -> bool {
             core::ptr::eq(self, other)
@@ -320,12 +317,15 @@ mod _contextvars {
 
     impl ContextVar {
         fn delete(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<()> {
-            zelf.cached.store(None);
+            let cached = zelf.cached.lock().take();
+            drop(cached);
 
             let ctx = PyContext::current(vm);
 
-            let mut vars = ctx.borrow_vars_mut();
-            if vars.swap_remove(zelf).is_none() {
+            let removed = ctx.borrow_vars_mut().swap_remove(zelf);
+            let existed = removed.is_some();
+            drop(removed);
+            if !existed {
                 // TODO:
                 // PyErr_SetObject(PyExc_LookupError, (PyObject *)var);
                 return Err(vm.new_lookup_error(zelf.as_object().repr(vm)?.as_wtf8().to_owned()));
@@ -338,16 +338,17 @@ mod _contextvars {
         fn set_inner(zelf: &Py<Self>, value: PyObjectRef, vm: &VirtualMachine) {
             let ctx = PyContext::current(vm);
 
-            let mut vars = ctx.borrow_vars_mut();
-            vars.insert(zelf.to_owned(), value.clone());
+            let replaced = ctx.borrow_vars_mut().insert(zelf.to_owned(), value.clone());
+            drop(replaced);
 
             zelf.cached_id.store(ctx.get_id(), Ordering::SeqCst);
 
             let cache = ContextVarCache {
                 object: value,
-                idx: ctx.inner.idx.get(),
+                idx: ctx.inner.idx.load(Ordering::Relaxed),
             };
-            zelf.cached.store(Some(cache));
+            let replaced = zelf.cached.lock().replace(cache);
+            drop(replaced);
         }
 
         fn generate_hash(zelf: &Py<Self>, vm: &VirtualMachine) -> PyHash {
@@ -370,28 +371,32 @@ mod _contextvars {
             default: OptionalArg<PyObjectRef>,
             vm: &VirtualMachine,
         ) -> PyResult<Option<PyObjectRef>> {
-            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 }
+            // 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 {
+                    return (None, None);
+                };
+                let mut cached = zelf.cached.lock();
+                if let Some(cached) = &*cached
                     && zelf.cached_id.load(Ordering::SeqCst) == ctx.get_id()
                     && cached.idx + 1 == ctxs.len()
                 {
-                    return Some(cached.object.clone());
+                    return (Some(cached.object.clone()), None);
                 }
-                let vars = ctx.borrow_vars();
-                let obj = vars.get(zelf)?;
+                let Some(obj) = ctx.borrow_vars().get(zelf).map(|obj| obj.to_owned()) else {
+                    return (None, None);
+                };
                 zelf.cached_id.store(ctx.get_id(), Ordering::SeqCst);
 
-                // TODO: ensure cached is not changed
-                let _removed = zelf.cached.swap(Some(ContextVarCache {
+                let replaced = cached.replace(ContextVarCache {
                     object: obj.clone(),
                     idx: ctxs.len() - 1,
-                }));
+                });
 
-                Some(obj.clone())
+                (Some(obj), replaced)
             });
+            drop(replaced);
 
             let value = if let Some(value) = found {
                 value
@@ -425,7 +430,7 @@ mod _contextvars {
 
         #[pymethod]
         fn reset(zelf: &Py<Self>, token: PyRef<ContextToken>, vm: &VirtualMachine) -> PyResult<()> {
-            if token.used.get() {
+            if token.used.load(Ordering::Acquire) {
                 return Err(vm.new_runtime_error(format!(
                     "{} has already been used once",
                     token.as_object().repr(vm)?
@@ -447,7 +452,7 @@ mod _contextvars {
                 )));
             }
 
-            token.used.set(true);
+            token.used.store(true, Ordering::Release);
 
             if let Some(old_value) = &token.old_value {
                 Self::set_inner(zelf, old_value.clone(), vm);
@@ -484,15 +489,13 @@ mod _contextvars {
                 name: args.name.to_string(),
                 default: args.default.into_option(),
                 cached_id: 0.into(),
-                cached: AtomicCell::new(None),
-                hash: UnsafeCell::new(0),
+                cached: PyMutex::new(None),
+                hash: AtomicI64::new(0),
             };
             let py_var = var.into_ref_with_type(vm, cls)?;
 
-            unsafe {
-                // SAFETY: py_var is not exposed to python memory model yet
-                *py_var.hash.get() = Self::generate_hash(&py_var, vm)
-            };
+            let hash = Self::generate_hash(&py_var, vm);
+            py_var.hash.store(hash, Ordering::Relaxed);
             Ok(py_var.into())
         }
 
@@ -504,14 +507,14 @@ mod _contextvars {
     impl core::hash::Hash for ContextVar {
         #[inline]
         fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
-            unsafe { *self.hash.get() }.hash(state)
+            self.hash.load(Ordering::Relaxed).hash(state)
         }
     }
 
     impl Hashable for ContextVar {
         #[inline]
         fn hash(zelf: &Py<Self>, _vm: &VirtualMachine) -> PyResult<PyHash> {
-            Ok(unsafe { *zelf.hash.get() })
+            Ok(zelf.hash.load(Ordering::Relaxed))
         }
     }
 
@@ -537,11 +540,9 @@ mod _contextvars {
         ctx: PyRef<PyContext>,          // tok_ctx in CPython
         var: PyRef<ContextVar>,         // tok_var in CPython
         old_value: Option<PyObjectRef>, // tok_oldval in CPython
-        used: Cell<bool>,
+        used: AtomicBool,
     }
 
-    unsafe impl Sync for ContextToken {}
-
     #[pyclass(with(Constructor, Representable))]
     impl ContextToken {
         #[pygetset]
@@ -598,7 +599,11 @@ mod _contextvars {
     impl Representable for ContextToken {
         #[inline]
         fn repr_wtf8(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<Wtf8Buf> {
-            let used = if zelf.used.get() { " used" } else { "" };
+            let used = if zelf.used.load(Ordering::Acquire) {
+                " used"
+            } else {
+                ""
+            };
             let var = Representable::repr_wtf8(&zelf.var, vm)?;
             let ptr = zelf.as_object().get_id() as *const u8;
             let mut result = Wtf8Buf::from(format!("<Token{used} var="));

crates/vm/src/coroutine.rs

@@ -51,6 +51,21 @@ unsafe impl Traverse for Coro {
     }
 }
 
+/// An exclusive claim on a generator's frame, released when dropped.
+///
+/// Only the holder may look at the frame or resume it. Resuming decides from
+/// the frame state whether the sent value goes on the value stack, so a state
+/// read taken before the claim can be answered by a frame that another thread
+/// then advances: resuming it leaves the stack short of what the code after
+/// the yield pops.
+struct RunningGuard<'a>(&'a Coro);
+
+impl Drop for RunningGuard<'_> {
+    fn drop(&mut self) {
+        self.0.running.store(false);
+    }
+}
+
 fn gen_name(jen: &PyObject, vm: &VirtualMachine) -> &'static str {
     let typ = jen.class();
     if typ.is(vm.ctx.types.coroutine_type) {
@@ -88,10 +103,10 @@ impl Coro {
         }
     }
 
-    fn maybe_close(&self, res: &PyResult<ExecutionResult>, entered_frame: bool) {
-        if !entered_frame {
-            return;
-        }
+    /// Retire the generator if the frame it just ran came to an end. The claim
+    /// is still held, so a thread waiting for it cannot resume a frame that has
+    /// already finished.
+    fn maybe_close(&self, res: &PyResult<ExecutionResult>, _claim: &RunningGuard<'_>) {
         match res {
             Ok(ExecutionResult::Return(_)) | Err(_) => {
                 self.closed.store(true);
@@ -109,45 +124,44 @@ impl Coro {
         }
     }
 
-    fn run_with_context<F>(
+    /// Take the frame for this thread, or report that another thread holds it.
+    ///
+    /// What the resume depends on -- whether the generator is closed, and
+    /// whether it has started -- has to be read from here onwards.
+    fn claim(&self, jen: &PyObject, vm: &VirtualMachine) -> PyResult<RunningGuard<'_>> {
+        if self.running.compare_exchange(false, true).is_err() {
+            return Err(vm.new_value_error(format!("{} already executing", gen_name(jen, vm))));
+        }
+        Ok(RunningGuard(self))
+    }
+
+    fn run_claimed<F>(
         &self,
-        jen: &PyObject,
+        _claim: &RunningGuard<'_>,
         vm: &VirtualMachine,
         func: F,
-    ) -> (PyResult<ExecutionResult>, bool)
+    ) -> PyResult<ExecutionResult>
     where
         F: FnOnce(&Py<FrameObject>) -> PyResult<ExecutionResult>,
     {
-        if self.running.compare_exchange(false, true).is_err() {
-            return (
-                Err(vm.new_value_error(format!("{} already executing", gen_name(jen, vm)))),
-                false,
-            );
-        }
-
-        // SAFETY: running.compare_exchange guarantees exclusive access
+        // SAFETY: the claim guarantees exclusive access
         let gen_exc = unsafe { self.exception.swap(None) };
         let exception_ptr = &self.exception as *const PyAtomicRef<Option<PyBaseException>>;
 
-        let result = vm.resume_gen_frame(&self.frame, gen_exc, |f| {
+        vm.resume_gen_frame(&self.frame, gen_exc, |f| {
             let result = func(f);
-            // SAFETY: exclusive access guaranteed by running flag
+            // SAFETY: exclusive access guaranteed by the claim
             let _old = unsafe { (*exception_ptr).swap(vm.current_exception()) };
             result
-        });
-
-        self.running.store(false);
-        (result, true)
+        })
     }
 
     fn finalize_send_result(
         &self,
         result: PyResult<ExecutionResult>,
-        entered_frame: bool,
         jen: &PyObject,
         vm: &VirtualMachine,
     ) -> PyResult<PyIterReturn> {
-        self.maybe_close(&result, entered_frame);
         match result {
             Ok(exec_res) => Ok(exec_res.into_iter_return(vm)),
             Err(e) => {
@@ -177,16 +191,20 @@ impl Coro {
         if self.closed.load() {
             return Ok(PyIterReturn::StopIteration(None));
         }
-        if self.running.load() {
-            return Err(vm.new_value_error(format!("{} already executing", gen_name(jen, vm))));
+        let claim = self.claim(jen, vm)?;
+        // The generator can have run to its end in the meantime.
+        if self.closed.load() {
+            return Ok(PyIterReturn::StopIteration(None));
         }
         let value = if self.frame.lasti() > 0 {
             Some(vm.ctx.none())
         } else {
             None
         };
-        let (result, entered_frame) = self.run_with_context(jen, vm, |f| f.resume(value, vm));
-        self.finalize_send_result(result, entered_frame, jen, vm)
+        let result = self.run_claimed(&claim, vm, |f| f.resume(value, vm));
+        self.maybe_close(&result, &claim);
+        drop(claim);
+        self.finalize_send_result(result, jen, vm)
     }
 
     pub fn send(
@@ -198,8 +216,10 @@ impl Coro {
         if self.closed.load() {
             return Ok(PyIterReturn::StopIteration(None));
         }
-        if self.running.load() {
-            return Err(vm.new_value_error(format!("{} already executing", gen_name(jen, vm))));
+        let claim = self.claim(jen, vm)?;
+        // The generator can have run to its end in the meantime.
+        if self.closed.load() {
+            return Ok(PyIterReturn::StopIteration(None));
         }
         let value = if self.frame.lasti() > 0 {
             Some(value)
@@ -211,8 +231,10 @@ impl Coro {
         } else {
             None
         };
-        let (result, entered_frame) = self.run_with_context(jen, vm, |f| f.resume(value, vm));
-        self.finalize_send_result(result, entered_frame, jen, vm)
+        let result = self.run_claimed(&claim, vm, |f| f.resume(value, vm));
+        self.maybe_close(&result, &claim);
+        drop(claim);
+        self.finalize_send_result(result, jen, vm)
     }
 
     pub fn throw(
@@ -237,13 +259,25 @@ impl Coro {
         // Validate exception type before entering generator context.
         // Invalid types propagate to caller without closing the generator.
         crate::exceptions::ExceptionCtor::try_from_object(vm, exc_type.clone())?;
-        let (result, entered_frame) =
-            self.run_with_context(jen, vm, |f| f.gen_throw(vm, exc_type, exc_val, exc_tb));
-        self.maybe_close(&result, entered_frame);
+        let claim = self.claim(jen, vm)?;
+        // The generator can have run to its end in the meantime. Normalizing
+        // runs the exception's constructor, so let the claim go first.
+        if self.closed.load() {
+            drop(claim);
+            return Err(vm.normalize_exception(exc_type, exc_val, exc_tb)?);
+        }
+        let result = self.run_claimed(&claim, vm, |f| f.gen_throw(vm, exc_type, exc_val, exc_tb));
+        self.maybe_close(&result, &claim);
+        drop(claim);
         Ok(result?.into_iter_return(vm))
     }
 
     pub fn close(&self, jen: &PyObject, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
+        if self.closed.load() {
+            return Ok(vm.ctx.none());
+        }
+        let claim = self.claim(jen, vm)?;
+        // The generator can have run to its end in the meantime.
         if self.closed.load() {
             return Ok(vm.ctx.none());
         }
@@ -252,7 +286,7 @@ impl Coro {
             self.closed.store(true);
             return Ok(vm.ctx.none());
         }
-        let (result, entered_frame) = self.run_with_context(jen, vm, |f| {
+        let result = self.run_claimed(&claim, vm, |f| {
             f.gen_throw(
                 vm,
                 vm.ctx.exceptions.generator_exit.to_owned().into(),
@@ -260,16 +294,11 @@ impl Coro {
                 vm.ctx.none(),
             )
         });
-        if !entered_frame {
-            return match result {
-                Err(err) => Err(err),
-                Ok(_) => unreachable!("run_with_context preflight returned without an error"),
-            };
-        }
         self.closed.store(true);
         // Release frame locals and stack to free references held by the
         // closed generator, matching gen_send_ex2 with close_on_completion.
         self.clear_frame_locals_on_close();
+        drop(claim);
         match result {
             Ok(ExecutionResult::Yield(_)) => {
                 Err(vm.new_runtime_error(format!("{} ignored GeneratorExit", gen_name(jen, vm))))

crates/vm/src/stdlib/_imp.rs

@@ -179,7 +179,7 @@ mod _imp {
         PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine,
         builtins::{PyBytesRef, PyCode, PyMemoryView, PyModule, PyStrRef, PyUtf8StrRef},
         convert::TryFromBorrowedObject,
-        function::{FuncArgs, OptionalArg},
+        function::OptionalArg,
         import, version,
     };
 
@@ -264,21 +264,20 @@ mod _imp {
         if let OptionalArg::Present(data) = data
             && !vm.is_none(&data)
         {
-            let buf = crate::protocol::PyBuffer::try_from_borrowed_object(vm, &data)?;
-            let contiguous = buf.as_contiguous().ok_or_else(|| {
-                vm.new_buffer_error("get_frozen_object() requires a contiguous buffer")
-            })?;
             let invalid_err = || {
                 vm.new_import_error(
                     format!("Frozen object named '{}' is invalid", name.as_str()),
                     name.clone().into_wtf8(),
                 )
             };
-            let bag = crate::builtins::code::PyVmBag(vm);
-            let code =
-                rustpython_compiler_core::marshal::deserialize_code(&mut &contiguous[..], bag)
-                    .map_err(|_| invalid_err())?;
-            return Ok(PyCode::new_ref_with_bag(vm, code));
+            // A non-buffer is a TypeError, not invalid frozen data.
+            crate::protocol::PyBuffer::try_from_borrowed_object(vm, &data)?;
+            // The data is a marshalled code object: a whole marshal value, which
+            // deserialize_code() does not read — it takes the code body alone,
+            // without the type byte the writer puts in front of it.
+            let loads = vm.import("marshal", 0)?.get_attr("loads", vm)?;
+            let code = loads.call((data,), vm).map_err(|_| invalid_err())?;
+            return code.downcast::<PyCode>().map_err(|_| invalid_err());
         }
         import::make_frozen(vm, name.as_str())
     }
@@ -317,19 +316,21 @@ mod _imp {
             .collect()
     }
 
+    #[derive(FromArgs)]
+    struct FindFrozenArgs {
+        #[pyarg(positional)]
+        name: PyUtf8StrRef,
+        #[pyarg(named, default = false)]
+        withdata: bool,
+    }
+
     #[allow(clippy::type_complexity)]
     #[pyfunction]
     fn find_frozen(
-        args: FuncArgs,
+        args: FindFrozenArgs,
         vm: &VirtualMachine,
     ) -> PyResult<Option<(Option<PyRef<PyMemoryView>>, bool, Option<PyStrRef>)>> {
-        if args.args.len() > 1 {
-            return Err(vm.new_type_error(format!(
-                "find_frozen() takes exactly 1 positional argument ({} given)",
-                args.args.len()
-            )));
-        }
-        let (name,): (PyUtf8StrRef,) = args.bind(vm)?;
+        let FindFrozenArgs { name, withdata } = args;
 
         let name_str = name.as_str();
         let info = match super::find_frozen(name_str, vm) {
@@ -340,6 +341,18 @@ mod _imp {
             Err(e) => return Err(e.to_pyexception(name_str, vm)),
         };
 
+        // The data is what get_frozen_object() takes back, i.e. marshalled code.
+        // Frozen modules are stored in their own encoding, so it has to be
+        // re-serialized rather than handed out as a view of the stored bytes.
+        let data = if withdata {
+            let code = PyCode::new_ref_from_frozen(vm, info.code);
+            let dumps = vm.import("marshal", 0)?.get_attr("dumps", vm)?;
+            let bytes = dumps.call((code,), vm)?;
+            Some(PyMemoryView::from_object(&bytes, vm)?.into_ref(&vm.ctx))
+        } else {
+            None
+        };
+
         // When origname is empty (e.g. __hello_only__), return None.
         // Otherwise return the resolved alias name.
         let origname_str = super::resolve_frozen_alias(name_str);
@@ -348,7 +361,7 @@ mod _imp {
         } else {
             Some(vm.ctx.new_utf8_str(origname_str).into())
         };
-        Ok(Some((None, info.package, origname)))
+        Ok(Some((data, info.package, origname)))
     }
 
     #[pyfunction]

crates/vm/src/stdlib/itertools.rs

@@ -4,11 +4,10 @@ pub(crate) use decl::module_def;
 mod decl {
     use crate::{
         AsObject, Py, PyObjectRef, PyPayload, PyRef, PyResult, PyWeakRef, VirtualMachine,
-        builtins::{PyGenericAlias, PyInt, PyIntRef, PyList, PyTuple, PyType, PyTypeRef, int},
-        common::{
-            lock::{PyMutex, PyRwLock, PyRwLockWriteGuard},
-            rc::PyRc,
+        builtins::{
+            PyGenericAlias, PyInt, PyIntRef, PyList, PyTuple, PyTupleRef, PyType, PyTypeRef, int,
         },
+        common::lock::{PyMutex, PyRwLock, PyRwLockWriteGuard},
         convert::ToPyObject,
         function::{FuncArgs, OptionalArg, OptionalOption, PosArgs},
         protocol::{PyIter, PyIterReturn, PyNumber},
@@ -962,20 +961,25 @@ mod decl {
         }
     }
 
-    #[derive(Debug)]
+    #[pyattr]
+    #[pyclass(name = "_tee_dataobject", traverse)]
+    #[derive(Debug, PyPayload)]
     struct PyItertoolsTeeData {
         iterable: PyIter,
         values: PyMutex<Vec<PyObjectRef>>,
+        #[pytraverse(skip)]
         running: AtomicBool,
     }
 
+    #[pyclass(flags(DISALLOW_INSTANTIATION))]
     impl PyItertoolsTeeData {
-        fn new(iterable: PyIter, _vm: &VirtualMachine) -> PyRc<Self> {
-            PyRc::new(Self {
+        fn new(iterable: PyIter, vm: &VirtualMachine) -> PyRef<Self> {
+            Self {
                 iterable,
                 values: PyMutex::new(vec![]),
                 running: AtomicBool::new(false),
-            })
+            }
+            .into_ref(&vm.ctx)
         }
 
         fn get_item(&self, vm: &VirtualMachine, index: usize) -> PyResult<PyIterReturn> {
@@ -988,13 +992,15 @@ mod decl {
                     return Ok(PyIterReturn::Return(values[index].clone()));
                 }
             }
-            // Prevent concurrent/reentrant calls to iterable.next()
+            // Prevent concurrent/reentrant calls to iterable.next(). The claim
+            // covers caching the value as well: released any earlier, a second
+            // tee at the same index fetches a value of its own and one of the
+            // two is dropped without ever reaching a caller.
             if self.running.swap(true, Ordering::Acquire) {
                 return Err(vm.new_runtime_error("cannot re-enter the tee iterator"));
             }
-            let result = self.iterable.next(vm);
-            self.running.store(false, Ordering::Release);
-            let obj = raise_if_stop!(result?);
+            scopeguard::defer! { self.running.store(false, Ordering::Release) }
+            let obj = raise_if_stop!(self.iterable.next(vm)?);
             let Some(mut values) = self.values.try_lock() else {
                 return Err(vm.new_runtime_error("cannot re-enter the tee iterator"));
             };
@@ -1006,59 +1012,44 @@ mod decl {
     }
 
     #[pyattr]
-    #[pyclass(name = "tee")]
+    #[pyclass(name = "_tee", traverse)]
     #[derive(Debug, PyPayload)]
     struct PyItertoolsTee {
-        tee_data: PyRc<PyItertoolsTeeData>,
+        tee_data: PyRef<PyItertoolsTeeData>,
+        #[pytraverse(skip)]
         index: AtomicCell<usize>,
-    }
-
-    #[derive(FromArgs)]
-    struct TeeNewArgs {
-        #[pyarg(positional)]
-        iterable: PyIter,
-        #[pyarg(positional, optional)]
-        n: OptionalArg<usize>,
+        #[pytraverse(skip)]
+        advancing: AtomicBool,
     }
 
     impl Constructor for PyItertoolsTee {
-        type Args = TeeNewArgs;
+        type Args = PyIter;
 
-        // TODO: make tee() a function, rename this class to itertools._tee and make
-        // teedata a python class
-        fn slot_new(_cls: PyTypeRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult {
-            let TeeNewArgs { iterable, n } = args.bind(vm)?;
-            let n = n.unwrap_or(2);
-
-            let copyable = if iterable.class().has_attr(identifier!(vm, __copy__)) {
-                vm.call_special_method(iterable.as_object(), identifier!(vm, __copy__), ())?
-            } else {
-                Self::from_iter(iterable, vm)?
-            };
-
-            let mut tee_vec: Vec<PyObjectRef> = Vec::with_capacity(n);
-            for _ in 0..n {
-                tee_vec.push(vm.call_special_method(&copyable, identifier!(vm, __copy__), ())?);
+        fn py_new(_cls: &Py<PyType>, iterator: Self::Args, vm: &VirtualMachine) -> PyResult<Self> {
+            // An iterator that is already a tee shares its buffer rather than
+            // getting one of its own.
+            if let Some(tee) = iterator.as_object().downcast_ref::<Self>() {
+                return Ok(tee.__copy__());
             }
-
-            Ok(PyTuple::new_ref(tee_vec, &vm.ctx).into())
-        }
-
-        fn py_new(_cls: &Py<PyType>, _args: Self::Args, _vm: &VirtualMachine) -> PyResult<Self> {
-            unimplemented!("use slot_new")
+            Ok(Self {
+                tee_data: PyItertoolsTeeData::new(iterator, vm),
+                index: AtomicCell::new(0),
+                advancing: AtomicBool::new(false),
+            })
         }
     }
 
-    #[pyclass(with(IterNext, Iterable, Constructor))]
+    #[pyclass(with(IterNext, Iterable, Constructor), flags(HAS_WEAKREF))]
     impl PyItertoolsTee {
         fn from_iter(iterator: PyIter, vm: &VirtualMachine) -> PyResult {
             let class = Self::class(&vm.ctx);
-            if iterator.class().is(Self::class(&vm.ctx)) {
+            if iterator.class().is(class) {
                 return vm.call_special_method(&iterator, identifier!(vm, __copy__), ());
             }
             Ok(Self {
                 tee_data: PyItertoolsTeeData::new(iterator, vm),
                 index: AtomicCell::new(0),
+                advancing: AtomicBool::new(false),
             }
             .into_ref_with_type(vm, class.to_owned())?
             .into())
@@ -1069,14 +1060,49 @@ mod decl {
             Self {
                 tee_data: self.tee_data.clone(),
                 index: AtomicCell::new(self.index.load()),
+                advancing: AtomicBool::new(false),
             }
         }
     }
+
+    #[pyfunction]
+    fn tee(iterable: PyIter, n: OptionalArg<isize>, vm: &VirtualMachine) -> PyResult<PyTupleRef> {
+        let n = n.unwrap_or(2);
+        if n < 0 {
+            return Err(vm.new_value_error("n must be >= 0"));
+        }
+        let n = n as usize;
+
+        // Only an iterator that cannot copy itself needs a tee to buffer it.
+        let copyable = if iterable.class().has_attr(identifier!(vm, __copy__)) {
+            iterable.into()
+        } else {
+            PyItertoolsTee::from_iter(iterable, vm)?
+        };
+
+        let mut tee_vec: Vec<PyObjectRef> = Vec::new();
+        tee_vec
+            .try_reserve_exact(n)
+            .map_err(|_| vm.new_memory_error(""))?;
+        for _ in 0..n {
+            tee_vec.push(vm.call_special_method(&copyable, identifier!(vm, __copy__), ())?);
+        }
+
+        Ok(PyTuple::new_ref(tee_vec, &vm.ctx))
+    }
     impl SelfIter for PyItertoolsTee {}
     impl IterNext for PyItertoolsTee {
         fn next(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<PyIterReturn> {
-            let value = raise_if_stop!(zelf.tee_data.get_item(vm, zelf.index.load())?);
-            zelf.index.fetch_add(1);
+            // Reading the index and moving it on is one step: two callers that
+            // read the same index hand out the same value twice and leave the
+            // buffer to be filled out of order.
+            if zelf.advancing.swap(true, Ordering::Acquire) {
+                return Err(vm.new_runtime_error("cannot re-enter the tee iterator"));
+            }
+            scopeguard::defer! { zelf.advancing.store(false, Ordering::Release) }
+            let index = zelf.index.load();
+            let value = raise_if_stop!(zelf.tee_data.get_item(vm, index)?);
+            zelf.index.store(index + 1);
             Ok(PyIterReturn::Return(value))
         }
     }

extra_tests/snippets/stdlib_asyncio.py

@@ -49,4 +49,27 @@ future = _asyncio.Future(loop=object())
 with assert_raises(TypeError):
     future.__await__().throw(BadException)
 
+# InvalidStateError is looked up in asyncio.exceptions every time a pending
+# future is asked for its result, so what is found there need not be an
+# exception type at all.
+import asyncio  # noqa: E402
+import asyncio.exceptions  # noqa: E402
+
+pending = _asyncio.Future(loop=object())
+with assert_raises(asyncio.exceptions.InvalidStateError):
+    pending.result()
+
+saved_invalid_state_error = asyncio.exceptions.InvalidStateError
+try:
+    for replacement in (lambda *args: 42, None):
+        asyncio.InvalidStateError = replacement
+        asyncio.exceptions.InvalidStateError = replacement
+        with assert_raises(RuntimeError):
+            pending.result()
+        with assert_raises(RuntimeError):
+            pending.exception()
+finally:
+    asyncio.InvalidStateError = saved_invalid_state_error
+    asyncio.exceptions.InvalidStateError = saved_invalid_state_error
+
 print("ok")

extra_tests/snippets/stdlib_gc.py

@@ -62,5 +62,7 @@ assert collects(lambda c: itertools.filterfalse(None, c))
 assert collects(lambda c: itertools.compress(c, [1]))
 assert collects(lambda c: itertools.product(c))
 assert collects(lambda c: itertools.combinations(c, 1))
+# tee holds its buffer through a second object, which has to be walked too
+assert collects(lambda c: itertools.tee(c)[0])
 
 print("ok")

extra_tests/snippets/stdlib_imp.py

@@ -36,3 +36,9 @@ assert hello.initialized == True
 with assert_raises(TypeError):
     _imp.find_frozen("x", True)
 assert _imp.find_frozen("_this_module_does_not_exist_") is None
+
+# and it hands back the marshalled code that get_frozen_object() takes
+data, ispkg, origname = _imp.find_frozen("__hello__", withdata=True)
+assert ispkg is False
+assert origname == "__hello__"
+assert _imp.get_frozen_object("__hello__", data).co_name == "<module>"

extra_tests/snippets/stdlib_threading_contextvars.py

@@ -0,0 +1,72 @@
+"""Stress contextvars from several threads at once.
+
+A Context holds the variable map, and both the map and the per-variable cache
+are shared between every thread that touches the Context. Reading and writing
+them has to be done under a lock rather than a cell borrow.
+
+Dropping a value that a set() or reset() displaced can run a __del__ that comes
+straight back into the same Context, so the displaced value has to be released
+after the lock is, not while it is held.
+"""
+
+import contextvars
+import threading
+
+ROUNDS = 2000
+
+var = contextvars.ContextVar("v", default=0)
+shared = contextvars.Context()
+errors = []
+
+
+class Reentrant:
+    """__del__ runs while the variable that held this value is being replaced."""
+
+    def __del__(self):
+        try:
+            var.get()
+        except Exception:  # a different context, or no value: not what is tested
+            pass
+
+
+def churn():
+    try:
+        for i in range(ROUNDS):
+            token = var.set(Reentrant())
+            var.get()
+            var.reset(token)
+            var.set(i)
+            var.get()
+            contextvars.copy_context()
+    except Exception as exc:  # noqa: BLE001
+        errors.append(exc)
+
+
+def run_in_shared():
+    for i in range(ROUNDS):
+        try:
+            shared.run(var.set, i)
+        except RuntimeError:
+            # the Context is already entered by another thread
+            pass
+
+
+threads = [threading.Thread(target=churn) for _ in range(4)]
+threads += [threading.Thread(target=run_in_shared) for _ in range(4)]
+for t in threads:
+    t.start()
+for t in threads:
+    t.join()
+
+assert not errors, errors
+
+# the map itself still behaves
+ctx = contextvars.copy_context()
+ctx.run(var.set, 42)
+assert ctx[var] == 42
+assert var in ctx
+assert list(ctx) == [var]
+assert ctx.get(var) == 42
+assert len(ctx) == 1
+
+print("ok")

extra_tests/snippets/stdlib_threading_generator.py

@@ -0,0 +1,95 @@
+"""Resume one generator from several threads at once.
+
+A generator is resumed by one thread at a time, and whether the sent value is
+pushed onto the frame's value stack depends on whether the generator has
+already started. Deciding that before the generator is claimed reads a frame
+another thread can advance in the meantime, and resuming it then leaves the
+stack short of what the code after the yield expects.
+
+Every yielded value still has to reach exactly one caller: threads that lose
+the race get a ValueError instead of a value.
+"""
+
+import threading
+
+WORKERS = 4
+ROUNDS = 400
+
+
+def counter():
+    yield 1
+    yield 2
+    yield 3
+
+
+gens = [counter() for _ in range(ROUNDS)]
+received = [[] for _ in range(ROUNDS)]
+start = threading.Barrier(WORKERS)
+errors = []
+
+
+def worker():
+    try:
+        for index, gen in enumerate(gens):
+            start.wait()
+            for _ in range(3):
+                try:
+                    received[index].append(next(gen))
+                except StopIteration:
+                    break
+                except ValueError:
+                    # another thread is running this generator
+                    pass
+    except Exception as exc:  # noqa: BLE001
+        errors.append(exc)
+        # the other workers are waiting at the barrier for this one
+        start.abort()
+
+
+threads = [threading.Thread(target=worker) for _ in range(WORKERS)]
+for t in threads:
+    t.start()
+for t in threads:
+    t.join()
+
+assert not errors, errors
+for got in received:
+    # no value handed out twice, and none skipped
+    assert sorted(got) == list(range(1, len(got) + 1)), got
+
+
+# a generator that is closed while it is being resumed stays consistent
+def loop():
+    while True:
+        yield 1
+
+
+shared = loop()
+closed = threading.Barrier(2)
+
+
+def resumer():
+    closed.wait()
+    for _ in range(ROUNDS):
+        try:
+            next(shared)
+        except (StopIteration, ValueError):
+            pass
+
+
+def closer():
+    closed.wait()
+    try:
+        shared.close()
+    except ValueError:
+        # the generator was running
+        pass
+
+
+pair = [threading.Thread(target=resumer), threading.Thread(target=closer)]
+for t in pair:
+    t.start()
+for t in pair:
+    t.join()
+
+print("ok")

extra_tests/snippets/stdlib_threading_itertools_tee.py

@@ -0,0 +1,57 @@
+"""Advance the iterators of one tee() from several threads at once.
+
+Every tee iterator reads its position, asks the shared buffer for that item and
+then moves the position on. Reading and moving it on has to be one step, and
+the buffer has to stay claimed until the value it fetched from the source is
+cached: otherwise two callers work on the same index, a fetched value is
+dropped, and the buffer is left to be filled out of order.
+
+A caller that loses the race gets a RuntimeError, never a value another caller
+has already been handed.
+"""
+
+import itertools
+import threading
+
+ROUNDS = 200
+WORKERS = 4
+
+errors = []
+
+
+def drain(iterator, out):
+    for _ in range(ROUNDS):
+        try:
+            out.append(next(iterator))
+        except StopIteration:
+            break
+        except RuntimeError:
+            # another thread is advancing this tee
+            pass
+        except Exception as exc:  # noqa: BLE001
+            errors.append(exc)
+            break
+
+
+for _ in range(10):
+    first, second = itertools.tee(iter(range(ROUNDS * WORKERS)))
+    taken = [[] for _ in range(WORKERS)]
+    threads = [
+        threading.Thread(target=drain, args=(first if i % 2 else second, taken[i]))
+        for i in range(WORKERS)
+    ]
+    for t in threads:
+        t.start()
+    for t in threads:
+        t.join()
+
+    assert not errors, errors
+    for got in taken:
+        # one iterator hands out ascending values, each of them once
+        assert got == sorted(set(got)), got
+    for side in (taken[1], taken[3]), (taken[0], taken[2]):
+        # the two threads sharing an iterator split its values between them
+        shared = side[0] + side[1]
+        assert len(shared) == len(set(shared)), shared
+
+print("ok")