From 2c109b9f8ff3e9c58dc601b5e8da0a2a5fb75035 Mon Sep 17 00:00:00 2001 From: Jeong YunWon Date: Fri, 14 Aug 2026 02:39:58 +0900 Subject: [PATCH 1/7] _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 --- crates/vm/src/stdlib/_imp.rs | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/crates/vm/src/stdlib/_imp.rs b/crates/vm/src/stdlib/_imp.rs index fa979fcadbb..c8ff5e74e4a 100644 --- a/crates/vm/src/stdlib/_imp.rs +++ b/crates/vm/src/stdlib/_imp.rs @@ -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::().map_err(|_| invalid_err()); } import::make_frozen(vm, name.as_str()) } From 9b3095da0de4e61f2c20025329f0fbb73b2cfd2d Mon Sep 17 00:00:00 2001 From: Jeong YunWon Date: Fri, 14 Aug 2026 02:42:56 +0900 Subject: [PATCH 2/7] _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 --- crates/vm/src/stdlib/_imp.rs | 34 +++++++++++++++++++++--------- extra_tests/snippets/stdlib_imp.py | 6 ++++++ 2 files changed, 30 insertions(+), 10 deletions(-) diff --git a/crates/vm/src/stdlib/_imp.rs b/crates/vm/src/stdlib/_imp.rs index c8ff5e74e4a..50f0b0be8ab 100644 --- a/crates/vm/src/stdlib/_imp.rs +++ b/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, }; @@ -316,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>, bool, Option)>> { - 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) { @@ -339,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); @@ -347,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] diff --git a/extra_tests/snippets/stdlib_imp.py b/extra_tests/snippets/stdlib_imp.py index 64f1a0ad67e..9fd5f8a36fa 100644 --- a/extra_tests/snippets/stdlib_imp.py +++ b/extra_tests/snippets/stdlib_imp.py @@ -36,3 +36,9 @@ def __init__(self, name): 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 == "" From 68730fd016c72ee4960bd6fcf48f635b48908712 Mon Sep 17 00:00:00 2001 From: Jeong YunWon Date: Fri, 14 Aug 2026 03:09:57 +0900 Subject: [PATCH 3/7] itertools: give tee's shared buffer a Python type The buffer was a PyRc, 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 --- Lib/test/test_itertools.py | 1 - crates/vm/src/stdlib/itertools.rs | 95 +++++++++++++++++-------------- extra_tests/snippets/stdlib_gc.py | 2 + 3 files changed, 55 insertions(+), 43 deletions(-) diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py index b91e3735d94..c1695690b72 100644 --- a/Lib/test/test_itertools.py +++ b/Lib/test/test_itertools.py @@ -1270,7 +1270,6 @@ def test_dropwhile(self): 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 diff --git a/crates/vm/src/stdlib/itertools.rs b/crates/vm/src/stdlib/itertools.rs index e633404e803..44d63bfe582 100644 --- a/crates/vm/src/stdlib/itertools.rs +++ b/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>, + #[pytraverse(skip)] running: AtomicBool, } + #[pyclass(flags(DISALLOW_INSTANTIATION))] impl PyItertoolsTeeData { - fn new(iterable: PyIter, _vm: &VirtualMachine) -> PyRc { - PyRc::new(Self { + fn new(iterable: PyIter, vm: &VirtualMachine) -> PyRef { + Self { iterable, values: PyMutex::new(vec![]), running: AtomicBool::new(false), - }) + } + .into_ref(&vm.ctx) } fn get_item(&self, vm: &VirtualMachine, index: usize) -> PyResult { @@ -1006,54 +1010,35 @@ mod decl { } #[pyattr] - #[pyclass(name = "tee")] + #[pyclass(name = "_tee", traverse)] #[derive(Debug, PyPayload)] struct PyItertoolsTee { - tee_data: PyRc, + tee_data: PyRef, + #[pytraverse(skip)] index: AtomicCell, } - #[derive(FromArgs)] - struct TeeNewArgs { - #[pyarg(positional)] - iterable: PyIter, - #[pyarg(positional, optional)] - n: OptionalArg, - } - impl Constructor for PyItertoolsTee { - type Args = TeeNewArgs; - - // 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)? - }; + type Args = PyIter; - let mut tee_vec: Vec = Vec::with_capacity(n); - for _ in 0..n { - tee_vec.push(vm.call_special_method(©able, identifier!(vm, __copy__), ())?); + fn py_new(_cls: &Py, iterator: Self::Args, vm: &VirtualMachine) -> PyResult { + // 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::() { + return Ok(tee.__copy__()); } - - Ok(PyTuple::new_ref(tee_vec, &vm.ctx).into()) - } - - fn py_new(_cls: &Py, _args: Self::Args, _vm: &VirtualMachine) -> PyResult { - unimplemented!("use slot_new") + Ok(Self { + tee_data: PyItertoolsTeeData::new(iterator, vm), + index: AtomicCell::new(0), + }) } } - #[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 { @@ -1072,6 +1057,32 @@ mod decl { } } } + + #[pyfunction] + fn tee(iterable: PyIter, n: OptionalArg, vm: &VirtualMachine) -> PyResult { + 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 = 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(©able, identifier!(vm, __copy__), ())?); + } + + Ok(PyTuple::new_ref(tee_vec, &vm.ctx)) + } impl SelfIter for PyItertoolsTee {} impl IterNext for PyItertoolsTee { fn next(zelf: &Py, vm: &VirtualMachine) -> PyResult { diff --git a/extra_tests/snippets/stdlib_gc.py b/extra_tests/snippets/stdlib_gc.py index 134b1b9f458..6c3169beedb 100644 --- a/extra_tests/snippets/stdlib_gc.py +++ b/extra_tests/snippets/stdlib_gc.py @@ -62,5 +62,7 @@ def build(): 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") From 81f06903646dd2c3a333827c748db2ec898a2a4b Mon Sep 17 00:00:00 2001 From: Jeong YunWon Date: Fri, 14 Aug 2026 03:20:39 +0900 Subject: [PATCH 4/7] 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 --- crates/stdlib/src/contextvars.rs | 149 +++++++++--------- .../snippets/stdlib_threading_contextvars.py | 72 +++++++++ 2 files changed, 149 insertions(+), 72 deletions(-) create mode 100644 extra_tests/snippets/stdlib_threading_contextvars.py diff --git a/crates/stdlib/src/contextvars.rs b/crates/stdlib/src/contextvars.rs index 19fbcb8412f..e3823f6ac59 100644 --- a/crates/stdlib/src/contextvars.rs +++ b/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: PyMutex, } #[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, + idx: AtomicUsize, vars: PyRef, // PyObject *ctx_weakreflist; - entered: Cell, + 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 + '_ { - self.inner.vars.hamt.borrow() + fn borrow_vars(&self) -> impl core::ops::DerefMut + '_ { + self.inner.vars.hamt.lock() } fn borrow_vars_mut(&self) -> impl core::ops::DerefMut + '_ { - self.inner.vars.hamt.borrow_mut() + self.inner.vars.hamt.lock() } fn enter(zelf: &Py, 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, 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, vm: &VirtualMachine, ) -> PyResult { - 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, #[pytraverse(skip)] - cached: AtomicCell>, + cached: PyMutex>, #[pytraverse(skip)] - cached_id: core::sync::atomic::AtomicUsize, // cached_tsid in CPython + cached_id: AtomicUsize, // cached_tsid in CPython #[pytraverse(skip)] - hash: UnsafeCell, + 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, 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, 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, vm: &VirtualMachine) -> PyHash { @@ -370,28 +371,32 @@ mod _contextvars { default: OptionalArg, vm: &VirtualMachine, ) -> PyResult> { - 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, token: PyRef, 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(&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, _vm: &VirtualMachine) -> PyResult { - Ok(unsafe { *zelf.hash.get() }) + Ok(zelf.hash.load(Ordering::Relaxed)) } } @@ -537,11 +540,9 @@ mod _contextvars { ctx: PyRef, // tok_ctx in CPython var: PyRef, // tok_var in CPython old_value: Option, // tok_oldval in CPython - used: Cell, + 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, vm: &VirtualMachine) -> PyResult { - 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!(" Date: Fri, 14 Aug 2026 03:38:55 +0900 Subject: [PATCH 5/7] 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 --- crates/vm/src/coroutine.rs | 111 +++++++++++------- .../snippets/stdlib_threading_generator.py | 93 +++++++++++++++ 2 files changed, 163 insertions(+), 41 deletions(-) create mode 100644 extra_tests/snippets/stdlib_threading_generator.py diff --git a/crates/vm/src/coroutine.rs b/crates/vm/src/coroutine.rs index 43a28320e00..61431ea82e3 100644 --- a/crates/vm/src/coroutine.rs +++ b/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, 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, _claim: &RunningGuard<'_>) { match res { Ok(ExecutionResult::Return(_)) | Err(_) => { self.closed.store(true); @@ -109,45 +124,44 @@ impl Coro { } } - fn run_with_context( + /// 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> { + 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( &self, - jen: &PyObject, + _claim: &RunningGuard<'_>, vm: &VirtualMachine, func: F, - ) -> (PyResult, bool) + ) -> PyResult where F: FnOnce(&Py) -> PyResult, { - 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>; - 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, - entered_frame: bool, jen: &PyObject, vm: &VirtualMachine, ) -> PyResult { - 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 { + 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)))) diff --git a/extra_tests/snippets/stdlib_threading_generator.py b/extra_tests/snippets/stdlib_threading_generator.py new file mode 100644 index 00000000000..908f397e32d --- /dev/null +++ b/extra_tests/snippets/stdlib_threading_generator.py @@ -0,0 +1,93 @@ +"""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) + + +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") From fafb31700df58aadbdff3a40198da47a7cd3b4f8 Mon Sep 17 00:00:00 2001 From: Jeong YunWon Date: Fri, 14 Aug 2026 12:58:51 +0900 Subject: [PATCH 6/7] 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 --- crates/vm/src/stdlib/itertools.rs | 27 +++++++-- .../snippets/stdlib_threading_generator.py | 2 + .../stdlib_threading_itertools_tee.py | 57 +++++++++++++++++++ 3 files changed, 80 insertions(+), 6 deletions(-) create mode 100644 extra_tests/snippets/stdlib_threading_itertools_tee.py diff --git a/crates/vm/src/stdlib/itertools.rs b/crates/vm/src/stdlib/itertools.rs index 44d63bfe582..6eb268d94c1 100644 --- a/crates/vm/src/stdlib/itertools.rs +++ b/crates/vm/src/stdlib/itertools.rs @@ -992,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")); }; @@ -1016,6 +1018,8 @@ mod decl { tee_data: PyRef, #[pytraverse(skip)] index: AtomicCell, + #[pytraverse(skip)] + advancing: AtomicBool, } impl Constructor for PyItertoolsTee { @@ -1030,6 +1034,7 @@ mod decl { Ok(Self { tee_data: PyItertoolsTeeData::new(iterator, vm), index: AtomicCell::new(0), + advancing: AtomicBool::new(false), }) } } @@ -1044,6 +1049,7 @@ mod decl { 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()) @@ -1054,6 +1060,7 @@ mod decl { Self { tee_data: self.tee_data.clone(), index: AtomicCell::new(self.index.load()), + advancing: AtomicBool::new(false), } } } @@ -1086,8 +1093,16 @@ mod decl { impl SelfIter for PyItertoolsTee {} impl IterNext for PyItertoolsTee { fn next(zelf: &Py, vm: &VirtualMachine) -> PyResult { - 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)) } } diff --git a/extra_tests/snippets/stdlib_threading_generator.py b/extra_tests/snippets/stdlib_threading_generator.py index 908f397e32d..e606ef11cec 100644 --- a/extra_tests/snippets/stdlib_threading_generator.py +++ b/extra_tests/snippets/stdlib_threading_generator.py @@ -42,6 +42,8 @@ def worker(): 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)] diff --git a/extra_tests/snippets/stdlib_threading_itertools_tee.py b/extra_tests/snippets/stdlib_threading_itertools_tee.py new file mode 100644 index 00000000000..e13e20731f8 --- /dev/null +++ b/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") From 2a42430705d58198ef4777fef64d6ce2f48b0830 Mon Sep 17 00:00:00 2001 From: Jeong YunWon Date: Fri, 14 Aug 2026 13:05:16 +0900 Subject: [PATCH 7/7] _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 --- crates/stdlib/src/_asyncio.rs | 22 +++++++++++++--------- extra_tests/snippets/stdlib_asyncio.py | 23 +++++++++++++++++++++++ 2 files changed, 36 insertions(+), 9 deletions(-) diff --git a/crates/stdlib/src/_asyncio.rs b/crates/stdlib/src/_asyncio.rs index b311db4a315..9ad75fb8d69 100644 --- a/crates/stdlib/src/_asyncio.rs +++ b/crates/stdlib/src/_asyncio.rs @@ -2745,16 +2745,20 @@ pub(crate) mod _asyncio { } } + fn get_invalid_state_error_type(vm: &VirtualMachine) -> PyResult { + 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()), } diff --git a/extra_tests/snippets/stdlib_asyncio.py b/extra_tests/snippets/stdlib_asyncio.py index 7f03aeb436b..d54f84564a3 100644 --- a/extra_tests/snippets/stdlib_asyncio.py +++ b/extra_tests/snippets/stdlib_asyncio.py @@ -49,4 +49,27 @@ def __new__(cls, *args): 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")