Stop appends, sync multipart copies, and commits at a failed part

docs/filesystem.md

@@ -91,11 +91,12 @@ this section apply to the compressed data.
 A multipart write keeps at most `max_workers` parts of the written data in flight and
 waits for one of them to finish before it uploads another. The parts take about
 `max_workers` × the block size of memory, in addition to the buffered data of the file.
-If one of these parts has failed when the write waits, the write raises its error
-instead of uploading the rest of the data, and aborts the multipart upload after the
-parts being uploaded finish. A part that fails after the last wait raises its error
-when the file is closed, or when the transaction commits, and the multipart upload is
-aborted.
+If one of these parts has failed when the write waits, or a part copied from the
+existing object in an append has failed before the write uploads another part, the
+write raises its error instead of uploading the rest of the data, and aborts the
+multipart upload after the parts being uploaded or copied finish. A part that fails
+later raises its error when the file is closed, or when the transaction commits, and
+the multipart upload is aborted.
 `max_workers` is an argument of `S3FileSystem`, `AioS3FileSystem`, `open`, `put`, and
 `pipe`.
 

pyathena/filesystem/s3.py

@@ -9,7 +9,7 @@
 import os.path
 import time
 from collections.abc import Callable, Mapping
-from concurrent.futures import FIRST_COMPLETED, Future, as_completed, wait
+from concurrent.futures import FIRST_COMPLETED, FIRST_EXCEPTION, Future, as_completed, wait
 from copy import deepcopy
 from datetime import datetime
 from io import BytesIO
