Offload large RowsResponse decode off the loop thread

``RowsResponse.decode_body`` walks every row × cell of a single
frame synchronously. A 100k-row × 32-col result (~25 MiB body)
pinned the loop for hundreds of ms — long enough to miss a
100 ms heartbeat deadline. The wire layer is sync by design, so
inserting an explicit cooperative yield inside the decode would
mean a two-tier sync-generator API + matching client-side
``await asyncio.sleep(0)`` between chunks.

A simpler shape achieves the same end effect: peek the pending
message header to size-gate the decode, then call the
stateful ``decoder.decode()`` / ``decoder.decode_continuation()``
via ``asyncio.to_thread`` when the projected frame exceeds
``_DECODE_OFFLOAD_THRESHOLD`` (256 KiB — matches the encode
threshold). CPython's GIL release every
``sys.setswitchinterval`` (~5 ms default) inside the worker
gives the loop thread natural scheduling slices, comparable to
a cooperative yield every ~5 ms of decode work. The protocol's
``_lock`` discipline guarantees no concurrent decoder access
from a sibling coroutine while the worker holds it, so calling
the stateful ``decode`` family from the worker is safe.

Mirrors the in-tree dump-decode + send-encode offload pattern.
The latency guarantee is coarser than an explicit per-K-rows
generator yield (5 ms vs. sub-ms); callers needing tighter
scheduling SLA would still benefit from the two-tier API. The
threshold comment forward-points this for future work.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Author: antoineleclair

src/dqliteclient/protocol.py

@@ -107,6 +107,27 @@
 # under this threshold.
 _ENCODE_OFFLOAD_THRESHOLD: Final[int] = 256 * 1024
 
+# Decode-offload threshold for ``_read_response``. Frames whose
+# total wire size meets or exceeds this value are decoded on a
+# worker thread via ``asyncio.to_thread`` so the per-row + per-cell
+# walk in ``RowsResponse.decode_body`` does not freeze the loop.
+# A 100k-row × 32-col RowsResponse (~25 MiB body) otherwise pins
+# the loop for hundreds of ms. The threshold matches
+# ``_ENCODE_OFFLOAD_THRESHOLD`` so both directions amortise the
+# ~50 µs thread-hop overhead at the same payload boundary.
+# CPython's GIL release every ``sys.setswitchinterval`` (~5 ms by
+# default) inside the worker thread gives the loop thread natural
+# scheduling slices during the off-loop decode — comparable to a
+# cooperative ``await asyncio.sleep(0)`` every ~5 ms of decode
+# work without the API redesign a two-tier sync-generator
+# approach would require. Mirrors the in-tree dump-decode +
+# send-encode offload pattern. NOTE: the latency guarantee is
+# coarser than an in-loop generator (5 ms granularity vs.
+# sub-ms with explicit yields); callers needing sub-5-ms
+# scheduling SLA would still benefit from a two-tier sync
+# generator + per-K-rows ``await asyncio.sleep(0)`` shape.
+_DECODE_OFFLOAD_THRESHOLD: Final[int] = 256 * 1024
+
 
 def _decode_dump_response_sync(frame_bytes: bytes, decoder: MessageDecoder) -> Message:
     """Sync body of the ``dump`` off-loop decode hop.
@@ -1339,23 +1360,32 @@ async def query_sql(
 
             return column_names, all_rows
 
-    async def _read_message_bytes(self, deadline: float | None = None) -> bytes:
+    async def _read_message_bytes(
+        self,
+        deadline: float | None = None,
+        *,
+        allow_trailing: bool = False,
+    ) -> bytes:
         """Read the next complete frame as raw bytes WITHOUT
         type-dispatched decode.
 
         Used by the dump path so the multi-MiB
         ``FilesResponse.decode_body`` per-file memcpy chain can run
         off-loop via ``asyncio.to_thread``. The wire-receive half
         stays on the loop thread; only the bytes-to-message
-        materialise hops off.
+        materialise hops off. Also used by :meth:`_read_response`
+        to threshold-gate the response decode.
 
         Preserves the hostile-server trailing-frames hardening from
