perf: stream s3 downloads via ranged GetObject request

goosebit/storage/s3.py

@@ -1,4 +1,5 @@
 import asyncio
+from functools import partial
 from typing import AsyncIterable
 from urllib.parse import urlparse
 
@@ -10,6 +11,10 @@
 from .base import StorageProtocol
 
 DOWNLOAD_CHUNK_SIZE = 64 * 1024
+# bytes fetched per ranged GetObject; request count and per-stream memory both scale with this
+RANGE_REQUEST_SIZE = 1024 * 1024
+# retries per range for transient read failures; botocore only retries get_object, not body.read()
+RANGE_READ_ATTEMPTS = 3
 
 
 class S3StorageBackend(StorageProtocol):
@@ -53,22 +58,65 @@ async def store_file(self, source_path: Path, dest_path: Path) -> str:
 
     async def get_file_stream(self, uri: str) -> AsyncIterable[bytes]:  # type: ignore[override]
         key = self._extract_key_from_uri(uri)
+        loop = asyncio.get_running_loop()
+
+        offset = 0
+        total_size: int | None = None
+        etag: str | None = None
 
         try:
-            loop = asyncio.get_running_loop()
-            response = await loop.run_in_executor(None, lambda: self.s3_client.get_object(Bucket=self.bucket, Key=key))
-
-            body = response["Body"]
-            try:
-                while True:
-                    chunk = await loop.run_in_executor(None, body.read, DOWNLOAD_CHUNK_SIZE)
-                    if not chunk:
+            while total_size is None or offset < total_size:
+                request = {
+                    "Bucket": self.bucket,
+                    "Key": key,
+                    "Range": f"bytes={offset}-{offset + RANGE_REQUEST_SIZE - 1}",
+                }
+                if etag is not None:
+                    # fail with 412 instead of splicing two versions if the artifact is replaced mid-download
+                    request["IfMatch"] = etag
+
+                # retry the whole range fetch on transient failures; ClientErrors are terminal
+                data = None
+                for attempt in range(RANGE_READ_ATTEMPTS):
+                    try:
+                        response = await loop.run_in_executor(None, partial(self.s3_client.get_object, **request))
+
+                        if total_size is None:
+                            etag = response.get("ETag")
+                            # ContentRange is "bytes 0-1048575/104857600"; fall back to
+                            # ContentLength if a backend ignored Range (200) or reports "*" total
+                            content_range = response.get("ContentRange")
+                            if content_range and content_range.rsplit("/", 1)[1] != "*":
+                                total_size = int(content_range.rsplit("/", 1)[1])
+                            else:
+                                total_size = response["ContentLength"]
+                            if etag is not None:
+                                request["IfMatch"] = etag  # pin version across ranges and retries
+
+                        body = response["Body"]
+                        try:
+                            data = await loop.run_in_executor(None, body.read)
+                        finally:
+                            await loop.run_in_executor(None, body.close)
                         break
-                    yield chunk
-            finally:
-                await loop.run_in_executor(None, body.close)
-
-        # BotoCoreError covers mid-stream failures (e.g. ResponseStreamingError
+                    except ClientError as e:
+                        if offset == 0 and e.response["Error"]["Code"] == "InvalidRange":
+                            return  # zero-byte object: any range request returns 416
+                        raise
+                    except BotoCoreError:
+                        if attempt + 1 == RANGE_READ_ATTEMPTS:
+                            raise
+                        await asyncio.sleep(0.5 * (attempt + 1))  # back off, then re-fetch this range
+
+                if not data:
+                    raise ValueError(f"S3 returned empty range at offset {offset} for {uri}")
+                offset += len(data)
+
+                # device-paced waits happen here, with no S3 request in flight to time out
+                for i in range(0, len(data), DOWNLOAD_CHUNK_SIZE):
+                    yield data[i : i + DOWNLOAD_CHUNK_SIZE]
+
+        # BotoCoreError covers mid-transfer failures (e.g. ResponseStreamingError
         # when the connection drops), which are not ClientError subclasses.
         except (BotoCoreError, ClientError) as e:
             raise ValueError(f"S3 download failed: {e}")

tests/e2e/s3/tests/test_e2e_s3.py

@@ -1,14 +1,17 @@
+import math
 import os
 import sys
 import time
 from pathlib import Path
 from typing import Any, Generator
+from unittest.mock import patch
 
 import boto3
 import httpx
 import pytest
 from botocore.exceptions import ClientError
 
+from goosebit.storage.s3 import RANGE_REQUEST_SIZE, S3StorageBackend
 from tests.e2e.utils import auth_token, compose_down, compose_up_build, wait_for_service
 
 BASE_URL = os.getenv("E2E_BASE_URL", "http://localhost:60053")
@@ -295,3 +298,53 @@ def test_e2e_artifact_delete_removes_from_minio(ensure_services_ready: bool) ->
                     break
                 time.sleep(1.0)
         assert deleted, f"S3 object still present after delete bucket={MINIO_BUCKET}, key={key}. Last error: {last_exc}"
+
+
+# ---------------------
+# Ranged streaming (S3StorageBackend.get_file_stream)
+# ---------------------
+
+
+def _minio_client() -> Any:
+    return boto3.client(
+        "s3",
+        endpoint_url=MINIO_URL,
+        aws_access_key_id=MINIO_ACCESS_KEY,
+        aws_secret_access_key=MINIO_SECRET_KEY,
+    )
+
+
+@pytest.mark.parametrize(
+    "label, size",
+    [
+        ("empty", 0),  # zero-byte object -> 416 InvalidRange -> empty stream
+        ("sub_range", 64 * 1024),  # smaller than one range -> single GetObject
+        ("multi_range", 2_500_000),  # larger than RANGE_REQUEST_SIZE -> offset-continuation loop
+    ],
+)
+async def test_e2e_s3_stream_reassembles_ranges(ensure_services_ready: bool, label: str, size: int) -> None:
+    """Ranged streaming reassembles bytes identically to the source, across object sizes."""
+    source = os.urandom(size)
+    key = f"e2e-stream/{label}.bin"
+    _minio_client().put_object(Bucket=MINIO_BUCKET, Key=key, Body=source)
+
+    backend = S3StorageBackend(
+        bucket=MINIO_BUCKET,
+        endpoint_url=MINIO_URL,
+        access_key_id=MINIO_ACCESS_KEY,
+        secret_access_key=MINIO_SECRET_KEY,
+    )
+    uri = f"s3://{MINIO_BUCKET}/{key}"
+
+    with patch.object(backend.s3_client, "get_object", wraps=backend.s3_client.get_object) as spy:
+        streamed = b"".join([chunk async for chunk in backend.get_file_stream(uri)])
+
+    assert streamed == source, f"{label}: streamed {len(streamed)} bytes != source {len(source)} bytes"
+
+    # one GetObject per range (empty object still costs the single 416 probe)
+    expected_calls = 1 if size == 0 else math.ceil(size / RANGE_REQUEST_SIZE)
+    assert (
+        spy.call_count == expected_calls
+    ), f"{label}: expected {expected_calls} ranged GetObject call(s), got {spy.call_count}"
+    if size > RANGE_REQUEST_SIZE:
+        assert spy.call_count > 1, f"{label}: object larger than a range must span multiple GetObject calls"