@@ -1831,7 +1831,8 @@ def _finish_multipart_upload(
         interrupted, the parts that have not started are cancelled, the
         running ones are waited for, and the multipart upload is aborted so
         that no incomplete upload or part is left behind, unless ``abort`` is
-        false. The original error is then re-raised.
+        false. The original error is then re-raised. A failed part is
+        handled as soon as it fails, without waiting for the parts before it.
 
         Args:
             upload: The multipart upload returned by creation.
@@ -1847,6 +1848,12 @@ def _finish_multipart_upload(
         """
         request_kwargs = request_kwargs or {}
         try:
+            done, _ = wait(futures, return_when=FIRST_EXCEPTION)
+            for future in done:
+                if not future.cancelled() and future.exception() is not None:
+                    # Raised without waiting for the parts before it, so that
+                    # the parts queued after it are cancelled.
+                    future.result()
             # The futures are in part-number order.
             parts = [future.result() for future in futures]
             return self.core.complete_multipart_upload(
@@ -2898,6 +2905,8 @@ def __init__(
         # The parts uploaded from the buffer that may not have finished. The
         # part copies of an append hold no data and are not counted.
         self._uploads_in_flight: set[Future[S3MultipartUploadPart]] = set()
+        # The part copies of an append that may not have finished.
+        self._copies_in_flight: set[Future[S3MultipartUploadPart]] = set()
         if append_info is not None:
             if append_data is not None:
                 self.write(append_data)
@@ -3056,6 +3065,7 @@ def _initiate_upload(self) -> None:
                     # cancel it or wait for it before the abort.
                     part: Future[S3MultipartUploadPart] = Future()
                     self.multipart_upload_parts.append(part)
+                    self._copies_in_flight.add(part)
                     self._executor.submit_to(
                         part,
                         writer.upload_part_copy,
@@ -3078,7 +3088,8 @@ def _upload_chunk(self, final: bool = False) -> bool:
         which uploads it with PutObject. Otherwise, the buffered data is
         uploaded as parts, and a part is submitted only while fewer than
         ``max_workers`` uploaded parts are in flight; the write waits for
-        one of them to finish first. A part found failed by that wait stops
+        one of them to finish first. A part found failed by that wait, or a
+        part copy of an append found failed before a part is submitted, stops
         the write with its error; one that fails afterwards is reported by
         :meth:`commit`. If the upload of the chunk fails or is interrupted,
         the file is closed without committing it.
@@ -3093,7 +3104,8 @@ def _upload_chunk(self, final: bool = False) -> bool:
 
         Raises:
             BaseException: The error of an uploaded part that the write found
-                failed while it waited for a part in flight.
+                failed while it waited for a part in flight, or of a part copy
+                of an append found failed before a part is submitted.
             RuntimeError: If no multipart upload is initialized.
             ValueError: If the data needs more parts than a multipart upload
                 allows.
@@ -3121,6 +3133,13 @@ def _upload_chunk(self, final: bool = False) -> bool:
             for part_number, data in writer.iter_parts(
                 buffer, first_part_number=len(self.multipart_upload_parts) + 1
             ):
+                if self._copies_in_flight:
+                    copied = {copy for copy in self._copies_in_flight if copy.done()}
+                    self._copies_in_flight -= copied
+                    for copy in copied:
+                        # Otherwise the rest of the data is uploaded after a
+                        # failed part copy.
+                        copy.result()
                 if len(self._uploads_in_flight) >= self.max_workers:
                     # Otherwise the parts waiting for a worker pile up in memory.
                     done, self._uploads_in_flight = wait(
@@ -3156,11 +3175,12 @@ def commit(self) -> None:
 
         Creates an empty object if nothing was written, uploads the buffered
         data with PutObject if no multipart upload part was submitted, and
-        otherwise completes the multipart upload, which is aborted if the
-        completion fails or is interrupted. If the abort also fails, the
-        upload is kept so that :meth:`discard` can abort it. Invalidates the
-        cache of the path afterwards. Does nothing for a file whose failed
-        write dropped the written data.
+        otherwise completes the multipart upload, which is aborted if a part
+        or the completion fails or is interrupted. A failed part is handled
+        as soon as it fails, without waiting for the parts before it. If the
+        abort also fails, the upload is kept so that :meth:`discard` can
+        abort it. Invalidates the cache of the path afterwards. Does nothing
+        for a file whose failed write dropped the written data.
 
         Raises:
             FileExistsError: If an object was created at the path after the
@@ -3189,6 +3209,12 @@ def commit(self) -> None:
 
             upload_id = cast(str, self.multipart_upload.upload_id)
             try:
+                done, _ = wait(self.multipart_upload_parts, return_when=FIRST_EXCEPTION)
+                for future in done:
+                    if not future.cancelled() and future.exception() is not None:
+                        # Raised without waiting for the parts before it, so
+                        # that the parts queued after it are cancelled.
+                        future.result()
                 parts = [future.result() for future in self.multipart_upload_parts]
                 self._get_multipart_writer().complete(
                     parts, **self._get_request_kwargs("complete_multipart_upload")
@@ -3225,6 +3251,7 @@ def discard(self) -> None:
         self.multipart_upload = None
         self.multipart_upload_parts = []
         self._uploads_in_flight = set()
+        self._copies_in_flight = set()
 
     def url(self, expiration: int = 3600, **kwargs) -> str:
         """Generate a presigned HTTP URL to read this file (if it already exists).

tests/pyathena/filesystem/test_s3.py

@@ -18,7 +18,7 @@
 import uuid
 import weakref
 from base64 import b64encode
-from concurrent.futures import Future, ThreadPoolExecutor, wait
+from concurrent.futures import FIRST_COMPLETED, FIRST_EXCEPTION, Future, ThreadPoolExecutor, wait
 from datetime import UTC, datetime
 from itertools import chain
 from pathlib import Path
@@ -3532,9 +3532,11 @@ def upload_part():
             release.wait(5)
             events.append("part 2 stored")
 
-        def wait_parts(futures):
-            release.set()
-            return wait(futures)
+        def wait_parts(futures, **kwargs):
+            if not kwargs:
+                # The cleanup waits for the running parts.
+                release.set()
+            return wait(futures, **kwargs)
 
         with (
             ThreadPoolExecutor(max_workers=1) as executor,
@@ -3550,9 +3552,55 @@ def wait_parts(futures):
                 )
 
         assert events == ["part 2 stored", "abort"]
-        waited.assert_called_once_with([failed, running])
+        assert waited.call_args_list == [
+            mock.call([failed, running, pending], return_when=FIRST_EXCEPTION),
+            mock.call([failed, running]),
+        ]
         assert pending.cancelled()
 
+    def test_finish_multipart_upload_stops_at_failed_part(self):
+        # GH-1178: a failed part is handled without waiting for the parts
+        # before it, so that the parts queued after it are cancelled. The
+        # parts used to be collected in part-number order.
+        fs = self._make_fs()
+        upload = S3MultipartUpload({"Bucket": "bucket", "Key": "key", "UploadId": "uploadid"})
+        fs.core.complete_multipart_upload = mock.MagicMock()
+        started = threading.Event()
+        release = threading.Event()
+        cleanup = []
+
+        def copy_part():
+            started.set()
+            # Copying until the cleanup waits for it.
+            release.wait(5)
+
+        def wait_parts(futures, **kwargs):
+            if not kwargs:
+                # Whether the earlier part was still copying when the cleanup
+                # started.
+                cleanup.append(running.done())
+                release.set()
+            return wait(futures, **kwargs)
+
+        failed: Future[SimpleNamespace] = Future()
+        failed.set_exception(PermissionError("Access Denied"))
+        with (
+            ThreadPoolExecutor(max_workers=1) as executor,
+            mock.patch("pyathena.filesystem.s3.wait", side_effect=wait_parts),
+        ):
+            running = executor.submit(copy_part)
+            pending = executor.submit(SimpleNamespace)
+            assert started.wait(5)
+            with pytest.raises(PermissionError, match="Access Denied"):
+                fs._finish_multipart_upload(upload=upload, futures=[running, failed, pending])
+
+        assert cleanup == [False]
+        assert pending.cancelled()
+        fs.core.complete_multipart_upload.assert_not_called()
+        fs._call.assert_called_once_with(
+            fs._client.abort_multipart_upload, Bucket="bucket", Key="key", UploadId="uploadid"
+        )
+
     def test_finish_multipart_upload_does_not_wait_for_cancelled_parts(self):
         # GH-976: a cancelled part is not waited for, as nothing may
         # acknowledge its cancellation, e.g., an event loop blocked by the
@@ -6668,6 +6716,7 @@ def _make_write_file(data: bytes, autocommit: bool):
         # a part.
         file.max_workers = S3Core.MULTIPART_UPLOAD_MAX_PARTS
         file._uploads_in_flight = set()
+        file._copies_in_flight = set()
         file._executor = mock.MagicMock()
         return file
 
@@ -7027,6 +7076,9 @@ def record_upload(future, fn, *args, **kwargs):
         executor.submit_to = record_upload  # type: ignore[method-assign]
 
         def release_newest(futures, **kwargs):
+            if kwargs.get("return_when") != FIRST_COMPLETED:
+                # The wait of commit() for the parts.
+                return wait(futures, **kwargs)
             # Only the newest part in flight finishes, so a write that waits
             # for the oldest one would time out.
             gates[max(n for n, upload in uploads.items() if not upload.done())].set()
@@ -7050,7 +7102,9 @@ def release_newest(futures, **kwargs):
         assert self._uploaded_object(fs, existing) == existing + b"".join(writes)
         # Each part was submitted once fewer than max_workers were in flight.
         assert in_flight == [min(n, max_workers) for n in range(1, len(uploads) + 1)]
-        assert waited.call_count == len(uploads) - max_workers
+        assert [c.kwargs.get("return_when") for c in waited.call_args_list].count(
+            FIRST_COMPLETED
+        ) == len(uploads) - max_workers
 
     def test_write_interrupted_while_waiting_for_part(self):
         # GH-1158: an interrupt while a write waits for a part closes the file
@@ -7190,6 +7244,51 @@ def upload_part(**kwargs):
         fs.core.complete_multipart_upload.assert_not_called()
         assert f.multipart_upload is None
 
+    @pytest.mark.parametrize("max_workers", [1, 4])
+    def test_append_failed_part_copy(self, max_workers):
+        # GH-1178: a failed part copy of an append stops the write before it
+        # uploads another part, whether or not the write waits for a part.
+        # The part copies are not counted toward max_workers, and the write
+        # used to upload the rest of the data.
+        fs = self._make_append_fs(b"a" * 6)
+        release = threading.Event()
+
+        def upload_part_copy(**kwargs):
+            # Copying until the first write has returned. The test always
+            # sets the event, so the wait needs no timeout.
+            release.wait()
+            raise PermissionError("Access Denied")
+
+        fs.core.upload_part_copy.side_effect = upload_part_copy
+        executor = mock.MagicMock(wraps=S3ThreadPoolExecutor(max_workers=max_workers))
+        f = S3File(
+            fs,
+            "s3://bucket/key.txt",
+            mode="ab",
+            block_size=4,
+            max_workers=max_workers,
+            executor=executor,
+        )
+        try:
+            f.write(b"b" * 4)
+            release.set()
+            # The part copy has failed before the next write.
+            assert wait(f.multipart_upload_parts[:1], timeout=5).done
+            with pytest.raises(PermissionError, match="Access Denied"):
+                f.write(b"b" * 4)
+        finally:
+            release.set()
+            executor.shutdown()
+
+        assert [
+            c.kwargs["part_number"] for c in executor.submit_to.call_args_list if "body" in c.kwargs
+        ] == [2]
+        fs._call.assert_called_once_with(
+            S3_CLIENT.abort_multipart_upload, Bucket="bucket", Key="key.txt", UploadId="uploadid"
+        )
+        fs.core.complete_multipart_upload.assert_not_called()
+        assert f.closed
+
     def test_multipart_write_request_parameters(self):
         # GH-946: the parts receive the parameters of the file that they
         # accept, such as RequestPayer and SSE-C, as does the completion.
@@ -7424,6 +7523,50 @@ def test_commit_failure_and_discard(self, caplog, error, abort_fails):
         assert file.multipart_upload is None
         assert file.multipart_upload_parts == []
 
+    def test_commit_stops_at_failed_part(self):
+        # GH-1178: commit() handles a failed part without waiting for the
+        # parts before it, so that the parts queued after it, such as the
+        # part copies of a large append, are cancelled. The parts used to be
+        # collected in part-number order.
+        file = self._make_multipart_write_file(b"x" * 16, autocommit=False)
+        started = threading.Event()
+        release = threading.Event()
+        cleanup = []
+
+        def copy_part():
+            started.set()
+            # Copying until the cleanup waits for it.
+            release.wait(5)
+
+        def wait_parts(futures, **kwargs):
+            if not kwargs:
+                # Whether the earlier part was still copying when the cleanup
+                # started.
+                cleanup.append(running.done())
+                release.set()
+            return wait(futures, **kwargs)
+
+        failed: Future[SimpleNamespace] = Future()
+        failed.set_exception(PermissionError("Access Denied"))
+        with (
+            ThreadPoolExecutor(max_workers=1) as executor,
+            mock.patch("pyathena.filesystem.s3.wait", side_effect=wait_parts),
+        ):
+            running = executor.submit(copy_part)
+            pending = executor.submit(SimpleNamespace)
+            assert started.wait(5)
+            file.multipart_upload_parts = [running, failed, pending]
+            with pytest.raises(PermissionError, match="Access Denied"):
+                file.commit()
+
+        assert cleanup == [False]
+        assert pending.cancelled()
+        file.fs.core.complete_multipart_upload.assert_not_called()
+        file.fs._call.assert_called_once_with(
+            S3_CLIENT.abort_multipart_upload, Bucket="bucket", Key="key.txt", UploadId="uploadid"
+        )
+        assert file.multipart_upload is None
+
     def test_commit_failure_and_interrupted_abort(self):
         # GH-945: if the abort after a failed completion is interrupted, the
         # interrupt propagates and the upload is kept so that discard()

tests/pyathena/filesystem/test_s3_async.py

@@ -8,7 +8,7 @@
 import urllib.parse
 import urllib.request
 import uuid
-from concurrent.futures import wait
+from concurrent.futures import FIRST_COMPLETED, wait
 from datetime import UTC, datetime
 from itertools import chain
 from pathlib import Path
@@ -2405,6 +2405,9 @@ def record_upload(executor, future, fn, *args, **kwargs):
             submit_to(executor, future, fn, *args, **kwargs)
 
         def release_newest(futures, **kwargs):
+            if kwargs.get("return_when") != FIRST_COMPLETED:
+                # The wait of commit() for the parts.
+                return wait(futures, **kwargs)
             # Only the newest part in flight finishes, so a write that waits
             # for the oldest one would time out.
             gates[max(n for n, upload in uploads.items() if not upload.done())].set()