Skip to content

Commit 259ce8c

Browse files
youknowoneyouknowone
authored
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
1 parent 00461a0 commit 259ce8c

12 files changed

Lines changed: 521 additions & 190 deletions

‎Lib/test/test_itertools.py‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1270,7 +1270,6 @@ def test_dropwhile(self):
12701270
self.assertRaises(TypeError, next, dropwhile(10, [(4,5)]))
12711271
self.assertRaises(ValueError, next, dropwhile(errfunc, [(4,5)]))
12721272

1273-
@unittest.expectedFailure # TODO: RUSTPYTHON
12741273
def test_tee(self):
12751274
n = 200
12761275

‎crates/stdlib/src/_asyncio.rs‎

Lines changed: 13 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2745,16 +2745,20 @@ pub(crate) mod _asyncio {
27452745
}
27462746
}
27472747

2748+
fn get_invalid_state_error_type(vm: &VirtualMachine) -> PyResult<PyTypeRef> {
2749+
let module = vm.import("asyncio.exceptions", 0)?;
2750+
let exc_type = vm
2751+
.get_attribute_opt(module, vm.ctx.intern_str("InvalidStateError"))?
2752+
.ok_or_else(|| vm.new_attribute_error("InvalidStateError not found"))?;
2753+
exc_type
2754+
.downcast()
2755+
.map_err(|_| vm.new_type_error("InvalidStateError is not a type"))
2756+
}
2757+
27482758
fn new_invalid_state_error(vm: &VirtualMachine, msg: &str) -> PyBaseExceptionRef {
2749-
match vm.import("asyncio.exceptions", 0) {
2750-
Ok(module) => {
2751-
match vm.get_attribute_opt(module, vm.ctx.intern_str("InvalidStateError")) {
2752-
Ok(Some(exc_type)) => match exc_type.call((msg,), vm) {
2753-
Ok(exc) => exc.downcast().unwrap(),
2754-
Err(_) => vm.new_runtime_error(msg.to_string()),
2755-
},
2756-
_ => vm.new_runtime_error(msg.to_string()),
2757-
}
2759+
match get_invalid_state_error_type(vm) {
2760+
Ok(invalid_state_error) => {
2761+
vm.new_exception_msg(invalid_state_error, msg.to_string().into())
27582762
}
27592763
Err(_) => vm.new_runtime_error(msg.to_string()),
27602764
}

‎crates/stdlib/src/contextvars.rs‎

Lines changed: 77 additions & 72 deletions
Original file line numberDiff line numberDiff line change
@@ -15,16 +15,16 @@ mod _contextvars {
1515
AsObject, Py, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, atomic_func,
1616
builtins::{PyGenericAlias, PyList, PyStrRef, PyType, PyTypeRef},
1717
class::StaticType,
18-
common::{hash::PyHash, lock::LazyLock, wtf8::Wtf8Buf},
18+
common::{
19+
hash::PyHash,
20+
lock::{LazyLock, PyMutex},
21+
wtf8::Wtf8Buf,
22+
},
1923
function::{ArgCallable, FuncArgs, OptionalArg},
2024
protocol::{PyMappingMethods, PySequenceMethods},
2125
types::{AsMapping, AsSequence, Constructor, Hashable, Iterable, Representable},
2226
};
23-
use core::{
24-
cell::{Cell, RefCell, UnsafeCell},
25-
sync::atomic::Ordering,
26-
};
27-
use crossbeam_utils::atomic::AtomicCell;
27+
use core::sync::atomic::{AtomicBool, AtomicI64, AtomicUsize, Ordering};
2828
use indexmap::IndexMap;
2929

3030
// TODO: Real hamt implementation
@@ -33,7 +33,7 @@ mod _contextvars {
3333
#[pyclass(no_attr, name = "Hamt", module = "contextvars")]
3434
#[derive(Debug, PyPayload)]
3535
pub(crate) struct HamtObject {
36-
hamt: RefCell<Hamt>,
36+
hamt: PyMutex<Hamt>,
3737
}
3838

3939
#[pyclass]
@@ -42,23 +42,19 @@ mod _contextvars {
4242
impl Default for HamtObject {
4343
fn default() -> Self {
4444
Self {
45-
hamt: RefCell::new(Hamt::default()),
45+
hamt: PyMutex::new(Hamt::default()),
4646
}
4747
}
4848
}
4949

50-
unsafe impl Sync for HamtObject {}
51-
5250
#[derive(Debug)]
5351
struct ContextInner {
54-
idx: Cell<usize>,
52+
idx: AtomicUsize,
5553
vars: PyRef<HamtObject>,
5654
// PyObject *ctx_weakreflist;
57-
entered: Cell<bool>,
55+
entered: AtomicBool,
5856
}
5957

60-
unsafe impl Sync for ContextInner {}
61-
6258
#[pyattr]
6359
#[pyclass(name = "Context")]
6460
#[derive(Debug, PyPayload)]
@@ -71,40 +67,46 @@ mod _contextvars {
7167
fn empty(vm: &VirtualMachine) -> Self {
7268
Self {
7369
inner: ContextInner {
74-
idx: Cell::new(usize::MAX),
70+
idx: AtomicUsize::new(usize::MAX),
7571
vars: HamtObject::default().into_ref(&vm.ctx),
76-
entered: Cell::new(false),
72+
entered: AtomicBool::new(false),
7773
},
7874
}
7975
}
8076

81-
fn borrow_vars(&self) -> impl core::ops::Deref<Target = Hamt> + '_ {
82-
self.inner.vars.hamt.borrow()
77+
fn borrow_vars(&self) -> impl core::ops::DerefMut<Target = Hamt> + '_ {
78+
self.inner.vars.hamt.lock()
8379
}
8480

8581
fn borrow_vars_mut(&self) -> impl core::ops::DerefMut<Target = Hamt> + '_ {
86-
self.inner.vars.hamt.borrow_mut()
82+
self.inner.vars.hamt.lock()
8783
}
8884

8985
fn enter(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<()> {
90-
if zelf.inner.entered.get() {
86+
// A context is entered by one thread at a time, so the check and the
87+
// claim have to be a single step.
88+
if zelf
89+
.inner
90+
.entered
91+
.compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
92+
.is_err()
93+
{
9194
return Err(vm.new_runtime_error(format!(
9295
"cannot enter context: {} is already entered",
9396
zelf.as_object().repr(vm)?
9497
)));
9598
}
9699

97100
super::CONTEXTS.with_borrow_mut(|ctxs| {
98-
zelf.inner.idx.set(ctxs.len());
101+
zelf.inner.idx.store(ctxs.len(), Ordering::Relaxed);
99102
ctxs.push(zelf.to_owned());
100103
});
101-
zelf.inner.entered.set(true);
102104

103105
Ok(())
104106
}
105107

106108
fn exit(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<()> {
107-
if !zelf.inner.entered.get() {
109+
if !zelf.inner.entered.load(Ordering::Acquire) {
108110
return Err(vm.new_runtime_error(format!(
109111
"cannot exit context: {} is not entered",
110112
zelf.as_object().repr(vm)?
@@ -120,7 +122,7 @@ mod _contextvars {
120122
)
121123
})
122124
})?;
123-
zelf.inner.entered.set(false);
125+
zelf.inner.entered.store(false, Ordering::Release);
124126

125127
Ok(())
126128
}
@@ -131,8 +133,8 @@ mod _contextvars {
131133
ctx.clone()
132134
} else {
133135
let ctx = Self::empty(vm);
134-
ctx.inner.idx.set(0);
135-
ctx.inner.entered.set(true);
136+
ctx.inner.idx.store(0, Ordering::Relaxed);
137+
ctx.inner.entered.store(true, Ordering::Release);
136138
let ctx = ctx.into_ref(&vm.ctx);
137139
ctxs.push(ctx);
138140
ctxs[0].clone()
@@ -170,13 +172,13 @@ mod _contextvars {
170172
fn copy(&self, vm: &VirtualMachine) -> Self {
171173
// Deep copy the vars - clone the underlying Hamt data, not just the PyRef
172174
let vars_copy = HamtObject {
173-
hamt: RefCell::new(self.inner.vars.hamt.borrow().clone()),
175+
hamt: PyMutex::new(self.inner.vars.hamt.lock().clone()),
174176
};
175177
Self {
176178
inner: ContextInner {
177-
idx: Cell::new(usize::MAX),
179+
idx: AtomicUsize::new(usize::MAX),
178180
vars: vars_copy.into_ref(&vm.ctx),
179-
entered: Cell::new(false),
181+
entered: AtomicBool::new(false),
180182
},
181183
}
182184
}
@@ -186,11 +188,8 @@ mod _contextvars {
186188
var: PyRef<ContextVar>,
187189
vm: &VirtualMachine,
188190
) -> PyResult<PyObjectRef> {
189-
let vars = self.borrow_vars();
190-
let item = vars
191-
.get(&*var)
192-
.ok_or_else(|| vm.new_key_error(var.into()))?;
193-
Ok(item.to_owned())
191+
let item = self.borrow_vars().get(&*var).map(|item| item.to_owned());
192+
item.ok_or_else(|| vm.new_key_error(var.into()))
194193
}
195194

196195
fn __len__(&self) -> usize {
@@ -290,11 +289,11 @@ mod _contextvars {
290289
name: String,
291290
default: Option<PyObjectRef>,
292291
#[pytraverse(skip)]
293-
cached: AtomicCell<Option<ContextVarCache>>,
292+
cached: PyMutex<Option<ContextVarCache>>,
294293
#[pytraverse(skip)]
295-
cached_id: core::sync::atomic::AtomicUsize, // cached_tsid in CPython
294+
cached_id: AtomicUsize, // cached_tsid in CPython
296295
#[pytraverse(skip)]
297-
hash: UnsafeCell<PyHash>,
296+
hash: AtomicI64,
298297
}
299298

300299
impl core::fmt::Debug for ContextVar {
@@ -303,8 +302,6 @@ mod _contextvars {
303302
}
304303
}
305304

306-
unsafe impl Sync for ContextVar {}
307-
308305
impl PartialEq for ContextVar {
309306
fn eq(&self, other: &Self) -> bool {
310307
core::ptr::eq(self, other)
@@ -320,12 +317,15 @@ mod _contextvars {
320317

321318
impl ContextVar {
322319
fn delete(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<()> {
323-
zelf.cached.store(None);
320+
let cached = zelf.cached.lock().take();
321+
drop(cached);
324322

325323
let ctx = PyContext::current(vm);
326324

327-
let mut vars = ctx.borrow_vars_mut();
328-
if vars.swap_remove(zelf).is_none() {
325+
let removed = ctx.borrow_vars_mut().swap_remove(zelf);
326+
let existed = removed.is_some();
327+
drop(removed);
328+
if !existed {
329329
// TODO:
330330
// PyErr_SetObject(PyExc_LookupError, (PyObject *)var);
331331
return Err(vm.new_lookup_error(zelf.as_object().repr(vm)?.as_wtf8().to_owned()));
@@ -338,16 +338,17 @@ mod _contextvars {
338338
fn set_inner(zelf: &Py<Self>, value: PyObjectRef, vm: &VirtualMachine) {
339339
let ctx = PyContext::current(vm);
340340

341-
let mut vars = ctx.borrow_vars_mut();
342-
vars.insert(zelf.to_owned(), value.clone());
341+
let replaced = ctx.borrow_vars_mut().insert(zelf.to_owned(), value.clone());
342+
drop(replaced);
343343

344344
zelf.cached_id.store(ctx.get_id(), Ordering::SeqCst);
345345

346346
let cache = ContextVarCache {
347347
object: value,
348-
idx: ctx.inner.idx.get(),
348+
idx: ctx.inner.idx.load(Ordering::Relaxed),
349349
};
350-
zelf.cached.store(Some(cache));
350+
let replaced = zelf.cached.lock().replace(cache);
351+
drop(replaced);
351352
}
352353

353354
fn generate_hash(zelf: &Py<Self>, vm: &VirtualMachine) -> PyHash {
@@ -370,28 +371,32 @@ mod _contextvars {
370371
default: OptionalArg<PyObjectRef>,
371372
vm: &VirtualMachine,
372373
) -> PyResult<Option<PyObjectRef>> {
373-
let found = super::CONTEXTS.with_borrow(|ctxs| {
374-
let ctx = ctxs.last()?;
375-
let cached_ptr = zelf.cached.as_ptr();
376-
debug_assert!(!cached_ptr.is_null());
377-
if let Some(cached) = unsafe { &*cached_ptr }
374+
// The replaced cache entry comes back out so that dropping it, which
375+
// can run a __del__ that calls back in, happens with no lock held.
376+
let (found, replaced) = super::CONTEXTS.with_borrow(|ctxs| {
377+
let Some(ctx) = ctxs.last() else {
378+
return (None, None);
379+
};
380+
let mut cached = zelf.cached.lock();
381+
if let Some(cached) = &*cached
378382
&& zelf.cached_id.load(Ordering::SeqCst) == ctx.get_id()
379383
&& cached.idx + 1 == ctxs.len()
380384
{
381-
return Some(cached.object.clone());
385+
return (Some(cached.object.clone()), None);
382386
}
383-
let vars = ctx.borrow_vars();
384-
let obj = vars.get(zelf)?;
387+
let Some(obj) = ctx.borrow_vars().get(zelf).map(|obj| obj.to_owned()) else {
388+
return (None, None);
389+
};
385390
zelf.cached_id.store(ctx.get_id(), Ordering::SeqCst);
386391

387-
// TODO: ensure cached is not changed
388-
let _removed = zelf.cached.swap(Some(ContextVarCache {
392+
let replaced = cached.replace(ContextVarCache {
389393
object: obj.clone(),
390394
idx: ctxs.len() - 1,
391-
}));
395+
});
392396

393-
Some(obj.clone())
397+
(Some(obj), replaced)
394398
});
399+
drop(replaced);
395400

396401
let value = if let Some(value) = found {
397402
value
@@ -425,7 +430,7 @@ mod _contextvars {
425430

426431
#[pymethod]
427432
fn reset(zelf: &Py<Self>, token: PyRef<ContextToken>, vm: &VirtualMachine) -> PyResult<()> {
428-
if token.used.get() {
433+
if token.used.load(Ordering::Acquire) {
429434
return Err(vm.new_runtime_error(format!(
430435
"{} has already been used once",
431436
token.as_object().repr(vm)?
@@ -447,7 +452,7 @@ mod _contextvars {
447452
)));
448453
}
449454

450-
token.used.set(true);
455+
token.used.store(true, Ordering::Release);
451456

452457
if let Some(old_value) = &token.old_value {
453458
Self::set_inner(zelf, old_value.clone(), vm);
@@ -484,15 +489,13 @@ mod _contextvars {
484489
name: args.name.to_string(),
485490
default: args.default.into_option(),
486491
cached_id: 0.into(),
487-
cached: AtomicCell::new(None),
488-
hash: UnsafeCell::new(0),
492+
cached: PyMutex::new(None),
493+
hash: AtomicI64::new(0),
489494
};
490495
let py_var = var.into_ref_with_type(vm, cls)?;
491496

492-
unsafe {
493-
// SAFETY: py_var is not exposed to python memory model yet
494-
*py_var.hash.get() = Self::generate_hash(&py_var, vm)
495-
};
497+
let hash = Self::generate_hash(&py_var, vm);
498+
py_var.hash.store(hash, Ordering::Relaxed);
496499
Ok(py_var.into())
497500
}
498501

@@ -504,14 +507,14 @@ mod _contextvars {
504507
impl core::hash::Hash for ContextVar {
505508
#[inline]
506509
fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
507-
unsafe { *self.hash.get() }.hash(state)
510+
self.hash.load(Ordering::Relaxed).hash(state)
508511
}
509512
}
510513

511514
impl Hashable for ContextVar {
512515
#[inline]
513516
fn hash(zelf: &Py<Self>, _vm: &VirtualMachine) -> PyResult<PyHash> {
514-
Ok(unsafe { *zelf.hash.get() })
517+
Ok(zelf.hash.load(Ordering::Relaxed))
515518
}
516519
}
517520

@@ -537,11 +540,9 @@ mod _contextvars {
537540
ctx: PyRef<PyContext>, // tok_ctx in CPython
538541
var: PyRef<ContextVar>, // tok_var in CPython
539542
old_value: Option<PyObjectRef>, // tok_oldval in CPython
540-
used: Cell<bool>,
543+
used: AtomicBool,
541544
}
542545

543-
unsafe impl Sync for ContextToken {}
544-
545546
#[pyclass(with(Constructor, Representable))]
546547
impl ContextToken {
547548
#[pygetset]
@@ -598,7 +599,11 @@ mod _contextvars {
598599
impl Representable for ContextToken {
599600
#[inline]
600601
fn repr_wtf8(zelf: &Py<Self>, vm: &VirtualMachine) -> PyResult<Wtf8Buf> {
601-
let used = if zelf.used.get() { " used" } else { "" };
602+
let used = if zelf.used.load(Ordering::Acquire) {
603+
" used"
604+
} else {
605+
""
606+
};
602607
let var = Representable::repr_wtf8(&zelf.var, vm)?;
603608
let ptr = zelf.as_object().get_id() as *const u8;
604609
let mut result = Wtf8Buf::from(format!("<Token{used} var="));

0 commit comments

Comments
 (0)