diff --git a/goosebit/storage/s3.py b/goosebit/storage/s3.py index 131eeb41..6720fc1d 100644 --- a/goosebit/storage/s3.py +++ b/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}") diff --git a/tests/e2e/s3/tests/test_e2e_s3.py b/tests/e2e/s3/tests/test_e2e_s3.py index 9940a9ae..6df48c9d 100644 --- a/tests/e2e/s3/tests/test_e2e_s3.py +++ b/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"