Fix multipart upload part retries missing connection exceptions

synapseclient/core/upload/multipart_upload.py

@@ -27,7 +27,7 @@
     SynapseUploadFailedException,
     _raise_for_status,
 )
-from synapseclient.core.retry import with_retry
+from synapseclient.core.retry import RETRYABLE_CONNECTION_EXCEPTIONS, with_retry
 from synapseclient.core.upload.upload_utils import (
     copy_md5_fn,
     copy_part_request_body_provider_fn,
@@ -295,7 +295,7 @@ def put_fn():
                 try:
                     # use our backoff mechanism here, we have encountered 500s on puts to AWS signed urls
                     response = with_retry(
-                        put_fn, retry_exceptions=[requests.exceptions.ConnectionError]
+                        put_fn, retry_exceptions=RETRYABLE_CONNECTION_EXCEPTIONS
                     )
                     _raise_for_status(response)
 

synapseclient/core/upload/multipart_upload_async.py

@@ -84,7 +84,6 @@
 
 import httpx
 import psutil
-import requests
 from opentelemetry import trace
 from tqdm.contrib.logging import logging_redirect_tqdm
 
@@ -104,7 +103,10 @@
     _raise_for_status_httpx,
 )
 from synapseclient.core.otel_config import get_tracer
-from synapseclient.core.retry import with_retry_time_based
+from synapseclient.core.retry import (
+    RETRYABLE_CONNECTION_EXCEPTIONS,
+    with_retry_time_based,
+)
 from synapseclient.core.transfer_bar import create_progress_bar
 from synapseclient.core.typing_utils import DataFrame as DATA_FRAME_TYPE
 from synapseclient.core.upload.upload_utils import (
@@ -606,7 +608,7 @@ def _put_part_with_retry(
                         content=body,  # noqa: F821
                         headers=signed_headers,
                     ),
-                    retry_exceptions=[requests.exceptions.ConnectionError],
+                    retry_exceptions=RETRYABLE_CONNECTION_EXCEPTIONS,
                 )
 
                 _raise_for_status_httpx(response=response, logger=self._syn.logger)

tests/unit/synapseclient/core/upload/test_multipart_upload_async.py

