diff --git a/crates/stdlib/src/pickle.rs b/crates/stdlib/src/pickle.rs index 07f29f2a1..a1ea76eb9 100644 --- a/crates/stdlib/src/pickle.rs +++ b/crates/stdlib/src/pickle.rs @@ -15,6 +15,7 @@ mod _pickle { PyInt, PyList, PySet, PyStr, PyTuple, PyTupleRef, PyType, PyTypeRef, }, function::{FuncArgs, OptionalArg, PySetterValue}, + object::{Traverse, TraverseFn}, protocol::{PyBuffer, PyIter, PyIterReturn}, types::{AsBuffer, Constructor, Initializer, Representable}, }; @@ -250,6 +251,15 @@ mod _pickle { frame_end: Option, } + // SAFETY: these are all the Python references owned by the input state. + unsafe impl Traverse for ReadState { + fn traverse(&self, traverse_fn: &mut TraverseFn<'_>) { + self.read.traverse(traverse_fn); + self.readline.traverse(traverse_fn); + self.peek.traverse(traverse_fn); + } + } + fn to_byte_vec(obj: &PyObject, what: &str, vm: &VirtualMachine) -> PyResult> { let bytes = obj .downcast_ref::() @@ -418,7 +428,7 @@ mod _pickle { // memo proxies #[pyattr] - #[pyclass(module = "_pickle", name = "UnpicklerMemoProxy")] + #[pyclass(module = "_pickle", name = "UnpicklerMemoProxy", traverse)] #[derive(Debug, PyPayload)] struct UnpicklerMemoProxy { unpickler: PyRef, @@ -484,6 +494,13 @@ mod _pickle { buffers: Option, } + // SAFETY: buffers is the only Python reference owned by the configuration. + unsafe impl Traverse for UnpicklerConfig { + fn traverse(&self, traverse_fn: &mut TraverseFn<'_>) { + self.buffers.traverse(traverse_fn); + } + } + impl Default for UnpicklerConfig { fn default() -> Self { Self { @@ -498,7 +515,7 @@ mod _pickle { } #[pyattr] - #[pyclass(module = "_pickle", name = "Unpickler")] + #[pyclass(module = "_pickle", name = "Unpickler", traverse = "manual")] #[derive(Debug, PyPayload)] pub(super) struct PyUnpickler { read_state: PyMutex, @@ -506,6 +523,34 @@ mod _pickle { config: PyRwLock, } + // SAFETY: visit each owned reference once, including sparse memo entries. + unsafe impl Traverse for PyUnpickler { + fn traverse(&self, traverse_fn: &mut TraverseFn<'_>) { + self.read_state.traverse(traverse_fn); + self.config.traverse(traverse_fn); + if let Some(memo) = self.memo.try_read_recursive() { + #[expect( + clippy::iter_over_hash_type, + reason = "GC tracing visits every memo reference regardless of index order" + )] + for value in memo.values() { + value.traverse(traverse_fn); + } + } + } + + fn clear(&mut self, out: &mut Vec) { + let state = self.read_state.get_mut(); + out.extend(state.read.take()); + out.extend(state.readline.take()); + out.extend(state.peek.take()); + out.extend(self.memo.get_mut().drain().map(|(_, value)| value)); + let config = self.config.get_mut(); + out.extend(config.buffers.take()); + config.initialized = false; + } + } + #[derive(FromArgs)] pub(super) struct UnpicklerNewArgs { #[pyarg(any)] @@ -576,23 +621,29 @@ mod _pickle { return Err(vm.new_type_error("file must have 'read' and 'readline' attributes")); } - *state = ReadState { - read, - readline, - peek, - ..ReadState::default() - }; + let old_state = core::mem::replace( + &mut *state, + ReadState { + read, + readline, + peek, + ..ReadState::default() + }, + ); let old_memo = core::mem::take(&mut *zelf.memo.write()); - *zelf.config.write() = UnpicklerConfig { - initialized: true, - proto: 0, - fix_imports: args.fix_imports, - encoding, - errors, - buffers, - }; + let old_config = core::mem::replace( + &mut *zelf.config.write(), + UnpicklerConfig { + initialized: true, + proto: 0, + fix_imports: args.fix_imports, + encoding, + errors, + buffers, + }, + ); drop(state); - drop(old_memo); + drop((old_state, old_memo, old_config)); Ok(()) } } diff --git a/extra_tests/snippets/stdlib_pickle.py b/extra_tests/snippets/stdlib_pickle.py index 999106d42..ecafea1b1 100644 --- a/extra_tests/snippets/stdlib_pickle.py +++ b/extra_tests/snippets/stdlib_pickle.py @@ -1,6 +1,8 @@ import _pickle +import gc import io import sys +import weakref from types import MappingProxyType, MethodType from testutils import assert_raises @@ -116,3 +118,42 @@ namespace = { with assert_raises(RuntimeError) as caught: exec("method.__reduce__()", namespace) assert caught.exception is name_error + + +# Unpickler payloads can retain a cycle through their memo proxy. +class UnpicklerCycleStream: + def __init__(self): + self.source = io.BytesIO(_pickle.dumps([1, 2, 3], protocol=4)) + self.read = self.source.read + self.readline = self.source.readline + self.peek = self.peek_bytes + + def peek_bytes(self, size=0): + position = self.source.tell() + data = self.source.read(size) + self.source.seek(position) + return data + + +class UnpicklerCycleBuffers: + def __iter__(self): + return self + + def __next__(self): + raise StopIteration + + +def unpickler_payload_cycle(kind): + stream = UnpicklerCycleStream() + buffers = UnpicklerCycleBuffers() + unpickler = _pickle.Unpickler(stream, buffers=buffers) + holder = stream if kind == "stream" else buffers + holder.proxy = unpickler.memo + return weakref.ref(holder) + + +for kind in ("stream", "buffers"): + reference = unpickler_payload_cycle(kind) + for _ in range(3): + gc.collect() + assert reference() is None, kind