diff --git a/crates/vm/src/builtins/bytearray.rs b/crates/vm/src/builtins/bytearray.rs index fc046460b70..ac7e071f511 100644 --- a/crates/vm/src/builtins/bytearray.rs +++ b/crates/vm/src/builtins/bytearray.rs @@ -15,15 +15,12 @@ use crate::{ ByteInnerTranslateOptions, ByteOr, BytesJoin, DecodeArgs, PyBytesInner, bytes_decode, }, class::PyClassImpl, - common::{ - atomic::{AtomicUsize, Ordering}, - lock::{ - PyDetachingRwLock, PyDetachingRwLockReadGuard, PyDetachingRwLockWriteGuard, - PyMappedDetachingRwLockReadGuard, PyMappedDetachingRwLockWriteGuard, PyMutex, - }, + common::lock::{ + PyDetachingRwLock, PyDetachingRwLockReadGuard, PyDetachingRwLockWriteGuard, + PyMappedDetachingRwLockReadGuard, PyMappedDetachingRwLockWriteGuard, PyMutex, }, convert::{ToPyObject, ToPyResult}, - function::{ArgBytesLike, FuncArgs, PyComparisonValue, PySsize}, + function::{ArgBytesLike, FuncArgs, OptionalArg, PyComparisonValue, PySsize}, protocol::{ BufferDescriptor, BufferFlags, BufferMethods, BufferResizeGuard, PyBuffer, PyIterReturn, PyMappingMethods, PyNumberMethods, PySequenceMethods, @@ -37,12 +34,62 @@ use crate::{ }; use bstr::ByteSlice; use core::mem::size_of; +use num_traits::ToPrimitive; #[pyclass(module = false, name = "bytearray", unhashable = true)] #[derive(Debug, Default)] pub struct PyByteArray { inner: PyDetachingRwLock, - exports: AtomicUsize, + exports: PyMutex, +} + +// While count is nonzero, len is the stable logical length of the payload. +// Publish both together; the length is meaningless between export epochs. +#[derive(Debug, Default)] +struct ExportState { + count: usize, + len: usize, +} + +impl ExportState { + fn snapshot(&self) -> Option { + (self.count != 0).then_some(self.len) + } + + fn retain_existing(&mut self) -> Option { + let len = self.snapshot()?; + self.count += 1; + Some(len) + } + + // The caller holds a payload read guard, so first publication cannot + // overlap a resize and an existing epoch has the same length. + fn retain(&mut self, len: usize) { + if self.count == 0 { + self.len = len; + } + self.count += 1; + } + + fn release(&mut self) { + self.count -= 1; + } +} + +struct ByteArrayExportGuard<'a> { + owner: &'a PyByteArray, + len: usize, +} + +impl Drop for ByteArrayExportGuard<'_> { + fn drop(&mut self) { + self.owner.exports.lock().release(); + } +} + +enum ResizeAccess<'a> { + Exclusive(PyDetachingRwLockWriteGuard<'a, PyBytesInner>), + Exported { len: usize }, } pub(crate) type PyByteArrayRef = PyRef; @@ -95,7 +142,7 @@ impl PyByteArray { const fn from_inner(inner: PyBytesInner) -> Self { Self { inner: PyDetachingRwLock::new(inner), - exports: AtomicUsize::new(0), + exports: PyMutex::new(ExportState { count: 0, len: 0 }), } } @@ -103,15 +150,54 @@ impl PyByteArray { PyDetachingRwLockReadGuard::map(self.inner.read(), |inner| &*inner.elements) } + /// Borrow for in-place byte replacement only. Do not change the length, + /// replace or reallocate storage, or hold this guard across Python calls. + /// Operations that may resize or reallocate must use export-checked access. pub fn borrow_buf_mut(&self) -> PyMappedDetachingRwLockWriteGuard<'_, Vec> { PyDetachingRwLockWriteGuard::map(self.inner.write(), |inner| &mut inner.elements) } + fn begin_export(&self) -> ByteArrayExportGuard<'_> { + // Existing epochs need no payload lock, including in codec callbacks + // that already hold a payload read guard. + let existing_len = { self.exports.lock().retain_existing() }; + let len = match existing_len { + Some(len) => len, + None => { + // Never wait for the payload with the metadata lock held. + // An earlier resize must finish before first publication. + let inner = self.inner.read(); + let len = inner.elements.len(); + self.exports.lock().retain(len); + len + } + }; + ByteArrayExportGuard { owner: self, len } + } + + fn resize_access(&self) -> ResizeAccess<'_> { + // An export is a borrow someone else still holds, so it is answered + // before the lock rather than by waiting on it. + let exported_len = { self.exports.lock().snapshot() }; + if let Some(len) = exported_len { + return ResizeAccess::Exported { len }; + } + let inner = self.inner.write(); + // An export may have started while we waited. Capture its length in + // the same decision, without borrowing the payload a second time. + let exported_len = { self.exports.lock().snapshot() }; + match exported_len { + Some(len) => { + drop(inner); + ResizeAccess::Exported { len } + } + None => ResizeAccess::Exclusive(inner), + } + } + fn with_export(&self, f: impl FnOnce(usize) -> R) -> R { - self.exports.fetch_add(1, Ordering::Release); - scopeguard::defer! { self.exports.fetch_sub(1, Ordering::Release); } - // Reading the length also waits for any earlier resize to finish. - f(self.__len__()) + let export = self.begin_export(); + f(export.len) } fn buffer_snapshot(buffer: &PyBuffer, vm: &VirtualMachine) -> PyResult> { @@ -333,8 +419,7 @@ impl PyByteArray { return Ok(self.borrow_buf().contains(&byte)); } // Converting the needle runs Python, which must not resize this bytearray. - self.exports.fetch_add(1, Ordering::Release); - scopeguard::defer! { self.exports.fetch_sub(1, Ordering::Release); } + let _export = self.begin_export(); let needle = ByteInnerSub::from_contains_arg(needle, vm)?; self.inner().contains(needle, vm) } @@ -502,6 +587,64 @@ impl Py { Ok(()) } + #[pymethod] + fn take_bytes(&self, n: OptionalArg, vm: &VirtualMachine) -> PyResult { + // __index__ can resize self or acquire/release exports. Convert before + // observing either the payload or the export epoch. + let n = match n.into_option().filter(|n| !vm.is_none(n)) { + None => None, + Some(n) => { + let index = n + .try_index_opt(vm) + .ok_or_else(|| vm.new_type_error("n must be an integer or None"))??; + Some(index.as_bigint().to_isize().ok_or_else(|| { + vm.new_index_error(format!( + "cannot fit '{}' into an index-sized integer", + n.class().slot_name() + )) + })?) + } + }; + let access = self.resize_access(); + let len = match &access { + ResizeAccess::Exclusive(inner) => inner.elements.len(), + ResizeAccess::Exported { len } => *len, + }; + // Vec lengths fit isize. Adding a nonnegative length to a negative + // index cannot overflow; preserve the adjusted index in the error. + let to_take = match n { + None => len as isize, + Some(n) if n < 0 => n + len as isize, + Some(n) => n, + }; + if to_take < 0 || to_take as usize > len { + drop(access); + return Err( + vm.new_index_error(format!("can't take {to_take} bytes outside size {len}")) + ); + } + let ResizeAccess::Exclusive(mut inner) = access else { + return Err(vm.new_buffer_error("Existing exports of data: object cannot be re-sized")); + }; + let to_take = to_take as usize; + let result = if to_take == 0 { + Vec::new() + } else if to_take == len { + core::mem::take(&mut inner.elements) + } else { + let mut result = Vec::new(); + if result.try_reserve_exact(to_take).is_err() { + drop(inner); + return Err(vm.no_memory_error()); + } + result.extend_from_slice(&inner.elements[..to_take]); + inner.elements.drain(..to_take); + result + }; + drop(inner); + Ok(PyBytes::from(result)) + } + #[pymethod] fn __reduce_ex__( &self, @@ -525,7 +668,7 @@ impl Py { #[cfg(debug_assertions)] #[pygetset] fn exports(&self) -> usize { - self.exports.load(Ordering::Relaxed) + self.exports.lock().count } #[pymethod] @@ -620,8 +763,7 @@ impl Py { // is borrowed, with the buffer exported so that resizing this // bytearray from there raises BufferError. let (sep, bytes_per_sep) = { - self.exports.fetch_add(1, Ordering::Release); - scopeguard::defer! { self.exports.fetch_sub(1, Ordering::Release); } + let _export = self.begin_export(); options.resolve(vm)? }; Ok(self.inner().hex(sep, bytes_per_sep)) @@ -687,9 +829,8 @@ impl Py { fn join(&self, iterable_of_bytes: PyObjectRef, vm: &VirtualMachine) -> PyResult { // Export before driving the iterable: resizing is forbidden, but // same-size mutations must be reflected in the joined result. - self.exports.fetch_add(1, Ordering::Release); - scopeguard::defer! { self.exports.fetch_sub(1, Ordering::Release); } - let separator_len = self.__len__(); + let _export = self.begin_export(); + let separator_len = _export.len; let items = BytesJoin::new(iterable_of_bytes, vm)?; Ok(items .join(separator_len, || self.borrow_buf().into(), vm)? @@ -801,8 +942,7 @@ impl Py { fn split(&self, args: FuncArgs, vm: &VirtualMachine) -> PyResult> { // Converting the separator runs Python, which must not resize this // bytearray, so the arguments are bound with the buffer exported. - self.exports.fetch_add(1, Ordering::Release); - scopeguard::defer! { self.exports.fetch_sub(1, Ordering::Release); } + let _export = self.begin_export(); let options: ByteInnerSplitOptions = args.bind_for(vm, "split")?; self.inner() .split(options, |s, vm| vm.ctx.new_bytearray(s.to_vec()).into(), vm) @@ -810,8 +950,7 @@ impl Py { #[pymethod] fn rsplit(&self, args: FuncArgs, vm: &VirtualMachine) -> PyResult> { - self.exports.fetch_add(1, Ordering::Release); - scopeguard::defer! { self.exports.fetch_sub(1, Ordering::Release); } + let _export = self.begin_export(); let options: ByteInnerSplitOptions = args.bind_for(vm, "rsplit")?; self.inner() .rsplit(options, |s, vm| vm.ctx.new_bytearray(s.to_vec()).into(), vm) @@ -1001,17 +1140,14 @@ static BUFFER_METHODS: BufferMethods = BufferMethods { ) .into() }, - release: |buffer| { - buffer - .obj_as::() - .exports - .fetch_sub(1, Ordering::Release); - }, + release: |buffer| buffer.obj_as::().exports.lock().release(), retain: |buffer| { + // as_buffer holds the payload read guard until PyBuffer::new returns. buffer .obj_as::() .exports - .fetch_add(1, Ordering::Release); + .lock() + .retain(buffer.desc.len); }, }; @@ -1045,14 +1181,10 @@ impl BufferResizeGuard for PyByteArray { type Resizable<'a> = PyDetachingRwLockWriteGuard<'a, PyBytesInner>; fn try_resizable_opt(&self) -> Option> { - // An export is a borrow someone else still holds, so it is answered - // before the lock rather than by waiting on it. - if self.exports.load(Ordering::SeqCst) != 0 { - return None; + match self.resize_access() { + ResizeAccess::Exclusive(inner) => Some(inner), + ResizeAccess::Exported { .. } => None, } - let inner = self.inner.write(); - // An export may have been acquired while we waited for the lock. - (self.exports.load(Ordering::SeqCst) == 0).then_some(inner) } } @@ -1241,3 +1373,44 @@ impl IterNext for PyByteArrayIterator { }) } } + +#[cfg(test)] +mod export_state_tests { + use super::ExportState; + + #[test] + fn export_epoch_length_is_stable_until_last_release() { + let mut state = ExportState::default(); + assert_eq!(state.snapshot(), None); + assert_eq!(state.retain_existing(), None); + state.retain(3); + assert_eq!(state.snapshot(), Some(3)); + assert_eq!(state.retain_existing(), Some(3)); + state.retain(3); + assert_eq!(state.count, 3); + state.release(); + state.release(); + assert_eq!(state.count, 1); + assert_eq!(state.snapshot(), Some(3)); + let rejected_length = state.snapshot(); + state.release(); + assert_eq!(state.snapshot(), None); + state.retain(7); + assert_eq!(state.snapshot(), Some(7)); + assert_eq!(rejected_length, Some(3)); + state.release(); + assert_eq!(state.snapshot(), None); + } + + #[test] + fn empty_export_is_distinct_from_no_export() { + let mut state = ExportState::default(); + state.retain(0); + assert_eq!(state.snapshot(), Some(0)); + assert_eq!(state.retain_existing(), Some(0)); + state.release(); + assert_eq!(state.snapshot(), Some(0)); + state.release(); + assert_eq!(state.snapshot(), None); + } +} diff --git a/extra_tests/snippets/builtin_bytearray.py b/extra_tests/snippets/builtin_bytearray.py index 1ea800b64db..dbee9a87a7a 100644 --- a/extra_tests/snippets/builtin_bytearray.py +++ b/extra_tests/snippets/builtin_bytearray.py @@ -1587,3 +1587,167 @@ def __release_buffer__(self, view): test_bytearray_contains_builtin_and_callback_exports() + + +def test_bytearray_take_bytes_bounded(): + for n, expected, remaining in ( + (None, b"abcd", b""), + (0, b"", b"abcd"), + (1, b"a", b"bcd"), + (2, b"ab", b"cd"), + (4, b"abcd", b""), + (-1, b"abc", b"d"), + (-4, b"", b"abcd"), + (True, b"a", b"bcd"), + ): + target = bytearray(b"abcd") + result = target.take_bytes(n) + assert type(result) is bytes + assert result == expected + assert target == remaining + target[:] = b"z" + assert result == expected + target = bytearray(b"abc") + assert target.take_bytes() == b"abc" + assert target == b"" + for n in (None, 0): + assert bytearray().take_bytes(n) == b"" + for n in (5, -5, sys.maxsize + 1, -sys.maxsize - 2): + target = bytearray(b"abcd") + with assert_raises(IndexError): + target.take_bytes(n) + assert target == b"abcd" + for n in (1.0, "1", b"1", object()): + with assert_raises(TypeError): + bytearray(b"a").take_bytes(n) + with assert_raises(TypeError): + bytearray(b"a").take_bytes(n=1) + with assert_raises(TypeError): + bytearray(b"a").take_bytes(1, 2) + for content in (b"", b"abcd"): + target = bytearray(content) + view = memoryview(target) + try: + for n in (None, 0, len(content)): + with assert_raises(BufferError): + target.take_bytes(n) + assert target == content + for n in (len(content) + 1, -len(content) - 1): + with assert_raises(IndexError): + target.take_bytes(n) + assert target == content + finally: + view.release() + assert target.take_bytes() == content + target = bytearray(b"abcd") + parent = memoryview(target) + child = parent[1:3] + clone = memoryview(parent) + parent.release() + try: + with assert_raises(BufferError): + target.take_bytes(0) + child.release() + with assert_raises(BufferError): + target.take_bytes(1) + finally: + child.release() + clone.release() + assert target.take_bytes() == b"abcd" + + +def test_bytearray_take_bytes_reentrant_bounded(): + class Index: + def __init__(self, action, result): + self.action = action + self.result = result + + def __index__(self): + self.action() + return self.result + + for replacement, n, expected, remaining in ( + (b"ab", 2, b"ab", b""), + (b"", 0, b"", b""), + (b"abcdef", 5, b"abcde", b"f"), + (b"ab", -1, b"a", b"b"), + ): + target = bytearray(b"abcd") + index = Index(lambda: target.__setitem__(slice(None), replacement), n) + assert target.take_bytes(index) == expected + assert target == remaining + target = bytearray(b"abcd") + with assert_raises(IndexError): + target.take_bytes(Index(target.clear, 1)) + assert target == b"" + for n, error in ((1, BufferError), (5, IndexError)): + target = bytearray(b"abcd") + views = [] + try: + with assert_raises(error): + target.take_bytes(Index(lambda: views.append(memoryview(target)), n)) + assert target == b"abcd" + finally: + for view in views: + view.release() + for n, error in ((2, None), (5, IndexError)): + target = bytearray(b"abcd") + view = memoryview(target) + try: + index = Index(view.release, n) + if error is None: + assert target.take_bytes(index) == b"ab" + assert target == b"cd" + else: + with assert_raises(error): + target.take_bytes(index) + assert target == b"abcd" + finally: + view.release() + target = bytearray(b"abcd") + with assert_raises(TypeError): + target.take_bytes(Index(lambda: None, "wrong")) + assert target == b"abcd" + + def raises(): + raise ValueError("index callback") + + with assert_raises(ValueError): + target.take_bytes(Index(raises, 1)) + assert target == b"abcd" + + class MembershipIndex: + def __index__(self): + with assert_raises(IndexError): + target.take_bytes(5) + with assert_raises(BufferError): + target.take_bytes(0) + return ord("a") + + assert MembershipIndex() in target + assert target.take_bytes() == b"abcd" + + +def test_bytearray_take_bytes_codec_callback_bounded(): + import codecs + + target = bytearray(b"\xff") + calls = [] + + def handler(error): + with assert_raises(IndexError): + target.take_bytes(2) + with assert_raises(BufferError): + target.take_bytes(0) + calls.append(error.start) + return ("?", error.end) + + codecs.register_error("bytearray_take_bytes_bounded", handler) + assert target.decode("ascii", "bytearray_take_bytes_bounded") == "?" + assert calls == [0] + assert target.take_bytes() == b"\xff" + + +test_bytearray_take_bytes_bounded() +test_bytearray_take_bytes_reentrant_bounded() +test_bytearray_take_bytes_codec_callback_bounded()