From f482e0c2cffdb3971e9a3553b874b6831bc0bc15 Mon Sep 17 00:00:00 2001 From: youknowdot Date: Wed, 21 Oct 2026 18:17:06 -0700 Subject: [PATCH] Track Unpickler payload references during cyclic GC (#9288) Traverse input callbacks, buffer iterators, memo slots and memo proxy owners. Clear outgoing payload edges through deferred drops and retire successful reinitialization state outside the state lock. Retain stream and buffers payload-cycle regressions in the existing pickle snippet. Assisted-by: OpenAI Codex:model-version-unavailable --- crates/stdlib/src/pickle.rs | 85 +++++++++++++++++++++------ extra_tests/snippets/stdlib_pickle.py | 41 +++++++++++++ 2 files changed, 109 insertions(+), 17 deletions(-) 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