|
| 1 | +"""Pin: ``decode_text``'s memoryview branch peeks a SHORT probe |
| 2 | +window first; only the long-cell case pays the 64 KiB |
| 3 | +materialise. |
| 4 | +
|
| 5 | +The prior shape always materialised |
| 6 | +``min(data_len, max_size + 1, _TEXT_ONE_SHOT_MAX)`` bytes per |
| 7 | +cell — which collapsed to 64 KiB for any cell sitting in a |
| 8 | +multi-MiB body buffer (a 1M-row RowsResponse is the canonical |
| 9 | +case). A 16-byte TEXT cell at offset 100 in a 1 MiB body paid |
| 10 | +~64 KiB of transient allocation per cell × 1M cells = ~64 GiB |
| 11 | +of allocator churn on the loop thread purely to support a |
| 12 | +``bytes.find`` that would have terminated in 16 bytes. |
| 13 | +
|
| 14 | +The fix peeks a short window (a few hundred bytes) first; when |
| 15 | +NUL is found there — the common case for short TEXT cells — |
| 16 | +the function returns without ever materialising the 64 KiB |
| 17 | +probe. Cells that genuinely exceed the short probe escalate to |
| 18 | +the existing 64 KiB path. The C reference uses ``strnlen`` on |
| 19 | +the cursor's in-place buffer (zero allocation); this is the |
| 20 | +Python adaptive-probe equivalent that mirrors the property |
| 21 | +without requiring a callsite-API change. |
| 22 | +""" |
| 23 | + |
| 24 | +from __future__ import annotations |
| 25 | + |
| 26 | +import pytest |
| 27 | + |
| 28 | +from dqlitewire.exceptions import DecodeError |
| 29 | +from dqlitewire.types import decode_text |
| 30 | + |
| 31 | + |
| 32 | +def test_decode_text_short_cell_in_long_buffer_does_not_materialise_64kib() -> None: |
| 33 | + """A 16-byte TEXT cell inside a 1 MiB memoryview must be |
| 34 | + decoded without materialising the full 64 KiB probe window. |
| 35 | +
|
| 36 | + We assert this structurally by patching ``bytes`` in the |
| 37 | + module's namespace to track the largest argument it received. |
| 38 | + Pre-fix: ``bytes(data[:65536])`` runs every call, so the max |
| 39 | + materialise size is ~64 KiB. Post-fix: the short probe stays |
| 40 | + well under 1 KiB. |
| 41 | + """ |
| 42 | + from dqlitewire import types as types_mod |
| 43 | + |
| 44 | + # 16-byte TEXT cell at the head of a 1 MiB buffer. Decode with |
| 45 | + # the DEFAULT ``max_size`` (the wire's 64 MiB row-cell ceiling) |
| 46 | + # so the existing probe-then-escalate path picks |
| 47 | + # ``min(data_len, max_size + 1, _TEXT_ONE_SHOT_MAX) = 65536`` |
| 48 | + # — the 64 KiB materialise we want to avoid for short cells. |
| 49 | + payload = b"short-text\x00\x00\x00\x00\x00\x00" + b"\xff" * (1 << 20) |
| 50 | + view = memoryview(payload) |
| 51 | + |
| 52 | + largest_bytes_size: list[int] = [0] |
| 53 | + real_bytes = bytes |
| 54 | + |
| 55 | + def _tracking_bytes(arg: object = b"") -> bytes: |
| 56 | + if isinstance(arg, memoryview): |
| 57 | + largest_bytes_size[0] = max(largest_bytes_size[0], len(arg)) |
| 58 | + return real_bytes(arg) # type: ignore[call-overload, no-any-return] |
| 59 | + |
| 60 | + monkeypatch_token = types_mod.__dict__.get("bytes", real_bytes) |
| 61 | + types_mod.__dict__["bytes"] = _tracking_bytes |
| 62 | + try: |
| 63 | + text, consumed = decode_text(view, label="TEXT") |
| 64 | + finally: |
| 65 | + if monkeypatch_token is real_bytes: |
| 66 | + types_mod.__dict__.pop("bytes", None) |
| 67 | + else: |
| 68 | + types_mod.__dict__["bytes"] = monkeypatch_token |
| 69 | + |
| 70 | + assert text == "short-text" |
| 71 | + assert consumed == 16 # 10 chars + NUL + 5 pad = 16 |
| 72 | + |
| 73 | + # Pre-fix: 65536 bytes materialised. Post-fix: short probe far |
| 74 | + # smaller (a few hundred bytes). Pin at 4 KiB as the upper |
| 75 | + # acceptable bound for the short probe — well under the 64 KiB |
| 76 | + # pre-fix value, well above any sensible short-probe choice. |
| 77 | + assert largest_bytes_size[0] <= 4096, ( |
| 78 | + f"decode_text materialised {largest_bytes_size[0]} bytes for a " |
| 79 | + f"16-byte cell in a 1 MiB buffer; the short-probe optimisation " |
| 80 | + f"should keep per-cell materialise well under the 64 KiB " |
| 81 | + f"one-shot ceiling for short cells." |
| 82 | + ) |
| 83 | + |
| 84 | + |
| 85 | +def test_decode_text_long_cell_still_decodes_correctly() -> None: |
| 86 | + """Regression guard: a cell that EXCEEDS the short probe but |
| 87 | + fits in the 64 KiB one-shot ceiling must still decode |
| 88 | + correctly via the escalated path. |
| 89 | + """ |
| 90 | + long_text = "x" * 8192 |
| 91 | + payload = long_text.encode("utf-8") + b"\x00" + b"\x00" * 7 |
| 92 | + view = memoryview(payload) |
| 93 | + text, consumed = decode_text(view, label="TEXT", max_size=16_384) |
| 94 | + assert text == long_text |
| 95 | + |
| 96 | + |
| 97 | +def test_decode_text_very_long_cell_uses_chunked_path() -> None: |
| 98 | + """Regression guard: a cell that exceeds the 64 KiB one-shot |
| 99 | + ceiling escalates to the chunked path and still decodes.""" |
| 100 | + from dqlitewire.types import _TEXT_ONE_SHOT_MAX |
| 101 | + |
| 102 | + very_long = "x" * (_TEXT_ONE_SHOT_MAX + 100) |
| 103 | + payload = very_long.encode("utf-8") + b"\x00" + b"\x00" * 7 |
| 104 | + view = memoryview(payload) |
| 105 | + text, consumed = decode_text(view, label="TEXT", max_size=_TEXT_ONE_SHOT_MAX * 2) |
| 106 | + assert text == very_long |
| 107 | + |
| 108 | + |
| 109 | +def test_decode_text_rejects_cell_exceeding_max_size_in_long_buffer() -> None: |
| 110 | + """A cell that fits in the buffer but exceeds ``max_size`` |
| 111 | + raises ``DecodeError`` — verify the short-probe optimisation |
| 112 | + does not swallow the max-size check. |
| 113 | + """ |
| 114 | + payload = b"a" * 100 + b"\x00" + b"\x00" * 7 + b"\xff" * 1024 |
| 115 | + view = memoryview(payload) |
| 116 | + with pytest.raises(DecodeError, match="exceeds maximum"): |
| 117 | + decode_text(view, label="TEXT", max_size=50) |
| 118 | + |
| 119 | + |
| 120 | +def test_decode_text_rejects_unterminated_cell() -> None: |
| 121 | + """A buffer with no NUL terminator at all raises |
| 122 | + ``DecodeError("...not null-terminated")``. The short-probe |
| 123 | + optimisation must escalate to find the missing-NUL case. |
| 124 | + """ |
| 125 | + payload = b"a" * 100 # no NUL anywhere |
| 126 | + view = memoryview(payload) |
| 127 | + with pytest.raises(DecodeError): |
| 128 | + decode_text(view, label="TEXT", max_size=1024) |
0 commit comments