Abort a cancelled async multipart copy after its running parts

[laughingman7743]

WHAT

AioS3FileSystem._copy_object_with_multipart_upload() now cleans up a cancelled copy. It waits for the part copies and the completion that are still running in threads, aborts the multipart upload unless it has completed, and re-raises the cancellation. This matches what S3FileSystem._finish_multipart_upload() does on an interrupt.

• The parts are awaited with asyncio.wait(..., return_when=FIRST_EXCEPTION) instead of asyncio.gather(). When the awaiting task is cancelled, gather() cancels the part tasks while their asyncio.to_thread() copies keep running, so they could no longer be waited for. wait() leaves the part tasks running. A failed part's error is raised, as before.
• The completion runs as a task awaited through asyncio.shield(), so a cancellation does not detach it from its thread.
• The cleanup catches BaseException instead of Exception, so it also runs on asyncio.CancelledError (for example task.cancel(), an asyncio.wait_for() timeout, or a cancelled gather() or TaskGroup of copies). It first waits for the running parts. Parts that have not started do not start after the cancellation, as after a failed part. It then waits for the completion, if it was sent, and aborts the upload unless the completion succeeded.
• The cleanup runs in its own task, awaited through asyncio.shield(). A repeated cancellation of the copy therefore returns at once, while the cleanup still runs to the abort. A module-level set keeps the cleanup task referenced while it runs, because the event loop keeps only weak references to tasks.

Behavior changes for the release notes, for AioS3FileSystem copies of objects over 5 GiB (_cp_file()/cp_file()/copy() and _mv()/mv()):

• A cancellation while the parts are copied now returns only after the running part copies finish and the upload is aborted. That can take as long as the slowest running part copy (max_workers parts of up to block_size each). Before, the cancellation returned at once and left the multipart upload, and the parts stored after the cancellation, in the destination bucket.
• A cancellation during CompleteMultipartUpload now waits for the completion. If the completion succeeded, the object exists and is not aborted. Its annotations are not copied, and the cancellation is re-raised. If the completion failed, the upload is aborted. Before, the cancellation returned at once and the completion went on in the background without an abort, even if it failed.
• A second cancellation while the cleanup runs returns at once, and the cleanup still aborts the upload.

Not handled:

• A cancellation during the CreateMultipartUpload request, before the upload ID is known. The sync path does not handle an interrupt there either.
• Event-loop shutdown, for example asyncio.run() cancelling the remaining tasks at exit. Shutdown cancels the part, completion, and cleanup tasks directly, and shield() does not protect against that. A cleanup that is cancelled skips the abort, as master always did. A cleanup that is still running can abort while a part thread is still copying, so that part may be stored after the abort; master did not abort at all in this case.
• A cleanup that a repeated cancellation detached does not report its own failure, for example when a loop's default executor is already shut down and the abort cannot be submitted. asyncio then logs "Task exception was never retrieved" with the error.

WHY

Fixes #1046. #973 (PR #1036) added the abort for failed parts and completions and deliberately left cancellation out.

TEST

Tested head: cd0c4a4 (rebased on master a4cf604). The S3 run below was at 50fe7cb; the later commits change only a docstring and the cancellation tests, whose offline runs below are at cd0c4a4.

• just format and just lint: passed.
• TestAioS3FileSystem.test_copy_object_with_multipart_upload_cancelled[1|2] (offline, mocked requests) holds two of three parts with an event and cancels the copy once or twice. It checks that both parts finish before the abort, that part 3 never starts, that the completion is not sent, and that CancelledError is raised. With two cancellations, it also checks that the copy returns while the parts still run and that the abort follows.
• TestAioS3FileSystem.test_copy_object_with_multipart_upload_cancelled_completion[False|True] holds the completion and cancels the copy. It keeps the completion held for 0.1 s after the cancellation and checks that the copy neither returns nor aborts meanwhile, that a successful completion is not aborted, and that a failed one is aborted after it finishes.
• Without the pyathena/ change (master behavior), all 4 new cases fail. With the first version of this PR (single cancellation only), cancelled[2] and both cancelled_completion cases fail.
• uv run --env-file .env pytest -n 1 tests/pyathena/filesystem/test_s3_async.py: 116 passed (includes the S3 integration tests).
• Offline -k multipart_upload tests with -W error: 15 passed on Python 3.13.1 and 3.11.11 (the requires-python floor). The cancelled tests passed 20 consecutive runs.
• The reproduction from Cancelling an async multipart copy leaves its multipart upload behind #1046 in a local script with asyncio.run(..., debug=True), for three cases: task.cancel(), task.cancel() with a part that then fails, and an asyncio.wait_for() timeout. In each case the abort ran after parts 1 and 2 finished, no part finished afterwards, and nothing was logged.

