diff --git a/docs/performance.md b/docs/performance.md index 9c0d0685..8c50cadb 100644 --- a/docs/performance.md +++ b/docs/performance.md @@ -73,7 +73,7 @@ With `backend=None`, cachekit keeps the returned object itself in memory and han With an L2 backend configured, L1 holds the serialized bytes instead, and a hit deserializes them, so that hit costs more as the payload grows. For the 100-user dict it took 284 ± 2μs on the same runner, measured before the decode changes described below, about 910 times an L1-only hit of the same dict. Nearly all of the difference is deserializing 23.5KB of MessagePack back into about 400 nested dicts. It is still about half of a Redis L2 hit there (539 ± 3μs), which deserializes the same bytes after the round trip. -Every read with the default serializer decodes its MessagePack in the Rust extension, in one pass that also enforces the bounds on untrusted MessagePack: nesting at most 1,024 levels deep, and no header declaring more elements or bytes than the rest of the input can back. That covers L2 reads, L1 hits and the decrypted plaintext of encrypted reads. The decode builds the Python objects directly. It shares most repeated short strings, map keys up to 64 bytes and values up to 32, as one object within a value (a small table per decode; two strings that land in one slot replace each other). It runs the Python hook that restores datetimes only on maps holding a `__datetime__`, `__date__` or `__time__` key, because the hook hands back any other map unchanged. It returns the same values and refuses the same inputs with the same exception classes as the decode it replaced: cachekit-core's structure walk, then msgpack-python's C unpacker calling the hook once per map. On the 100-user dict (23,547 bytes of MessagePack, about 3,400 objects), the decode fell from 2,580,104 to 1,238,204 instructions on CPython 3.14 (−52.0%, callgrind main-thread instructions per call, the [Instruction Budgets](#instruction-budgets) method), level with `pickle.loads` of the same value (1,203,437). The 1,000-user dict (239 KB) fell 51.0% and a 1,090-byte one 45.3%. A list of 200 records each holding a datetime fell 33.2%, since every datetime still runs the hook. On the abi3t extension that CPython 3.15 and later load, the 100-user dict fell 50.5%. +Every read with the default serializer decodes its MessagePack in the Rust extension, in one pass that also enforces the bounds on untrusted MessagePack: nesting at most 1,024 levels deep, and no header declaring more elements or bytes than the rest of the input can back. That covers L2 reads, L1 hits and the decrypted plaintext of encrypted reads. The decode builds the Python objects directly. It shares most repeated short strings, map keys up to 64 bytes and values up to 32, as one object within a value (a small table per decode; two strings that land in one slot replace each other). It runs the Python hook that restores datetimes only on maps holding a `__datetime__`, `__date__` or `__time__` key, because the hook hands back any other map unchanged. It returns the same values and refuses the same inputs with the same exception classes as the decode it replaced: cachekit-core's structure walk, then msgpack-python's C unpacker calling the hook once per map. Running out of memory is the one difference. A large string, bin, ext or list still raises `MemoryError`. If CPython cannot allocate a dict, a list of up to 64 items, a float or an int outside -5 to 256, the process aborts. With Linux's default overcommit, the OOM killer usually ends the process before either outcome. On the 100-user dict (23,547 bytes of MessagePack, about 3,400 objects), the decode fell from 2,577,967 to 1,220,541 instructions on CPython 3.14 (−52.7%, callgrind main-thread instructions per call, the [Instruction Budgets](#instruction-budgets) method), level with `pickle.loads` of the same value (1,199,580). The 1,000-user dict (239 KB) fell 51.5% and a 1,090-byte one 45.8%. A list of 200 records each holding a datetime fell 33.9%, since every datetime still runs the hook. On the abi3t extension that CPython 3.15 and later load, the 100-user dict fell 51.0%. For a plaintext cache with the default serializer and integrity checking on, an entry's first L1 hit decodes it as every read does, and then, when its verified MessagePack is no longer than the envelope, which is true of small values and incompressible ones, L1 holds that MessagePack in the envelope's place. Later hits on it skip the envelope's LZ4 decompress and xxHash3-64 checksum. Every hit still runs the bounded decode, which costs no more than the unbounded one did. The 100-user dict compresses to a 3,171-byte envelope from 23.5KB of MessagePack, so L1 keeps its envelope; no budget gates this large-value hit, and `l1_hit_backed` in [Instruction Budgets](#instruction-budgets) gates the small-value hit on held MessagePack. Encrypted caches always keep the ciphertext envelope, and every hit on one decrypts and decodes. diff --git a/rust/src/msgpack_decode.rs b/rust/src/msgpack_decode.rs index 79d5a776..4b77e6a2 100644 --- a/rust/src/msgpack_decode.rs +++ b/rust/src/msgpack_decode.rs @@ -12,16 +12,26 @@ //! //! The hook runs only on a map holding a `str` key equal to a marker key. Any other map it would //! return unchanged, because each marker test reads its key with `dict.get`. +//! +//! Safe Rust only (`SECURITY.md`): every object is a PyO3 `Bound`, so each reference is released +//! by its owner on any exit. PyO3's plain constructors abort the process when CPython is out of +//! memory (the release profile sets `panic = "abort"`). A single allocation the input sizes, a +//! long array's list or the bytes of a str, bin, ext or `ExtraData`, goes through a constructor +//! that raises MemoryError instead (`nones`, `bytes_of`, `PyString::from_bytes`). A dict, a list +//! of at most `SHORT_LIST` items, a float or an int outside -5..=256 does not: running out of +//! memory while building one aborts the process, where msgpack-python raises MemoryError. + +// cargo-geiger checks the crate only after a merge; this fails any build with the `python` +// feature (a default one), which every PR's clippy job and wheel builds compile. +#![forbid(unsafe_code)] -use std::os::raw::{c_char, c_long}; -use std::ptr; +use std::convert::Infallible; use cachekit_core::check_msgpack_structure; use pyo3::exceptions::PyValueError; -use pyo3::ffi; use pyo3::prelude::*; use pyo3::sync::PyOnceLock; -use pyo3::types::PyBytes; +use pyo3::types::{PyBool, PyBytes, PyDict, PyFloat, PyInt, PyList, PyString}; /// The keys `_standard_object_hook` tests (`standard_serializer.py`). const MARKER_KEYS: [&[u8]; 3] = [b"__datetime__", b"__date__", b"__time__"]; @@ -33,64 +43,55 @@ const VALUE_CACHE_MAX_LEN: usize = 32; const CACHE_BITS: u32 = 8; const CACHE_SLOTS: usize = 1 << CACHE_BITS; -/// A Python error is already set (`PyErr::fetch` takes it), or the input broke a bound. +/// Lists at most this long (512 bytes of slots) are made by PyO3's plain constructor (`nones`). +const SHORT_LIST: u64 = 64; + +/// A Python error is set (`restore` put it there, `PyErr::fetch` takes it), or the input broke a +/// bound. A fieldless enum, so a decode step's result stays two words. enum Fail { Py, Bound, } -/// An open array or map, holding a strong reference to its container. -struct Frame { - obj: *mut ffi::PyObject, +/// An open container, filled as each element arrives. +enum Open<'py> { + /// An array: its list, made with `None` in every slot, and the next index to fill. + List(Bound<'py, PyList>, usize), + /// A map: its dict, and the key read, awaiting its value. + Map(Bound<'py, PyDict>, Option>), +} + +/// An open array or map. +struct Frame<'py> { + open: Open<'py>, /// Elements still to read: items for an array, keys plus values for a map. owed: u64, - /// Array: next index. Map: unused. - index: ffi::Py_ssize_t, - /// Map: the key read, awaiting its value (strong reference), or null. - key: *mut ffi::PyObject, - is_map: bool, /// Map: a `str` key equal to a marker key was read. marker: bool, } -/// One cached string: its bytes' position in the input, and the `str` (a strong reference). -#[derive(Clone, Copy)] -struct StrSlot { +/// One cached string: its bytes' position in the input, and the `str`. +struct StrSlot<'py> { start: usize, len: usize, - obj: *mut ffi::PyObject, + obj: Bound<'py, PyAny>, } -struct Decoder<'a> { +struct Decoder<'a, 'py> { + py: Python<'py>, data: &'a [u8], pos: usize, max_depth: usize, /// Elements owed across every open container, the root included (the walk's `pending`). pending: u64, - stack: Vec, + stack: Vec>, /// Cached strings; a table slot holds an index into it plus one, or 0 when empty. A slot holds /// one string at a time, a colliding one replacing it, so this never exceeds `2 * CACHE_SLOTS`. - cached: Vec, + cached: Vec>, keys: [u16; CACHE_SLOTS], values: [u16; CACHE_SLOTS], } -impl Drop for Decoder<'_> { - fn drop(&mut self) { - // SAFETY: every pointer here is a strong reference this decoder owns, or null; the - // thread is attached (the decoder never outlives the pyfunction call). - unsafe { - for f in self.stack.drain(..) { - ffi::Py_XDECREF(f.key); - ffi::Py_DECREF(f.obj); - } - for slot in &self.cached { - ffi::Py_DECREF(slot.obj); - } - } - } -} - /// The cache slot for `bytes`: its first and last eight bytes, multiplied and shifted. #[inline] fn slot_of(bytes: &[u8]) -> usize { @@ -106,18 +107,10 @@ fn slot_of(bytes: &[u8]) -> usize { (h >> (64 - CACHE_BITS)) as usize } -#[inline] -unsafe fn decode_utf8(bytes: &[u8]) -> *mut ffi::PyObject { - ffi::PyUnicode_DecodeUTF8( - bytes.as_ptr().cast::(), - bytes.len() as ffi::Py_ssize_t, - ptr::null(), - ) -} - -impl<'a> Decoder<'a> { - fn new(data: &'a [u8], max_depth: usize) -> Self { +impl<'a, 'py> Decoder<'a, 'py> { + fn new(py: Python<'py>, data: &'a [u8], max_depth: usize) -> Self { Self { + py, data, pos: 0, max_depth, @@ -156,14 +149,17 @@ impl<'a> Decoder<'a> { /// Whether the next element read is a map key. #[inline] fn at_key(&self) -> bool { - matches!(self.stack.last(), Some(f) if f.is_map && f.key.is_null()) + matches!( + self.stack.last(), + Some(Frame { + open: Open::Map(_, None), + .. + }) + ) } - fn decode( - &mut self, - py: Python<'_>, - hook: &Bound<'_, PyAny>, - ) -> Result<*mut ffi::PyObject, Fail> { + fn decode(&mut self, hook: &Bound<'py, PyAny>) -> Result, Fail> { + let py = self.py; loop { let marker = *self.data.get(self.pos).ok_or(Fail::Bound)?; self.pos += 1; @@ -171,82 +167,72 @@ impl<'a> Decoder<'a> { if let Some(f) = self.stack.last_mut() { f.owed -= 1; } - // SAFETY: the thread is attached; every constructor returns a new reference or null - // with an exception set, which `check` turns into Fail::Py. - let mut value = unsafe { - match marker { - 0x00..=0x7f => self.scalar(ffi::PyLong_FromLong(c_long::from(marker)))?, - 0xe0..=0xff => self.scalar(ffi::PyLong_FromLong(c_long::from(marker as i8)))?, - 0xc0 => self.scalar(ffi::Py_NewRef(ffi::Py_None()))?, - 0xc2 => self.scalar(ffi::Py_NewRef(ffi::Py_False()))?, - 0xc3 => self.scalar(ffi::Py_NewRef(ffi::Py_True()))?, - 0xa0..=0xbf => self.string(u64::from(marker & 0x1f))?, - 0xd9 => self.prefixed(1).and_then(|n| self.string(n))?, - 0xda => self.prefixed(2).and_then(|n| self.string(n))?, - 0xdb => self.prefixed(4).and_then(|n| self.string(n))?, - 0xc4 => self.prefixed(1).and_then(|n| self.binary(n))?, - 0xc5 => self.prefixed(2).and_then(|n| self.binary(n))?, - 0xc6 => self.prefixed(4).and_then(|n| self.binary(n))?, - 0xc7 => self.prefixed(1).and_then(|n| self.ext(py, n))?, - 0xc8 => self.prefixed(2).and_then(|n| self.ext(py, n))?, - 0xc9 => self.prefixed(4).and_then(|n| self.ext(py, n))?, - 0xd4..=0xd8 => self.ext(py, 1u64 << (marker - 0xd4))?, - 0xca..=0xd3 => self.number(marker)?, - 0x90..=0x9f => self.open(false, u64::from(marker & 0x0f))?, - 0xdc => self.prefixed(2).and_then(|n| self.open(false, n))?, - 0xdd => self.prefixed(4).and_then(|n| self.open(false, n))?, - 0x80..=0x8f => self.open(true, u64::from(marker & 0x0f))?, - 0xde => self.prefixed(2).and_then(|n| self.open(true, n))?, - 0xdf => self.prefixed(4).and_then(|n| self.open(true, n))?, - 0xc1 => return Err(Fail::Bound), - } + let value = match marker { + 0x00..=0x7f => self.scalar(int(py, i64::from(marker)))?, + 0xe0..=0xff => self.scalar(int(py, i64::from(marker as i8)))?, + 0xc0 => self.scalar(py.None().into_bound(py))?, + 0xc2 => self.scalar(PyBool::new(py, false).to_owned().into_any())?, + 0xc3 => self.scalar(PyBool::new(py, true).to_owned().into_any())?, + 0xa0..=0xbf => self.string(u64::from(marker & 0x1f))?, + 0xd9 => self.prefixed(1).and_then(|n| self.string(n))?, + 0xda => self.prefixed(2).and_then(|n| self.string(n))?, + 0xdb => self.prefixed(4).and_then(|n| self.string(n))?, + 0xc4 => self.prefixed(1).and_then(|n| self.binary(n))?, + 0xc5 => self.prefixed(2).and_then(|n| self.binary(n))?, + 0xc6 => self.prefixed(4).and_then(|n| self.binary(n))?, + 0xc7 => self.prefixed(1).and_then(|n| self.ext(n))?, + 0xc8 => self.prefixed(2).and_then(|n| self.ext(n))?, + 0xc9 => self.prefixed(4).and_then(|n| self.ext(n))?, + 0xd4..=0xd8 => self.ext(1u64 << (marker - 0xd4))?, + 0xca..=0xd3 => self.number(marker)?, + 0x90..=0x9f => self.open(false, u64::from(marker & 0x0f))?, + 0xdc => self.prefixed(2).and_then(|n| self.open(false, n))?, + 0xdd => self.prefixed(4).and_then(|n| self.open(false, n))?, + 0x80..=0x8f => self.open(true, u64::from(marker & 0x0f))?, + 0xde => self.prefixed(2).and_then(|n| self.open(true, n))?, + 0xdf => self.prefixed(4).and_then(|n| self.open(true, n))?, + 0xc1 => return Err(Fail::Bound), }; - // A non-empty container was pushed: read its first element. - if value.is_null() { + // None: a non-empty container was pushed, so read its first element. + let Some(mut value) = value else { continue; - } + }; // Hand the value to its container, closing every container it completes. loop { let Some(f) = self.stack.last_mut() else { return Ok(value); }; - // SAFETY: `f.obj` is a strong reference to a list or dict this decoder built; - // `value` and `f.key` are strong references the set calls consume. - unsafe { - if !f.is_map { - if ffi::PyList_SetItem(f.obj, f.index, value) != 0 { - return Err(Fail::Py); - } - f.index += 1; - } else if f.key.is_null() { - f.key = value; - break; - } else { - let key = std::mem::replace(&mut f.key, ptr::null_mut()); - let ok = key_admitted(key); - let rc = if ok { - ffi::PyDict_SetItem(f.obj, key, value) - } else { - 0 - }; - if !ok { - set_key_error(py, key); + match &mut f.open { + Open::List(list, index) => { + list.set_item(*index, value).map_err(|e| restore(py, e))?; + *index += 1; + } + Open::Map(dict, awaiting_key) => match awaiting_key.take() { + None => { + *awaiting_key = Some(value); + break; } - ffi::Py_DECREF(key); - ffi::Py_DECREF(value); - if !ok || rc != 0 { - return Err(Fail::Py); + Some(key) => { + // strict_map_key=True: a map key must be exactly `str` or `bytes`. + if !(key.is_exact_instance_of::() + || key.is_exact_instance_of::()) + { + return Err(restore(py, key_error(&key))); + } + dict.set_item(key, value).map_err(|e| restore(py, e))?; } - } - if f.owed > 0 { - break; - } - let done = self.stack.pop().expect("a frame was just read"); - value = done.obj; - if done.marker { - let obj = Bound::from_owned_ptr(py, value); - value = hook.call1((obj,)).map_err(|e| restore(py, e))?.into_ptr(); - } + }, + } + if f.owed > 0 { + break; + } + let done = self.stack.pop().expect("a frame was just read"); + value = match done.open { + Open::List(list, _) => list.into_any(), + Open::Map(dict, _) => dict.into_any(), + }; + if done.marker { + value = hook.call1((value,)).map_err(|e| restore(py, e))?; } } } @@ -265,19 +251,14 @@ impl<'a> Decoder<'a> { Ok(n) } - /// A one-byte element, already built as `obj`. + /// A one-byte element, already built as `obj` (dropped on a bound failure). #[inline] - fn scalar(&mut self, obj: *mut ffi::PyObject) -> Result<*mut ffi::PyObject, Fail> { - // Accounting first: on a bound failure the object is dropped, not leaked. - if let Err(e) = self.account(0, 0) { - // SAFETY: obj is a new reference or null. - unsafe { ffi::Py_XDECREF(obj) }; - return Err(e); - } - check(obj) + fn scalar(&mut self, obj: Bound<'py, PyAny>) -> Result>, Fail> { + self.account(0, 0)?; + Ok(Some(obj)) } - unsafe fn string(&mut self, len: u64) -> Result<*mut ffi::PyObject, Fail> { + fn string(&mut self, len: u64) -> Result>, Fail> { self.account(len, 0)?; let at_key = self.at_key(); let start = self.pos; @@ -293,22 +274,22 @@ impl<'a> Decoder<'a> { VALUE_CACHE_MAX_LEN }; if bytes.len() > max { - return check(decode_utf8(bytes)); + return self.decode_utf8(bytes).map(Some); } let i = slot_of(bytes); let held = if at_key { self.keys[i] } else { self.values[i] }; - if let Some(s) = (held as usize).checked_sub(1).map(|j| self.cached[j]) { + if let Some(s) = (held as usize).checked_sub(1).map(|j| &self.cached[j]) { if s.len == bytes.len() && &self.data[s.start..s.start + s.len] == bytes { - return Ok(ffi::Py_NewRef(s.obj)); + return Ok(Some(s.obj.clone())); } } // Not interned, unlike msgpack-python's keys: equal either way, and interning a distinct key // costs more than every later lookup of it saves. - let obj = check(decode_utf8(bytes))?; + let obj = self.decode_utf8(bytes)?; let slot = StrSlot { start, len: bytes.len(), - obj: ffi::Py_NewRef(obj), + obj: obj.clone(), }; if held == 0 { self.cached.push(slot); @@ -319,110 +300,142 @@ impl<'a> Decoder<'a> { &mut self.values[i] }) = self.cached.len() as u16; } else { - let old = std::mem::replace(&mut self.cached[held as usize - 1], slot); - ffi::Py_DECREF(old.obj); + self.cached[held as usize - 1] = slot; } - Ok(obj) + Ok(Some(obj)) } - unsafe fn binary(&mut self, len: u64) -> Result<*mut ffi::PyObject, Fail> { + /// `bytes` as a `str`, or the UnicodeDecodeError of the strict UTF-8 decode msgpack-python runs. + #[inline] + fn decode_utf8(&self, bytes: &[u8]) -> Result, Fail> { + PyString::from_bytes(self.py, bytes) + .map(Bound::into_any) + .map_err(|e| restore(self.py, e)) + } + + fn binary(&mut self, len: u64) -> Result>, Fail> { self.account(len, 0)?; let bytes = self.take(len); - check(ffi::PyBytes_FromStringAndSize( - bytes.as_ptr().cast::(), - bytes.len() as ffi::Py_ssize_t, - )) + bytes_of(self.py, bytes) + .map(|b| Some(b.into_any())) + .map_err(|e| restore(self.py, e)) } - unsafe fn number(&mut self, marker: u8) -> Result<*mut ffi::PyObject, Fail> { + fn number(&mut self, marker: u8) -> Result>, Fail> { + let py = self.py; let width = 1u64 << (marker & 0x03); self.account(width, 0)?; let raw = big_endian(self.take(width)); - check(match marker { - 0xca => ffi::PyFloat_FromDouble(f64::from(f32::from_bits(raw as u32))), - 0xcb => ffi::PyFloat_FromDouble(f64::from_bits(raw)), - 0xcc..=0xce => ffi::PyLong_FromLongLong(raw as i64), - 0xcf => ffi::PyLong_FromUnsignedLongLong(raw), - 0xd0 => ffi::PyLong_FromLong(c_long::from(raw as u8 as i8)), - 0xd1 => ffi::PyLong_FromLong(c_long::from(raw as u16 as i16)), - 0xd2 => ffi::PyLong_FromLong(c_long::from(raw as u32 as i32)), - _ => ffi::PyLong_FromLongLong(raw as i64), - }) + Ok(Some(match marker { + 0xca => PyFloat::new(py, f64::from(f32::from_bits(raw as u32))).into_any(), + 0xcb => PyFloat::new(py, f64::from_bits(raw)).into_any(), + 0xcf => int(py, raw), + 0xd0 => int(py, i64::from(raw as u8 as i8)), + 0xd1 => int(py, i64::from(raw as u16 as i16)), + 0xd2 => int(py, i64::from(raw as u32 as i32)), + // 0xcc..=0xce (unsigned, at most 32 bits) and 0xd3 (signed 64 bits). + _ => int(py, raw as i64), + })) } /// An ext of `len` data bytes after its type byte, decoded as msgpack-python's defaults do: /// type -1 to `msgpack.Timestamp` (through its own `from_bytes`), any other to /// `msgpack.ExtType(code, data)`. cachekit never writes an ext, so this path is not hot. - unsafe fn ext(&mut self, py: Python<'_>, len: u64) -> Result<*mut ffi::PyObject, Fail> { + fn ext(&mut self, len: u64) -> Result>, Fail> { + let py = self.py; self.account(len + 1, 0)?; let code = self.take(1)[0] as i8; let data = self.take(len); - let made = if code == -1 { - msgpack_attr(py, "Timestamp") - .and_then(|t| t.call_method1("from_bytes", (PyBytes::new(py, data),))) - } else { - msgpack_attr(py, "ExtType").and_then(|t| t.call1((code, PyBytes::new(py, data)))) - }; - made.map(Bound::into_ptr).map_err(|e| restore(py, e)) + let made = bytes_of(py, data).and_then(|data| { + if code == -1 { + msgpack_attr(py, "Timestamp").and_then(|t| t.call_method1("from_bytes", (data,))) + } else { + msgpack_attr(py, "ExtType").and_then(|t| t.call1((code, data))) + } + }); + made.map(Some).map_err(|e| restore(py, e)) } - /// Open a container of `n` elements (`2n` slots for a map). Returns null after pushing a + /// Open a container of `n` elements (`2n` slots for a map). Returns None after pushing a /// non-empty one, or the empty container. - unsafe fn open(&mut self, is_map: bool, n: u64) -> Result<*mut ffi::PyObject, Fail> { + fn open(&mut self, is_map: bool, n: u64) -> Result>, Fail> { // An empty container is still a level, so depth is checked before the n > 0 push. if self.stack.len() >= self.max_depth { return Err(Fail::Bound); } let owed = if is_map { 2 * n } else { n }; self.account(0, owed)?; - // n <= the input length now, so it fits Py_ssize_t. - let obj = check(if is_map { - ffi::PyDict_New() - } else { - ffi::PyList_New(n as ffi::Py_ssize_t) - })?; if n == 0 { - return Ok(obj); + return Ok(Some(if is_map { + PyDict::new(self.py).into_any() + } else { + PyList::empty(self.py).into_any() + })); } + let open = if is_map { + Open::Map(PyDict::new(self.py), None) + } else { + Open::List(self.nones(n).map_err(|e| restore(self.py, e))?, 0) + }; self.stack.push(Frame { - obj, + open, owed, - index: 0, - key: ptr::null_mut(), - is_map, marker: false, }); - Ok(ptr::null_mut()) + Ok(None) + } + + /// A list of `n` `None`s, which the decode replaces in order. `n` is at most the input length, + /// so a list longer than `SHORT_LIST` is made as `[None] * n`, which raises MemoryError where + /// PyO3's list constructor would abort. A shorter one, at most 512 bytes of slots, comes from + /// that constructor: making it by repeat would add about 1% to a small document's decode. + fn nones(&self, n: u64) -> PyResult> { + let py = self.py; + let none = py.None().into_bound(py); + if n <= SHORT_LIST { + return PyList::new(py, std::iter::repeat_n(none, n as usize)); + } + let seed = PyList::new(py, [none])?; + Ok(seed + .as_sequence() + .repeat(n as usize)? + .cast_into::()?) } } +/// A Python `int`. #[inline] -fn big_endian(bytes: &[u8]) -> u64 { - bytes.iter().fold(0u64, |acc, b| (acc << 8) | u64::from(*b)) +fn int<'py>( + py: Python<'py>, + n: impl IntoPyObject<'py, Output = Bound<'py, PyInt>, Error = Infallible>, +) -> Bound<'py, PyAny> { + let Ok(n) = n.into_pyobject(py); + n.into_any() +} + +/// `bytes` holding `data`. `new_with`, not `new`: it raises MemoryError where `new` would abort, +/// at the cost of zeroing the buffer before the copy. +fn bytes_of<'py>(py: Python<'py>, data: &[u8]) -> PyResult> { + PyBytes::new_with(py, data.len(), |b| { + b.copy_from_slice(data); + Ok(()) + }) } #[inline] -fn check(obj: *mut ffi::PyObject) -> Result<*mut ffi::PyObject, Fail> { - if obj.is_null() { - Err(Fail::Py) - } else { - Ok(obj) - } +fn big_endian(bytes: &[u8]) -> u64 { + bytes.iter().fold(0u64, |acc, b| (acc << 8) | u64::from(*b)) } -/// Put a fetched error back so the `Fail::Py` path finds it. +/// Set `e` as the current Python error, for the `Fail::Py` path to fetch. fn restore(py: Python<'_>, e: PyErr) -> Fail { e.restore(py); Fail::Py } -/// `strict_map_key=True`: a map key must be exactly `str` or `bytes`. -unsafe fn key_admitted(key: *mut ffi::PyObject) -> bool { - ffi::PyUnicode_CheckExact(key) != 0 || ffi::PyBytes_CheckExact(key) != 0 -} - -unsafe fn set_key_error(py: Python<'_>, key: *mut ffi::PyObject) { - let name = Bound::from_borrowed_ptr(py, key) +/// msgpack-python's `strict_map_key=True` refusal of `key`. +fn key_error(key: &Bound<'_, PyAny>) -> PyErr { + let name = key .get_type() .name() .map(|n| n.to_string()) @@ -430,7 +443,6 @@ unsafe fn set_key_error(py: Python<'_>, key: *mut ffi::PyObject) { PyValueError::new_err(format!( "{name:.100} is not allowed for map key when strict_map_key=True" )) - .restore(py); } fn msgpack_attr<'py>(py: Python<'py>, name: &str) -> PyResult> { @@ -449,20 +461,18 @@ pub fn unpackb_standard( max_depth: usize, hook: &Bound<'_, PyAny>, ) -> PyResult> { - let mut decoder = Decoder::new(data, max_depth); - let outcome = decoder.decode(py, hook); + let mut decoder = Decoder::new(py, data, max_depth); + let outcome = decoder.decode(hook); let consumed = decoder.pos; drop(decoder); let failure = match outcome { Ok(value) => { - // SAFETY: a strong reference the decode returned. - let value = unsafe { Bound::from_owned_ptr(py, value) }; if consumed == data.len() { return Ok(value.unbind()); } let extra = msgpack_attr(py, "ExtraData")?; return Err(PyErr::from_value( - extra.call1((value, PyBytes::new(py, &data[consumed..])))?, + extra.call1((value, bytes_of(py, &data[consumed..])?))?, )); } Err(Fail::Py) => PyErr::fetch(py), diff --git a/tests/unit/test_standard_decode.py b/tests/unit/test_standard_decode.py index bdeacaaa..1b266ef5 100644 --- a/tests/unit/test_standard_decode.py +++ b/tests/unit/test_standard_decode.py @@ -11,7 +11,8 @@ read paths' error handling see what they saw before; - the hook runs only on maps holding a marker key (the rule the speed rests on), and a heap amplification payload costs no more than it did; -- no path leaks a reference, and the decode is safe to run from many threads. +- no path leaks an object it allocates, nor a reference to its input, the input's buffer or the hook, + and the decode is safe to run from many threads. The inputs are values packed as the serializer packs them, the same with ``bytes`` packed as str (invalid UTF-8), and both after random byte edits, so the bounds, truncation, reserved-marker and @@ -21,6 +22,10 @@ from __future__ import annotations import gc +import os +import subprocess +import sys +import sysconfig import threading from collections.abc import Callable from datetime import UTC, datetime, timedelta, timezone @@ -46,10 +51,18 @@ def reference(data: bytes) -> Any: return unpackb_bounded(data, raw=False, use_list=True, object_hook=_standard_object_hook) -def rust(data: bytes) -> Any: +def rust(data: bytes | bytearray | memoryview) -> Any: return unpackb_standard(data, MSGPACK_MAX_NESTING, _standard_object_hook) +def _decode_or_refuse(data: bytes | bytearray | memoryview) -> None: + """Decode ``data``, swallowing only the refusals CASES pins; any other error fails the caller.""" + try: + rust(data) + except (ValueError, SerializationError, TypeError): # noqa: S110 - the refusal is the expected outcome + pass + + def outcome(decode: Callable[[bytes], Any], data: bytes) -> tuple[Any, ...]: """What ``decode`` returned, as a repr, or the class it raised. A bound's message is kept: it names the bound.""" try: @@ -302,15 +315,51 @@ def test_an_amplifying_payload_costs_no_more_heap_than_before() -> None: assert now <= before, f"{name}: {now} bytes at peak, {before} before" +_OUT_OF_MEMORY = """ +import resource, sys +from cachekit._rust_serializer import unpackb_standard +n = 1 << 27 +data = {shape} +with open("/proc/self/statm") as f: + mapped = int(f.read().split()[0]) * resource.getpagesize() +resource.setrlimit(resource.RLIMIT_AS, (mapped + (64 << 20), resource.getrlimit(resource.RLIMIT_AS)[1])) +try: + unpackb_standard(data, 1024, None) +except MemoryError: + sys.exit(0) +sys.exit(3) +""" +# Each single allocation the input sizes (n is 128 MiB), as the `data` expression the script decodes. Many small +# objects (dicts, short lists, floats, ints past the small-int cache) still abort: see msgpack_decode.rs's module doc. +_OUT_OF_MEMORY_SHAPES = { + "str": 'b"\\xdb" + n.to_bytes(4, "big") + b"a" * n', + "bin": 'b"\\xc6" + n.to_bytes(4, "big") + bytes(n)', + "ext": 'b"\\xc9" + n.to_bytes(4, "big") + b"\\x05" + bytes(n)', + "trailing bytes": 'b"\\xc0" + bytes(n)', + "array of n / 8 items": 'b"\\xdd" + (n >> 3).to_bytes(4, "big") + b"\\xc0" * (n >> 3)', +} + + +@pytest.mark.skipif(not sys.platform.startswith("linux"), reason="RLIMIT_AS is enforced on Linux") +@pytest.mark.skipif( + bool(sysconfig.get_config_var("Py_GIL_DISABLED")), + reason="mimalloc's reserved arenas hold the allocation, so RLIMIT_AS never refuses it; the GIL legs run the same Rust", +) +@pytest.mark.parametrize("shape", _OUT_OF_MEMORY_SHAPES.values(), ids=_OUT_OF_MEMORY_SHAPES.keys()) +def test_an_allocation_the_input_sizes_raises_memory_error(shape: str) -> None: + """With panic = "abort", a PyO3 constructor that meets a NULL would kill the process, not raise.""" + script = _OUT_OF_MEMORY.format(shape=shape) + env = {**os.environ, "PYTHONPATH": os.pathsep.join(sys.path)} # import this checkout's cachekit, as pytest does + run = subprocess.run([sys.executable, "-c", script], capture_output=True, check=False, env=env) # noqa: S603 - this module's own script + assert run.returncode == 0, run.stderr.decode()[-2000:] + + def _measure_leaks() -> list[int]: import tracemalloc def every_path() -> None: for data in (*CASES.values(), _pack({"k": ["v" * 40, datetime(2024, 1, 15), b"b"] * 3}, True)): - try: - rust(data) - except (ValueError, SerializationError, TypeError): # noqa: S110 - every refusal CASES pins; any other error fails - pass + _decode_or_refuse(data) for _ in range(50): # warm caches and lazily imported msgpack classes every_path() @@ -326,22 +375,44 @@ def every_path() -> None: return sizes -# A leak on any path keeps at least one object per call, and the smallest object (an int) is 28 bytes, so it adds at -# least 500 * 28 = 13.7 KiB to EVERY round. The interpreter's own growth (caches, free lists, a dict resizing) comes -# in a few one-off steps, and one such step can exceed any fixed budget on the total. So the gate bounds the quietest -# of five rounds, which grows by under 1 KiB when nothing leaks. +# This gate catches per-call allocations only: a path that keeps an object it allocated keeps at least one per call, +# and the smallest object (an int) is 28 bytes, so it adds at least 500 * 28 = 13.7 KiB to EVERY round. A leaked +# reference to an object that already exists (the input, its buffer, the hook) allocates nothing, so tracemalloc +# cannot see it: test_no_decode_keeps_a_reference_to_its_input counts those. The interpreter's own growth (caches, +# free lists, a dict resizing) comes in a few one-off steps, and one such step can exceed any fixed budget on the +# total. So the gate bounds the quietest of five rounds, which grows by under 1 KiB when nothing leaks. LEAK_ROUNDS = 5 LEAK_CALLS_PER_ROUND = 500 LEAK_ROUND_BUDGET = 8 * 1024 -def test_no_path_leaks_a_reference() -> None: - """A leaked reference keeps its object alive, so traced memory would climb in every round.""" +def test_no_path_leaks_an_object() -> None: + """An object a decode allocates and never frees climbs traced memory every round (per-call allocations only).""" sizes = measure_in_subprocess(_measure_leaks) rounds = [after - before for before, after in zip(sizes, sizes[1:], strict=False)] assert min(rounds) < LEAK_ROUND_BUDGET, sizes +# Empty and one-byte bytes are shared cached singletons (immortal from 3.12), so their counts are not the decode's alone. +_MULTI_BYTE = {name: data for name, data in CASES.items() if len(data) > 1} + + +@pytest.mark.parametrize("data", _MULTI_BYTE.values(), ids=_MULTI_BYTE.keys()) +def test_no_decode_keeps_a_reference_to_its_input(data: bytes) -> None: + """Every path, from every buffer type, gives back each reference it took to the input, its buffer and the hook.""" + owned = bytes(bytearray(data)) # a fresh object: a code constant may be immortal too + view, buffer = memoryview(owned), bytearray(data) + assert view.obj is owned + watched = (owned, view, buffer, _standard_object_hook) + before = [sys.getrefcount(o) for o in watched] + for source in (owned, view, buffer): + for _ in range(100): + _decode_or_refuse(source) + del source # still bound to the last input + assert [sys.getrefcount(o) for o in watched] == before + buffer.append(0) # a buffer export never released refuses the resize with BufferError + + def test_threads_decode_one_document_concurrently() -> None: doc = _pack({"users": [{"id": i, "name": f"user {i}", "at": datetime(2024, 1, 15)} for i in range(50)]}, True) expected, failures, barrier = reference(doc), [], threading.Barrier(8)