fix: retry transient source-read and segment-copy failures in UploadPartCopy

s3proxy/handlers/base.py

@@ -2,15 +2,25 @@
 
 import asyncio
 import base64
+import os
 import re
 from collections.abc import AsyncIterator
 from datetime import datetime
 from typing import NoReturn
 from urllib.parse import parse_qs, unquote
 
+import aiohttp
 import httpx
 import structlog
-from botocore.exceptions import ClientError
+from botocore.exceptions import (
+    ClientError,
+    ConnectionClosedError,
+    ConnectTimeoutError,
+    EndpointConnectionError,
+    IncompleteReadError,
+    ReadTimeoutError,
+    ResponseStreamingError,
+)
 from fastapi import Request, Response
 from structlog.stdlib import BoundLogger
 
@@ -26,6 +36,69 @@
 PATH_RE = re.compile(r"^/([^/]+)/(.+)$")
 BUCKET_RE = re.compile(r"^/([^/]+)/?$")  # Handles /bucket and /bucket/
 
+# botocore's retry layer only covers a request up to the response headers; a
+# body that dies mid-read (backend drops a long-lived connection, observed as
+# ClientPayloadError at a fixed offset after minutes of sequential frame GETs)
+# surfaces here with no retry and used to fail whole UploadPartCopy parts.
+# 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"))
+
+_RETRYABLE_S3_ERROR_CODES = frozenset(
+    {"InternalError", "SlowDown", "RequestTimeout", "ServiceUnavailable", "BadGateway"}
+    | {"500", "502", "503"}
+)
+
+_RETRYABLE_TRANSPORT_ERRORS = (
+    aiohttp.ClientPayloadError,
+    aiohttp.ClientConnectionError,
+    TimeoutError,
+    ConnectionClosedError,
+    ConnectTimeoutError,
+    EndpointConnectionError,
+    IncompleteReadError,
+    ReadTimeoutError,
+    ResponseStreamingError,
+)
+
+
+def is_retryable_source_error(exc: BaseException) -> bool:
+    """Transient transport/backend failures where re-issuing the request is safe."""
+    if isinstance(exc, _RETRYABLE_TRANSPORT_ERRORS):
+        return True
+    if isinstance(exc, ClientError):
+        code = str(exc.response.get("Error", {}).get("Code", ""))
+        return code in _RETRYABLE_S3_ERROR_CODES
+    return False
+
+
+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."""
+    for attempt in range(1, SOURCE_READ_ATTEMPTS + 1):
+        try:
+            resp = await client.get_object(bucket, key, range_header=range_header)
+            async with resp["Body"] as body:
+                return await body.read()
+        except Exception as exc:
+            if not is_retryable_source_error(exc) or attempt == SOURCE_READ_ATTEMPTS:
+                raise
+            logger.warning(
+                "SOURCE_READ_RETRY",
+                bucket=bucket,
+                key=key,
+                range=range_header,
+                attempt=attempt,
+                error=str(exc),
+            )
+            await asyncio.sleep(SOURCE_READ_BACKOFF_SEC * (2 ** (attempt - 1)))
+    raise AssertionError("unreachable")
+
+
 # Shared httpx client for connection reuse
 _http_client: httpx.AsyncClient | None = None
 _http_client_lock = asyncio.Lock()
@@ -229,9 +302,7 @@ def _check_conditional_headers(
     async def _download_encrypted_single(
         self, client: S3Client, bucket: str, key: str, wrapped_dek_b64: str, kid: str = ""
     ) -> bytes:
-        resp = await client.get_object(bucket, key)
-        async with resp["Body"] as body:
-            ciphertext = await body.read()
+        ciphertext = await read_source_bytes(client, bucket, key)
         wrapped_dek = base64.b64decode(wrapped_dek_b64)
         return crypto.decrypt_object(ciphertext, wrapped_dek, self.keyring.key_by_id(kid))
 
@@ -283,9 +354,9 @@ async def _iter_multipart_plaintext(
 
                     if in_range:
                         ct_end = ct_offset + fsize - 1
-                        resp = await client.get_object(bucket, key, f"bytes={ct_offset}-{ct_end}")
-                        async with resp["Body"] as body:
-                            ciphertext = await body.read()
+                        ciphertext = await read_source_bytes(
+                            client, bucket, key, f"bytes={ct_offset}-{ct_end}"
+                        )
                         plaintext = crypto.decrypt(ciphertext, dek)
 
                         if range_start is not None:

s3proxy/handlers/multipart/copy.py

@@ -31,7 +31,8 @@
     persist_upload_state,
 )
 from ...utils import format_iso8601
-from ..base import BaseHandler
+from .. import base
+from ..base import BaseHandler, is_retryable_source_error, read_source_bytes
 from .upload_part import _PlaintextReader
 
 logger: BoundLogger = structlog.get_logger(__name__)
@@ -876,14 +877,36 @@ async def copy_segment(idx: int, seg: _CiphertextSegment) -> InternalPartMetadat
             copy_range = f"bytes={seg.ct_offset}-{ct_end}"
             part_reserve = crypto.copy_passthrough_segment_peak(seg.plaintext_size)
             async with segment_sem, concurrency.reserve_copy_memory(part_reserve):
-                resp = await client.upload_part_copy(
-                    bucket,
-                    key,
-                    upload_id,
-                    internal_num,
-                    copy_source_path,
-                    copy_source_range=copy_range,
-                )
+                # The backend can accept the copy, stream a 200, and still fail
+                # ("InternalError: The server did not respond in time" embedded
+                # in the body). Same part number + range makes a retry safe.
+                for attempt in range(1, base.SOURCE_READ_ATTEMPTS + 1):
+                    try:
+                        resp = await client.upload_part_copy(
+                            bucket,
+                            key,
+                            upload_id,
+                            internal_num,
+                            copy_source_path,
+                            copy_source_range=copy_range,
+                        )
+                        break
+                    except Exception as exc:
+                        if (
+                            not is_retryable_source_error(exc)
+                            or attempt == base.SOURCE_READ_ATTEMPTS
+                        ):
+                            raise
+                        logger.warning(
+                            "UPLOAD_PART_COPY_SEGMENT_RETRY",
+                            bucket=bucket,
+                            key=key,
+                            part_num=part_num,
+                            internal_part=internal_num,
+                            attempt=attempt,
+                            error=str(exc),
+                        )
+                        await asyncio.sleep(base.SOURCE_READ_BACKOFF_SEC * (2 ** (attempt - 1)))
             return InternalPartMetadata(
                 internal_part_number=internal_num,
                 plaintext_size=seg.plaintext_size,
@@ -1103,9 +1126,7 @@ async def _fetch_copy_plaintext(
     ) -> bytes:
         """Download the full plaintext of a copy source (small sources only)."""
         if not src_wrapped_dek and not src_multipart_meta:
-            resp = await client.get_object(src_bucket, src_key, range_header=copy_source_range)
-            async with resp["Body"] as body:
-                return await body.read()
+            return await read_source_bytes(client, src_bucket, src_key, copy_source_range)
         elif src_multipart_meta:
             range_start, range_end = self._parse_copy_source_range(
                 copy_source_range, src_multipart_meta.total_plaintext_size
@@ -1308,16 +1329,60 @@ async def _iter_copy_source(
                 plaintext = plaintext[start : end + 1]
             yield plaintext
         else:
-            resp = await client.get_object(src_bucket, src_key, range_header=copy_source_range)
-            async with resp["Body"] as body:
-                # resp["Body"] enters as an aiohttp ClientResponse, whose read()
-                # takes no size arg; stream via its StreamReader in bounded chunks
-                # (body.read(n) raised TypeError and 500'd every passthrough copy).
-                while True:
-                    chunk = await body.content.read(crypto.MAX_BUFFER_SIZE)
-                    if not chunk:
-                        break
-                    yield chunk
+            async for chunk in self._stream_raw_source_with_resume(
+                client, src_bucket, src_key, copy_source_range
+            ):
+                yield chunk
+
+    async def _stream_raw_source_with_resume(
+        self,
+        client: S3Client,
+        src_bucket: str,
+        src_key: str,
+        copy_source_range: str | None,
+    ) -> AsyncIterator[bytes]:
+        """Stream an unencrypted copy source, resuming with a ranged GET if the
+        body dies mid-read (backend dropping a long-lived connection). Bytes
+        already yielded are never re-fetched: the resume range starts at the
+        exact offset the previous attempt reached.
+        """
+        if copy_source_range:
+            range_start, range_end = self._parse_raw_copy_source_range(copy_source_range)
+        else:
+            range_start, range_end = 0, None
+        sent = 0
+        attempt = 1
+        while True:
+            if sent == 0 and copy_source_range is None:
+                range_header = None
+            else:
+                end_part = "" if range_end is None else str(range_end)
+                range_header = f"bytes={range_start + sent}-{end_part}"
+            try:
+                resp = await client.get_object(src_bucket, src_key, range_header=range_header)
+                async with resp["Body"] as body:
+                    # resp["Body"] enters as an aiohttp ClientResponse, whose read()
+                    # takes no size arg; stream via its StreamReader in bounded chunks
+                    # (body.read(n) raised TypeError and 500'd every passthrough copy).
+                    while True:
+                        chunk = await body.content.read(crypto.MAX_BUFFER_SIZE)
+                        if not chunk:
+                            return
+                        sent += len(chunk)
+                        yield chunk
+            except Exception as exc:
+                if not is_retryable_source_error(exc) or attempt >= base.SOURCE_READ_ATTEMPTS:
+                    raise
+                logger.warning(
+                    "COPY_SOURCE_STREAM_RESUME",
+                    bucket=src_bucket,
+                    key=src_key,
+                    resumed_at=range_start + sent,
+                    attempt=attempt,
+                    error=str(exc),
+                )
+                await asyncio.sleep(base.SOURCE_READ_BACKOFF_SEC * (2 ** (attempt - 1)))
+                attempt += 1
 
     async def _pump_copy_chunks(
         self,