Not run: cancellation against real S3, because a multipart copy needs a source object over 5 GiB.

🤖 Generated with Claude Code

[laughingman7743]

Self-review round one (behavior and implementation): FINDINGS, repaired.

Scope: base e49d576ce1e8cb60d74e1fde56de9c72ff3cd8e1 (merge-base with master), head 43b3fe471fdcf4ccf657a2383ea6df284e18e8c4, both changed files: pyathena/filesystem/s3_async.py and tests/pyathena/filesystem/test_s3_async.py.

Finding: this docstring still said that the part copies run in parallel with asyncio.gather. After this PR they run as tasks that are awaited with asyncio.wait, so the docstring no longer matched the code. It also did not mention the new behavior on cancellation.

Repair in 185115d: the docstring now says the copies run as asyncio tasks with asyncio.to_thread. It also states that a cancellation while the parts are copied or the upload is completed waits for the running part copies, aborts the upload, and is re-raised. The wording leaves out a cancellation during CreateMultipartUpload, which this PR does not handle (see the PR description). just lint passed.

[laughingman7743]

Self-review round one, checked: CLEAN.

• asyncio.wait() neither cancels the part tasks when the awaiting task is cancelled nor creates an outer future that could hold an unretrieved exception. An earlier asyncio.shield(asyncio.gather(...)) version logged _GatheringFuture exception was never retrieved when a cancellation and a failed part combined. This version logs nothing, checked with asyncio.run(..., debug=True).
• Failure path: FIRST_EXCEPTION returns at the first failed part. Its error is raised, and the existing cleanup then waits for the running parts and aborts. test_copy_object_with_multipart_upload_failed_part and test_copy_object_with_multipart_upload_waits_for_running_parts pass unchanged. One difference: if two parts fail within the same event-loop step, either error may be raised, whereas gather raised the first one to complete. Both are part-copy errors, so this is not actionable.
• An empty ranges would make asyncio.wait([]) raise ValueError where gather() returned []. This cannot happen: the only caller, _copy_file, passes size1 > MULTIPART_UPLOAD_MAX_PART_SIZE, and _get_copy_ranges(0, ...) already raises before this point.
• asyncio.wait_for() on Python 3.11, the floor in requires-python, waits for the cancelled inner task to finish. 3.12+ implements it with asyncio.timeout(), which lets the cleanup run before the timeout is raised. The 3.13 local check covers wait_for; 3.11 was not run locally.

[laughingman7743]

Self-review round one, test quality: CLEAN.

The test cancels the copy only after both running parts have signalled that they started. The CancelledError reaches the task on the next loop step, about 200 ms before the part threads finish. The test then checks the observable order (both end events before abort), that part 3 never starts, that no CompleteMultipartUpload is sent, and that CancelledError is raised. With the pyathena/ change reverted, it fails: no abort is called. It reuses the mocking style of test_copy_object_with_multipart_upload_waits_for_running_parts and adds no new fake framework.

[laughingman7743]

Self-review round two (claims, callers, and operations): FINDINGS, repaired in the PR description.

Scope: full pass over base e49d576ce1e8cb60d74e1fde56de9c72ff3cd8e1, head 185115db8963dabe413396d3c71a883e0516acb6. It covers the PR description, both commit messages, the changed docstring, the code and test comments, and the premises of #1046.