@@ -0,0 +1,116 @@
+from unittest import mock
+
+import httpx
+import pytest
+
+from synapseclient.core.upload.multipart_upload_async import (
+    HandlePartResult,
+    UploadAttemptAsync,
+)
+from synapseclient.core.utils import md5_fn
+
+PART_SIZE = 256
+PART_NUMBER = 1
+
+
+def _init_upload_attempt(syn):
+    upload_request_payload = {
+        "concreteType": "org.sagebionetworks.repo.model.file.MultipartUploadRequest",
+        "contentMD5Hex": "abc",
+        "contentType": "application/text",
+        "fileName": "target.txt",
+        "fileSizeBytes": 1024,
+        "generatePreview": False,
+        "storageLocationId": "1234",
+        "partSizeBytes": PART_SIZE,
+    }
+
+    def part_request_body_provider_fn(part_number):
+        return (f"{part_number}" * PART_SIZE).encode("utf-8")
+
+    upload = UploadAttemptAsync(
+        syn,
+        "target.txt",
+        upload_request_payload,
+        part_request_body_provider_fn,
+        md5_fn,
+        False,
+    )
+    upload._upload_id = "123"
+    upload._pre_signed_part_urls = {
+        PART_NUMBER: ("https://foo.com/1", {"a": 1}),
+    }
+    return upload
+
+
+class TestPutPartWithRetry:
+    """Regression coverage for the retry paths inside
+    UploadAttemptAsync._put_part_with_retry. multipart_upload_async.py previously
+    passed requests.exceptions.ConnectionError as the retryable exception allowlist,
+    but the session performing the PUT is an httpx.Client, which never raises that
+    exception type -- so a transient httpx connection error skipped the retry
+    window entirely and failed the part immediately.
+    """
+
+    @pytest.mark.parametrize(
+        "exception",
+        [
+            httpx.ConnectError("connection refused"),
+            httpx.ReadError("broken"),
+            httpx.ReadTimeout("timed out"),
+            httpx.ConnectTimeout("timed out"),
+            httpx.RemoteProtocolError("disconnected"),
+        ],
+    )
+    def test_handle_part__httpx_connection_error_then_success(self, syn, exception):
+        upload = _init_upload_attempt(syn)
+        mock_session = mock.Mock()
+        mock_session.put.side_effect = [exception, mock.Mock(status_code=200)]
+
+        with mock.patch.object(syn, "_requests_session_storage", mock_session):
+            result = upload._handle_part(PART_NUMBER)
+
+        assert mock_session.put.call_count == 2
+        body = (f"{PART_NUMBER}" * PART_SIZE).encode("utf-8")
+        assert result == HandlePartResult(PART_NUMBER, PART_SIZE, md5_fn(body, None))
+
+    def test_handle_part__retryable_status_then_success(self, syn):
+        upload = _init_upload_attempt(syn)
+        mock_session = mock.Mock()
+        mock_503 = mock.Mock(status_code=503, headers={}, text="")
+        mock_session.put.side_effect = [mock_503, mock.Mock(status_code=200)]
+
+        with mock.patch.object(syn, "_requests_session_storage", mock_session):
+            upload._handle_part(PART_NUMBER)
+
+        assert mock_session.put.call_count == 2
+
+    def test_handle_part__expired_url_then_success(self, syn):
+        upload = _init_upload_attempt(syn)
+        mock_session = mock.Mock()
+        mock_403 = mock.Mock(status_code=403, headers={}, text="")
+        mock_session.put.side_effect = [mock_403, mock.Mock(status_code=200)]
+
+        with (
+            mock.patch.object(syn, "_requests_session_storage", mock_session),
+            mock.patch.object(
+                upload,
+                "_refresh_pre_signed_part_urls",
+                return_value=("https://bar.com/1", {"a": 2}),
+            ) as refresh_urls,
+        ):
+            upload._handle_part(PART_NUMBER)
+
+        refresh_urls.assert_called_once_with(PART_NUMBER, "https://foo.com/1")
+        assert mock_session.put.call_count == 2
+
+    def test_handle_part__non_retryable_exception_fails_immediately(self, syn):
+        upload = _init_upload_attempt(syn)
+        mock_session = mock.Mock()
+        mock_session.put.side_effect = ValueError("boom")
+
+        with mock.patch.object(syn, "_requests_session_storage", mock_session):
+            with pytest.raises(ValueError):
+                upload._handle_part(PART_NUMBER)
+
+        assert mock_session.put.call_count == 1

tests/unit/synapseclient/core/upload/unit_test_multipart_upload.py

@@ -259,8 +259,20 @@ def test_handle_part__500(self, syn):
             None,
         )
 
-    def test_handle_part__connection_error(self, syn):
-        """Test that we retry if we encounter a ConnectionError on a reqeust to PUT to an AWS presigend url"""
+    @pytest.mark.parametrize(
+        "exception",
+        [
+            requests.exceptions.ConnectionError("aborted"),
+            ConnectionResetError("reset"),
+            requests.exceptions.Timeout("timed out"),
+            requests.exceptions.ChunkedEncodingError("truncated"),
+            requests.exceptions.ReadTimeout("read timed out"),
+            requests.exceptions.ConnectTimeout("connect timed out"),
+        ],
+    )
+    def test_handle_part__retryable_connection_exception(self, syn, exception):
+        """Test that we retry if we encounter a retryable connection exception (per
+        RETRYABLE_CONNECTION_EXCEPTIONS) on a request to PUT to an AWS presigned url."""
 
         upload = self._init_upload_attempt(syn)
         upload._upload_id = "123"
@@ -277,12 +289,10 @@ def test_handle_part__connection_error(self, syn):
             upload,
             part_number,
             pre_signed_url,
-            # initial call is expired and results in a 500
-            # second call is successful
             [
                 (
                     mock.call(pre_signed_url, chunk, headers=signed_headers),
-                    requests.exceptions.ConnectionError("aborted"),
+                    exception,
                 ),
                 (
                     mock.call(pre_signed_url, chunk, headers=signed_headers),