perf(file): eliminate two full-payload copies on the non-mmap read path (LAB-770)

rust/src/python_bindings.rs

@@ -52,13 +52,45 @@ 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<u8>, String)> {
+    pub fn retrieve(
+        &self,
+        py: Python,
+        envelope_bytes: PyBuffer<u8>,
+    ) -> PyResult<(Vec<u8>, String)> {
+        let owned: Vec<u8>;
+        let data: &[u8] = if envelope_bytes.readonly() && envelope_bytes.is_c_contiguous() {
+            if envelope_bytes.item_count() == 0 {
+                // buf_ptr may be NULL for an empty buffer; from_raw_parts requires non-null.
+                &[]
+            } else {
+                // SAFETY: readonly + C-contiguous checked above, and `envelope_bytes` holds
+                // the Py_buffer view alive for the whole call (resizing an exported bytearray
+                // raises BufferError in the mutator, so the pointer cannot dangle). Residual:
+                // a thread mutating memory behind a readonly view over a still-mutable
+                // exporter during the detached read is a data race on this slice — UB —
+                // accepted per CPython hashlib's own GIL-release idiom; the slice is parsed
+                // once into an owned envelope in safe Rust, so a torn read fails checksum
+                // rather than corrupting memory.
+                unsafe {
+                    std::slice::from_raw_parts(
+                        envelope_bytes.buf_ptr() as *const u8,
+                        envelope_bytes.item_count(),
+                    )
+                }
+            }
+        } else {
+            // Writable or non-contiguous exporter: copy to owned bytes (fail-safe fallback).
+            owned = envelope_bytes.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)))
     }
 

src/cachekit/backends/file/backend.py

@@ -167,22 +167,25 @@ 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
                             self._safe_unlink(file_path)
                             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)

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
 
 

tests/critical/test_byte_storage_error_injection.py

@@ -375,3 +375,121 @@ 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_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):
+        """The readonly path BORROWS — a bytes() coercion or to_vec creeping back fails here.
+
+        tracemalloc sees only Python-heap allocations: retrieve's output bytes (~1x
+        payload for incompressible input). A revert to copy-the-envelope adds another
+        ~1x. This is the non-slow guard; the end-to-end bound lives in
+        tests/performance/test_large_object_memory.py.
+        """
+        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))

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)"
     )