Finding: the description listed "a cancellation during CompleteMultipartUpload sends the abort while the completion may still run" under Unchanged. That is false. Before this PR, except Exception did not catch CancelledError, so the completion thread finished and created the object without an abort. With except BaseException here, the abort is now sent while the completion request may still run in its thread. Either the abort wins and no object is created, or the completion wins and the failed abort is logged by S3FileSystem._abort_multipart_upload() (pyathena/filesystem/s3.py:2342). The sync _finish_multipart_upload() also aborts after an interrupt during completion, so this matches the issue's expectation. The description now lists it as a behavior change for the release notes.

Also added to the description: a cancelled copy now returns only after the running part copies finish, which can take up to the slowest running part copy (max_workers parts of up to block_size). This matters for asyncio.wait_for() callers. "The first part error found is raised" became "A failed part's error is raised"; see the round-one note on simultaneous failures.

Claims checked and held:

• gather() cancels its children when cancelled, while the to_thread() copies keep running (the issue's reproduction, rerun on master behavior); wait() does not cancel them (new test).
• Affected entry points: _copy_file() reaches the multipart path only for size1 > MULTIPART_UPLOAD_MAX_PART_SIZE, and it is called from _cp_file() (s3_async.py:409) and _mv() (s3_async.py:370).
• The cleanup runs under asyncio.wait_for() on Python 3.11.11, the requires-python floor, and on 3.13.1: the local script and the 12 offline -k multipart_upload tests on 3.11.
• "Second cancellation" and "cancellation during CreateMultipartUpload" are left unhandled in both the async and the sync paths.

Evidence limits, as stated in the PR: the AWS test run belongs to 43b3fe4 (185115d changes only the docstring). No real-S3 cancellation and no test for cancellation during completion were run.

[laughingman7743]

Independent review (relayed), static review only. Reviewer: Codex CLI 0.160.0, model gpt-6-astra, reasoning effort max, --sandbox read-only, session 01a10457-ec39-7471-aaef-d09f81f11567. Scope: base e49d576ce1e8cb60d74e1fde56de9c72ff3cd8e1, head 185115db8963dabe413396d3c71a883e0516acb6, detached snapshot plus the literal diff. The prompt left out the PR number, the description, the commit messages, and the self-review results. The reviewer ran no edits, builds, or tests and had no GitHub access, and the snapshot and the PR worktree were unchanged afterwards. Covered: async multipart copy and its _copy_file/_cp_file/_mv callers, sync completion/abort handling, tests; success, failures, cancellation propagation, exception retrieval, docs, Python 3.11-3.14. Verdict: FINDINGS (3, all P2).

Finding 1 (P2): a repeated cancellation can skip the abort. Reviewer's scenario: two part copies are running. Cancel the copy, then cancel it again while the cleanup awaits gather. The second cancellation cancels the part tasks and interrupts the cleanup before _abort_multipart_upload, the threads continue, and the upload is never aborted. The same happens for a cancellation during the cleanup after an ordinary part failure. Reviewer's classification: an incomplete repair of the pre-existing leak, not a new one.

Verified. The 185115d version fails the new test_copy_object_with_multipart_upload_cancelled[2]. Repaired in 50fe7cb (maintainer chose to fix it in this PR): the cleanup is the _abort() task, awaited through asyncio.shield() here and kept referenced in _cleanup_tasks. A repeated cancellation returns at once, and the cleanup still waits for the parts and aborts. Remaining limit, stated in the PR: an event-loop shutdown that cancels all remaining tasks also cancels the cleanup.

[laughingman7743]

Independent review (relayed), static review only. Reviewer: Codex CLI 0.160.0, model gpt-6-astra, reasoning effort max, --sandbox read-only, session 01a10457-ec39-7471-aaef-d09f81f11567. Scope: base e49d576ce1e8cb60d74e1fde56de9c72ff3cd8e1, head 185115db8963dabe413396d3c71a883e0516acb6, detached snapshot plus the literal diff. The prompt left out the PR number, the description, the commit messages, and the self-review results. The reviewer ran no edits, builds, or tests and had no GitHub access, and the snapshot and the PR worktree were unchanged afterwards. Covered: async multipart copy and its _copy_file/_cp_file/_mv callers, sync completion/abort handling, tests; success, failures, cancellation propagation, exception retrieval, docs, Python 3.11-3.14. Verdict: FINDINGS (3, all P2).

Finding 2 (P2): the completion-cancellation guarantee was too strong. Reviewer's scenario: cancel while _complete_multipart_upload runs in its thread. The cleanup waited only for the parts and then raced the abort against the completion. If the completion won, the object existed, the abort could fail with NoSuchUpload, and the docstring's "aborts the upload" did not hold. Classification: background completion predates the diff; the concurrent abort and the docstring claim were new.

Verified. The 185115d version fails the new test_copy_object_with_multipart_upload_cancelled_completion[False]. Repaired in 50fe7cb (maintainer's choice: wait for the completion, then decide). The completion is a task awaited through shield() here. _abort() (line 600) waits for it and returns without aborting if it succeeded. If it failed, the upload is aborted. The docstring now says "the upload is aborted unless it has completed".

[laughingman7743]

Independent review (relayed), static review only. Reviewer: Codex CLI 0.160.0, model gpt-6-astra, reasoning effort max, --sandbox read-only, session 01a10457-ec39-7471-aaef-d09f81f11567. Scope: base e49d576ce1e8cb60d74e1fde56de9c72ff3cd8e1, head 185115db8963dabe413396d3c71a883e0516acb6, detached snapshot plus the literal diff. The prompt left out the PR number, the description, the commit messages, and the self-review results. The reviewer ran no edits, builds, or tests and had no GitHub access, and the snapshot and the PR worktree were unchanged afterwards. Covered: async multipart copy and its _copy_file/_cp_file/_mv callers, sync completion/abort handling, tests; success, failures, cancellation propagation, exception retrieval, docs, Python 3.11-3.14. Verdict: FINDINGS (3, all P2).

Finding 3 (P2): the new test could fail with correct code. The semaphore proved that two parts had started, but a 200 ms sleep was all that kept them running until the cancellation. Under a scheduling delay, the parts could finish, part 3 could start, or the copy could complete before task.cancel(). The reviewer asked for explicit start/release synchronization and draining in finally.

Verified. Repaired in 50fe7cb. The parts block on a threading.Event that is released only after the cancellation, and an except cancels the copy if the setup fails. A finally always releases the threads. No timing assumption is left beyond asyncio's FIFO scheduling of the sleep(0) that lets the copy enter its cleanup. The test passed 20 consecutive runs with -W error, and it still fails on master (assert [] == ['abort']).

[laughingman7743]

Self-review of the repair (both perspectives), base a4cf60427f3027db084a51883a59cfea871475f4 (merge-base after rebasing on master), head 50fe7cb0fc2676fc02af422935ce37ae51b87f8b. Full pass, because the repair expands the contract (completion and repeated cancellation). The rebase brought #1052 (S3Path, _copy_file path parsing only) and #1047 (commit() abort handling). Neither touches this method or _abort_multipart_upload().

Round one (behavior): CLEAN.

• _abort() reads completion through its closure when it runs, so it sees the task assigned in try.
• On a part failure, completion is None and the cleanup aborts, as before.
• When the completion raises, shield() re-raises the error, completion.exception() is not None, and the cleanup aborts and re-raises.
• completion.exception() and gather(return_exceptions=True) retrieve every task exception, so no "never retrieved" warning is logged (-W error runs).
• _abort_multipart_upload() logs and swallows Exception, so the background cleanup does not raise after a repeated cancellation.
• _cleanup_tasks may hold tasks of several loops (fsspec's IO thread and the caller's loop). set.add and set.discard are atomic under the GIL.

Round two (claims and operations): the docstring claims hold within the stated limits, and the PR description is rewritten for the final behavior. A cancelled completion that succeeded leaves the object without its annotations; the description now states this. The cancellation can take as long as the slowest running part or the completion. The loop-shutdown and CreateMultipartUpload limits are listed. Evidence: just lint; 116 passed in test_s3_async.py with S3, on 50fe7cb; offline multipart tests with -W error on Python 3.13.1 and 3.11.11. Both the master and 185115d behavior fail the new tests, as described in the PR.

[laughingman7743]

Independent follow-up review (relayed), static review only. Reviewer: Codex CLI 0.160.0, model gpt-6-astra, reasoning effort max, --sandbox read-only, session 01a1046c-4083-7dc1-9a09-e5542d4bd974. This was a full review of base a4cf60427f3027db084a51883a59cfea871475f4 to head 50fe7cb0fc2676fc02af422935ce37ae51b87f8b, because the repair expanded the contract. The reviewer saw a detached snapshot and the literal diff, without the PR text, the commit messages, or prior findings. The snapshot and the PR worktree were unchanged afterwards. Verdict: FINDINGS (2 P2, 2 P3).

Finding 1 (P2): shutdown can abort before the worker threads finish. At event-loop shutdown, the part and completion tasks are cancelled directly, and shield() does not prevent that. The cleanup's gather (line 603) then sees them as finished and may abort while an UploadPartCopy thread is still running, so that part can be stored after the abort. Shutdown can also cancel the cleanup itself and skip the abort. Reviewer's classification: the premature abort is new, and the shutdown leak was pre-existing.

Deferred, documented. On master, shutdown during a multipart copy never aborts, so the whole upload and all its parts are left behind. With this PR, the cleanup that runs aborts the upload, and only a part still copying at that moment can remain. That is no worse than master. Tracking the worker threads apart from their asyncio wrappers would need a different executor design, which this cancellation fix does not call for. The PR description's Not handled list now describes both shutdown outcomes.

Finding 2 (P2): exceptions of a detached cleanup are never retrieved. After a repeated cancellation detaches the cleanup, a manually managed loop may shut down its default executor. The abort's to_thread then raises RuntimeError on submission, and nothing retrieves it.

Deferred, documented. This needs a loop that shuts down its default executor while a detached cleanup still runs. Even then, the error is not silent: asyncio logs it as "Task exception was never retrieved" with its traceback. The PR description lists it under Not handled.

[laughingman7743]

Independent follow-up review (relayed), static review only. Reviewer: Codex CLI 0.160.0, model gpt-6-astra, reasoning effort max, --sandbox read-only, session 01a1046c-4083-7dc1-9a09-e5542d4bd974. This was a full review of base a4cf60427f3027db084a51883a59cfea871475f4 to head 50fe7cb0fc2676fc02af422935ce37ae51b87f8b, because the repair expanded the contract. The reviewer saw a detached snapshot and the literal diff, without the PR text, the commit messages, or prior findings. The snapshot and the PR worktree were unchanged afterwards. Verdict: FINDINGS (2 P2, 2 P3).

Finding 3 (P3): the successful-completion case did not prove that the copy waits. The test released the completion right after cancel(). A version that returned at once could still see "complete" appended before the assertion, so [False] could pass without the fix.

Repaired in 28f0b6a. The completion stays held while the test cancels the copy, yields once so that the cancellation is processed, and asserts not task.done() and events == []. Only then is the completion released. Checked by swapping in older pyathena/ sources: with master (a4cf604), all 4 cancelled* cases fail; with 185115d, cancelled[2] and both cancelled_completion cases fail. At 28f0b6a the cases passed 20 consecutive runs with -W error, and the offline -k multipart_upload tests passed on Python 3.13.1 and 3.11.11.

[laughingman7743]

Independent follow-up review (relayed), static review only. Reviewer: Codex CLI 0.160.0, model gpt-6-astra, reasoning effort max, --sandbox read-only, session 01a1046c-4083-7dc1-9a09-e5542d4bd974. This was a full review of base a4cf60427f3027db084a51883a59cfea871475f4 to head 50fe7cb0fc2676fc02af422935ce37ae51b87f8b, because the repair expanded the contract. The reviewer saw a detached snapshot and the literal diff, without the PR text, the commit messages, or prior findings. The snapshot and the PR worktree were unchanged afterwards. Verdict: FINDINGS (2 P2, 2 P3).

Finding 4 (P3): the docstring overstated the guarantee. A cancellation during the unshielded CreateMultipartUpload await can let the worker create an upload whose ID is never kept, and no cleanup runs. That leak is pre-existing, but the blanket claim was new.

Repaired in 28f0b6a: the docstring now says "On a cancellation after the upload is created". The PR description already lists the creation window under Not handled.

Self-review of these repairs: they change one test and one docstring, and add no new behavior. Both rounds: CLEAN. The new assertions observe the documented wait, and the docstring and the PR text now state the same scope.

[laughingman7743]

Independent narrow follow-up (relayed), static review only. Reviewer: Codex CLI 0.160.0, model gpt-6-astra, reasoning effort max, --sandbox read-only, session 01a10477-5f3b-70f2-9166-62387628c16c. Scope: patch 50fe7cb0fc2676fc02af422935ce37ae51b87f8b..28f0b6af9d25b2f667a424146491de6dfbdb7d74 (same base a4cf6042): the docstring and test_copy_object_with_multipart_upload_cancelled_completion. The snapshot and the PR worktree were unchanged afterwards. Verdict: FINDINGS. The docstring now matches the protected phase, and the task.done() assertion catches an implementation that returns at once without aborting, for both parameters.

Finding (P2): the cleanup was not synchronized before the completion was released. A single sleep(0) lets the cancelled copy schedule _abort(), but the test can resume before _abort() runs. A broken cleanup could then pass both parameters if the completion won that race. Such a cleanup would skip the wait and abort unless the completion had already succeeded.

Repaired in cd0c4a4. The completion stays held for asyncio.sleep(0.1) after the cancellation, and the test asserts that the copy has neither returned nor aborted (not task.done(), events == []). A cleanup that aborts early or returns would do so during that window. Correct code cannot fail because of the window's length, since the completion is held until it is released. Checked by swapping in older pyathena/ sources: master fails all 4 cancelled* cases, and 185115d fails 3.

[laughingman7743]

Independent narrow follow-up of the repair (relayed), static review only: CLEAN. Reviewer: Codex CLI 0.160.0, model gpt-6-astra, reasoning effort max, --sandbox read-only, session 01a1047e-aa1e-78c1-a298-919900d318a0. Scope: patch 28f0b6af9d25b2f667a424146491de6dfbdb7d74..cd0c4a42ccad8f2ff64df135845208ce3f0e5241 (base a4cf6042), both cancellation tests and all their parameters. The snapshot was unchanged afterwards.

Reviewer's notes: both finally blocks release the blocked workers on every failing path. Once the workers have started, a scheduling delay cannot release them early. For both completion parameters, an early abort fails the event assertions, and returning while the completion is held fails assert not task.done(). The 100 ms wait is a bounded observation window, not deterministic synchronization with the cleanup. The existing 5 s start and abort waits can still time out under severe scheduling delays. These waits only bound a hang, and they fail loudly; not actionable.

[laughingman7743]

Independent narrow follow-up (relayed), static review only. Reviewer: Codex CLI 0.160.0, model gpt-6-astra, reasoning effort max, --sandbox read-only, session 01a10477-5f3b-70f2-9166-62387628c16c. Scope: patch 50fe7cb0fc2676fc02af422935ce37ae51b87f8b..28f0b6af9d25b2f667a424146491de6dfbdb7d74 (same base a4cf6042): the docstring and test_copy_object_with_multipart_upload_cancelled_completion. The snapshot and the PR worktree were unchanged afterwards. Verdict: FINDINGS. The docstring now matches the protected phase, and the task.done() assertion catches an implementation that returns at once without aborting, for both parameters.

Finding (P3): the hold depended on elapsed time. The mock ignored the result of release.wait(5). If the loop thread were delayed by 5 s, the completion would go ahead unreleased and fail the assertion with correct code.

Repaired in cd0c4a4, in both cancellation tests. The mocks call release.wait() with no timeout. Every test path reaches release.set() in a finally, so a failing test cannot leave a thread blocked. Self-review of the repair, both perspectives: CLEAN. The change is test-only; it passed 20 consecutive runs with -W error, and the offline multipart tests passed on Python 3.13.1 and 3.11.11. With the older sources the tests failed as described instead of hanging.