diff --git a/.secrets.baseline b/.secrets.baseline index 2debecd..595fd63 100644 --- a/.secrets.baseline +++ b/.secrets.baseline @@ -213,7 +213,7 @@ "filename": "src/cachekit/cache_handler.py", "hashed_secret": "5baa61e4c9b93f3f0682250b6cf8331b7ee68fd8", "is_verified": false, - "line_number": 411 + "line_number": 427 } ], "src/cachekit/config/decorator.py": [ @@ -425,5 +425,5 @@ } ] }, - "generated_at": "2026-07-23T09:21:58Z" + "generated_at": "2026-07-25T12:35:03Z" } diff --git a/docs/backends/file.md b/docs/backends/file.md index f6c90a3..549ff62 100644 --- a/docs/backends/file.md +++ b/docs/backends/file.md @@ -85,6 +85,27 @@ backend = FileBackend(config) - Cross-process: No (single-process only) - Platform support: Full on Linux/macOS, limited on Windows (no O_NOFOLLOW) +## Bounded-Memory Large Values (Arrow) + +FileBackend is the only backend that implements **both** large-value capability protocols, +making it the right choice for caching DataFrames bigger than your RAM headroom: + +- **Streaming writes** (`BufferWritableBackend.set_streaming`): plaintext Arrow values are + serialized record batch by record batch straight into the cache file (temp file + atomic + rename), so the full serialized payload never exists in memory. Write peak RSS is ~2.3x the + DataFrame's logical size, versus ~5.6x on the buffered path every other backend uses. +- **Zero-copy reads** (`BufferReadableBackend.get_buffer`): with `compression="none"`, cached + entries are served back through a POSIX mmap without materializing the payload on the heap. + +Both engage automatically — no API change — whenever the serializer is plaintext Arrow +(`serializer_name="arrow"`, no encryption). **Encrypted values are excluded from both paths**: +AES-256-GCM needs the whole ciphertext to produce/verify its auth tag, so the secure path +stays on buffered `set()`/`get()`. Streamed values also intentionally skip the L1 in-memory +cache — holding a multi-GB envelope in-process would defeat the point. + +Other backends (Redis, CachekitIO, Memcached) fall back transparently to buffered +`set(bytes)` with no behaviour change. + ## TTL Inspection & Sliding Expiration FileBackend implements the `TTLInspectableBackend` protocol (`get_ttl` / `refresh_ttl`) by diff --git a/docs/configuration.md b/docs/configuration.md index c58e1d6..e54051c 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -587,7 +587,7 @@ export CACHEKIT_ARROW_COMPRESSION=zstd ``` > [!TIP] -> Compression saves network bandwidth for large values. `none` can enable zero-copy mmap reads on the File backend — eligibility also requires a plaintext (unencrypted) Arrow payload read back as pandas (`return_format="pandas"`) and a backend that supports buffer reads. +> Compression saves network bandwidth for large values. `none` can enable zero-copy mmap reads on the File backend — eligibility also requires a plaintext (unencrypted) Arrow payload read back as pandas (`return_format="pandas"`) and a backend that supports buffer reads. Writes are independent of this setting: plaintext Arrow values stream to the File backend record batch by record batch (compressed or not), never materializing the full payload in memory; encrypted values and other backends use the buffered write path. ### Connection Pooling diff --git a/docs/serializers/arrow.md b/docs/serializers/arrow.md index ec0e066..ecd7a48 100644 --- a/docs/serializers/arrow.md +++ b/docs/serializers/arrow.md @@ -158,6 +158,13 @@ ArrowSerializer uses memory-mapped deserialization, which means: - Default deserialization: +15 MB memory allocation - Arrow deserialization: +2 MB memory allocation +**Writes stream on the File backend.** When the cache backend supports streaming writes +(File backend only today) and the value is plaintext (no encryption), serialization streams +~8 MiB record batches directly into the cache file instead of materializing the whole Arrow +IPC payload in memory first — cutting write peak RSS from ~5.6x to ~2.3x the DataFrame's +logical size. Wire backends (Redis, CachekitIO, Memcached) and encrypted values use the +buffered path unchanged. See [File Backend](../backends/file.md#bounded-memory-large-values-arrow). + ## Polars Support Polars DataFrames are supported via the `__arrow_c_stream__` interface (Arrow C Data Interface). This means zero-copy interchange between polars and Arrow — no intermediate conversion. diff --git a/src/cachekit/backends/base.py b/src/cachekit/backends/base.py index 61da726..42b3316 100644 --- a/src/cachekit/backends/base.py +++ b/src/cachekit/backends/base.py @@ -4,13 +4,14 @@ All L2 backends (Redis, HTTP, etc.) must implement BaseBackend protocol. Optional capability protocols (TTLInspectableBackend, LockableBackend, -TimeoutConfigurableBackend) enable advanced features with graceful degradation. +TimeoutConfigurableBackend, BufferReadableBackend, BufferWritableBackend) +enable advanced features with graceful degradation. """ from __future__ import annotations -from collections.abc import AsyncIterator -from typing import Any, Optional, Protocol, runtime_checkable +from collections.abc import AsyncIterator, Callable +from typing import Any, BinaryIO, Optional, Protocol, runtime_checkable # Re-export BackendError for convenience (public API) from cachekit.backends.errors import BackendError # noqa: F401 @@ -212,6 +213,53 @@ def get_buffer(self, key: str) -> Optional[BufferHandle]: ... +@runtime_checkable +class BufferWritableBackend(Protocol): + """Optional protocol for backends that can accept a value written incrementally through a + seekable sink, instead of one materialized ``bytes`` blob (LAB-766, write-side twin of + ``BufferReadableBackend``). + + Lets large values (e.g. multi-GB Arrow IPC) reach storage without the whole serialized + payload ever existing in memory. Only the File backend implements this today; backends that + don't are simply written via ``set`` as usual (the caller falls back transparently). + + Why a writer callback and not an ``Iterable[bytes]``: the Arrow envelope stores an + xxHash3-64 checksum BEFORE the payload it covers, so the producer must patch bytes it has + already emitted once the stream is complete. A forward-only chunk iterator cannot express + that without serializing twice; a seekable sink can. This deliberately scopes the protocol + to random-access backends (File) — a forward-only transport (chunked HTTP upload) needs a + different contract and is out of scope until a backend demands it. + + Contract: + - ``write_payload`` receives a writable, seekable binary sink positioned at the START of + the payload region. Producers MUST anchor seeks with ``tell()`` (never absolute offset 0): + backends may keep private framing before the payload. + - The value becomes visible under ``key`` only if ``write_payload`` returns successfully; + on any exception the backend discards the partial write, leaves any previous value for + the key intact, and re-raises the producer's exception unwrapped (so serialization errors + keep their type for the caller's error handling). + - ``ttl`` semantics are identical to ``set``. + - A successful ``set_streaming(key, w, ttl)`` is observably identical to + ``set(key, full_bytes, ttl)`` for the same payload bytes — same readback via + ``get``/``get_buffer``, same expiry, same size-limit enforcement (violations raise + ``BackendError`` and nothing is stored). + """ + + def set_streaming(self, key: str, write_payload: Callable[[BinaryIO], None], ttl: Optional[int] = None) -> None: + """Store the payload produced by ``write_payload`` under ``key``, never holding it whole. + + Args: + key: Cache key to store + write_payload: Callable that writes the complete payload to the provided sink + ttl: Time-to-live in seconds (None = no expiry), same semantics as ``set`` + + Raises: + BackendError: If the backend operation fails (I/O, size limits, TTL bounds) + Exception: Whatever ``write_payload`` raises, unwrapped (partial write discarded) + """ + ... + + @runtime_checkable class LockableBackend(Protocol): """Optional protocol for backends supporting distributed locking. diff --git a/src/cachekit/backends/file/backend.py b/src/cachekit/backends/file/backend.py index 504bb7d..14373c9 100644 --- a/src/cachekit/backends/file/backend.py +++ b/src/cachekit/backends/file/backend.py @@ -20,8 +20,9 @@ import struct import threading import time +from collections.abc import Callable from pathlib import Path -from typing import Any +from typing import Any, BinaryIO from cachekit.backends.errors import BackendError, BackendErrorType @@ -330,29 +331,8 @@ def set(self, key: str, value: bytes, ttl: int | None = None) -> None: file_path = self._key_to_path(key) - # Calculate expiry timestamp (0 = never expire) - if ttl is None or ttl == 0: - expiry_timestamp = 0 - else: - # Validate TTL bounds (security: prevent integer overflow/underflow) - if ttl < 0 or ttl > MAX_TTL_SECONDS: - raise BackendError( - f"TTL {ttl} out of range [0, {MAX_TTL_SECONDS}] (max 10 years)", - BackendErrorType.PERMANENT, - ) - expiry_timestamp = int(time.time() + ttl) - - # Build header (14 bytes) - header = ( - MAGIC # [0:2] Magic bytes - + bytes([FORMAT_VERSION]) # [2:3] Version - + bytes([RESERVED]) # [3:4] Reserved - + struct.pack(">H", 0) # [4:6] Flags (no compression/encryption yet) - + struct.pack(">Q", expiry_timestamp) # [6:14] Expiry timestamp - ) - - # Combine header + payload - file_data = header + value + # Combine header (14 bytes) + payload + file_data = self._build_header(self._expiry_from_ttl(ttl)) + value # Generate temp file name temp_path = self._generate_temp_path(file_path) @@ -360,15 +340,7 @@ def set(self, key: str, value: bytes, ttl: int | None = None) -> None: with self._lock: try: # Check entry count BEFORE write (security: prevent file persisting on error) - # Allow overwrites (existing key doesn't increase count) - if self.config.max_entry_count > 0: - _, entry_count = self._calculate_cache_size() - # Only check if this is a NEW entry (not overwriting existing) - if not os.path.exists(file_path) and entry_count >= self.config.max_entry_count: - raise BackendError( - f"Entry count {entry_count} would exceed max_entry_count ({self.config.max_entry_count})", - BackendErrorType.PERMANENT, - ) + self._check_entry_capacity(file_path) # Write to temp file with O_NOFOLLOW for security fd = os.open( @@ -410,6 +382,105 @@ def set(self, key: str, value: bytes, ttl: int | None = None) -> None: key=key, ) from exc + def set_streaming(self, key: str, write_payload: Callable[[BinaryIO], None], ttl: int | None = None) -> None: + """Store a streamed value with the same atomicity, locking, and limits as ``set`` (LAB-766). + + ``write_payload`` receives a writable, seekable buffered binary file positioned at the + start of the payload region (just past this backend's 14-byte header) and must write the + complete payload — record batch by record batch, so the full value never exists in + memory. The temp-file + atomic-rename pattern from ``set`` is preserved: a failed stream + never leaves a partial value visible under ``key``, and any previous value for the key + survives. Producer exceptions propagate unwrapped (after the temp file is discarded) so + serialization errors keep their type for the caller. + + ``max_value_mb`` is enforced AFTER the stream completes (``fstat`` on the temp file, + before the rename): a streaming write cannot know its size upfront. The caller-side byte + budget (CACHEKIT_MAX_VALUE_SIZE, enforced incrementally by the serializer) bounds the + transient disk usage of an over-limit attempt. + + Args: + key: Cache key to store + write_payload: Callable that writes the complete payload to the provided file + ttl: Time-to-live in seconds (None or 0 = never expire), same semantics as ``set`` + + Raises: + BackendError: If the write fails (disk full, permissions, size/TTL limits, etc.) + """ + header = self._build_header(self._expiry_from_ttl(ttl)) + file_path = self._key_to_path(key) + temp_path = self._generate_temp_path(file_path) + + committed = False + try: + # Check entry count BEFORE write (security: prevent file persisting on error) + with self._lock: + self._check_entry_capacity(file_path) + + # Stream OUTSIDE self._lock: unlike set()'s bounded os.write, the producer runs + # the whole serialization here (minutes for multi-GB frames), and holding the + # process-wide lock would block every other FileBackend op for that window. The + # in-process lock isn't what makes this safe anyway — the temp path is unique + # per pid+ns, the flock gives cross-process exclusion, and the rename commit + # below is atomic. Write to temp file with O_NOFOLLOW for security. + fd = os.open( + temp_path, + os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW, + self.config.permissions, + ) + f: BinaryIO | None = None + try: + # Acquire exclusive write lock + self._acquire_file_lock(fd, exclusive=True) + try: + f = os.fdopen(fd, "wb") # buffered writer; takes ownership of fd + f.write(header) + write_payload(f) + f.flush() + + # Enforce max_value_mb on the real on-disk size (fstat, not tell(): + # the producer may have seeked back to patch bytes it already wrote). + payload_size = os.fstat(fd).st_size - HEADER_SIZE + max_bytes = self.config.max_value_mb * 1024 * 1024 + if payload_size > max_bytes: + raise BackendError( + f"Value size {payload_size} exceeds max_value_mb ({self.config.max_value_mb}MB)", + BackendErrorType.PERMANENT, + ) + + # fsync to ensure data is on disk + os.fsync(fd) + finally: + self._release_file_lock(fd) + finally: + if f is not None: + f.close() + else: + os.close(fd) + + # Commit under the process lock: atomic rename + eviction bookkeeping only. + with self._lock: + os.rename(temp_path, file_path) + committed = True + + # Trigger eviction if over threshold + self._maybe_evict() + + except OSError as exc: + raise BackendError( + f"Failed to write cache file: {exc}", + error_type=self._classify_os_error(exc, is_directory=False), + original_exception=exc, + operation="set_streaming", + key=key, + ) from exc + finally: + # Any failure — backend I/O, size limit, or a producer exception (which + # propagates unwrapped per the BufferWritableBackend contract, so callers keep + # its type and log it once) — discards the partial write. finally (not a + # BaseException catch) so even KeyboardInterrupt can't leak a temp file. + if not committed: + self._safe_unlink(temp_path) + def delete(self, key: str) -> bool: """Delete key from file storage. @@ -640,16 +711,8 @@ async def refresh_ttl(self, key: str, ttl: int) -> bool: only bytes [6:14] are rewritten, so the payload and all other header fields are untouched (and cross-SDK File readers stay compatible). """ - # Same TTL bounds as set() (security: prevent integer overflow/underflow). - if ttl == 0: - new_expiry = 0 - elif ttl < 0 or ttl > MAX_TTL_SECONDS: - raise BackendError( - f"TTL {ttl} out of range [0, {MAX_TTL_SECONDS}] (max 10 years)", - BackendErrorType.PERMANENT, - ) - else: - new_expiry = int(time.time() + ttl) + # Same TTL bounds as set()/set_streaming (security: prevent integer overflow/underflow). + new_expiry = self._expiry_from_ttl(ttl) file_path = self._key_to_path(key) @@ -704,6 +767,50 @@ async def refresh_ttl(self, key: str, ttl: int) -> bool: # Private helper methods + @staticmethod + def _expiry_from_ttl(ttl: int | None) -> int: + """Expiry timestamp for a TTL (0 = never expire), with TTL bounds validation. + + Raises: + BackendError: If TTL is negative or exceeds MAX_TTL_SECONDS (security: + prevent integer overflow/underflow in the packed uint64 timestamp) + """ + if ttl is None or ttl == 0: + return 0 + if ttl < 0 or ttl > MAX_TTL_SECONDS: + raise BackendError( + f"TTL {ttl} out of range [0, {MAX_TTL_SECONDS}] (max 10 years)", + BackendErrorType.PERMANENT, + ) + return int(time.time() + ttl) + + @staticmethod + def _build_header(expiry_timestamp: int) -> bytes: + """Build the 14-byte on-disk entry header.""" + return ( + MAGIC # [0:2] Magic bytes + + bytes([FORMAT_VERSION]) # [2:3] Version + + bytes([RESERVED]) # [3:4] Reserved + + struct.pack(">H", 0) # [4:6] Flags (no compression/encryption yet) + + struct.pack(">Q", expiry_timestamp) # [6:14] Expiry timestamp + ) + + def _check_entry_capacity(self, file_path: str) -> None: + """Reject a NEW entry when max_entry_count is reached (overwrites always pass). + + Raises: + BackendError: If storing a new entry would exceed max_entry_count + """ + if self.config.max_entry_count <= 0: + return + _, entry_count = self._calculate_cache_size() + # Only check if this is a NEW entry (not overwriting existing) + if not os.path.exists(file_path) and entry_count >= self.config.max_entry_count: + raise BackendError( + f"Entry count {entry_count} would exceed max_entry_count ({self.config.max_entry_count})", + BackendErrorType.PERMANENT, + ) + def _key_to_path(self, key: str) -> str: """Convert cache key to file path using blake2b hash. diff --git a/src/cachekit/cache_handler.py b/src/cachekit/cache_handler.py index 1d217a8..0f0ab17 100644 --- a/src/cachekit/cache_handler.py +++ b/src/cachekit/cache_handler.py @@ -11,9 +11,16 @@ import threading import warnings from collections.abc import Callable -from typing import TYPE_CHECKING, Any, Optional, Protocol, TypeGuard, Union, runtime_checkable - -from cachekit.backends.base import BackendError, BaseBackend, BufferHandle, BufferReadableBackend, TTLInspectableBackend +from typing import TYPE_CHECKING, Any, BinaryIO, Optional, Protocol, TypeGuard, Union, runtime_checkable + +from cachekit.backends.base import ( + BackendError, + BaseBackend, + BufferHandle, + BufferReadableBackend, + BufferWritableBackend, + TTLInspectableBackend, +) from cachekit.backends.provider import ( BackendProviderInterface, DefaultBackendProvider, @@ -221,6 +228,15 @@ def supports_buffer_read(backend: BaseBackend) -> TypeGuard[BufferReadableBacken return hasattr(backend, "get_buffer") +def supports_streaming_write(backend: BaseBackend) -> TypeGuard[BufferWritableBackend]: + """Type guard: backend can accept an incrementally-streamed value via set_streaming (LAB-766, File only). + + Returns: + True if backend implements BufferWritableBackend (used for the streaming Arrow write path). + """ + return hasattr(backend, "set_streaming") + + class SWRCapableBackend(Protocol): """Backend with server-signaled stale-while-revalidate reads (LAB-381). @@ -932,6 +948,54 @@ def supports_mmap_read(self) -> bool: and getattr(self._base_serializer, "return_format", None) == "pandas" ) + def supports_streaming_write(self) -> bool: + """True iff writes can stream the envelope to a capable backend (LAB-766). + + Eligible only for PLAINTEXT Arrow: + - encrypted values can never stream (AES-256-GCM emits its auth tag only after the + whole ciphertext exists, so the secure path stays on buffered ``set(bytes)``); + - interop/v1 stores raw documents with no CK envelope (different serializer, no benefit); + - non-Arrow serializers materialize their output in one shot anyway. + + The backend must also support streaming writes (File only today); that is checked + separately at the cache-handler layer, so True here on a non-streaming backend simply + means the write falls back to the buffered path. + """ + return not self.encryption and not self.interop_mode and self._serializer_string_name == "arrow" + + def write_serialized_to(self, sink: BinaryIO, data: Any) -> None: + """Streaming twin of :meth:`serialize_data` (LAB-766): write the SAME envelope bytes into + ``sink`` without ever materializing the payload in memory. + + Only valid when :meth:`supports_streaming_write` is True (plaintext Arrow — no + encryption wrapper, so no tenant extraction / AAD / cache_key binding applies). The + envelope is byte-identical to a buffered ``serialize_data`` of the same value: frame + prefix first (metadata is deterministic per serializer config), then the Arrow + serializer streams ``[checksum][IPC record batches]`` directly into the sink. + + The L2 oversized-entry ceiling (``max_value_size``, issue #163) is enforced + incrementally mid-stream instead of on the finished blob — same limit, same ValueError, + but an over-limit stream aborts early instead of materializing first. + + Raises: + TypeError: If data is not a DataFrame or dict of arrays + ValueError: If the envelope would exceed max_value_size + SerializationError: If Arrow serialization fails + """ + from cachekit.serializers.arrow_serializer import ArrowSerializer # already loaded on the arrow path + + serializer = self._base_serializer + if not isinstance(serializer, ArrowSerializer): + # supports_streaming_write() gates callers to the arrow serializer; anything else + # here is a caller bug — fail loud rather than stream a wrong envelope. + raise SerializationError(f"write_serialized_to requires the arrow serializer, got {type(serializer).__name__}") + + metadata = serializer.serialization_metadata() + prefix = SerializationWrapper.wrap_prefix(metadata.to_dict(), self._serializer_string_name) + sink.write(prefix) + budget = get_settings().max_value_size - len(prefix) + serializer.serialize_to_sink(data, sink, max_bytes=budget) + def deserialize_data(self, data: str | bytes | memoryview, cache_key: str = "") -> Any: """Deserialize data from cache storage with cache_key verification. @@ -1413,7 +1477,11 @@ def store_result( kwargs: Function kwargs (for tenant extraction in encryption) Returns: - Serialized bytes (for L1 cache storage), or None if serialization failed + Serialized bytes when the buffered path stored the value (eligible for L1 + backfill), or None when there is nothing for L1: the value was STREAMED to the + backend (success — L1 intentionally skipped, LAB-766), the streaming attempt + failed (logged, never retried buffered), or serialization failed. None is NOT + a failure signal. Note: Requires cache_handler to be set via set_cache_handler() before calling. @@ -1423,6 +1491,28 @@ def store_result( if self._cache_handler is None: raise RuntimeError("Cache handler must be set before calling store_result") + # Streaming write fast path (LAB-766): plaintext Arrow to a streaming-writable + # backend (File) writes the envelope without ever materializing it. SWR + # (stale_ttl) stays buffered — no streaming backend signals freshness today. + # getattr guards custom CacheHandlerStrategy impls that predate set_streaming. + if stale_ttl is None and self.serialization_handler.supports_streaming_write(): + set_streaming = getattr(self._cache_handler, "set_streaming", None) + if set_streaming is not None: + streamed = set_streaming( + cache_key, + lambda sink: self.serialization_handler.write_serialized_to(sink, result), + ttl, + ) + if streamed is not None: + # Streamed (True) or attempt failed and was logged (False — do NOT + # retry buffered: that re-materializes the payload this path exists + # to avoid). Either way return None: a multi-GB envelope must never + # be copied into L1 (mirrors the mmap read path's L1 exclusion, #171). + if streamed: + get_logger().cache_stored(cache_key, ttl) + return None + # None: backend can't stream — fall through to the buffered path. + # Pass cache_key for AAD binding (required for encrypted data) serialized_data = self.serialization_handler.serialize_data(result, args, kwargs, cache_key) # Only thread the SWR kwarg when set: strategy implementations without @@ -1461,7 +1551,11 @@ async def store_result_async( kwargs: Function kwargs (for tenant extraction in encryption) Returns: - Serialized bytes (for L1 cache storage), or None if serialization failed + Serialized bytes when the buffered path stored the value (eligible for L1 + backfill), or None when there is nothing for L1: the value was STREAMED to the + backend (success — L1 intentionally skipped, LAB-766), the streaming attempt + failed (logged, never retried buffered), or serialization failed. None is NOT + a failure signal. Note: Requires cache_handler to be set via set_cache_handler() before calling. @@ -1471,6 +1565,25 @@ async def store_result_async( if self._cache_handler is None: raise RuntimeError("Cache handler must be set before calling store_result_async") + # Streaming write fast path (LAB-766) — async twin of store_result's gate. The + # whole stream (pyarrow serialization included) runs in the handler's thread + # pool, which also moves the blocking serialize OFF the event loop (the buffered + # path below serializes on the loop thread). + if self.serialization_handler.supports_streaming_write(): + set_streaming_async = getattr(self._cache_handler, "set_streaming_async", None) + if set_streaming_async is not None: + streamed = await set_streaming_async( + cache_key, + lambda sink: self.serialization_handler.write_serialized_to(sink, result), + ttl, + ) + if streamed is not None: + # See store_result: never retry buffered on failure, never hand the + # envelope to L1 on success. + if streamed: + get_logger().cache_stored(cache_key, ttl) + return None + # Pass cache_key for AAD binding (required for encrypted data) serialized_data = self.serialization_handler.serialize_data(result, args, kwargs, cache_key) await self._cache_handler.set_async(cache_key, serialized_data, ttl) @@ -1607,6 +1720,16 @@ def set(self, key: str, value: Union[str, bytes], ttl: Optional[int] = None, **m """Set value in cache with TTL and optional metadata.""" ... + def set_streaming(self, key: str, write_payload: Callable[[BinaryIO], None], ttl: Optional[int] = None) -> Optional[bool]: + """Stream a value to a streaming-writable backend (LAB-766); None when unsupported.""" + ... + + async def set_streaming_async( + self, key: str, write_payload: Callable[[BinaryIO], None], ttl: Optional[int] = None + ) -> Optional[bool]: + """Async variant of set_streaming.""" + ... + def delete(self, key: str) -> bool: """Delete key from cache.""" ... @@ -1860,6 +1983,47 @@ def set( get_logger().error(f"Unexpected error setting key {key}: {e}") return False + def set_streaming(self, key: str, write_payload: Callable[[BinaryIO], None], ttl: Optional[int] = None) -> Optional[bool]: + """Stream a value to a streaming-writable backend (LAB-766). + + Returns: + None when the backend doesn't implement BufferWritableBackend — the caller falls + back to the buffered ``set``. True on success. False when the streaming attempt + failed (logged, mirroring ``set``'s degradation contract); the caller must NOT + retry with the buffered path — that would re-materialize the very payload this + path exists to keep out of memory. + """ + if not supports_streaming_write(self.backend): + return None + try: + self._with_backpressure_and_timeout(self.backend.set_streaming, key, write_payload, ttl) + return True + except BackendError as e: + get_logger().error(f"Backend error streaming key {redact_cache_key(key)}: {e}") + return False + except Exception as e: + # Producer-side failure (serialization error, max_value_size budget): the backend + # already discarded its partial write; surface the real cause, not a backend error. + get_logger().error(f"Streaming serialization failed for key {redact_cache_key(key)}: {e}") + return False + + async def set_streaming_async( + self, key: str, write_payload: Callable[[BinaryIO], None], ttl: Optional[int] = None + ) -> Optional[bool]: + """Async variant of :meth:`set_streaming`: the whole stream — pyarrow serialization + included — runs in the thread pool, keeping blocking IPC writes off the event loop.""" + if not supports_streaming_write(self.backend): + return None + try: + await self._with_backpressure_and_timeout_async(self.backend.set_streaming, key, write_payload, ttl) + return True + except BackendError as e: + get_logger().error(f"Backend error streaming key {redact_cache_key(key)}: {e}") + return False + except Exception as e: + get_logger().error(f"Streaming serialization failed for key {redact_cache_key(key)}: {e}") + return False + def delete(self, key: str) -> bool: """Delete key from cache using backend. diff --git a/src/cachekit/config/settings.py b/src/cachekit/config/settings.py index b29a7ee..e8ca5d6 100644 --- a/src/cachekit/config/settings.py +++ b/src/cachekit/config/settings.py @@ -105,7 +105,9 @@ class CachekitConfig(BaseSettings): "'zstd'/'lz4' shrink the stored payload but must be decompressed into the heap on read. " "'none' stores uncompressed Arrow IPC, which lets the File backend serve plaintext " "DataFrame reads via a zero-copy mmap (low steady-state read RSS; peak transiently " - "higher) at the cost of a larger payload. Env: CACHEKIT_ARROW_COMPRESSION." + "higher) at the cost of a larger payload. Writes are codec-independent: plaintext " + "Arrow streams to the File backend without materializing the payload (encrypted " + "values stay buffered). Env: CACHEKIT_ARROW_COMPRESSION." ), ) retry_on_timeout: bool = Field( diff --git a/src/cachekit/serializers/arrow_serializer.py b/src/cachekit/serializers/arrow_serializer.py index 3a70856..95effa7 100644 --- a/src/cachekit/serializers/arrow_serializer.py +++ b/src/cachekit/serializers/arrow_serializer.py @@ -22,7 +22,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, ClassVar +from typing import TYPE_CHECKING, Any, BinaryIO, ClassVar from .base import SerializationError, SerializationFormat, SerializationMetadata @@ -82,6 +82,56 @@ def _write_bounded_batches(writer: pa.ipc.RecordBatchFileWriter, table: pa.Table offset += rows +class _HashingSink: + """Forward-only adapter pyarrow's IPC writer streams into (LAB-766): forwards each write to + the underlying sink while hashing it (xxHash3-64) and enforcing an optional byte budget, so + the checksum over the streamed IPC bytes is known the moment streaming completes — without + ever holding those bytes together in memory. + """ + + # pyarrow's PythonFile wrapper probes these before writing. + closed = False + + def __init__(self, raw: BinaryIO, budget: int | None) -> None: + self._raw = raw + self._hasher = xxhash.xxh3_64() + self._budget = budget + self._written = 0 + + def writable(self) -> bool: + return True + + def readable(self) -> bool: + return False + + def seekable(self) -> bool: + return False + + def write(self, data: bytes) -> int: + n = data.nbytes if isinstance(data, memoryview) else len(data) + self._written += n + if self._budget is not None and self._written > self._budget: + # Reject BEFORE forwarding: earlier chunks may be on disk (the backend discards + # its partial write), but no bytes past the budget ever leave this process. + raise ValueError( + f"Streaming serialization exceeded its {self._budget}-byte payload budget " + f"(CACHEKIT_MAX_VALUE_SIZE); refusing to cache. Increase CACHEKIT_MAX_VALUE_SIZE " + f"to cache larger values." + ) + self._hasher.update(data) + self._raw.write(data) + return n + + def tell(self) -> int: + return self._written + + def flush(self) -> None: + self._raw.flush() + + def digest(self) -> bytes: + return self._hasher.digest() + + class ArrowSerializer: """Apache Arrow IPC serializer for memory-efficient DataFrame caching with xxHash3-64 integrity. @@ -98,7 +148,9 @@ class ArrowSerializer: Memory profile (bounded, low-copy): - Serialize: builds the compressed IPC once and prepends the checksum; the source Arrow - table is freed before the IPC bytes are materialized. + table is freed before the IPC bytes are materialized. On a streaming-writable backend + (File), ``serialize_to_sink`` skips even that: record batches stream straight into the + cache file and the full payload never exists in memory (LAB-766). - Deserialize: the envelope is sliced with a memoryview (no full-body copy), wrapped via pa.py_buffer (zero-copy), and Arrow->pandas conversion uses self_destruct + split_blocks to free Arrow buffers during conversion. zstd is decompressed transparently by the reader. @@ -190,6 +242,56 @@ def _resolve_compression(compression: str | None) -> str | None: raise ValueError(f"Invalid compression: {compression!r}. Valid options: 'auto', 'zstd', 'lz4', None ('none').") return compression + @staticmethod + def _to_table(obj: Any) -> pa.Table: # type: ignore[name-defined] + """Convert a supported object to an Arrow Table (shared by the buffered and streaming paths). + + preserve_index=None (pyarrow default): a RangeIndex is stored as cheap + schema metadata (no materialized column / extra copy) and restored as a + RangeIndex; named/MultiIndex are still preserved as columns. preserve_index=True + would force even a RangeIndex into a materialized column. + + Raises: + TypeError: If obj is not a DataFrame (pandas, polars) or dict of arrays + """ + if HAS_PANDAS and isinstance(obj, pd.DataFrame): + return pa.Table.from_pandas(obj, preserve_index=None) + if hasattr(obj, "__arrow_c_stream__"): # polars DataFrame (zero-copy C Stream) + return pa.table(obj) + if isinstance(obj, dict): + # dict of arrays (columnar). Normalize pyarrow's raw conversion errors + # (e.g. dict-of-scalars -> "'int' object is not iterable") into the + # documented TypeError so callers get a consistent, actionable message. + try: + return pa.table(obj) + except (pa.ArrowInvalid, pa.ArrowTypeError, TypeError, ValueError) as e: + raise TypeError( + f"ArrowSerializer only supports DataFrames " + f"(pandas.DataFrame, polars.DataFrame) or dict of arrays (columnar). " + f"Got a dict that is not convertible to an Arrow table: {e}. " + f"For scalar values or nested dicts, use AutoSerializer." + ) from e + raise TypeError( + f"ArrowSerializer only supports DataFrames " + f"(pandas.DataFrame, polars.DataFrame) or dict of arrays. " + f"Got: {type(obj).__name__}. " + f"For scalar values or nested dicts, use AutoSerializer." + ) + + def serialization_metadata(self) -> SerializationMetadata: + """Metadata a serialize() call will produce, computable BEFORE any payload byte exists. + + The streaming write path (LAB-766) needs this upfront: the cache envelope's metadata + header is written before the payload streams. Deterministic per serializer config — + never payload-dependent. ``serialize`` returns exactly this object. + """ + return SerializationMetadata( + serialization_format=SerializationFormat.ARROW, + compressed=self.compression is not None, # reflects the configured codec (None = uncompressed) + encrypted=False, # Encryption is EncryptionWrapper's responsibility + original_type="arrow", + ) + def serialize(self, obj: Any) -> tuple[bytes, SerializationMetadata]: # type: ignore[name-defined] """Serialize DataFrame to Arrow IPC format bytes with optional xxHash3-64 integrity protection. @@ -206,37 +308,7 @@ def serialize(self, obj: Any) -> tuple[bytes, SerializationMetadata]: # type: i SerializationError: If Arrow conversion fails """ try: - # Convert to Arrow Table (supports pandas, polars, dict of arrays). - # preserve_index=None (pyarrow default): a RangeIndex is stored as cheap - # schema metadata (no materialized column / extra copy) and restored as a - # RangeIndex; named/MultiIndex are still preserved as columns. preserve_index=True - # would force even a RangeIndex into a materialized column. - table = None - if HAS_PANDAS and isinstance(obj, pd.DataFrame): - table = pa.Table.from_pandas(obj, preserve_index=None) - elif hasattr(obj, "__arrow_c_stream__"): # polars DataFrame (zero-copy C Stream) - table = pa.table(obj) - elif isinstance(obj, dict): - # dict of arrays (columnar). Normalize pyarrow's raw conversion errors - # (e.g. dict-of-scalars -> "'int' object is not iterable") into the - # documented TypeError so callers get a consistent, actionable message. - try: - table = pa.table(obj) - except (pa.ArrowInvalid, pa.ArrowTypeError, TypeError, ValueError) as e: - raise TypeError( - f"ArrowSerializer only supports DataFrames " - f"(pandas.DataFrame, polars.DataFrame) or dict of arrays (columnar). " - f"Got a dict that is not convertible to an Arrow table: {e}. " - f"For scalar values or nested dicts, use AutoSerializer." - ) from e - - if table is None: - raise TypeError( - f"ArrowSerializer only supports DataFrames " - f"(pandas.DataFrame, polars.DataFrame) or dict of arrays. " - f"Got: {type(obj).__name__}. " - f"For scalar values or nested dicts, use AutoSerializer." - ) + table = self._to_table(obj) # Serialize to Arrow IPC. Compression (when enabled) runs per record-batch, so # writing in bounded batches keeps the compressor's working set bounded (one big @@ -266,15 +338,57 @@ def serialize(self, obj: Any) -> tuple[bytes, SerializationMetadata]: # type: i del buf pa.default_memory_pool().release_unused() - return envelope, SerializationMetadata( - serialization_format=SerializationFormat.ARROW, - compressed=self.compression is not None, # reflects the configured codec (None = uncompressed) - encrypted=False, # Encryption is EncryptionWrapper's responsibility - original_type="arrow", - ) + return envelope, self.serialization_metadata() except (pa.ArrowInvalid, pa.ArrowTypeError, ValueError) as e: raise SerializationError(f"Failed to serialize DataFrame to Arrow IPC format: {e}") from e + def serialize_to_sink(self, obj: Any, sink: BinaryIO, max_bytes: int | None = None) -> SerializationMetadata: + """Stream-serialize a DataFrame into ``sink``, never materializing the full payload (LAB-766). + + Writes the exact bytes ``serialize`` returns — [8-byte xxHash3-64 checksum][Arrow IPC] — + but the IPC record batches (~8 MiB actual bytes each, see ``_write_bounded_batches``) go + straight to the sink instead of accumulating in a ``pa.BufferOutputStream``, which is what + held ~3.75x the payload in memory on the buffered path. The checksum covers all IPC bytes yet is + stored FIRST, so it is hashed incrementally while streaming and patched in place + afterwards: ``sink`` must therefore be seekable (``tell``/``seek``). Seeks are anchored + with ``tell()``, never absolute offsets — the sink may sit inside backend framing. + + Args: + obj: DataFrame (pandas, polars) or dict of arrays (columnar) + sink: Writable, seekable binary sink positioned where the envelope should start + max_bytes: Optional envelope byte budget; exceeding it raises ValueError mid-stream + (bounds disk churn on values that would be rejected as oversized anyway) + + Returns: + The same SerializationMetadata ``serialize`` would return. + + Raises: + TypeError: If obj is not a DataFrame or dict of arrays + ValueError: If the envelope would exceed ``max_bytes`` + SerializationError: If Arrow conversion fails + """ + try: + table = self._to_table(obj) + checksum_pos = sink.tell() + sink.write(b"\x00" * 8) # checksum placeholder, patched below + hashing_sink = _HashingSink(sink, max_bytes - 8 if max_bytes is not None else None) + write_options = pa.ipc.IpcWriteOptions(compression=self.compression) if self.compression else None + with pa.ipc.new_file(hashing_sink, table.schema, options=write_options) as writer: + _write_bounded_batches(writer, table) + del table # free the Arrow table promptly (mirrors serialize) + + end_pos = sink.tell() + sink.seek(checksum_pos) + sink.write(hashing_sink.digest()) + sink.seek(end_pos) + pa.default_memory_pool().release_unused() # return per-batch working memory to the OS + return self.serialization_metadata() + # NOTE: unlike serialize(), ValueError is deliberately NOT wrapped here — the + # _HashingSink budget ValueError must reach the caller with its original type + # (the max_value_size mid-stream abort contract; the handler and tests match on it). + except (pa.ArrowInvalid, pa.ArrowTypeError) as e: + raise SerializationError(f"Failed to serialize DataFrame to Arrow IPC format: {e}") from e + def deserialize(self, data: bytes | memoryview, metadata: SerializationMetadata | None = None) -> Any: """Deserialize Arrow IPC bytes with optional xxHash3-64 integrity validation. diff --git a/src/cachekit/serializers/wrapper.py b/src/cachekit/serializers/wrapper.py index 3350d72..1376a7f 100644 --- a/src/cachekit/serializers/wrapper.py +++ b/src/cachekit/serializers/wrapper.py @@ -70,34 +70,43 @@ class SerializationWrapper: """ @staticmethod - def wrap(data: bytes, metadata: dict[str, Any], serializer_name: str, version: str = "2.0") -> bytes: - """Frame serialized data with a metadata header for cache storage. - - Args: - data: Serialized bytes to wrap (stored raw — no base64). - metadata: Serialization metadata dict (must include "format" key). - serializer_name: Name of serializer used (e.g., "default", "arrow"). - version: Logical serializer-envelope version (carried in the header for - downstream compatibility checks; distinct from the binary frame version). + def wrap_prefix(metadata: dict[str, Any], serializer_name: str, version: str = "2.0") -> bytes: + """Build the v3 frame prefix (everything before the payload): MAGIC | VERSION | HDR_LEN | HEADER. - Returns: - v3 binary frame bytes: MAGIC | VERSION | HDR_LEN | HEADER(json) | PAYLOAD(raw). + The prefix depends only on the metadata/serializer name, never on the payload bytes, so + the streaming write path (LAB-766) can emit it before a single payload byte exists and + the stored frame stays byte-identical to a buffered ``wrap`` of the same payload. """ header = json.dumps( {"s": serializer_name, "m": metadata, "v": version}, ensure_ascii=False, ).encode("utf-8") - # Single allocation; the payload is copied exactly once. return b"".join( ( _MAGIC, bytes((_FRAME_VERSION,)), len(header).to_bytes(_HEADER_LEN_BYTES, "big"), header, - data, ) ) + @staticmethod + def wrap(data: bytes, metadata: dict[str, Any], serializer_name: str, version: str = "2.0") -> bytes: + """Frame serialized data with a metadata header for cache storage. + + Args: + data: Serialized bytes to wrap (stored raw — no base64). + metadata: Serialization metadata dict (must include "format" key). + serializer_name: Name of serializer used (e.g., "default", "arrow"). + version: Logical serializer-envelope version (carried in the header for + downstream compatibility checks; distinct from the binary frame version). + + Returns: + v3 binary frame bytes: MAGIC | VERSION | HDR_LEN | HEADER(json) | PAYLOAD(raw). + """ + # Single allocation; the payload is copied exactly once. + return SerializationWrapper.wrap_prefix(metadata, serializer_name, version) + data + @staticmethod def unwrap( wrapped_data: Union[str, bytes, bytearray, memoryview], diff --git a/tests/performance/test_large_object_memory.py b/tests/performance/test_large_object_memory.py index f0cbea4..a1b0df3 100644 --- a/tests/performance/test_large_object_memory.py +++ b/tests/performance/test_large_object_memory.py @@ -128,6 +128,73 @@ def test_full_roundtrip_through_cache_handler_is_correct_and_compact(): pd.testing.assert_frame_equal(out, df) +@pytest.mark.slow +@pytest.mark.performance +def test_streaming_write_path_peak_rss_bounded(): + """LAB-766: the File-backend streaming write path never materializes the envelope. + + Peak RSS in a dedicated subprocess (tracemalloc cannot see pyarrow's C++ pools, + which is exactly where the buffered path's residual lives). A 400MB incompressible + frame is stored through the REAL write flow (store_result -> set_streaming -> + FileBackend), record batch by record batch. + + Measured on the reference box: streaming peaks at ~2.3x logical (frame + the Arrow + table conversion + interpreter); the buffered path peaks at ~5.6x (BufferOutputStream + doubling-growth + memory-pool retention + envelope + wrap copies). The 3.2x bound + sits between them with margin on both sides: a full-payload materialization + returning to this path re-adds ~2-3x and trips it loudly. + """ + frame_mb = 400 + script = textwrap.dedent( + f""" + import resource, tempfile + + import numpy as np + import pandas as pd + + from cachekit.backends.file import FileBackend + from cachekit.backends.file.config import FileBackendConfig + from cachekit.cache_handler import CacheOperationHandler, CacheSerializationHandler, StandardCacheHandler + from cachekit.key_generator import CacheKeyGenerator + from cachekit.serializers.arrow_serializer import ArrowSerializer + + _MB = 1024 * 1024 + cols = 20 + rows = {frame_mb} * _MB // (8 * cols) + rng = np.random.default_rng(0) + df = pd.DataFrame({{f"c{{i}}": rng.standard_normal(rows) for i in range(cols)}}) + logical = int(df.memory_usage(deep=True, index=True).sum()) + + with tempfile.TemporaryDirectory() as d: + backend = FileBackend(FileBackendConfig(cache_dir=d, max_size_mb=4096, max_value_mb=2048)) + sh = CacheSerializationHandler(serializer_name="arrow") + sh._base_serializer = ArrowSerializer(compression=None) # uncompressed: on-disk ~1x, mmap-readable + op = CacheOperationHandler(sh, CacheKeyGenerator()) + op.set_cache_handler(StandardCacheHandler(backend)) + + ret = op.store_result("k", df, ttl=60) + assert ret is None, "streaming path did not engage (fell back to buffered set)" + # Sanity via on-disk size, NOT backend.get(): a full readback would materialize + # the envelope and pollute the ru_maxrss high-water mark this test measures. + import os as _os + + (entry,) = [f for f in _os.scandir(d) if ".tmp." not in f.name] + assert entry.stat().st_size > logical * 0.95, "streamed entry missing/truncated" + + peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss * 1024 # KiB on Linux + print(logical, peak) + """ + ) + env = dict(os.environ, CACHEKIT_MAX_VALUE_SIZE=str(2 * 1024**3)) + result = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True, env=env) # noqa: S603 (trusted: sys.executable + literal code) + assert result.returncode == 0, f"streaming-write subprocess failed: {result.stderr}" + logical, peak = (int(v) for v in result.stdout.split()) + assert peak < logical * 3.2, ( + f"streaming write peak RSS {peak / logical:.2f}x logical — the envelope is being " + f"materialized again on the write path (streaming ~2.3x, buffered ~5.6x)" + ) + + @pytest.mark.slow @pytest.mark.performance @pytest.mark.skipif(sys.platform != "linux", reason="peak-RSS via /proc/self/status VmHWM is Linux-only") diff --git a/tests/unit/test_streaming_write_path.py b/tests/unit/test_streaming_write_path.py new file mode 100644 index 0000000..1b76522 --- /dev/null +++ b/tests/unit/test_streaming_write_path.py @@ -0,0 +1,317 @@ +"""LAB-766 streaming write-path wiring (write-side twin of test_mmap_read_path.py). + +Four units compose the streaming Arrow write path: +- ``CacheSerializationHandler.supports_streaming_write()`` — eligibility (plaintext Arrow). +- ``CacheSerializationHandler.write_serialized_to()`` — streams the SAME envelope bytes + ``serialize_data`` produces (frame prefix, then [checksum][IPC]) into a seekable sink. +- ``StandardCacheHandler.set_streaming(_async)`` — delegates to a backend that exposes + ``set_streaming``; returns None (fall back to buffered set) when the backend doesn't. +- ``CacheOperationHandler.store_result(_async)`` — when eligible AND the backend streams, + writes via the callback and returns None so the multi-GB envelope never reaches L1; + otherwise the buffered path runs unchanged and still returns bytes for L1. + +FileBackend.set_streaming itself (atomicity, TTL, size limits) is covered here too — it is +the first and load-bearing BufferWritableBackend implementation. +""" + +from __future__ import annotations + +import io +import os +from unittest.mock import MagicMock + +import pytest + +from cachekit.backends.errors import BackendError +from cachekit.cache_handler import ( + CacheOperationHandler, + CacheSerializationHandler, + StandardCacheHandler, + supports_streaming_write, +) +from cachekit.key_generator import CacheKeyGenerator + +pd = pytest.importorskip("pandas") +pytest.importorskip("pyarrow") + + +@pytest.fixture +def df() -> pd.DataFrame: + return pd.DataFrame({"a": range(1000), "b": [float(i) for i in range(1000)]}) + + +@pytest.fixture +def file_backend(tmp_path): + from cachekit.backends.file import FileBackend + from cachekit.backends.file.config import FileBackendConfig + + return FileBackend(FileBackendConfig(cache_dir=str(tmp_path), max_size_mb=64, max_value_mb=32)) + + +def _operation_handler(backend) -> CacheOperationHandler: + handler = CacheOperationHandler(CacheSerializationHandler(serializer_name="arrow"), CacheKeyGenerator()) + handler.set_cache_handler(StandardCacheHandler(backend)) + return handler + + +class _PlainBackend: + """Minimal BaseBackend without set_streaming (the fallback target).""" + + def __init__(self) -> None: + self.store: dict[str, bytes] = {} + + def get(self, key): + return self.store.get(key) + + def set(self, key, value, ttl=None): + self.store[key] = bytes(value) + + def delete(self, key): + return self.store.pop(key, None) is not None + + def exists(self, key): + return key in self.store + + def health_check(self): + return True, {} + + +@pytest.mark.unit +class TestSupportsStreamingWrite: + def test_arrow_plaintext_is_eligible(self) -> None: + sh = CacheSerializationHandler(serializer_name="arrow") + assert sh.supports_streaming_write() is True + + def test_default_serializer_not_eligible(self) -> None: + sh = CacheSerializationHandler(serializer_name="default") + assert sh.supports_streaming_write() is False + + def test_encrypted_arrow_not_eligible(self) -> None: + """AES-256-GCM emits its tag only after the whole ciphertext — the secure path stays buffered.""" + sh = CacheSerializationHandler(serializer_name="arrow") + sh.encryption = True + assert sh.supports_streaming_write() is False + + def test_backend_type_guard(self, file_backend) -> None: + assert supports_streaming_write(file_backend) is True + assert supports_streaming_write(_PlainBackend()) is False + + +@pytest.mark.unit +class TestWriteSerializedTo: + def test_envelope_byte_identical_to_buffered(self, df) -> None: + """The streamed frame must be indistinguishable from serialize_data's output — + same prefix, same [checksum][IPC] payload — so every existing read path Just Works.""" + sh = CacheSerializationHandler(serializer_name="arrow") + buffered = sh.serialize_data(df, cache_key="k") + sink = io.BytesIO() + sh.write_serialized_to(sink, df) + assert sink.getvalue() == buffered + + def test_streamed_envelope_deserializes(self, df) -> None: + sh = CacheSerializationHandler(serializer_name="arrow") + sink = io.BytesIO() + sh.write_serialized_to(sink, df) + out = sh.deserialize_data(sink.getvalue(), cache_key="k") + pd.testing.assert_frame_equal(out, df) + + def test_max_value_size_budget_enforced_mid_stream(self, df, monkeypatch) -> None: + """The L2 oversized-entry ceiling (issue #163) applies to streamed writes too, + and aborts mid-stream instead of after materializing.""" + from cachekit.config.singleton import reset_settings + + monkeypatch.setenv("CACHEKIT_MAX_VALUE_SIZE", "1024") + reset_settings() + try: + sh = CacheSerializationHandler(serializer_name="arrow") + with pytest.raises(ValueError, match="budget"): + sh.write_serialized_to(io.BytesIO(), df) + finally: + reset_settings() + + def test_non_dataframe_raises_type_error(self) -> None: + sh = CacheSerializationHandler(serializer_name="arrow") + with pytest.raises(TypeError, match="only supports DataFrames"): + sh.write_serialized_to(io.BytesIO(), 42) + + +@pytest.mark.unit +class TestStandardCacheHandlerSetStreaming: + def test_returns_none_when_backend_lacks_set_streaming(self) -> None: + ch = StandardCacheHandler(_PlainBackend()) # type: ignore[arg-type] + assert ch.set_streaming("k", lambda sink: None) is None + + def test_delegates_to_backend_when_supported(self) -> None: + backend = MagicMock() + ch = StandardCacheHandler(backend) + cb = MagicMock() + assert ch.set_streaming("k", cb, ttl=60) is True + backend.set_streaming.assert_called_once_with("k", cb, 60) + + def test_returns_false_on_backend_error(self) -> None: + backend = MagicMock() + backend.set_streaming.side_effect = BackendError("boom") + ch = StandardCacheHandler(backend) + assert ch.set_streaming("k", lambda sink: None) is False + + def test_returns_false_on_producer_error(self) -> None: + """A serialization failure inside the callback degrades to False — never raises + out of the cache write (a cache must not break the wrapped function).""" + backend = MagicMock() + backend.set_streaming.side_effect = ValueError("payload budget") + ch = StandardCacheHandler(backend) + assert ch.set_streaming("k", lambda sink: None) is False + + @pytest.mark.asyncio + async def test_async_returns_none_when_unsupported(self) -> None: + ch = StandardCacheHandler(_PlainBackend()) # type: ignore[arg-type] + assert await ch.set_streaming_async("k", lambda sink: None) is None + + +@pytest.mark.unit +class TestStoreResultRouting: + def test_streams_to_file_backend_and_skips_l1(self, df, file_backend) -> None: + """Streaming success returns None: the envelope must never be copied into L1 + (mirrors the mmap read path's L1 exclusion, #171).""" + op = _operation_handler(file_backend) + assert op.store_result("k1", df, ttl=60) is None + hit = op.get_cached_value("k1") + assert hit is not None + pd.testing.assert_frame_equal(hit[1], df) + + def test_streamed_bytes_identical_to_buffered_set(self, df, file_backend) -> None: + op = _operation_handler(file_backend) + op.store_result("k1", df, ttl=60) + stored = file_backend.get("k1") + assert stored == op.serialization_handler.serialize_data(df, cache_key="k1") + + def test_falls_back_to_buffered_set_and_returns_bytes_for_l1(self, df) -> None: + backend = _PlainBackend() + op = _operation_handler(backend) + ret = op.store_result("k1", df, ttl=60) + assert ret is not None # L1 backfill contract unchanged on the buffered path + assert backend.store["k1"] == ret + + def test_non_arrow_serializer_stays_buffered(self, file_backend) -> None: + op = CacheOperationHandler(CacheSerializationHandler(serializer_name="default"), CacheKeyGenerator()) + op.set_cache_handler(StandardCacheHandler(file_backend)) + ret = op.store_result("k2", {"a": 1}, ttl=60) + assert ret is not None + assert op.get_cached_value("k2")[1] == {"a": 1} + + def test_stale_ttl_stays_buffered(self, df, file_backend, monkeypatch) -> None: + """SWR writes carry stale_ttl, which set_streaming has no channel for — buffered path.""" + op = _operation_handler(file_backend) + called = {"streaming": False} + original = file_backend.set_streaming + monkeypatch.setattr( + file_backend, "set_streaming", lambda *a, **k: called.__setitem__("streaming", True) or original(*a, **k) + ) + ret = op.store_result("k3", df, ttl=60, stale_ttl=30) + assert called["streaming"] is False + assert ret is not None # buffered path returns bytes + + def test_streaming_failure_does_not_retry_buffered(self, df, file_backend, monkeypatch) -> None: + """A failed stream must NOT re-materialize the payload via set(bytes).""" + op = _operation_handler(file_backend) + + def explode(key, cb, ttl=None): + raise BackendError("disk full") + + monkeypatch.setattr(file_backend, "set_streaming", explode) + set_spy = MagicMock(wraps=file_backend.set) + monkeypatch.setattr(file_backend, "set", set_spy) + assert op.store_result("k4", df, ttl=60) is None + set_spy.assert_not_called() + assert file_backend.get("k4") is None + + @pytest.mark.asyncio + async def test_async_streams_and_roundtrips(self, df, file_backend) -> None: + op = _operation_handler(file_backend) + assert await op.store_result_async("k5", df, ttl=60) is None + hit = await op.get_cached_value_async("k5") + assert hit is not None + pd.testing.assert_frame_equal(hit[1], df) + + +@pytest.mark.unit +class TestFileBackendSetStreaming: + def test_roundtrip_matches_set(self, file_backend) -> None: + payload = os.urandom(64 * 1024) + file_backend.set_streaming("k", lambda f: f.write(payload), ttl=60) + assert file_backend.get("k") == payload + + def test_seek_patch_within_payload(self, file_backend) -> None: + """Producers patch already-written bytes via tell()-anchored seeks (the checksum + pattern); the backend's 14-byte header must be untouched by it.""" + + def writer(f) -> None: + start = f.tell() + f.write(b"\x00" * 8) + f.write(b"body-bytes") + end = f.tell() + f.seek(start) + f.write(b"CHECKSUM") + f.seek(end) + + file_backend.set_streaming("k", writer, ttl=60) + assert file_backend.get("k") == b"CHECKSUM" + b"body-bytes" + + def test_ttl_zero_and_none_mean_no_expiry(self, file_backend) -> None: + file_backend.set_streaming("k", lambda f: f.write(b"x"), ttl=None) + assert file_backend.get("k") == b"x" + file_backend.set_streaming("k2", lambda f: f.write(b"y"), ttl=0) + assert file_backend.get("k2") == b"y" + + def test_invalid_ttl_rejected_before_writing(self, file_backend, tmp_path) -> None: + with pytest.raises(BackendError, match="out of range"): + file_backend.set_streaming("k", lambda f: f.write(b"x"), ttl=-1) + assert file_backend.get("k") is None + assert not [f for f in os.listdir(tmp_path) if ".tmp." in f] + + def test_producer_exception_discards_partial_and_keeps_old_value(self, file_backend, tmp_path) -> None: + class BoomError(Exception): + pass + + file_backend.set("k", b"old", ttl=60) + + def exploding(f) -> None: + f.write(b"partial") + raise BoomError("mid-stream") + + with pytest.raises(BoomError): # unwrapped, per BufferWritableBackend contract + file_backend.set_streaming("k", exploding, ttl=60) + assert file_backend.get("k") == b"old" + assert not [f for f in os.listdir(tmp_path) if ".tmp." in f] + + def test_max_value_mb_enforced(self, tmp_path) -> None: + from cachekit.backends.file import FileBackend + from cachekit.backends.file.config import FileBackendConfig + + backend = FileBackend(FileBackendConfig(cache_dir=str(tmp_path), max_size_mb=4, max_value_mb=1)) + oversized = b"x" * (1024 * 1024 + 1) + with pytest.raises(BackendError, match="max_value_mb"): + backend.set_streaming("k", lambda f: f.write(oversized), ttl=60) + assert backend.get("k") is None + assert not [f for f in os.listdir(tmp_path) if ".tmp." in f] + + def test_mmap_read_path_reads_streamed_entry(self, tmp_path) -> None: + """Write-side streaming and read-side mmap (#171) compose: an uncompressed Arrow + entry streamed to disk is served back through get_buffer zero-copy.""" + from cachekit.backends.file import FileBackend + from cachekit.backends.file.config import FileBackendConfig + from cachekit.serializers.arrow_serializer import ArrowSerializer + + backend = FileBackend(FileBackendConfig(cache_dir=str(tmp_path), max_size_mb=64, max_value_mb=32)) + sh = CacheSerializationHandler(serializer_name="arrow") + sh._base_serializer = ArrowSerializer(compression=None) # mmap needs plaintext, uncompressed + frame = pd.DataFrame({"a": range(100)}) + + backend.set_streaming("k", lambda sink: sh.write_serialized_to(sink, frame), ttl=60) + handle = backend.get_buffer("k") + assert handle is not None + try: + out = sh.deserialize_data(handle.view, cache_key="k") + finally: + handle.close() + pd.testing.assert_frame_equal(out, frame)