diff --git a/clients/python/README.md b/clients/python/README.md index a37e0c68..d8462fc3 100644 --- a/clients/python/README.md +++ b/clients/python/README.md @@ -100,6 +100,14 @@ session.put(video_data, compress="none") session.put(zstd_data, precompressed="zstd") ``` +### Resumable Uploads + +`Session.put()` automatically uses resumable uploads for eligible sources of at +least 32 MiB. Set `resumable_threshold_bytes` on `Usecase` or override it on an +individual `put()`; `None` disables resumable uploads. Configure recovery with +`Usecase(resumable_retries=ResumableRetryPolicy(...))`; zero retries also disables +resumable uploads. See `Session.put()` and `ResumableRetryPolicy` for details. + ### Custom Metadata Arbitrary key-value pairs can be attached to objects and retrieved on download. diff --git a/clients/python/src/objectstore_client/__init__.py b/clients/python/src/objectstore_client/__init__.py index d5c9955f..08fb580c 100644 --- a/clients/python/src/objectstore_client/__init__.py +++ b/clients/python/src/objectstore_client/__init__.py @@ -2,6 +2,7 @@ from objectstore_client.client import ( Client, GetResponse, + ResumableRetryPolicy, Session, Usecase, ) @@ -23,6 +24,7 @@ __all__ = [ "Client", "Usecase", + "ResumableRetryPolicy", "Session", "GetResponse", "RequestError", diff --git a/clients/python/src/objectstore_client/_resumable.py b/clients/python/src/objectstore_client/_resumable.py index 8c467ade..6a5934f4 100644 --- a/clients/python/src/objectstore_client/_resumable.py +++ b/clients/python/src/objectstore_client/_resumable.py @@ -1,12 +1,16 @@ from __future__ import annotations +import random +import time from dataclasses import dataclass +from io import SEEK_END from typing import IO, TYPE_CHECKING from urllib.parse import urlencode import urllib3 from objectstore_client.errors import RequestError, raise_for_status +from objectstore_client.metadata import Compression, ExpirationPolicy from objectstore_client.metrics import measure_storage_operation from objectstore_client.tracing import storage_span @@ -156,6 +160,12 @@ def put( headers = session._make_headers() headers["Upload-Offset"] = str(offset) headers["Content-Length"] = str(length) + # Disable pool retries that replay the body. The resumable upload loop + # uses the Usecase's recovery policy and probes the persisted offset + # before resending data; connection retries retain the pool's policy. + retries = urllib3.Retry.from_int(session._pool.retries).new( + read=0, status=0, other=0, raise_on_status=False + ) with ( storage_span( "resumable.put", @@ -174,6 +184,7 @@ def put( f"{session._make_url(self.key)}?{query}", headers=headers, body=body, + retries=retries, preload_content=True, decode_content=True, ) @@ -207,3 +218,107 @@ def cancel(self) -> None: else RequestError ) raise_for_status(response, error_type=error_type) + + +def get_size(contents: bytes | IO[bytes]) -> int | None: + if isinstance(contents, bytes): + return len(contents) + try: + if not contents.seekable(): + return None + start = contents.tell() + except (OSError, ValueError): + return None + try: + end = contents.seek(0, SEEK_END) + except (OSError, ValueError): + return None + finally: + contents.seek(start) + return max(0, end - start) + + +def is_transient(error: Exception) -> bool: + if isinstance(error, urllib3.exceptions.MaxRetryError): + # Exhausted status retries carry ResponseError rather than the response. + return isinstance(error.reason, urllib3.exceptions.ResponseError) or ( + isinstance(error.reason, Exception) and is_transient(error.reason) + ) + if isinstance(error, RequestError): + return error.status in (408, 429, 502, 503, 504) + return isinstance( + error, + ( + urllib3.exceptions.ReadTimeoutError, + urllib3.exceptions.ProtocolError, + ), + ) + + +def upload( + session: Session, + body: IO[bytes], + encoded_size: int, + key: str | None = None, + compression: Compression | None = None, + content_type: str | None = None, + metadata: dict[str, str] | None = None, + expiration_policy: ExpirationPolicy | None = None, + origin: str | None = None, + filename: str | None = None, +) -> str | None: + policy = session._usecase._resumable_retries + start = body.tell() + try: + handle = session._create_upload( + encoded_size, + key=key, + compression=compression, + content_type=content_type, + metadata=metadata, + expiration_policy=expiration_policy, + origin=origin, + filename=filename, + ) + except Exception: + handle = None + if handle is None: + return None + + try: + offset = 0 + retries = 0 + probe = False + while True: + try: + if probe: + result = handle.progress() + else: + body.seek(start + offset) + result = handle.put(offset, (body, encoded_size - offset)) + except UploadOffsetMismatch as error: + result = UploadIncomplete(error.offset) + except Exception as error: + if not is_transient(error) or retries >= policy.retries: + raise + time.sleep(policy.delay * 2**retries + random.uniform(0, policy.jitter)) + retries += 1 + probe = True + continue + + if isinstance(result, UploadComplete): + return handle.key + if not offset <= result.offset <= encoded_size: + raise ValueError("Invalid upload offset") + if result.offset == offset and not probe: + raise ValueError("Upload made no progress") + offset = result.offset + probe = False + except Exception as error: + try: + handle.cancel() + except Exception: + pass + status = error.status if isinstance(error, RequestError) else None + response = error.response if isinstance(error, RequestError) else None + raise RequestError("upload failed", status, response) from error diff --git a/clients/python/src/objectstore_client/client.py b/clients/python/src/objectstore_client/client.py index 579e9e3b..810a1dc0 100644 --- a/clients/python/src/objectstore_client/client.py +++ b/clients/python/src/objectstore_client/client.py @@ -37,6 +37,7 @@ from objectstore_client.metrics import ( MetricsBackend, NoOpMetricsBackend, + StorageMetricEmitter, measure_storage_operation, ) from objectstore_client.multipart import MultipartUpload @@ -55,6 +56,30 @@ class GetResponse(NamedTuple): payload: IO[bytes] +@dataclass(frozen=True) +class ResumableRetryPolicy: + """Recovery limits for one resumable upload, including progress queries. + + ``retries`` counts recovery retries across the upload; zero disables resumable + uploads. + + ``delay`` is the initial backoff in seconds, doubling on each retry. + + ``jitter`` is the maximum random delay added in seconds. + + These limits are separate from the pool's per-request retries; retryable + failures are determined internally. + """ + + retries: int = 2 + delay: float = 2.0 + jitter: float = 1.0 + + def __post_init__(self) -> None: + if self.retries < 0 or self.delay < 0 or self.jitter < 0: + raise ValueError("Resumable upload retry settings must be non-negative") + + class Usecase: """ An identifier for a workload in Objectstore, along with defaults to use for all @@ -63,6 +88,10 @@ class Usecase: Usecases need to be statically defined in Objectstore's configuration server-side. Objectstore can make decisions based on the Usecase. For example, choosing the most suitable storage backend. + + ``resumable_threshold_bytes`` defaults to 32 MiB of remaining source bytes, + before any compression. ``None`` disables resumable uploads. + ``resumable_retries`` configures recovery via `ResumableRetryPolicy`. """ name: str @@ -74,10 +103,16 @@ def __init__( name: str, compression: Compression = "zstd", expiration_policy: ExpirationPolicy | None = None, + resumable_threshold_bytes: int | None = 32 * 1024 * 1024, + resumable_retries: ResumableRetryPolicy = ResumableRetryPolicy(), ): + if resumable_threshold_bytes is not None and resumable_threshold_bytes < 0: + raise ValueError("resumable_threshold_bytes must be non-negative") self.name = name self._compression = compression self._expiration_policy = expiration_policy + self._resumable_threshold_bytes = resumable_threshold_bytes + self._resumable_retries = resumable_retries # Connect timeout used unless overridden in connection parameters. @@ -393,6 +428,7 @@ def put( expiration_policy: ExpirationPolicy | None = None, origin: str | None = None, filename: str | None = None, + resumable_threshold_bytes: int | None | Literal["unset"] = "unset", ) -> str: """ Uploads the given `contents` to blob storage. @@ -412,6 +448,10 @@ def put( You can use the utility function `objectstore_client.utils.guess_mime_type` to attempt to guess a `content_type` based on magic bytes. + ``resumable_threshold_bytes`` overrides the Usecase threshold when supplied; + ``None`` disables resumable uploads. Eligible uploads use the Usecase's + ``resumable_retries`` policy; a retry count of zero disables resumable uploads. + `compression` is deprecated in favor of `compress`. """ if compression is not None: @@ -431,16 +471,24 @@ def put( if precompressed and precompressed != "zstd": raise ValueError(f"Invalid compression: {precompressed}") - body = BytesIO(contents) if isinstance(contents, bytes) else contents - original_body: IO[bytes] = body + if resumable_threshold_bytes == "unset": + resumable_threshold_bytes = self._usecase._resumable_threshold_bytes + if resumable_threshold_bytes is not None and resumable_threshold_bytes < 0: + raise ValueError("resumable_threshold_bytes must be non-negative") encoding = precompressed or compress or self._usecase._compression compress_with = encoding if precompressed is None else "none" - if compress_with == "zstd": - cctx = zstandard.ZstdCompressor() - body = cctx.stream_reader(original_body) - body = cast(IO[bytes], utils._ZstdCompressionReaderWrapper(body)) + + replayable = isinstance(contents, bytes) or compress_with == "none" + body_size = _resumable.get_size(contents) + use_resumable = ( + body_size is not None + and resumable_threshold_bytes is not None + and self._usecase._resumable_retries.retries > 0 + and replayable + and body_size >= resumable_threshold_bytes + ) headers = self._metadata_headers( compression=encoding, @@ -460,8 +508,87 @@ def put( self._metrics_backend, "put", self._usecase.name ) as metrics, ): - retries = None # by default use the pool's value, set by the Client - if compress_with != "none": + if use_resumable: + assert body_size is not None + if isinstance(contents, bytes): + encoded = ( + zstandard.ZstdCompressor().compress(contents) + if compress_with == "zstd" + else contents + ) + body: IO[bytes] = BytesIO(encoded) + encoded_size = len(encoded) + else: + body = contents + encoded_size = body_size + try: + result_key = _resumable.upload( + self, + body, + encoded_size, + key=key, + compression=encoding, + content_type=content_type, + metadata=metadata, + expiration_policy=expiration_policy, + origin=origin, + filename=filename, + ) + if result_key is None: + headers["Content-Length"] = str(encoded_size) + result_key = self._put_direct( + body, key, headers, compress=False + ) + if precompressed is None: + metrics.record_uncompressed_size(body_size) + if encoding != "none": + metrics.record_compressed_size(encoded_size, encoding) + finally: + if isinstance(contents, bytes): + body.close() + else: + result_key = self._put_direct( + contents, + key, + headers, + compress=compress_with == "zstd", + metrics=metrics, + record_source=precompressed is None, + encoding=encoding, + ) + + # Set after the response, since the key may be server-generated. + span.set_attribute("objectstore.key", result_key) + span.set_attribute("objectstore.compression", encoding) + if metrics.uncompressed_size is not None: + span.set_attribute( + "objectstore.uncompressed_size", metrics.uncompressed_size + ) + if metrics.compressed_size is not None: + span.set_attribute( + "objectstore.compressed_size", metrics.compressed_size + ) + return result_key + + def _put_direct( + self, + contents: bytes | IO[bytes], + key: str | None, + headers: dict[str, str], + compress: bool, + metrics: StorageMetricEmitter | None = None, + record_source: bool = True, + encoding: Compression = "none", + ) -> str: + """Stream a direct upload, optionally compressing on the fly.""" + body = BytesIO(contents) if isinstance(contents, bytes) else contents + original_body: IO[bytes] = body + retries = None # by default use the pool's value, set by the Client + try: + if compress: + cctx = zstandard.ZstdCompressor() + body = cctx.stream_reader(original_body, closefd=False) + body = cast(IO[bytes], utils._ZstdCompressionReaderWrapper(body)) # For on-the-fly compression, don't attempt read retries, # as the stream cannot be rewound after data has been consumed. pool_retries = self._pool.retries @@ -484,23 +611,17 @@ def put( # Must do this after streaming `body` as that's what is responsible # for advancing the seek position in both streams - if precompressed is None: - metrics.record_uncompressed_size(original_body.tell()) - if encoding != "none": - metrics.record_compressed_size(body.tell(), encoding) - - # Set after the response, since the key may be server-generated. - span.set_attribute("objectstore.key", res["key"]) - span.set_attribute("objectstore.compression", encoding) - if metrics.uncompressed_size is not None: - span.set_attribute( - "objectstore.uncompressed_size", metrics.uncompressed_size - ) - if metrics.compressed_size is not None: - span.set_attribute( - "objectstore.compressed_size", metrics.compressed_size - ) + if metrics is not None: + if record_source: + metrics.record_uncompressed_size(original_body.tell()) + if encoding != "none": + metrics.record_compressed_size(body.tell(), encoding) return res["key"] + finally: + if body is not original_body: + body.close() + if isinstance(contents, bytes): + original_body.close() def get( self, diff --git a/clients/python/src/objectstore_client/errors.py b/clients/python/src/objectstore_client/errors.py index 2f94eff6..2f623cbd 100644 --- a/clients/python/src/objectstore_client/errors.py +++ b/clients/python/src/objectstore_client/errors.py @@ -6,7 +6,7 @@ class RequestError(Exception): """Exception raised if an API call to Objectstore fails.""" - def __init__(self, message: str, status: int, response: str): + def __init__(self, message: str, status: int | None, response: str | None): super().__init__(message) self.status = status self.response = response diff --git a/clients/python/tests/test_e2e.py b/clients/python/tests/test_e2e.py index 8529d57c..51d0ae17 100644 --- a/clients/python/tests/test_e2e.py +++ b/clients/python/tests/test_e2e.py @@ -12,6 +12,7 @@ from datetime import timedelta from io import BytesIO from pathlib import Path +from unittest.mock import Mock, patch import pytest import urllib3 @@ -24,6 +25,7 @@ Usecase, ) from objectstore_client._resumable import ( + ResumableUpload, ResumableUploadUnavailable, UploadComplete, UploadIncomplete, @@ -1153,3 +1155,70 @@ def test_put_stores_under_literal_key(server_url: str) -> None: status, body = _fetch(url) assert status == 200 assert body == payload + + +@pytest.mark.parametrize("source_kind", ["bytes", "precompressed_bytes", "stream"]) +def test_automatic_resumable_upload(server_url: str, source_kind: str) -> None: + session = Client(server_url, token=TestSecretKey.get()).session( + Usecase( + "test-usecase", + compression="none", + expiration_policy=TimeToLive(timedelta(days=1)), + resumable_threshold_bytes=1, + ), + org=42, + ) + contents = b"file contents\n" * 100 + precompressed = source_kind == "precompressed_bytes" + encoded = ( + zstandard.ZstdCompressor().compress(contents) if precompressed else contents + ) + with ( + tempfile.TemporaryFile() as source, + patch.object( + ResumableUpload, "put", autospec=True, side_effect=ResumableUpload.put + ) as put, + ): + source.write(b"skip this prefix" + encoded) + source.seek(len(b"skip this prefix")) + key = session.put( + source if source_kind == "stream" else encoded, + precompressed="zstd" if precompressed else None, + content_type="text/plain", + metadata={"source": "file"}, + origin="203.0.113.42", + filename="example.txt", + ) + put.assert_called_once() + assert not source.closed + stored = session.head(key) + assert stored is not None + assert stored.compression == ("zstd" if precompressed else None) + assert stored.content_type == "text/plain" + assert stored.filename == "example.txt" + assert stored.origin == "203.0.113.42" + assert stored.custom == {"source": "file"} + retrieved = session.get(key) + assert retrieved is not None + assert retrieved.payload.read() == contents + + +@pytest.mark.parametrize( + "error", [None, RequestError("creation rejected", 403, "forbidden")] +) +def test_resumable_creation_fallback( + server_url: str, monkeypatch: pytest.MonkeyPatch, error: RequestError | None +) -> None: + session = Client(server_url, token=TestSecretKey.get()).session( + Usecase("test-usecase", resumable_threshold_bytes=1) + ) + create = Mock(return_value=None, side_effect=error) + monkeypatch.setattr(session, "_create_upload", create) + source = BytesIO(b"prefixpayload") + source.seek(len(b"prefix")) + key = session.put(source, compress="none") + create.assert_called_once() + assert not source.closed + retrieved = session.get(key) + assert retrieved is not None + assert retrieved.payload.read() == b"payload" diff --git a/clients/python/tests/test_upload.py b/clients/python/tests/test_upload.py new file mode 100644 index 00000000..349c2570 --- /dev/null +++ b/clients/python/tests/test_upload.py @@ -0,0 +1,117 @@ +from io import BytesIO +from typing import Any +from unittest.mock import Mock + +import pytest +import urllib3 +from objectstore_client import Client, RequestError, Session, Usecase + + +@pytest.fixture +def session(monkeypatch: pytest.MonkeyPatch) -> Session: + monkeypatch.setattr("objectstore_client._resumable.time.sleep", Mock()) + return Client("http://localhost:8888").session( + Usecase("test", compression="none", resumable_threshold_bytes=4) + ) + + +@pytest.mark.parametrize("complete", [False, True]) +def test_failure_recovery( + session: Session, monkeypatch: pytest.MonkeyPatch, complete: bool +) -> None: + failure = urllib3.exceptions.ReadTimeoutError(session._pool, "/", "lost response") + outcomes = Mock( + side_effect=[ + urllib3.HTTPResponse(status=201, body=b'{"key":"key","session":"token"}'), + failure, + urllib3.HTTPResponse( + status=201 if complete else 204, headers={"Upload-Offset": "3"} + ), + urllib3.HTTPResponse(status=201), + ] + ) + sent = [] + + def request(*args: Any, **kwargs: Any) -> urllib3.HTTPResponse: + if body := kwargs.get("body"): + sent.append((kwargs["headers"]["Upload-Offset"], body.read())) + return outcomes() + + monkeypatch.setattr(session._pool, "_make_request", request) + source = BytesIO(b"prefixabcdefgh") + source.seek(len(b"prefix")) + assert session.put(source) == "key" + assert not source.closed + assert sent == ( + [("0", b"abcdefgh")] if complete else [("0", b"abcdefgh"), ("3", b"defgh")] + ) + + +@pytest.mark.parametrize("status_failure", [False, True]) +def test_retry_exhaustion( + session: Session, + monkeypatch: pytest.MonkeyPatch, + status_failure: bool, +) -> None: + policy = urllib3.Retry( + total=2, + read=2, + status=2, + status_forcelist=[503], + ) + session._pool.retries = policy + failure = urllib3.exceptions.ReadTimeoutError(session._pool, "/", "lost response") + sleep = Mock() + monkeypatch.setattr("objectstore_client._resumable.time.sleep", sleep) + + def request(*args: Any, **kwargs: Any) -> urllib3.HTTPResponse: + headers = kwargs["headers"] + if args[1] == "DELETE": + return urllib3.HTTPResponse(status=204) + if "Upload-Length" in headers: + return urllib3.HTTPResponse( + status=201, body=b'{"key":"key","session":"token"}' + ) + if headers["Upload-Offset"] == "*": + return urllib3.HTTPResponse(status=204, headers={"Upload-Offset": "0"}) + if status_failure: + return urllib3.HTTPResponse(status=503, body=b"unavailable") + raise failure + + make_request = Mock(side_effect=request) + monkeypatch.setattr(session._pool, "_make_request", make_request) + with pytest.raises(RequestError, match="^upload failed$") as raised: + session.put(b"payload") + if status_failure: + assert isinstance(raised.value.__cause__, RequestError) + assert raised.value.__cause__.status == 503 + assert raised.value.response == "unavailable" + else: + assert isinstance(raised.value.__cause__, urllib3.exceptions.MaxRetryError) + assert raised.value.__cause__.reason is failure + assert raised.value.response is None + # Creation, two progress queries, three writes, and cancellation. + assert make_request.call_count == 1 + 2 + 3 + 1 + assert session._pool.retries is policy + assert sleep.call_count == 2 + for call, delay in zip(sleep.call_args_list, [2, 4], strict=True): + assert delay <= call.args[0] <= delay + 1 + + +def test_connection_retries_are_not_multiplied( + session: Session, monkeypatch: pytest.MonkeyPatch +) -> None: + policy = urllib3.Retry(total=2, connect=2) + session._pool.retries = policy + handle = session._resume_upload("key", "token") + monkeypatch.setattr(session, "_create_upload", Mock(return_value=handle)) + failure = urllib3.exceptions.ConnectTimeoutError("connection timed out") + request = Mock(side_effect=failure) + monkeypatch.setattr(session._pool, "_make_request", request) + with pytest.raises(RequestError) as raised: + session.put(b"payload") + assert isinstance(raised.value.__cause__, urllib3.exceptions.MaxRetryError) + assert raised.value.__cause__.reason is failure + # Three connection attempts each for the upload and cancellation. + assert request.call_count == 6 + assert session._pool.retries is policy