diff --git a/pyproject.toml b/pyproject.toml index 1e4283c3bcc1..8f2ec0154704 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,7 +61,7 @@ dependencies = [ "mmh3>=4.0.0", "msgspec>=0.19.0", "msgpack>=1.2.1", - "objectstore-client>=0.2.5", + "objectstore-client>=0.2.10", "orjson>=3.10.10", "packaging>=24.1", "parsimonious>=0.10.0", diff --git a/src/sentry/preprod/snapshots/storage.py b/src/sentry/preprod/snapshots/storage.py index 7b8b3fb02041..99c5afe64aa4 100644 --- a/src/sentry/preprod/snapshots/storage.py +++ b/src/sentry/preprod/snapshots/storage.py @@ -1,10 +1,9 @@ from __future__ import annotations import logging -from typing import IO, Literal +from typing import IO from objectstore_client import Compression, GetResponse, Metadata, RequestError, Session -from objectstore_client.multipart import MultipartUpload from urllib3.exceptions import HTTPError from sentry.models.project import Project @@ -48,14 +47,21 @@ def _record_fallback(self, op: str, found: bool) -> None: }, ) - def put(self, contents: bytes | IO[bytes], *, key: str, content_type: str | None = None) -> str: - return self._primary.put(contents, key=key, content_type=content_type) - - def initiate_multipart_upload( - self, *, key: str, compression: Compression | Literal["none"], content_type: str - ) -> MultipartUpload: - return self._primary.initiate_multipart_upload( - key=key, compression=compression, content_type=content_type + def put( + self, + contents: bytes | IO[bytes], + *, + key: str, + content_type: str | None = None, + compress: Compression | None = None, + resumable_threshold_bytes: int | None = None, + ) -> str: + return self._primary.put( + contents, + key=key, + content_type=content_type, + compress=compress, + resumable_threshold_bytes=resumable_threshold_bytes, ) def delete(self, key: str) -> None: diff --git a/src/sentry/preprod/snapshots/zip_tasks.py b/src/sentry/preprod/snapshots/zip_tasks.py index 3a7e6978b50b..1638ceec3b8d 100644 --- a/src/sentry/preprod/snapshots/zip_tasks.py +++ b/src/sentry/preprod/snapshots/zip_tasks.py @@ -2,14 +2,10 @@ import logging import os -import time from tempfile import NamedTemporaryFile -from typing import IO, Any +from typing import Any import orjson -from objectstore_client import RequestError -from objectstore_client.multipart import CompletePart, MultipartUpload -from urllib3.exceptions import HTTPError from sentry.models.organization import Organization from sentry.preprod.snapshots.manifest import SnapshotManifest @@ -29,9 +25,6 @@ logger = logging.getLogger(__name__) -MULTIPART_PART_SIZE = 32 * 1024 * 1024 -MULTIPART_MAX_RETRIES = 3 - def _snapshot_page_url(organization: Organization, artifact_id: int) -> str: # Link to the snapshots UI rather than the raw download endpoint: the page's @@ -66,17 +59,6 @@ def _send_archive_email( ).send_async([email]) -def _put_part_with_retry(upload: MultipartUpload, chunk: bytes, part_number: int) -> CompletePart: - for attempt in range(MULTIPART_MAX_RETRIES): - try: - return upload.put_part(chunk, part_number=part_number, content_length=len(chunk)) - except (RequestError, HTTPError): - if attempt == MULTIPART_MAX_RETRIES - 1: - raise - time.sleep(2**attempt) - raise AssertionError("unreachable") - - def _archive_available(org_id: int, project_id: int, artifact_id: int) -> bool: try: session = get_snapshot_storage(project_id, org=org_id) @@ -85,28 +67,6 @@ def _archive_available(org_id: int, project_id: int, artifact_id: int) -> bool: return False -def _upload_archive_multipart(session: SnapshotStorage, key: str, tmp: IO[bytes]) -> None: - upload = session.initiate_multipart_upload( - key=key, compression="none", content_type="application/zip" - ) - try: - parts: list[CompletePart] = [] - part_number = 1 - while True: - chunk = tmp.read(MULTIPART_PART_SIZE) - if not chunk: - break - parts.append(_put_part_with_retry(upload, chunk, part_number)) - part_number += 1 - upload.complete(parts) - except Exception: - try: - upload.abort() - except (RequestError, HTTPError): - logger.warning("preprod_snapshot_zip.archive_abort_failed", extra={"key": key}) - raise - - @instrumented_task( name="sentry.preprod.tasks.build_snapshot_images_zip", namespace=preprod_snapshots_tasks, @@ -196,7 +156,17 @@ def build_snapshot_images_zip( "archive_size_bytes": archive_size_bytes, }, ) - _upload_archive_multipart(session, key, tmp) + session.put( + tmp, + key=key, + content_type="application/zip", + # A `zip` is already compressed, so it would likely be not beneficial to compress it again. + compress="none", + # Force this `put` to use a resumable upload, overriding the default threshold. + # This already incorporates retries with exponential backoffs and allows the client to retry + # from the persisted offset when the upload fails midway through. + resumable_threshold_bytes=0, + ) logger.info( "preprod_snapshot_zip.upload_completed", extra={ diff --git a/tests/sentry/preprod/snapshots/test_storage.py b/tests/sentry/preprod/snapshots/test_storage.py index 177d55c7c29f..a24ed5c1577e 100644 --- a/tests/sentry/preprod/snapshots/test_storage.py +++ b/tests/sentry/preprod/snapshots/test_storage.py @@ -52,8 +52,12 @@ def test_head_uses_fallback_when_primary_missing(sessions) -> None: def test_put_only_writes_primary(sessions) -> None: primary, fallback = sessions storage = SnapshotStorage(primary, fallback) - storage.put(b"data", key="k", content_type="image/png") - primary.put.assert_called_once_with(b"data", key="k", content_type="image/png") + storage.put( + b"data", key="k", content_type="image/png", compress="none", resumable_threshold_bytes=0 + ) + primary.put.assert_called_once_with( + b"data", key="k", content_type="image/png", compress="none", resumable_threshold_bytes=0 + ) fallback.put.assert_not_called() diff --git a/tests/sentry/preprod/snapshots/test_zip_tasks.py b/tests/sentry/preprod/snapshots/test_zip_tasks.py index 29d8ea27f1e0..6a78c80a221f 100644 --- a/tests/sentry/preprod/snapshots/test_zip_tasks.py +++ b/tests/sentry/preprod/snapshots/test_zip_tasks.py @@ -5,18 +5,11 @@ from unittest.mock import MagicMock, patch import orjson -import pytest from objectstore_client import RequestError -from objectstore_client.multipart import CompletePart from sentry.preprod.models import PreprodArtifact -from sentry.preprod.snapshots import zip_tasks from sentry.preprod.snapshots.models import PreprodSnapshotMetrics -from sentry.preprod.snapshots.zip_tasks import ( - _put_part_with_retry, - _upload_archive_multipart, - build_snapshot_images_zip, -) +from sentry.preprod.snapshots.zip_tasks import build_snapshot_images_zip from sentry.testutils.cases import TestCase SESSION_TARGET = "sentry.preprod.snapshots.zip_tasks.get_snapshot_storage" @@ -67,25 +60,30 @@ def _get(key): def test_build_includes_manifest(self, mock_session, mock_email): session = self._session() mock_session.return_value = session + archive_bytes = b"" + + def upload(contents, **kwargs): + nonlocal archive_bytes + archive_bytes = contents.read() + return kwargs["key"] + + session.put.side_effect = upload build_snapshot_images_zip( org_id=self.org.id, project_id=self.project.id, artifact_id=self.artifact.id, user_id=self.user.id, ) - session.initiate_multipart_upload.assert_called_once() - kwargs = session.initiate_multipart_upload.call_args.kwargs + session.put.assert_called_once() + kwargs = session.put.call_args.kwargs assert kwargs["key"] == f"snapshot_archives/{self.artifact.id}.zip" - assert kwargs["compression"] == "none" assert kwargs["content_type"] == "application/zip" + assert kwargs["compress"] == "none" + assert kwargs["resumable_threshold_bytes"] == 0 - upload = session.initiate_multipart_upload.return_value - assert upload.put_part.call_count >= 1 - archive_bytes = b"".join(call.args[0] for call in upload.put_part.call_args_list) with zipfile.ZipFile(io.BytesIO(archive_bytes)) as archive: assert sorted(archive.namelist()) == ["a.png", "manifest.json"] assert archive.read("manifest.json") == _manifest_bytes() - upload.complete.assert_called_once() mock_email.assert_called_once() assert mock_email.call_args.kwargs["ready"] is True @@ -100,7 +98,7 @@ def test_reuses_existing_archive(self, mock_session, mock_email): artifact_id=self.artifact.id, user_id=self.user.id, ) - session.initiate_multipart_upload.assert_not_called() + session.put.assert_not_called() mock_email.assert_called_once() assert mock_email.call_args.kwargs["ready"] is True @@ -108,7 +106,7 @@ def test_reuses_existing_archive(self, mock_session, mock_email): @patch(SESSION_TARGET) def test_emails_failure_when_upload_raises(self, mock_session, mock_email): session = self._session() - session.initiate_multipart_upload.side_effect = RuntimeError("boom") + session.put.side_effect = RuntimeError("boom") mock_session.return_value = session build_snapshot_images_zip( org_id=self.org.id, @@ -126,7 +124,7 @@ def test_emails_ready_when_failure_superseded_by_existing_archive( self, mock_session, mock_email, mock_available ): session = self._session() - session.initiate_multipart_upload.side_effect = RuntimeError("boom") + session.put.side_effect = RuntimeError("boom") mock_session.return_value = session build_snapshot_images_zip( org_id=self.org.id, @@ -174,80 +172,3 @@ def test_emails_failure_when_objectstore_probe_errors(self, mock_session, mock_e ) mock_email.assert_called_once() assert mock_email.call_args.kwargs["ready"] is False - - -class MultipartUploadHelperTest(TestCase): - @patch("sentry.preprod.snapshots.zip_tasks.time.sleep") - def test_put_part_retries_then_succeeds(self, mock_sleep): - upload = MagicMock() - part = CompletePart(part_number=1, etag="e1") - upload.put_part.side_effect = [ - RequestError("transient", status=500, response=""), - part, - ] - - result = _put_part_with_retry(upload, b"data", 1) - - assert result is part - assert upload.put_part.call_count == 2 - - @patch("sentry.preprod.snapshots.zip_tasks.time.sleep") - def test_put_part_retries_on_transport_error(self, mock_sleep): - from urllib3.exceptions import ProtocolError - - upload = MagicMock() - part = CompletePart(part_number=1, etag="e1") - upload.put_part.side_effect = [ProtocolError("connection aborted"), part] - - result = _put_part_with_retry(upload, b"data", 1) - - assert result is part - assert upload.put_part.call_count == 2 - - @patch("sentry.preprod.snapshots.zip_tasks.time.sleep") - def test_put_part_raises_after_max_retries(self, mock_sleep): - upload = MagicMock() - upload.put_part.side_effect = RequestError("down", status=500, response="") - - with pytest.raises(RequestError): - _put_part_with_retry(upload, b"data", 1) - - assert upload.put_part.call_count == 3 - - def test_upload_archive_multipart_uploads_all_parts(self): - upload = MagicMock() - upload.put_part.side_effect = lambda chunk, *, part_number, content_length: CompletePart( - part_number=part_number, etag=f"e{part_number}" - ) - session = MagicMock() - session.initiate_multipart_upload.return_value = upload - - with patch.object(zip_tasks, "MULTIPART_PART_SIZE", 4): - _upload_archive_multipart(session, "snapshot_archives/1.zip", io.BytesIO(b"abcdefghij")) - - session.initiate_multipart_upload.assert_called_once_with( - key="snapshot_archives/1.zip", - compression="none", - content_type="application/zip", - ) - assert upload.put_part.call_count == 3 - upload.complete.assert_called_once() - parts = upload.complete.call_args.args[0] - assert [p.part_number for p in parts] == [1, 2, 3] - upload.abort.assert_not_called() - - @patch("sentry.preprod.snapshots.zip_tasks.time.sleep") - def test_upload_archive_multipart_aborts_on_failure(self, mock_sleep): - upload = MagicMock() - upload.put_part.side_effect = RequestError("down", status=500, response="") - session = MagicMock() - session.initiate_multipart_upload.return_value = upload - - with patch.object(zip_tasks, "MULTIPART_PART_SIZE", 4): - with pytest.raises(RequestError): - _upload_archive_multipart( - session, "snapshot_archives/1.zip", io.BytesIO(b"abcdefgh") - ) - - upload.abort.assert_called_once() - upload.complete.assert_not_called() diff --git a/uv.lock b/uv.lock index 337fbc6f1765..2fea691c4497 100644 --- a/uv.lock +++ b/uv.lock @@ -1423,7 +1423,7 @@ wheels = [ [[package]] name = "objectstore-client" -version = "0.2.5" +version = "0.2.10" source = { registry = "https://pypi.devinfra.sentry.io/simple" } dependencies = [ { name = "filetype" }, @@ -1433,7 +1433,7 @@ dependencies = [ { name = "zstandard" }, ] wheels = [ - { url = "https://pypi.devinfra.sentry.io/wheels/objectstore_client-0.2.5-py3-none-any.whl", hash = "sha256:d080ae54573d7b1eb10a6a282dee03de1efa459b54c3dc99c2569b253d0b6dd3" }, + { url = "https://pypi.devinfra.sentry.io/wheels/objectstore_client-0.2.10-py3-none-any.whl", hash = "sha256:62c953d720c96aa28351f38835d9a9628c9b45920d73abd7e8ff895793305c77" }, ] [[package]] @@ -2412,7 +2412,7 @@ requires-dist = [ { name = "mmh3", specifier = ">=4.0.0" }, { name = "msgpack", specifier = ">=1.2.1" }, { name = "msgspec", specifier = ">=0.19.0" }, - { name = "objectstore-client", specifier = ">=0.2.5" }, + { name = "objectstore-client", specifier = ">=0.2.10" }, { name = "orjson", specifier = ">=3.10.10" }, { name = "packaging", specifier = ">=24.1" }, { name = "parsimonious", specifier = ">=0.10.0" },