ref(preprod): Use regular PUTs for snapshot archive uploads

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",

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:

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={

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()
 
 

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,15 +98,15 @@ 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
 
     @patch("sentry.preprod.snapshots.zip_tasks._send_archive_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()