Skip to content

Fix the itertools.tee leak, the contextvars and generator races, and _imp's frozen data #8518

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
youknowone wants to merge 5 commits into
base: main
Choose a base branch
from
Draft

Fix the itertools.tee leak, the contextvars and generator races, and _imp's frozen data #8518

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
2c109b9
_imp: read get_frozen_object's data as a whole marshal value
youknowone Aug 13, 2026
9b3095d
_imp: bind find_frozen's arguments with FromArgs
youknowone Aug 13, 2026
68730fd
itertools: give tee's shared buffer a Python type
youknowone Aug 13, 2026
81f0690
contextvars: hold the shared context state under locks
youknowone Aug 13, 2026
59a8029
generator: read the frame state under the running claim
youknowone Aug 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion Lib/test/test_itertools.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
149 changes: 77 additions & 72 deletions crates/stdlib/src/contextvars.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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]
Expand All @@ -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)]
Expand All @@ -71,40 +67,46 @@ 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)?
)));
}

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)?
Expand All @@ -120,7 +122,7 @@ mod _contextvars {
)
})
})?;
zelf.inner.entered.set(false);
zelf.inner.entered.store(false, Ordering::Release);

Ok(())
}
Expand All @@ -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()
Expand Down Expand Up @@ -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),
},
}
}
Expand All @@ -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 {
Expand Down Expand Up @@ -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 {
Expand All @@ -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)
Expand All @@ -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()));
Expand All @@ -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 {
Expand All @@ -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
Expand Down Expand Up @@ -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)?
Expand All @@ -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);
Expand Down Expand Up @@ -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())
}

Expand All @@ -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))
}
}

Expand All @@ -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]
Expand Down Expand Up @@ -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="));
Expand Down
Loading
Loading