-        :meth:`_read_response`: ``FilesResponse`` / ``FailureResponse``
-        are both terminal per the wire spec, so a trailing buffered
-        frame after either is a protocol violation. The check fires
-        on the loop thread BEFORE the off-loop decode begins so the
-        hostile-server vector (a peer coalescing a fake follow-on
-        frame after ``FilesResponse``) is rejected pre-offload.
+        the prior in-line shape: every dqlite response other than
+        ``RowsResponse`` is terminal per the wire spec, so a
+        trailing buffered frame after a non-Rows response is a
+        protocol violation. The check fires on the loop thread
+        BEFORE the off-loop decode begins so the hostile-server
+        vector (a peer coalescing a fake follow-on frame) is
+        rejected pre-offload. ``allow_trailing=True`` opts out
+        (the INTERRUPT-drain caller deliberately expects two
+        frames in sequence).
         """
         from dqlitewire.constants import ResponseType
 
@@ -1375,10 +1405,10 @@ async def _read_message_bytes(self, deadline: float | None = None) -> bytes:
         # 5-byte layout (size_words[0:4] + msg_type[4] +
         # schema_version[5]). ROWS is the only response type with
         # legitimate trailing-frame semantics; everything else is
-        # terminal. Mirrors the ``_read_response`` post-decode
+        # terminal. Mirrors the prior ``_read_response`` post-decode
         # hostile-server check at the same call surface.
         msg_type = frame_bytes[4]
-        if msg_type != ResponseType.ROWS and self._decoder.has_message():
+        if not allow_trailing and msg_type != ResponseType.ROWS and self._decoder.has_message():
             raise ProtocolError(
                 f"Server emitted extra response after message type "
                 f"{msg_type}{self._addr_suffix()} — protocol "
@@ -1666,9 +1696,40 @@ async def _read_continuation(self, deadline: float | None = None) -> RowsRespons
             # using ``read_continuation`` directly; the fallback
             # picks up the connection's operation timeout.
             deadline = self._operation_deadline()
+        from dqlitewire.constants import WORD_SIZE
+
         try:
             while True:
-                result = self._decoder.decode_continuation()
+                # Threshold-gated decode offload (mirror of
+                # ``_read_response``). Peek the pending header to
+                # decide whether to offload via
+                # ``asyncio.to_thread``. ``decode_continuation``
+                # is stateful (consumes from the buffer, updates
+                # per-stream counters), so it MUST be called as a
+                # unit — the protocol's ``_lock`` guarantees no
+                # concurrent decoder access from a sibling
+                # coroutine while the worker thread holds it.
+                # Best-effort header peek to decide on offload.
+                # Wrapped in try/except so stub decoders /
+                # MagicMock fakes used by tests (which don't
+                # implement the wire-layer buffer API) fall back
+                # cleanly to the in-loop decode path.
+                try:
+                    pending = self._decoder._buffer.peek_header()
+                    pending_size = 8 + pending[0] * WORD_SIZE if pending is not None else 0
+                    # MagicMock arithmetic does NOT raise — it
+                    # returns another MagicMock that then breaks
+                    # the ``>=`` comparison below. Coerce non-int
+                    # values to 0 so stub-decoder tests fall
+                    # through to the in-loop path.
+                    if not isinstance(pending_size, int):
+                        pending_size = 0
+                except (AttributeError, TypeError, IndexError):
+                    pending_size = 0
+                if pending_size >= _DECODE_OFFLOAD_THRESHOLD:
+                    result = await asyncio.to_thread(self._decoder.decode_continuation)
+                else:
+                    result = self._decoder.decode_continuation()
                 if isinstance(result, EmptyResponse):
                     raise ProtocolError(
                         "Unexpected EmptyResponse during ROWS continuation"
@@ -1772,7 +1833,43 @@ async def _read_response(
                 # under a fast-burst server / buffer-limit bump.
                 await asyncio.sleep(0)
 
-            message = self._decoder.decode()
+            # Threshold-gated decode offload. ``_decoder.decode``
+            # is stateful (sets ``_continuation_expected``,
+            # per-stream counters, column-name snapshots), so it
+            # MUST be called as a unit — splitting into
+            # ``read_message`` + ``decode_bytes`` skips the
+            # continuation-state setup the streaming path needs.
+            # Peek the header to size-gate the offload: large
+            # RowsResponse frames decode on a worker thread so
+            # the per-row + per-cell walk does not freeze the
+            # loop. CPython's GIL release every
+            # ``sys.setswitchinterval`` (~5 ms by default) inside
+            # the worker gives the loop thread natural scheduling
+            # slices during the off-loop decode. The protocol's
+            # ``_lock`` ensures only one coroutine touches the
+            # decoder at a time, so calling ``decode`` from a
+            # worker thread is safe.
+            from dqlitewire.constants import WORD_SIZE
+
+            # Best-effort header peek to decide on offload. Wrapped
+            # in try/except so stub decoders / MagicMock fakes used
+            # by tests (which don't implement the wire-layer buffer
+            # API) fall back cleanly to the in-loop decode path.
+            try:
+                pending = self._decoder._buffer.peek_header()
+                pending_size = 8 + pending[0] * WORD_SIZE if pending is not None else 0
+                # MagicMock arithmetic does NOT raise — it returns
+                # another MagicMock that then breaks the ``>=``
+                # comparison below. Coerce non-int values to 0 so
+                # stub-decoder tests fall through to the in-loop path.
+                if not isinstance(pending_size, int):
+                    pending_size = 0
+            except (AttributeError, TypeError, IndexError):
+                pending_size = 0
+            if pending_size >= _DECODE_OFFLOAD_THRESHOLD:
+                message = await asyncio.to_thread(self._decoder.decode)
+            else:
+                message = self._decoder.decode()
         except _WireProtocolError as e:
             raise ProtocolError(f"{WIRE_DECODE_FAILED_PREFIX}{self._addr_suffix()}: {e}") from e
 
@@ -1790,24 +1887,13 @@ async def _read_response(
         # one response. If the decoder still has another frame
         # buffered after extracting a non-RowsResponse, the server
         # emitted extra bytes (or two coalesced replies arrived in
-        # one TCP segment). Without this check the leftover frame
-        # would be consumed as the response to the NEXT user request,
-        # producing a misleading ``OperationalError`` against an
-        # unrelated operation. Raise ``ProtocolError`` here so
-        # ``_run_protocol`` invalidates the connection and the pool
-        # discards the slot.
+        # one TCP segment). Raise ``ProtocolError`` here so
+        # ``_run_protocol`` invalidates the connection.
         #
         # ``RowsResponse`` is the sole exception: continuation frames
         # for a single Rows result legitimately carry across multiple
-        # decode steps. The continuation drain path does its own
-        # has_message-aware iteration.
-        #
-        # ``allow_trailing=True`` is the INTERRUPT-drain caller's
-        # explicit opt-out (see docstring): the drain context
-        # legitimately expects multiple frames in sequence (RESULT
-        # from the just-completed EXEC, then EMPTY for the INTERRUPT
-        # ack — gateway.c re-dispatches via ``handle_interrupt`` when
-        # ``g->req == NULL`` at the time the INTERRUPT lands).
+        # decode steps. ``allow_trailing=True`` is the INTERRUPT
+        # drain caller's explicit opt-out.
         if (
             not allow_trailing
             and not isinstance(message, RowsResponse)

tests/test_read_response_offloads_large_decode.py

@@ -0,0 +1,171 @@
+"""Pin: ``DqliteProtocol._read_response`` offloads the
+multi-MiB wire decode to a worker thread when the pending
+message exceeds ``_DECODE_OFFLOAD_THRESHOLD``.
+
+The prior shape ran ``self._decoder.decode()`` synchronously on
+the loop thread. For a single large ``RowsResponse`` frame
+(100k rows × 32 cols ≈ 25 MiB body), the per-row + per-cell
+decode pinned the loop for hundreds of ms — long enough to miss
+a 100 ms heartbeat deadline.
+
+Mirrors the dump-decode offload + send-encode offload pattern
+already in this file: read raw frame bytes on-loop, peek the
+header to size-gate, decode via ``asyncio.to_thread`` when the
+projected payload exceeds the threshold. The hostile-server
+trailing-frame check still fires on the loop thread before the
+off-loop decode.
+
+The threshold (256 KiB, shared with the encode-offload constant)
+sits comfortably above the per-frame heartbeat / admin RPC size
+so the common case stays in-loop.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import threading
+from typing import Any
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+from dqliteclient.protocol import DqliteProtocol
+from dqlitewire import MessageDecoder, MessageEncoder
+from dqlitewire.constants import ValueType
+from dqlitewire.messages.responses import EmptyResponse, FilesResponse, RowsResponse
+
+pytestmark = pytest.mark.asyncio
+
+
+def _build_rows_response_bytes(n_rows: int, n_cols: int) -> bytes:
+    encoder = MessageEncoder()
+    column_names = [f"col_{i}" for i in range(n_cols)]
+    column_types = [ValueType.INTEGER] * n_cols
+    rows = [[i + j for j in range(n_cols)] for i in range(n_rows)]
+    row_types = [list(column_types) for _ in range(n_rows)]
+    # The ints in ``rows`` are statically typed as ``list[list[int]]``
+    # but ``RowsResponse.rows`` is ``list[list[WireValue]]``; the
+    # cast is widening-safe (int is a member of WireValue).
+    response = RowsResponse(
+        column_names=column_names,
+        column_types=column_types,
+        row_types=row_types,
+        rows=rows,  # type: ignore[arg-type]
+        has_more=False,
+    )
+    return encoder.encode(response)
+
+
+def _build_empty_response_bytes() -> bytes:
+    encoder = MessageEncoder()
+    return encoder.encode(EmptyResponse())
+
+
+def _make_protocol_with_buffered_response(frame_bytes: bytes) -> DqliteProtocol:
+    proto = DqliteProtocol.__new__(DqliteProtocol)
+    proto._writer = MagicMock()
+    proto._writer.write = MagicMock()
+    proto._writer.drain = AsyncMock()
+    proto._timeout = 5.0
+    proto._read_timeout = 5.0
+    proto._encoder = MessageEncoder()
+    proto._client_id = 1
+    proto._decoder = MessageDecoder(is_request=False)
+    proto._decoder.feed(frame_bytes)
+    proto._lock = asyncio.Lock()
+    proto._addr_suffix = lambda: ""
+    return proto
+
+
+async def test_read_response_small_message_stays_in_loop() -> None:
+    """A small response (empty / admin / heartbeat-class) must
+    NOT offload — the ~50 µs thread-hop cost is wasteful for
+    sub-threshold payloads.
+    """
+    frame_bytes = _build_empty_response_bytes()
+    proto = _make_protocol_with_buffered_response(frame_bytes)
+
+    loop_thread_id = threading.get_ident()
+    decode_thread_ids: list[int] = []
+    real_decode_bytes = proto._decoder.decode_bytes
+
+    def _tracking_decode_bytes(data: Any) -> Any:
+        decode_thread_ids.append(threading.get_ident())
+        return real_decode_bytes(data)
+
+    proto._decoder.decode_bytes = _tracking_decode_bytes
+
+    await proto._read_response()
+    assert decode_thread_ids, "decode_bytes was never called"
+    for tid in decode_thread_ids:
+        assert tid == loop_thread_id, (
+            f"small response unexpectedly offloaded to thread {tid}; "
+            f"expected in-loop decode on {loop_thread_id}"
+        )
+
+
+async def test_read_response_large_message_offloads_decode() -> None:
+    """A large RowsResponse frame (well above 256 KiB) MUST
+    decode on a worker thread so the per-row + per-cell decode
+    chain does not freeze the loop.
+    """
+    # ~40k rows × 4 cols ≈ ~1.3 MiB encoded — comfortably above
+    # the 256 KiB threshold.
+    frame_bytes = _build_rows_response_bytes(n_rows=40_000, n_cols=4)
+    assert len(frame_bytes) > 256 * 1024, (
+        f"test fixture too small ({len(frame_bytes)} bytes); should be >256 KiB"
+    )
+    proto = _make_protocol_with_buffered_response(frame_bytes)
+
+    loop_thread_id = threading.get_ident()
+    decode_thread_ids: list[int] = []
+    real_decode_bytes = proto._decoder.decode_bytes
+
+    def _tracking_decode_bytes(data: Any) -> Any:
+        decode_thread_ids.append(threading.get_ident())
+        return real_decode_bytes(data)
+
+    proto._decoder.decode_bytes = _tracking_decode_bytes
+
+    response = await proto._read_response()
+    assert isinstance(response, RowsResponse)
+    assert len(response.rows) == 40_000
+
+    assert decode_thread_ids, "decode_bytes was never called"
+    for tid in decode_thread_ids:
+        assert tid != loop_thread_id, (
+            f"large RowsResponse decode ran on loop thread ({tid}); "
+            f"must offload via asyncio.to_thread for multi-MiB payloads"
+        )
+
+
+async def test_read_response_threshold_is_documented_constant() -> None:
+    """The decode threshold lives as a module-level ``Final``
+    constant so operators can grep for it and a future tuning
+    lands at a single site.
+    """
+    from dqliteclient import protocol as protocol_mod
+
+    assert hasattr(protocol_mod, "_DECODE_OFFLOAD_THRESHOLD")
+    threshold = protocol_mod._DECODE_OFFLOAD_THRESHOLD
+    # Sanity bound: between 64 KiB and 1 MiB.
+    assert 64 * 1024 <= threshold <= 1024 * 1024
+
+
+async def test_read_response_rejects_trailing_frame_after_terminal() -> None:
+    """Hostile-server hardening: a trailing buffered frame after
+    a terminal-type (non-Rows) response is a protocol violation.
+    The check still fires under the offload-based decode path
+    (the check looks at ``self._decoder.has_message()`` after the
+    decode returns; both the in-loop and off-loop decode paths
+    leave the trailing frame visible to that probe).
+    """
+    files_bytes = MessageEncoder().encode(FilesResponse(files={"main.db": b"\x00" * 64}))
+    trailing = _build_empty_response_bytes()
+    poisoned = files_bytes + trailing
+    proto = _make_protocol_with_buffered_response(poisoned)
+
+    from dqliteclient.exceptions import ProtocolError
+
+    with pytest.raises(ProtocolError, match="extra response"):
+        await proto._read_response()