diff --git a/rust/src/python_bindings.rs b/rust/src/python_bindings.rs index 0a61232..699b86a 100644 --- a/rust/src/python_bindings.rs +++ b/rust/src/python_bindings.rs @@ -52,13 +52,64 @@ impl PyByteStorage { /// Retrieve and validate stored bytes /// /// Args: - /// envelope_bytes: Serialized StorageEnvelope bytes + /// envelope_bytes: Serialized StorageEnvelope — any buffer-protocol object + /// (`bytes`, `memoryview`, `bytearray`), so callers holding a zero-copy + /// `memoryview` (SerializationWrapper.unwrap) never re-coerce to `bytes` (LAB-770) /// /// Returns: /// Tuple[bytes, str]: (original_data, format_identifier) - pub fn retrieve(&self, py: Python, envelope_bytes: &[u8]) -> PyResult<(Vec, String)> { + pub fn retrieve( + &self, + py: Python, + envelope_bytes: &Bound<'_, PyAny>, + ) -> PyResult<(Vec, String)> { + let owned: Vec; + let buf: PyBuffer; + let base_bytes: Option>; + let data: &[u8] = if let Ok(b) = envelope_bytes.cast::() { + // `bytes` is immutable and kept alive by the Bound for the whole call: + // a safe zero-copy borrow with no data-race exposure. + b.as_bytes() + } else { + buf = PyBuffer::get(envelope_bytes)?; + // Zero-copy is only sound when the BACKING STORAGE is provably immutable — + // readonly() describes the view, not the exporter (memoryview(bytearray) + // .toreadonly() passes it while another thread can still mutate the bytearray + // during the detached read below: a data race, UB). The proof is a pointer- + // range check, not attribute trust: `.obj` naming a bytes object is spoofable + // by a PEP 688 __buffer__ exporter with a decoy attribute, so we borrow only + // when the buffer memory PROVABLY lies inside that immutable bytes object — + // true for the memoryview-over-bytes shape SerializationWrapper.unwrap + // produces, impossible to fake with memory the bytes doesn't own. `base_bytes` + // is held at function scope so the backing bytes outlives the detached read + // even if the exporter drops its own references mid-call. + base_bytes = envelope_bytes + .getattr("obj") + .ok() + .and_then(|base| base.cast_into::().ok()); + let provably_immutable = base_bytes.as_ref().is_some_and(|base| { + let start = base.as_bytes().as_ptr() as usize; + let ptr = buf.buf_ptr() as usize; + buf.readonly() + && buf.is_c_contiguous() + && buf.item_count() > 0 + && ptr >= start + && ptr + buf.item_count() <= start + base.as_bytes().len() + }); + if provably_immutable { + // SAFETY: non-null (len > 0), C-contiguous, element type validated u8 by + // PyBuffer extraction, and the range check above proves the memory sits + // inside an immutable `bytes` object kept alive by `base_bytes`. + unsafe { std::slice::from_raw_parts(buf.buf_ptr() as *const u8, buf.item_count()) } + } else { + // Mutable, non-bytes-backed, non-contiguous, or empty exporter: copy to + // owned bytes (fail-safe fallback; empty also sidesteps NULL buf_ptr). + owned = buf.to_vec(py)?; + &owned + } + }; // Detach from the GIL for decompression + checksum (see store()). - py.detach(|| self.inner.retrieve(envelope_bytes)) + py.detach(|| self.inner.retrieve(data)) .map_err(|e| PyValueError::new_err(format!("Retrieval failed: {}", e))) } diff --git a/src/cachekit/backends/file/backend.py b/src/cachekit/backends/file/backend.py index 14373c9..0e9a497 100644 --- a/src/cachekit/backends/file/backend.py +++ b/src/cachekit/backends/file/backend.py @@ -167,11 +167,14 @@ def get(self, key: str) -> bytes | None: self._acquire_file_lock(fd, exclusive=False) try: - # Read entire file - file_data = os.read(fd, os.fstat(fd).st_size) + # Header-first read (LAB-770): reading the 14-byte header separately, + # then the payload in one os.read, avoids the full-payload + # file_data[HEADER_SIZE:] slice copy (~1x payload off the read peak). + st_size = os.fstat(fd).st_size + header = os.read(fd, HEADER_SIZE) # Validate header - if len(file_data) < HEADER_SIZE: + if len(header) < HEADER_SIZE: # Corrupted file, delete it os.close(fd) fd_closed = True @@ -179,10 +182,10 @@ def get(self, key: str) -> bytes | None: return None # Parse header - magic = file_data[0:2] - version = file_data[2] - # flags = struct.unpack(">H", file_data[4:6])[0] # uint16 BE (reserved for future) - expiry_timestamp = struct.unpack(">Q", file_data[6:14])[0] # uint64 BE + magic = header[0:2] + version = header[2] + # flags = struct.unpack(">H", header[4:6])[0] # uint16 BE (reserved for future) + expiry_timestamp = struct.unpack(">Q", header[6:14])[0] # uint64 BE # Validate magic and version if magic != MAGIC or version != FORMAT_VERSION: @@ -200,9 +203,8 @@ def get(self, key: str) -> bytes | None: self._safe_unlink(file_path) return None - # Extract payload - payload = file_data[HEADER_SIZE:] - return payload + # Read payload directly — exactly the bytes after the header + return os.read(fd, st_size - HEADER_SIZE) finally: self._release_file_lock(fd) diff --git a/src/cachekit/serializers/standard_serializer.py b/src/cachekit/serializers/standard_serializer.py index 649819a..489f5ab 100644 --- a/src/cachekit/serializers/standard_serializer.py +++ b/src/cachekit/serializers/standard_serializer.py @@ -328,7 +328,8 @@ def deserialize(self, data: bytes | memoryview, metadata: SerializationMetadata >>> result == {"test": 123} True """ - data = bytes(data) # coerce unwrap's zero-copy memoryview; no-op when already bytes (Rust retrieve needs bytes) + # No bytes() coercion: Rust retrieve accepts the buffer protocol (LAB-770), so + # unwrap's zero-copy memoryview flows through without a full-payload copy. try: if self.enable_integrity_checking: # Unwrap ByteStorage envelope (decompress + validate integrity) @@ -342,7 +343,9 @@ def deserialize(self, data: bytes | memoryview, metadata: SerializationMetadata except SerializationError: # Re-raise SerializationError (integrity check failure) without swallowing raise - except (msgpack.exceptions.UnpackException, ValueError, TypeError) as e: + except (msgpack.exceptions.UnpackException, ValueError, TypeError, BufferError) as e: + # BufferError: a non-u8 buffer exporter (e.g. numpy float array) rejected at the + # PyO3 boundary — pre-LAB-770 the bytes() coercion surfaced these as ValueError. raise SerializationError(f"Failed to deserialize MessagePack data: {e}") from e diff --git a/tests/critical/test_byte_storage_error_injection.py b/tests/critical/test_byte_storage_error_injection.py index 33b481d..14d21d6 100644 --- a/tests/critical/test_byte_storage_error_injection.py +++ b/tests/critical/test_byte_storage_error_injection.py @@ -375,3 +375,164 @@ def test_final_envelope_size_security_check(self): error_msg = str(exc_info.value).lower() assert "exceeds maximum size" in error_msg or "too large" in error_msg + + +class TestByteStorageBufferProtocol: + """retrieve() accepts the buffer protocol (LAB-770) — no bytes() coercion needed. + + The zero-copy read path hands retrieve() a memoryview (SerializationWrapper.unwrap + slices one past the frame header); the PyO3 boundary must take it directly, plus + fall back to a copy for writable or non-contiguous exporters. + """ + + def test_retrieve_accepts_readonly_memoryview(self): + """Zero-copy path: memoryview over bytes round-trips identically to bytes.""" + from cachekit._rust_serializer import ByteStorage + + storage = ByteStorage("msgpack") + payload = b"buffer-protocol-roundtrip" * 1000 + envelope = storage.store(payload, None) + + data, fmt = storage.retrieve(memoryview(envelope)) + assert data == payload + assert fmt == "msgpack" + + def test_retrieve_accepts_offset_memoryview(self): + """The exact shape unwrap produces: a view sliced past a frame prefix.""" + from cachekit._rust_serializer import ByteStorage + + storage = ByteStorage("msgpack") + payload = b"offset-view" * 500 + envelope = storage.store(payload, None) + + framed = b"JUNKHDR" + envelope + data, _ = storage.retrieve(memoryview(framed)[7:]) + assert data == payload + + def test_retrieve_accepts_writable_buffer(self): + """Copy-fallback path: bytearray (writable exporter) still round-trips.""" + from cachekit._rust_serializer import ByteStorage + + storage = ByteStorage("msgpack") + payload = b"writable-exporter" * 500 + envelope = storage.store(payload, None) + + data, _ = storage.retrieve(bytearray(envelope)) + assert data == payload + data, _ = storage.retrieve(memoryview(bytearray(envelope))) + assert data == payload + + def test_retrieve_readonly_view_over_mutable_exporter_round_trips(self): + """A readonly VIEW whose backing storage is still mutable must not be borrowed + across the GIL release (data race). It takes the copy path — readonly() alone + is not the zero-copy gate; the exporter must be immutable bytes.""" + from cachekit._rust_serializer import ByteStorage + + storage = ByteStorage("msgpack") + payload = b"readonly-view-mutable-backing" * 500 + envelope = storage.store(payload, None) + + ro_view = memoryview(bytearray(envelope)).toreadonly() + assert ro_view.readonly + data, _ = storage.retrieve(ro_view) + assert data == payload + + def test_retrieve_spoofed_obj_attribute_round_trips(self): + """A PEP 688 exporter with a decoy bytes `.obj` attribute must not trick the + zero-copy gate: the pointer-range proof sees its memory is NOT inside the + decoy bytes and takes the copy path. Round-trip stays correct.""" + import sys + + if sys.version_info < (3, 12): + pytest.skip("__buffer__ protocol requires Python 3.12+") + from cachekit._rust_serializer import ByteStorage + + storage = ByteStorage("msgpack") + payload = b"spoofed-exporter" * 500 + envelope = storage.store(payload, None) + + class Spoof: + obj = b"decoy-bytes-not-the-buffer" + + def __init__(self, backing: bytearray) -> None: + self.backing = backing + + def __buffer__(self, flags: int) -> memoryview: + return memoryview(self.backing).toreadonly() + + data, _ = storage.retrieve(Spoof(bytearray(envelope))) + assert data == payload + + def test_retrieve_accepts_non_contiguous_view(self): + """Copy-fallback path: a strided view is copied, not misread.""" + from cachekit._rust_serializer import ByteStorage + + storage = ByteStorage("msgpack") + payload = b"strided-view" * 500 + envelope = storage.store(payload, None) + + interleaved = bytes(b for byte in envelope for b in (byte, 0xFF)) + data, _ = storage.retrieve(memoryview(interleaved)[::2]) + assert data == payload + + def test_retrieve_rejects_corrupt_memoryview(self): + """Error semantics are unchanged for buffer-protocol inputs.""" + from cachekit._rust_serializer import ByteStorage + + storage = ByteStorage("msgpack") + with pytest.raises(ValueError): + storage.retrieve(memoryview(b"not an envelope")) + + def test_retrieve_memoryview_is_zero_copy(self): + """No PYTHON-side full-envelope copy on the memoryview path (LAB-770). + + tracemalloc sees only Python-heap allocations: retrieve's output bytes (~1x + payload for incompressible input). A bytes(data)-style coercion creeping back + in front of retrieve adds another ~1x and fails here. Known blind spot: a + Rust-side copy (to_vec) is invisible to tracemalloc (measured: borrow and + copy paths both read 1.00x), so the borrow itself is pinned by review of + retrieve() in rust/src/python_bindings.rs, not by this suite. + """ + import gc + import os + import tracemalloc + + from cachekit._rust_serializer import ByteStorage + + storage = ByteStorage("msgpack") + payload = os.urandom(8 * 1024 * 1024) # incompressible: envelope ~= payload + envelope = storage.store(payload, None) + view = memoryview(envelope) + + gc.collect() + tracemalloc.start() + data, _ = storage.retrieve(view) + peak = tracemalloc.get_traced_memory()[1] + tracemalloc.stop() + + assert data == payload + assert peak / len(payload) < 1.5, ( + f"retrieve(memoryview) peak {peak / len(payload):.2f}x payload — the zero-copy borrow " + f"regressed to a full envelope copy (expected ~1x: just the output bytes)" + ) + + def test_standard_serializer_deserialize_memoryview(self): + """End-to-end: deserialize() takes unwrap's memoryview without re-coercing.""" + from cachekit.serializers.standard_serializer import StandardSerializer + + serializer = StandardSerializer() + obj = {"key": [1, 2, 3], "blob": b"x" * 4096} + data, _ = serializer.serialize(obj) + + assert serializer.deserialize(memoryview(data)) == obj + + def test_standard_serializer_deserialize_non_u8_buffer_raises_serialization_error(self): + """A non-u8 exporter (rejected as BufferError at the PyO3 boundary) keeps the + documented SerializationError contract — pre-LAB-770 the bytes() coercion + surfaced these as ValueError -> SerializationError.""" + np = pytest.importorskip("numpy") + from cachekit.serializers.base import SerializationError + from cachekit.serializers.standard_serializer import StandardSerializer + + with pytest.raises(SerializationError): + StandardSerializer().deserialize(np.zeros(4)) diff --git a/tests/performance/test_large_object_memory.py b/tests/performance/test_large_object_memory.py index a1b0df3..ceb22bf 100644 --- a/tests/performance/test_large_object_memory.py +++ b/tests/performance/test_large_object_memory.py @@ -312,13 +312,14 @@ def test_file_backend_bytes_read_python_allocations_bounded(tmp_path: Path) -> N """END-TO-END read through FileBackend.get() (default serializer, no mmap). This is the path most cached functions take (anything that isn't a plaintext - Arrow DataFrame). Measured cost today: ~5x payload on the Python heap — - FileBackend.get's two full-payload copies (os.read + the file_data[14:] slice, - the exact copies #169 calls out), StandardSerializer.deserialize's ``bytes(data)`` - re-coercion of the envelope's zero-copy memoryview (Rust retrieve needs bytes), - the decompressed msgpack document, and the unpacked output. The bound pins that: - one MORE full-payload copy (~6x) fails. Tightening below 5x means fixing those - copies (separate ticket per #169 — this test is the measurement). + Arrow DataFrame). Measured cost today: ~3x payload on the Python heap — + FileBackend.get's single payload os.read (header read separately, LAB-770), + the decompressed msgpack document, and the unpacked output. The two avoidable + copies #169 called out are gone: the file_data[14:] slice (header-first read) + and deserialize's ``bytes(data)`` coercion (Rust retrieve takes the buffer + protocol, so unwrap's zero-copy memoryview flows through). The bound pins + that: one full-payload copy creeping back (~4x) fails. ~3x is the floor — + decompress + unpack are inherent (retrieve returns owned bytes by construction). """ payload = np.random.default_rng(0).bytes(50 * _MB) # incompressible: envelope ~= payload size backend, operation = _file_read_stack(tmp_path / "cache", "default") @@ -333,9 +334,10 @@ def test_file_backend_bytes_read_python_allocations_bounded(tmp_path: Path) -> N assert hit is not None, "end-to-end File read missed (errors read as miss — check logs)" assert hit[1] == payload - assert peak / len(payload) < 5.7, ( - f"File-backend bytes read peak {peak / len(payload):.2f}x payload — an additional full-payload " - f"read-side copy crept in (known cost ~5x: os.read + slice + bytes() coercion + decode + output)" + assert peak / len(payload) < 3.5, ( + f"File-backend bytes read peak {peak / len(payload):.2f}x payload — a full-payload read-side " + f"copy crept back in (known cost ~3x: payload os.read + decode + output; LAB-770 removed " + f"the header slice and the bytes() coercion)" ) diff --git a/uv.lock b/uv.lock index 4f9df25..49b3a46 100644 --- a/uv.lock +++ b/uv.lock @@ -1283,11 +1283,11 @@ wheels = [ [[package]] name = "pip" -version = "26.1.2" +version = "26.2.1" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/01/91/47e7d486260f618783899587af63ccf7980fb60245c3e63dd4571c6b57ad/pip-26.1.2.tar.gz", hash = "sha256:f49cd134c61cf2fd75e0ce2676db03e4054504a5a4986d00f8299ae632dc4605", size = 1840799, upload-time = "2026-05-31T17:33:58.56Z" } +sdist = { url = "https://files.pythonhosted.org/packages/ae/15/4500e320e6b101ec3b719ae85b697d9940b6cda672bc555bd6016fc60c6f/pip-26.2.1.tar.gz", hash = "sha256:f6ad667e89a1fe78046c8f13232b247200f5258d7828f3f7883d660878e0813f", size = 1848877, upload-time = "2026-08-04T22:51:14.148Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/5d/95/6b5cb3461ea5673ba0995989746db58eb18b91b54dbf331e72f569540946/pip-26.1.2-py3-none-any.whl", hash = "sha256:382ff9f685ee3bc25864f820aa50505825f10f5458ffff07e30a6d96e5715cab", size = 1813144, upload-time = "2026-05-31T17:33:56.772Z" }, + { url = "https://files.pythonhosted.org/packages/f3/6e/1736e5b4ae2b778ef2f81c47d797de9f891d4d8acb047a24ca37a60294dd/pip-26.2.1-py3-none-any.whl", hash = "sha256:71138adf1f4ca900cdb7d289c21b7494329f2332b6d85f0e1c42108c0384ed3e", size = 1816632, upload-time = "2026-08-04T22:51:12.472Z" }, ] [[package]]