fix: retry 504 GatewayTimeout and resume truncated source reads

s3proxy/handlers/base.py

@@ -43,10 +43,27 @@
 # Re-issuing the same ranged GET on a fresh connection is always safe.
 SOURCE_READ_ATTEMPTS = int(os.environ.get("S3PROXY_SOURCE_READ_ATTEMPTS", "4"))
 SOURCE_READ_BACKOFF_SEC = float(os.environ.get("S3PROXY_SOURCE_READ_BACKOFF", "0.5"))
-
+# Chunk size for accumulating a source body so a mid-body truncation keeps what
+# already arrived (see read_source_bytes).
+SOURCE_READ_CHUNK_SIZE = int(os.environ.get("S3PROXY_SOURCE_READ_CHUNK_SIZE", str(1024 * 1024)))
+
+# GatewayTimeout/504 belongs here alongside 502/503: Hetzner returns it when its
+# own server-side copy exceeds an internal deadline. Copy latency is heavy-tailed
+# (measured p50 1.14s, p90 3.31s, max 61.86s over 572 UploadPartCopy calls), so a
+# 504 is transient congestion, not a permanent condition. botocore's own retries
+# all fire inside the same congestion window ("reached max retries: 3") and
+# exhaust; retrying here with exponential backoff gives the backend time to clear.
+# UploadPartCopy is idempotent for a given PartNumber+range, so the retry is safe.
 _RETRYABLE_S3_ERROR_CODES = frozenset(
-    {"InternalError", "SlowDown", "RequestTimeout", "ServiceUnavailable", "BadGateway"}
-    | {"500", "502", "503"}
+    {
+        "InternalError",
+        "SlowDown",
+        "RequestTimeout",
+        "ServiceUnavailable",
+        "BadGateway",
+        "GatewayTimeout",
+    }
+    | {"500", "502", "503", "504"}
 )
 
 _RETRYABLE_TRANSPORT_ERRORS = (
@@ -72,18 +89,60 @@ def is_retryable_source_error(exc: BaseException) -> bool:
     return False
 
 
+def _parse_byte_range(range_header: str) -> tuple[int, int | None]:
+    """Parse a `bytes=start-end` header into (start, end). end is None if open-ended."""
+    spec = range_header.removeprefix("bytes=").strip()
+    start_s, _, end_s = spec.partition("-")
+    return int(start_s), (int(end_s) if end_s else None)
+
+
 async def read_source_bytes(
     client: S3Client,
     bucket: str,
     key: str,
     range_header: str | None = None,
 ) -> bytes:
-    """GET an object (or range) and read the full body, retrying truncated reads."""
+    """GET an object (or range) and read the full body, resuming truncated reads.
+
+    Hetzner terminates a response mid-body with a clean TCP FIN while still
+    advertising the full Content-Length, so aiohttp raises ClientPayloadError
+    (ContentLengthError) with only part of the payload delivered. Observed at
+    ~0.5% per read, always ending on a page boundary (1017 or 2034 x 4096) --
+    the same vendor behaviour COPY_SOURCE_STREAM_RESUME already handles on the
+    upload side.
+
+    Re-requesting the identical range reproduces the identical truncation, so
+    plain retry cannot make progress. Instead keep what arrived and ask only for
+    the remainder: a fresh request gets a fresh response buffer. The stitched
+    result is byte-identical to an untruncated read.
+    """
+    start = end = None
+    if range_header is not None:
+        start, end = _parse_byte_range(range_header)
+
+    buf = bytearray()
     for attempt in range(1, SOURCE_READ_ATTEMPTS + 1):
+        if not buf:
+            resume_header = range_header
+        elif start is None:
+            # Whole-object read: no range was requested, so resume by offset.
+            resume_header = f"bytes={len(buf)}-"
+        else:
+            resume_header = f"bytes={start + len(buf)}-" + (str(end) if end is not None else "")
+
         try:
-            resp = await client.get_object(bucket, key, range_header=range_header)
-            async with resp["Body"] as body:
-                return await body.read()
+            resp = await client.get_object(bucket, key, range_header=resume_header)
+            stream = resp["Body"]
+            # Read in bounded chunks off the StreamingBody itself rather than
+            # calling read() once: read() collects blocks in a local list and
+            # drops all of them if the stream raises mid-body, which is exactly
+            # the case this function must survive. Note `async with` on the body
+            # unwraps to the raw ClientResponse, whose read() takes no size, so
+            # the chunk loop must not run inside that context.
+            async with stream:
+                while chunk := await stream.read(SOURCE_READ_CHUNK_SIZE):
+                    buf += chunk
+            return bytes(buf)
         except Exception as exc:
             if not is_retryable_source_error(exc) or attempt == SOURCE_READ_ATTEMPTS:
                 raise
@@ -92,6 +151,8 @@ async def read_source_bytes(
                 bucket=bucket,
                 key=key,
                 range=range_header,
+                resume_range=resume_header,
+                bytes_buffered=len(buf),
                 attempt=attempt,
                 error=str(exc),
             )

tests/conftest.py

@@ -128,12 +128,33 @@ async def read(self, n: int = -1) -> bytes:
         return result
 
     async def __aenter__(self):
-        return self
+        # Fidelity: aiobotocore's StreamingBody.__aenter__ returns the *wrapped*
+        # aiohttp.ClientResponse, not the StreamingBody, so the size-aware
+        # read(amt) is NOT available on the value `async with` binds --
+        # ClientResponse.read() takes no arguments. Code that needs bounded
+        # chunks must read off the StreamingBody itself. Returning `self` here
+        # would let a `async with body: body.read(n)` bug pass tests and fail in
+        # production (observed while fixing source-read resume).
+        return _MockEnteredResponse(self)
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
         return None
 
 
+class _MockEnteredResponse:
+    """What `async with StreamingBody` yields: no size-aware read()."""
+
+    def __init__(self, body: MockS3Response):
+        self._body = body
+
+    @property
+    def content(self):
+        return self._body
+
+    async def read(self) -> bytes:
+        return await self._body.read()
+
+
 class MockS3Client:
     """Mock S3 client for testing without real S3 backend."""
 

tests/unit/test_source_read_resume.py

@@ -0,0 +1,228 @@
+"""read_source_bytes resumes a truncated body instead of restarting it.
+
+Prod incident 2026-07-29: the Scylla backup of `main.companies` failed on 10 of
+13 racks. Hetzner terminates a response mid-body with a clean TCP FIN while
+still advertising the full Content-Length, so aiohttp raises ClientPayloadError
+after delivering only part of the payload. All 327 captured occurrences were
+clean closes (no ConnectionError suffix), i.e. the sender deliberately ended the
+response rather than a network fault.
+
+Measured at ~0.5% per read (2 of 400 in a sustained 16-way probe), and the
+truncation always lands on a page boundary -- only two values ever appeared:
+
+    190 x  received 4165632 of 8388636   (1017 x 4096)
+    137 x  received 8331264 of 8388636   (2034 x 4096)
+
+Re-requesting the identical range reproduces the identical truncation (9 of 12
+retried ranges truncated at the same offset every attempt), so plain retry
+cannot make progress -- which is why ~12% of reads never recovered no matter how
+many attempts they were given. At ~572 reads per 4.7GB copy a single large file
+had a ~94% chance of hitting at least one, so `companies` failed nearly every run
+while smaller tables passed.
+
+Resuming works because a fresh request gets a fresh response buffer. Verified
+against production: resuming 4 real failing ranges returned exactly the missing
+bytes, and the stitched result was SHA-256 identical to an untruncated read.
+"""
+
+from __future__ import annotations
+
+import aiohttp
+import pytest
+from botocore.exceptions import ClientError
+
+from s3proxy.handlers import base
+from s3proxy.handlers.base import read_source_bytes
+
+BODY = bytes(range(256)) * 512  # 131072 bytes, non-uniform so stitching is checked
+
+
+@pytest.fixture(autouse=True)
+def _no_backoff(monkeypatch):
+    monkeypatch.setattr(base, "SOURCE_READ_BACKOFF_SEC", 0.0)
+
+
+class _Body:
+    """Yields `cut` bytes then raises, mimicking a mid-body FIN."""
+
+    def __init__(self, data: bytes, cut: int | None):
+        self._data = data
+        self._cut = cut
+        self._sent = 0
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, *exc):
+        return False
+
+    async def read(self, n: int = -1) -> bytes:
+        limit = len(self._data) if self._cut is None else self._cut
+        if self._sent >= limit:
+            if self._cut is not None:
+                raise aiohttp.ClientPayloadError(
+                    "Response payload is not completed: ContentLengthError(400, "
+                    'message="Not enough data to satisfy content length header '
+                    f'(received {self._cut} of {len(self._data)} bytes)."))'
+                )
+            return b""
+        end = limit if n < 0 else min(self._sent + n, limit)
+        chunk = self._data[self._sent : end]
+        self._sent = end
+        return chunk
+
+
+class _Client:
+    """Serves BODY, truncating each successive request per `cuts`."""
+
+    def __init__(self, cuts: list[int | None]):
+        self._cuts = list(cuts)
+        self.ranges: list[str | None] = []
+
+    async def get_object(self, bucket, key, range_header=None):
+        self.ranges.append(range_header)
+        start, end = 0, len(BODY) - 1
+        if range_header:
+            start, end = base._parse_byte_range(range_header)
+            if end is None:
+                end = len(BODY) - 1
+        segment = BODY[start : end + 1]
+        cut = self._cuts.pop(0) if self._cuts else None
+        return {"Body": _Body(segment, cut)}
+
+
+@pytest.mark.asyncio
+async def test_clean_read_needs_one_request():
+    c = _Client([None])
+    assert await read_source_bytes(c, "b", "k") == BODY
+    assert c.ranges == [None]
+
+
+@pytest.mark.asyncio
+async def test_resumes_ranged_read_and_stitches_identically():
+    """The prod shape: full range truncated once, then resumed."""
+    c = _Client([50_000, None])
+    got = await read_source_bytes(c, "b", "k", "bytes=0-131071")
+    assert got == BODY, "stitched result must equal an untruncated read"
+    assert c.ranges == ["bytes=0-131071", "bytes=50000-131071"]
+
+
+@pytest.mark.asyncio
+async def test_resume_offsets_are_absolute_not_relative():
+    """A non-zero range start must not be dropped when resuming."""
+    c = _Client([1000, None])
+    got = await read_source_bytes(c, "b", "k", "bytes=4096-20479")
+    assert got == BODY[4096:20480]
+    assert c.ranges == ["bytes=4096-20479", "bytes=5096-20479"]
+
+
+@pytest.mark.asyncio
+async def test_repeated_truncation_still_converges():
+    """Several truncations in a row: each attempt keeps prior progress."""
+    c = _Client([40_000, 30_000, 20_000, None])
+    got = await read_source_bytes(c, "b", "k", "bytes=0-131071")
+    assert got == BODY
+    assert c.ranges == [
+        "bytes=0-131071",
+        "bytes=40000-131071",
+        "bytes=70000-131071",
+        "bytes=90000-131071",
+    ]
+
+
+@pytest.mark.asyncio
+async def test_whole_object_read_resumes_by_offset():
+    """No range requested: resume must still target the missing tail."""
+    c = _Client([70_000, None])
+    got = await read_source_bytes(c, "b", "k")
+    assert got == BODY
+    assert c.ranges == [None, "bytes=70000-"]
+
+
+@pytest.mark.asyncio
+async def test_open_ended_range_preserved_on_resume():
+    c = _Client([2048, None])
+    got = await read_source_bytes(c, "b", "k", "bytes=1024-")
+    assert got == BODY[1024:]
+    assert c.ranges == ["bytes=1024-", "bytes=3072-"]
+
+
+@pytest.mark.asyncio
+async def test_gives_up_after_attempt_cap():
+    c = _Client([10_000] * base.SOURCE_READ_ATTEMPTS)
+    with pytest.raises(aiohttp.ClientPayloadError):
+        await read_source_bytes(c, "b", "k", "bytes=0-131071")
+    assert len(c.ranges) == base.SOURCE_READ_ATTEMPTS
+
+
+@pytest.mark.asyncio
+async def test_non_retryable_error_is_not_resumed():
+    class _Boom:
+        async def get_object(self, *a, **k):
+            raise ValueError("permanent")
+
+    with pytest.raises(ValueError):
+        await read_source_bytes(_Boom(), "b", "k", "bytes=0-10")
+
+
+def test_parse_byte_range():
+    assert base._parse_byte_range("bytes=0-99") == (0, 99)
+    assert base._parse_byte_range("bytes=4096-") == (4096, None)
+    assert base._parse_byte_range("4096-8191") == (4096, 8191)
+
+
+# --- 504 GatewayTimeout retryability ---------------------------------------
+#
+# 273 of 288 fatal copy failures in the 2026-07-29 capture were HTTP 504
+# GatewayTimeout on UploadPartCopy. The segment loop in copy.py already retries
+# via is_retryable_source_error, but 504 was absent from the retryable set, so it
+# hit `raise` on the first occurrence -- which is why only 4 SEGMENT_RETRY events
+# were logged against 273 fatal 504s.
+#
+# Reproduced: 572 UploadPartCopy calls at 8-way concurrency against the real
+# object yielded 1 x 504 (0.17%), with latency p50 1.14s / p90 3.31s / max
+# 61.86s. Heavy-tailed latency means a 504 is transient congestion. At 0.17% per
+# copy a 4.7GB file (572 copies) has a ~63% chance of hitting one.
+
+
+@pytest.mark.parametrize(
+    "code",
+    ["GatewayTimeout", "504", "BadGateway", "502", "503", "500", "InternalError", "SlowDown"],
+)
+def test_transient_backend_codes_are_retryable(code):
+    exc = ClientError(
+        {"Error": {"Code": code, "Message": "x"}, "ResponseMetadata": {"HTTPStatusCode": 504}},
+        "UploadPartCopy",
+    )
+    assert base.is_retryable_source_error(exc) is True
+
+
+@pytest.mark.parametrize("code", ["AccessDenied", "NoSuchKey", "InvalidPart", "EntityTooSmall"])
+def test_permanent_codes_are_not_retryable(code):
+    exc = ClientError({"Error": {"Code": code, "Message": "x"}}, "UploadPartCopy")
+    assert base.is_retryable_source_error(exc) is False
+
+
+@pytest.mark.asyncio
+async def test_source_read_retries_gateway_timeout():
+    """A 504 mid-read must be retried, not raised on first sight."""
+
+    class _Flaky:
+        def __init__(self):
+            self.calls = 0
+
+        async def get_object(self, bucket, key, range_header=None):
+            self.calls += 1
+            if self.calls == 1:
+                raise ClientError(
+                    {
+                        "Error": {"Code": "GatewayTimeout", "Message": "did not respond in time"},
+                        "ResponseMetadata": {"HTTPStatusCode": 504},
+                    },
+                    "GetObject",
+                )
+            return {"Body": _Body(BODY, None)}
+
+    c = _Flaky()
+    assert await read_source_bytes(c, "b", "k", "bytes=0-131071") == BODY
+    assert c.calls == 2

tests/unit/test_source_read_retry.py

@@ -53,9 +53,15 @@ def _client_error(code, operation="GetObject"):
 class _Body:
     def __init__(self, data: bytes):
         self._data = data
+        self._sent = 0
 
-    async def read(self):
-        return self._data
+    async def read(self, n: int = -1):
+        if self._sent >= len(self._data):
+            return b""
+        end = len(self._data) if n < 0 else min(self._sent + n, len(self._data))
+        chunk = self._data[self._sent : end]
+        self._sent = end
+        return chunk
 
     async def __aenter__(self):
         return self
@@ -67,7 +73,7 @@ async def __aexit__(self, *exc):
 class _TruncatingBody:
     """read() dies mid-body, the exact prod failure shape."""
 
-    async def read(self):
+    async def read(self, n: int = -1):
         raise _payload_error()
 
     async def __aenter__(self):