diff --git a/crates/vm/src/stdlib/itertools.rs b/crates/vm/src/stdlib/itertools.rs index e72b4f08d..f09e556eb 100644 --- a/crates/vm/src/stdlib/itertools.rs +++ b/crates/vm/src/stdlib/itertools.rs @@ -1025,47 +1025,30 @@ mod decl { struct PyItertoolsTeeData { iterable: PyIter, values: PyMutex>, + next: PyMutex>>, #[pytraverse(skip)] running: AtomicBool, } #[pyclass(flags(DISALLOW_INSTANTIATION))] impl PyItertoolsTeeData { + const BLOCK_SIZE: usize = 64; + fn new(iterable: PyIter, vm: &VirtualMachine) -> PyRef { Self { iterable, - values: PyMutex::new(vec![]), + values: PyMutex::new(Vec::with_capacity(Self::BLOCK_SIZE)), + next: PyMutex::new(None), running: AtomicBool::new(false), } .into_ref(&vm.ctx) } - fn get_item(&self, vm: &VirtualMachine, index: usize) -> PyResult { - // Return cached value if available - { - let Some(values) = self.values.try_lock() else { - return Err(vm.new_runtime_error("cannot re-enter the tee iterator")); - }; - if index < values.len() { - return Ok(PyIterReturn::Return(values[index].clone())); - } - } - // 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")); - } - 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")); - }; - if values.len() == index { - values.push(obj); - } - Ok(PyIterReturn::Return(values[index].clone())) + fn next_block(&self, vm: &VirtualMachine) -> PyRef { + self.next + .lock() + .get_or_insert_with(|| Self::new(self.iterable.clone(), vm)) + .clone() } } @@ -1073,7 +1056,7 @@ mod decl { #[pyclass(name = "_tee", traverse)] #[derive(Debug, PyPayload)] struct PyItertoolsTee { - tee_data: PyRef, + tee_data: PyMutex>, #[pytraverse(skip)] index: AtomicCell, #[pytraverse(skip)] @@ -1091,7 +1074,7 @@ mod decl { return Ok(tee.__copy__()); } Ok(Self { - tee_data: PyItertoolsTeeData::new(iterator, vm), + tee_data: PyMutex::new(PyItertoolsTeeData::new(iterator, vm)), index: AtomicCell::new(0), advancing: AtomicBool::new(false), }) @@ -1105,7 +1088,7 @@ mod decl { return vm.call_special_method(&iterator, identifier!(vm, __copy__), ()); } Ok(Self { - tee_data: PyItertoolsTeeData::new(iterator, vm), + tee_data: PyMutex::new(PyItertoolsTeeData::new(iterator, vm)), index: AtomicCell::new(0), advancing: AtomicBool::new(false), } @@ -1118,8 +1101,11 @@ mod decl { impl Py { #[pymethod] fn __copy__(&self) -> PyItertoolsTee { + // A block transition resets the index under this lock, so the pair + // is a consistent snapshot even while another thread advances us. + let data = self.tee_data.lock(); PyItertoolsTee { - tee_data: self.tee_data.clone(), + tee_data: PyMutex::new(data.clone()), index: AtomicCell::new(self.index.load()), advancing: AtomicBool::new(false), } @@ -1162,15 +1148,65 @@ mod decl { impl SelfIter for PyItertoolsTee {} impl IterNext for PyItertoolsTee { fn next(zelf: &Py, vm: &VirtualMachine) -> PyResult { + let mut current = zelf.tee_data.lock(); // 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) { + // The block lock serializes claims; an atomic exchange is not needed. + if zelf.advancing.load(Ordering::Acquire) { + drop(current); return Err(vm.new_runtime_error("cannot re-enter the tee iterator")); } + zelf.advancing.store(true, Ordering::Relaxed); scopeguard::defer! { zelf.advancing.store(false, Ordering::Release) } + if zelf.index.load() == PyItertoolsTeeData::BLOCK_SIZE { + let data = current.clone(); + drop(current); + let next = data.next_block(vm); + current = zelf.tee_data.lock(); + let previous = core::mem::replace(&mut *current, next); + zelf.index.store(0); + drop(current); + // Releasing a consumed block may run item finalizers. No cursor + // or buffer lock may be held when the last owner lets it go. + drop(data); + drop(previous); + current = zelf.tee_data.lock(); + } + + // Cached reads neither call Python nor release the block, so they + // can update the cursor without cloning its data reference. let index = zelf.index.load(); - let value = raise_if_stop!(zelf.tee_data.get_item(vm, index)?); + let cached = { + let values = current.values.lock(); + if let Some(value) = values.get(index) { + Some(value.clone()) + } else { + // Claim the source while the cache is locked: another tee + // cannot fill this slot between the check and the claim. + if current.running.load(Ordering::Acquire) { + drop(values); + drop(current); + return Err(vm.new_runtime_error("cannot re-enter the tee iterator")); + } + current.running.store(true, Ordering::Relaxed); + None + } + }; + if let Some(value) = cached { + zelf.index.store(index + 1); + return Ok(PyIterReturn::Return(value)); + } + + let data = current.clone(); + drop(current); + // 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. + scopeguard::defer! { data.running.store(false, Ordering::Release) } + let value = raise_if_stop!(data.iterable.next(vm)?); + data.values.lock().push(value.clone()); zelf.index.store(index + 1); Ok(PyIterReturn::Return(value)) } diff --git a/extra_tests/snippets/stdlib_itertools.py b/extra_tests/snippets/stdlib_itertools.py index 38e213ea9..2eb7ca7bb 100644 --- a/extra_tests/snippets/stdlib_itertools.py +++ b/extra_tests/snippets/stdlib_itertools.py @@ -679,3 +679,22 @@ class PairwiseWithLabel(itertools.pairwise): pairs = PairwiseWithLabel([1, 2], label="pairs") assert pairs.label == "pairs" assert next(pairs) == (1, 2) + + +# Consumed tee values must not remain alive for the lifetime of the iterators. +import weakref + + +class TeeItem: + pass + + +tee_left, tee_right = itertools.tee(TeeItem() for _ in range(200)) +tee_first = next(tee_left) +tee_ref = weakref.ref(tee_first) +del tee_first +next(tee_right) +for _ in range(199): + next(tee_left) + next(tee_right) +assert tee_ref() is None