fix(storage): respect deadline in transfer_manager.upload_many

packages/google-cloud-storage/google/cloud/storage/transfer_manager.py

@@ -97,6 +97,13 @@ def convert_threads_or_raise(*args, **kwargs):
     return convert_threads_or_raise
 
 
+def _process_pool_terminate_workers(executor):
+    """Terminate a ProcessPoolExecutor's worker processes, if any."""
+    # executor.terminate_workers() is only supported on Python >= 3.14
+    for process in (getattr(executor, "_processes", None) or {}).values():
+        process.terminate()
+
+
 @_deprecate_threads_param
 def upload_many(
     file_blob_pairs,
@@ -213,7 +220,11 @@ def upload_many(
 
     pool_class, needs_pickling = _get_pool_class_and_requirements(worker_type)
 
-    with pool_class(max_workers=max_workers) as executor:
+    # The executor is managed explicitly rather than via a context manager.
+    # The context manager's implicit shutdown(wait=True) prevents correct implementation
+    # of deadline.
+    executor = pool_class(max_workers=max_workers)
+    try:
         futures = []
         for path_or_file, blob in file_blob_pairs:
             # File objects are only supported by the THREAD worker because they can't
@@ -236,9 +247,19 @@ def upload_many(
                     **upload_kwargs,
                 )
             )
-        concurrent.futures.wait(
+        _, not_done = concurrent.futures.wait(
             futures, timeout=deadline, return_when=concurrent.futures.ALL_COMPLETED
         )
+        if not_done:
+            # Deadline exceeded. If using process mode, kill the workers.
+            if isinstance(executor, concurrent.futures.ProcessPoolExecutor):
+                _process_pool_terminate_workers(executor)
+            raise concurrent.futures.TimeoutError(
+                "Deadline of {} second(s) exceeded while waiting for uploads "
+                "to complete.".format(deadline)
+            )
+    finally:
+        executor.shutdown(wait=False, cancel_futures=True)
 
     results = []
     for future in futures:

packages/google-cloud-storage/tests/unit/test_transfer_manager.py

@@ -12,9 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import concurrent.futures
+import io
 import os
 import pickle
 import tempfile
+import time
 
 import mock
 import pytest
@@ -114,6 +117,7 @@ def test_upload_many_passes_concurrency_options():
         mock.patch("concurrent.futures.ThreadPoolExecutor") as pool_patch,
         mock.patch("concurrent.futures.wait") as wait_patch,
     ):
+        wait_patch.return_value = ({pool_patch.return_value.submit.return_value}, set())
         transfer_manager.upload_many(
             FILE_BLOB_PAIRS,
             deadline=DEADLINE,
@@ -135,6 +139,7 @@ def test_threads_deprecation_with_upload():
         mock.patch("concurrent.futures.ThreadPoolExecutor") as pool_patch,
         mock.patch("concurrent.futures.wait") as wait_patch,
     ):
+        wait_patch.return_value = ({pool_patch.return_value.submit.return_value}, set())
         with pytest.warns():
             transfer_manager.upload_many(
                 FILE_BLOB_PAIRS, deadline=DEADLINE, threads=MAX_WORKERS
@@ -189,6 +194,57 @@ def test_upload_many_raises_exceptions():
         )
 
 
+def test_upload_many_raises_timeout_error_when_deadline_exceeded():
+    # Thread-mode: A stuck upload must not make upload_many hang past the deadline.
+    def blocking_upload(*args, **kwargs):
+        time.sleep(5)
+
+    mock_blob = mock.Mock(spec=Blob)
+    mock_blob._prep_and_do_upload.side_effect = blocking_upload
+
+    with pytest.raises(concurrent.futures.TimeoutError):
+        transfer_manager.upload_many(
+            [(io.BytesIO(b"data"), mock_blob)],
+            worker_type=transfer_manager.THREAD,
+            deadline=0.1,
+        )
+
+
+def test_upload_many_terminates_process_workers_on_deadline():
+    # Process-mode: A stuck upload must not make upload_many hang past the deadline.
+    # After the deadline, the worker processes should be terminated.
+    fake_processes = {1: mock.Mock(), 2: mock.Mock()}
+
+    class FakeProcessPoolExecutor(concurrent.futures.ProcessPoolExecutor):
+        def __init__(self, *args, **kwargs):
+            self._processes = fake_processes
+
+        def submit(self, *args, **kwargs):
+            return concurrent.futures.Future()
+
+        def shutdown(self, *args, **kwargs):
+            pass
+
+    with (
+        mock.patch(
+            "google.cloud.storage.transfer_manager._get_pool_class_and_requirements",
+            return_value=(FakeProcessPoolExecutor, False),
+        ),
+        mock.patch("concurrent.futures.wait") as wait_patch,
+    ):
+        # A non-empty not_done set signals the deadline was exceeded.
+        wait_patch.return_value = (set(), {concurrent.futures.Future()})
+        with pytest.raises(concurrent.futures.TimeoutError):
+            transfer_manager.upload_many(
+                [("file_a.txt", mock.Mock(spec=Blob))],
+                worker_type=transfer_manager.PROCESS,
+                deadline=0.1,
+            )
+
+    for process in fake_processes.values():
+        process.terminate.assert_called_once_with()
+
+
 def test_upload_many_suppresses_412_with_skip_if_exists():
     FILE_BLOB_PAIRS = [
         ("file_a.txt", mock.Mock(spec=Blob)),