Repository navigation
Align AioS3FileSystem with S3FileSystem for cancellation and arguments (GH-1337) - #1395
Conversation
GH-1337) A cancelled coroutine of the filesystem now waits for its blocking call, whose stop scope keeps a write from sending further parts or committing, before it re-raises. The batch coroutines of fsspec cancel their per-object tasks and wait for them when they are cancelled. A recursive _put() no longer creates a directory, and thereby looks up or creates the bucket, for each local directory. _expand_path() runs S3FileSystem.expand_path(), so that its keyword arguments reach the listing, and S3FileSystem.glob() lists only the keys that start with the stem before the first wildcard, as fsspec's asynchronous glob does. _rm() takes batch_size and the deprecated max_workers as S3FileSystem.rm() does, and _cp_file() no longer sends the unused recursive, maxdepth, and on_error arguments to CopyObject. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The calls of the cancelled and failed batch tests were released by a timer or after a fixed wait, which a slow event loop could reach before the stop. They are now held until their stop scope stops them, and a short time longer, so that a batch that does not wait for them returns first. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
| batch._tasks.add(task) | ||
| task.add_done_callback(batch._tasks.discard) | ||
|
|
||
| async def run(self, batch: Awaitable[T]) -> T: |
There was a problem hiding this comment.
Self-review round one (implementation behavior): base 6589eb6efeeddbd83e9faf07b879f6bd4270afa1, head df64393031089f3bb95cee1f5758d24eed52632d (initial pass at 1a08eb1f).
Covered: S3Executor.run() and every _thread_pool.run caller in s3_async.py (lookups, _cat_file, _pipe_file/_put_file/_get_file including the transaction paths, the simple copies, _delete_paths, _move_pairs, _copy_pairs, _put_pairs); _S3BatchTasks with the overrides of _cat, _cat_ranges, _pipe, _put, _get, _copy and the tracking in the five per-object coroutines, against fsspec 2026.9.0's _run_coros_in_chunks (asyncio.wait without cleanup; on failure, cancel + gather of the pending tasks); _mv with asyncio.gather; _put_pairs extracted from S3FileSystem.put; _expand_path delegation and S3FileSystem.glob; _delete_batch_limit for both rms; _cp_file arguments; docs/aio.md; the new and changed tests.
FINDINGS, all handled:
test_batch_cancelled_waits_for_calls(shutdown cases) released its calls with a 0.2 s timer, and the cancelled/failure cases after a fixed 0.1 s wait: a slow event loop could release a call before its stop arrived, makingstoppedflaky. Fixed in df64393:_held_until_stoppedholds each call until its stop scope stops it, then 0.2 s longer, so a batch that does not wait returns first. The five affected mutants (no drain inrun, no drain/re-cancel/no tracking in_S3BatchTasks) are still killed.- The first draft replaced
asyncio.gatherin_mvwith a tracked variant. The mutation check showed it unnecessary:Task.cancel()cancels the awaited gather synchronously, so atasyncio.run()shutdown both cancellations reach each copy before it runs and collapse into one, and a single cancellation already waits for the copies. Dropped before the PR was opened;_mvstays covered by the parametrized test. - Behavior changes to state in the PR body:
_rm(batch_size=0)now raisesValueError(it was ignored), and_putpasses its keyword arguments to the localexpand_path()asS3FileSystem.put()does (fsspec's_putdid not).
Limits (amended I1, repeated cancellation): a cancellation that reaches fsspec's own failure path while it gathers the cancelled tasks, or a copy of _mv cancelled separately before _mv, can be cancelled twice, which ends its wait; the call then ends in its thread, already stopped.
| return {f.name: f for f in files} | ||
| return [f.name for f in files] | ||
|
|
||
| def glob(self, path, maxdepth=None, **kwargs): |
There was a problem hiding this comment.
Self-review round two (claims, compatibility, operations): base 6589eb6efeeddbd83e9faf07b879f6bd4270afa1, head df64393031089f3bb95cee1f5758d24eed52632d.
Claims checked:
- "fsspec's asynchronous glob passes the stem as
prefix=, the synchronousglob()does not": read in fsspec 2026.9.0 (asyn.py_glob,spec.pyglob); 2026.9.0 is the floor inpyproject.toml, so it holds for the whole supported range. Matches unchanged:test_glob_matches_fsspec(incl.dir/s*) passes, and a bucket-level wildcard still raises in_find("")as before. - "A cancelled write sends no further part and does not commit":
test_put_file_cancelled_stops_upload(requests are exactly Create, one UploadPart, Abort, and the task is not done before the abort); the small-object and download checks are those of Run the batch operations of S3FileSystem concurrently (GH-1322) #1344, whichrun()now reaches through its stop scope. - "fsspec left the tasks running after a cancellation / its fail-fast path did not wait for them":
fsspec/asyn.py:380and:391-396; seen failing withrun()or_S3BatchTasksnot draining (mutants). docs/aio.md: the table row, the cancellation paragraph, and the batch paragraph match the code;_mv()runs all copies at once (gather), the othersbatch_sizeat a time.- Defect 5:
test_request_pool_size[None-...]checksmax_workers=Noneresolves to the client pool.
Existing callers: run() keeps its signature and return/exception; only the moment a cancellation is raised moves. _rm, _cp_file, _put, _expand_path accept a superset of their old arguments, except _rm(batch_size=0) and other invalid values, now ValueError. S3FileSystem.put()/rm() behavior is unchanged by the helper extraction (their #1344 tests pass).
Operator: fewer ListObjectsV2 keys for a glob with a stem, also for the synchronous filesystem. A cancellation, including an asyncio.wait_for() timeout, now returns only when the in-flight request ends, e.g. after the whole body read of a large _cat_file(). That is the decided B6 behavior; added to the PR body's release notes with the other observable changes.
Evidence: local AWS filesystem suite and just benchmark test at 1a08eb1; df64393 changes tests only (rerun with their mutants). AWS CI not run yet (Draft).
FINDINGS: the PR body lacked the observable changes; corrected.
- A _put_file() or _pipe_file() in a transaction that is cancelled is discarded and dropped from the transaction once its write has ended. A write cancelled while its last part ran had nothing left to stop, and the transaction committed it. - _put() and _pipe() return the values of their per-object calls, as fsspec's do. - The pairing of put() strips the protocol of a single path-like destination as fsspec does, instead of iterating it. - The batch tests hold their stopped calls until the test releases them, and check that the batch is still waiting, instead of relying on a fixed delay. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
| return | ||
| await self._thread_pool.run(self._sync_fs.pipe_file, path, value, mode=mode, **kwargs) | ||
|
|
||
| async def _write_in_transaction(self, write: Callable[[list[S3File]], None]) -> None: |
There was a problem hiding this comment.
Independent review (relayed): Codex CLI codex exec -s read-only (model gpt-6-astra, session 01a1cd67-1179-77f3-9c2f-21e1154dfab1), static review only, base 6589eb6efeeddbd83e9faf07b879f6bd4270afa1, head df64393031089f3bb95cee1f5758d24eed52632d, detached snapshot (unchanged after the review). Prompt: the diff, the intended contract, repository conventions and docs/testing.md; no PR text, commits, or prior findings.
Covered: exact diff and conventions; executor cancellation races, scheduling, shutdown, reference ownership; batch overrides against fsspec 2026.9.0 (pairing, callbacks, return values, errors, concurrency); multipart writes, transactions, range reads, glob prefix, delete limits, copy arguments; repository callers, documentation, and all added or rewritten tests. Result: FINDINGS (4). Not a finding: the rewritten test_run_cancelled_while_retrying keeps its ordinary-versus-shielded retry-count checks; no match-changing glob regression found.
1. P2 — Cancelled transaction uploads can still commit (reviewed at s3_async.py:333). Inside fs.transaction, upload a file exactly one block long and cancel while its final UploadPart is running and close() is waiting. If that request succeeds, the worker returns without discarding the upload: autocommit=False skips completion, and no subsequent stop check triggers an abort. After the caller catches CancelledError, normal transaction exit commits that cancelled upload. Retain the transaction file and discard it when cancellation stops the write before its worker returns.
There was a problem hiding this comment.
Verified and fixed in 511ffe4 (pre-existing path, but within the guarantee of this PR). _put_file/_pipe_file in a transaction run through _write_in_transaction(): the write appends the file it opens to a list, and on a cancellation, once run() has waited for the write, _discard_transaction_files() (shielded) removes the file from transaction.files and discards it; a failed discard is logged. test_transaction_put_file_cancelled_is_discarded (2 parts, cancelled while part 2 runs) asserts Create, 2 × UploadPart, Abort and no completion at transaction exit; killed without the discard. docs/aio.md states it.
| self._sync_fs._put_pairs, lpath, rpath, recursive=recursive, maxdepth=maxdepth, **kwargs | ||
| ) | ||
| files = [(lp, rp) for lp, rp in pairs if not os.path.isdir(lp)] | ||
| return await _S3BatchTasks().run( |
There was a problem hiding this comment.
Independent review 2. P2 — _put and _pipe discard fsspec's return values (reviewed at s3_async.py:697, :720). A successful single-object _pipe or _put previously returned [None]; these overrides now return None. The generated synchronous wrappers inherit this change, breaking callers that iterate or count results. Return the result of _S3BatchTasks.run().
There was a problem hiding this comment.
Verified and fixed in 511ffe4: _put() and _pipe() return the result of _S3BatchTasks.run() (annotated list[Any]). test_put_pipe_return_values asserts [None] for both; killed by either mutant.
| return [] | ||
| dest_is_dir = isinstance(rpath, str) and (trailing_sep(rpath) or self.isdir(rpath)) | ||
| # A list of destinations or one path, also path-like, as in fsspec. | ||
| rpath = self._strip_protocol(rpath) |
There was a problem hiding this comment.
Independent review 3. P2 — _put now rejects path-like destinations (reviewed at s3.py:2773). _put(str(local_file), PurePosixPath("bucket/key")) previously worked because fsspec normalizes the destination through _strip_protocol(). The newly used _put_pairs helper treats every non-string destination as iterable, raising TypeError for PurePosixPath before uploading. Normalize scalar path-like destinations before distinguishing them from destination lists.
There was a problem hiding this comment.
Verified and fixed in 511ffe4: _put_pairs() strips the destination with _strip_protocol() as fsspec does, which maps a list and stringifies a single path-like. This also fixes S3FileSystem.put(), which uses the same pairing since #1344. Covered by test_put_pipe_return_values (PurePosixPath("bucket/b")); killed with the old normalization.
| assert requests == ["create_multipart_upload", "upload_part", "abort_multipart_upload"] | ||
|
|
||
| @staticmethod | ||
| def _held_until_stopped(state, release, result=None): |
There was a problem hiding this comment.
Independent review 4. P3 — Batch drain tests can pass without the drain (reviewed at test_s3_async.py:1845). _held_until_stopped assumes its 0.2-second sleep lets the test observe an early return. If the observing thread is delayed beyond that interval, workers finish before the assertions, so cancellation/failure tests can pass even when waiting for workers is removed. Use an explicit release gate and assert that the batch remains pending while cleanup is held.
There was a problem hiding this comment.
Verified and fixed in 511ffe4: _held_until_stopped(state, release) holds each stopped call until the test releases it. The cancel and failure tests wait until the calls are stopped, assert the batch is still pending, and only then release. The shutdown case releases from a helper thread once both calls are stopped and records whether asyncio.run() returned while they were held (expected not). A slow observer can now only delay the test, not let a missing drain pass. Mutants (no drain in run, no drain/re-cancel/tracking in _S3BatchTasks) still killed.
…H-1337) The follow-up review found that a second cancellation could end the wait for a transaction write before its file was registered, let the transaction commit the file while its discard was queued, or start a new wait behind the running write. _S3TransactionWrite now drops the file from the transaction when the write is cancelled, or when it is opened after that, and the file is discarded once by whichever of the write and the cancelled coroutine ends last; a write that is still running discards its own file, so the coroutine does not wait for it again. The batch cancellation test records, in the event loop, how many calls had ended when the batch task finished, instead of observing the return of asyncio.run() from the test thread, and cancels the calls before the batch in a new case, the order in which the shutdown of the event loop can reach them. The moves of the mv tests are compared by destination, since the copies of S3FileSystem.mv() run in parallel since GH-1322. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
| except Exception: | ||
| _logger.exception(f"Failed to discard the cancelled write to s3://{file.path}.") | ||
|
|
||
| def _opened(self, file: S3File) -> None: |
There was a problem hiding this comment.
Independent follow-up review 1 (relayed): Codex CLI codex exec -s read-only (model gpt-6-astra, session 01a1cd79-38c2-7f52-a491-62db0d583416), static review only, repair df64393031089f3bb95cee1f5758d24eed52632d..511ffe40a7afc3f782a7e2bb6c4f321c26ca3d6d (same base 6589eb6efeeddbd83e9faf07b879f6bd4270afa1), detached snapshot unchanged.
Covered: the repair diff, fsspec 2026.9.0 transactions and sync wrappers, fs.transaction.files, multipart and small-object writes, repeated cancellation, shared _put_pairs, the docs sentence, and the named tests against adjacent coverage. The return-value and path-like repairs appear complete; their test and the final-part cancellation test detect the original defects; the cancellation and failure drain tests now check the pending state before releasing workers. Result: FINDINGS (4).
1. P2 — Cleanup can finish before the worker registers its file (reviewed at s3_async.py:310). With multiple pool workers, hold _put_file() in callback.set_size() before open(), then cancel twice. The second cancellation interrupts run()'s drain, and cleanup processes an empty files list. Once released, the worker can open and buffer a small object: this path performs no multipart stop check. If the transaction remains open, its later exit commits the cancelled write.
There was a problem hiding this comment.
Verified and fixed in 1bf31e1. _S3TransactionWrite records the file through _opened(), which drops it from the transaction at once when the write has already been cancelled; the worker discards it when the write ends. test_transaction_put_file_cancelled_twice[before_open] holds the write in set_size(), cancels twice, releases it while the transaction is open, and checks no PutObject at exit; killed without the detach in _opened().
| if cancelled: | ||
| self.discard() | ||
|
|
||
| def cancel(self) -> bool: |
There was a problem hiding this comment.
Independent follow-up 1, finding 2. P2 — Transaction exit can overtake queued cleanup (reviewed at s3_async.py:324). After the write finishes, discard can queue behind another blocking filesystem call. A second cancellation interrupts run_shielded()'s wait, allowing the caller to catch cancellation and exit the transaction while the file remains in its deque. Transaction completion commits it; the eventual discard does nothing to an already committed file. This also contradicts the added documentation guarantee.
There was a problem hiding this comment.
Verified and fixed in 1bf31e1. cancel() drops the files from the transaction synchronously, in the loop thread, before the coroutine awaits the discard, so a second cancellation of that wait leaves nothing for the transaction to commit. test_transaction_put_file_cancelled_twice[discard_queued] (one thread, the discard queued behind a held call, second cancellation once the file is dropped) checks no PutObject at exit; killed without the detach in cancel().
| return | ||
| await self._thread_pool.run(self._sync_fs.pipe_file, path, value, mode=mode, **kwargs) | ||
|
|
||
| async def _write_in_transaction( |
There was a problem hiding this comment.
Independent follow-up 1, finding 3. P2 — A second cancellation can start another blocking wait (reviewed at s3_async.py:308). With max_workers=1 and the last UploadPart held, the second cancellation interrupts the write drain but enters this handler and awaits discard behind the still-running write. The task remains pending until that request ends or a third cancellation arrives, regressing the documented second-cancellation escape.
There was a problem hiding this comment.
Verified and fixed in 1bf31e1. Exactly one side discards: the worker sets _ended and reads _cancelled, and cancel() sets _cancelled and reads _ended, under one lock. A write still running when the coroutine gives up discards its own file at its end, so the coroutine awaits a discard only when the write has ended, and a second cancellation adds no wait. test_transaction_put_file_cancelled_twice[last_part] requires the task to end within 2 s while part 2 is held, then checks the abort and no completion; killed when the coroutine always awaits the discard and when the worker does not discard.
| } | ||
| # The calls that had ended when the batch finished, recorded in the | ||
| # event loop. | ||
| ended_when_done = [] |
There was a problem hiding this comment.
Independent follow-up 1, finding 4. P3 — The shutdown drain test still permits a delayed-observer false pass (reviewed at test_s3_async.py:1897). returned is set by the test thread after run_bounded() returns. Without draining, asyncio.run() can finish early while that thread remains unscheduled. The helper's 0.1-second wait then records False and releases the workers; they finish before the test thread resumes, satisfying every assertion.
There was a problem hiding this comment.
Verified and fixed in 1bf31e1. The batch task records state["ended"] in a done callback in the event loop, and the shutdown case releases the held calls only once the batch has finished or after 0.5 s, so an early return is seen whatever the test thread does. A mutant that re-cancels draining calls could still pass the shutdown case when asyncio.run() happened to step the batch before its calls (the cancellations then merge). The new calls_first case cancels the calls, waits until they are stopped, then cancels the batch, which kills that mutant deterministically. _mv is excluded from it: asyncio.gather re-cancels synchronously within the parent cancellation, which at shutdown merges with the first.
The second follow-up review found that a write cancelled twice could still join the transaction between fsspec's open() and the record of its file, and that a second cancellation arriving just as the write ended started a new wait for the discard. _S3TransactionWrite now opens the file with _open(autocommit=False) and adds it to the transaction itself, under the lock that the cancellation takes, so a cancelled write never joins it. After a repeated cancellation, the discard runs without being waited for. The batch cancellation test releases the held calls only once the end of the batch has been recorded in the event loop, and the transaction test checks that the queued discard was still held when the task ended, with a new case in which the write ends just before the second cancellation. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
| if cancelled: | ||
| self.discard() | ||
|
|
||
| def open(self, path: str, mode: str, **kwargs) -> S3File: |
There was a problem hiding this comment.
Independent follow-up review 2 (relayed): Codex CLI codex exec -s read-only (model gpt-6-astra, session 01a1cd94-7711-7f40-b9cc-89ace6ae275e), static review only, repair 511ffe40a7afc3f782a7e2bb6c4f321c26ca3d6d..1bf31e1690c55d4e4d89ff3d0c2ae2868138d2f7 (same base 6589eb6efeeddbd83e9faf07b879f6bd4270afa1), detached snapshot unchanged.
Covered: the repair diff, cancellation/worker handoffs, fsspec 2026.9.0 transaction registration and exit, discard ownership, locks, pool shutdown, normal writes, the changed tests, _moves() callers, and both documentation contracts. Result: FINDINGS (5).
1. P2 — Transaction exit can still commit a cancelled file (reviewed at s3_async.py:193). After two cancellations return while the worker is before open(), the worker can resume and append its file to fsspec's transaction, then lose execution before _opened(). The caller can now exit the transaction and commit that file, even as an empty object, before _detach() runs. The unconditional exclusion claim in docs/aio.md:335 remains false.
There was a problem hiding this comment.
Verified and fixed in 727a1dd. The write opens its file through _S3TransactionWrite.open(), which calls _open(..., autocommit=False) and adds the file to transaction.files itself, under the lock that cancel() takes, and only if the write has not been cancelled. There is no longer a moment when fsspec has registered the file and the cancellation cannot see it. An explicit autocommit=True still commits on close without joining, as with open(). test_transaction_put_file_cancelled_twice[before_open] kills the mutant that registers regardless of the cancellation. test_put_file_in_transaction_open_parameters and test_transaction_pipe_put_file_exceeding_max_parts now patch _open, with the stripped path and autocommit=False.
| await self._thread_pool.run(transaction_write.run) | ||
| except asyncio.CancelledError: | ||
| if transaction_write.cancel(): | ||
| if task is not None and task.cancelling() > cancellations + 1: |
There was a problem hiding this comment.
Independent follow-up 2, finding 2. P2 — A second cancellation can still start another blocking wait (reviewed at s3_async.py:409). While the first cancellation drains S3Executor.run(), the worker can set _ended=True immediately before the second cancellation interrupts that drain. cancel() then returns true and starts a fresh run_shielded() wait. With the sole pool thread occupied by another blocking call, this requires a third cancellation to escape. The new tests hold the worker on either side of this boundary and miss the race.
There was a problem hiding this comment.
Verified and fixed in 727a1dd. _write_in_transaction records Task.cancelling() when it starts. When run() raised after a repeated cancellation, it submits the discard without awaiting it, so a second cancellation adds no wait. Reproduced deterministically by the new write_ended case of test_transaction_put_file_cancelled_twice: one event-loop callback releases the write, waits until the only thread has taken a queued call (so the write has ended), and cancels again before the task resumes. It requires the task to end within 2 s while that call is held. Killed by the mutant that always awaits.
| self._thread_pool.submit(transaction_write.discard) | ||
| else: | ||
| # A repeated cancellation does not stop the discard. | ||
| await self._thread_pool.run_shielded(transaction_write.discard) |
There was a problem hiding this comment.
Independent follow-up 2, finding 3. P2 — Pool shutdown can prevent discard and replace cancellation with RuntimeError (reviewed at s3_async.py:411). If pool.shutdown(wait=False) occurs while the first cancellation drains a multipart write's last part, the worker finishes with _cancelled=False and performs no discard. The coroutine detaches the file, but submitting its discard fails because the pool rejects new work: no discard runs, the upload remains, and RuntimeError replaces CancelledError.
There was a problem hiding this comment.
Not taken: not reachable through the API. AioS3FileSystem never shuts its _thread_pool down. close() and the with blocks call S3ThreadPool.release(), after which the pool still accepts calls (s3_async.py close(); S3ThreadPool.release). Only direct code on the private attribute can call shutdown(), and then every coroutine of the filesystem fails with RuntimeError anyway.
| # While the write or the queued discard is still held. | ||
| with pytest.raises(asyncio.CancelledError): | ||
| await await_bounded(task, timeout=2) | ||
| # The queued discard was not what ended the task. |
There was a problem hiding this comment.
Independent follow-up 2, finding 4. P3 — The queued-discard test can pass without the repair (reviewed at test_s3_async.py:2106). With the old implementation, the transaction remains nonempty until the queued discard runs, so this wait can continue until the blocker's unchecked EVENT_TIMEOUT expires; discard then runs and satisfies the predicate, the second cancel() may target an already finished task, and all assertions pass. Verify that the blocker remains held when cancellation and transaction exit occur.
There was a problem hiding this comment.
Verified and fixed in 727a1dd. The case asserts not blocker_ended.is_set() right after the task has ended, so the task must end while the queued discard is still held behind the blocker. The wait for the file to leave the transaction is bounded at 2 s, well below the blocker's 10 s.
| # event loop. | ||
| ended_when_done = [] | ||
|
|
||
| def release_once_stopped(): |
There was a problem hiding this comment.
Independent follow-up 2, finding 5. P3 — The shutdown observation still has a release-before-observation race (reviewed at test_s3_async.py:1907). A prematurely completed batch sets task.done() before its scheduled done callback executes. The helper thread can observe that, release both workers, and let them increment ended before the callback records it, so the test records [2] despite the batch finishing early.
There was a problem hiding this comment.
Verified and fixed in 727a1dd. The helper thread now releases the held calls only once the done callback has recorded state["ended"] in the event loop (or after 0.5 s), not when task.done() turns true, so the record always precedes the release. The re-cancellation mutant is killed deterministically by the calls_first case regardless.
The third follow-up review found that the count of cancellations at the start of the write misses a cancellation still pending then, and that dropping a file from the transaction could fail while another write appended to it, since the search compares files with their __eq__(). A write that ends stopped now drops and discards its file itself, so that the wait of a single cancellation for the write includes the discard, and the coroutine never waits for a discard: it drops the file of a write that ended before it was stopped and discards it in a thread. The removal from the transaction is retried when another write changed it during the search. pyathena.util.retries_stopped() reports whether a stop scope has stopped. The transaction tests cover a cancellation after the write returned and hold a write at the moment its file would join the transaction. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
| except Exception: | ||
| _logger.exception(f"Failed to discard the cancelled write to s3://{file.path}.") | ||
|
|
||
| def _detach_all(self) -> None: |
There was a problem hiding this comment.
Independent follow-up review 3 (relayed): Codex CLI codex exec -s read-only (model gpt-6-astra, session 01a1cda4-9755-7a33-8eae-af417f0586fd), static review only, repair 1bf31e1690c55d4e4d89ff3d0c2ae2868138d2f7..727a1dd612a34752ed12d6199122dc0807ad0c4e (same base 6589eb6efeeddbd83e9faf07b879f6bd4270afa1), detached snapshot unchanged.
Covered: the repair diff; fsspec 2026.9.0 path handling, wb/xb, kwargs and autocommit; captured/current transactions, restarts and nesting; transaction deque concurrency; executor draining and discard; Python 3.11+ cancellation counts, timeout()/wait_for() unwinding; the changed tests' synchronization and regression sensitivity. Result: FINDINGS (3).
1. P2 — Per-write locking does not protect the shared transaction deque (reviewed at s3_async.py:198). While cancellation removes a non-first file, deque.remove() invokes fsspec's Python __eq__. Another worker can append to the same deque under its own lock; on Python 3.11, mutation during that comparison raises IndexError: deque mutated during iteration, which the ValueError suppression misses. If the cancelled write already ended, discard is never scheduled and its file remains registered.
There was a problem hiding this comment.
Verified and fixed in 6894974. _detach_all() retries the removal when it raises IndexError or RuntimeError, which is what happens when another write appends to the deque during the __eq__ search; ValueError (not there) ends it. Not covered by a deterministic test: the window is inside deque.remove().
| try: | ||
| await self._thread_pool.run(transaction_write.run) | ||
| except asyncio.CancelledError: | ||
| if transaction_write.cancel(): |
There was a problem hiding this comment.
Independent follow-up 3, finding 2. P2 — An undelivered cancellation at entry defeats repeated-cancellation detection (reviewed at s3_async.py:413). A caller can invoke current_task().cancel() immediately before awaiting the write. The baseline is then 1, although that cancellation is delivered inside S3Executor.run(). If the worker finishes during draining and a second cancellation interrupts that drain, 2 > 1 + 1 is false, and the handler starts another awaited discard, potentially behind an indefinitely blocked pool call.
There was a problem hiding this comment.
Verified and fixed in 6894974 by removing the count. A write that ends stopped (retries_stopped(), the new bool form of raise_if_stopped() in pyathena.util) drops and discards its file itself, in its own thread, so a single cancellation still waits for the discard through the drain of run(). The coroutine never waits for a discard: only for a write that ended before it was stopped does cancel() return true, and then the coroutine drops the file and submits the discard without awaiting it. test_transaction_put_file_cancelled_late covers twice before open, last part, write ended just before the second cancellation, and a single cancellation after the write returned (discard queued behind a held call). It requires the task to end within 2 s while the thread is held, and checks the abort and that nothing is committed. Mutants killed: ungated registration, no self-discard, no coroutine discard, no detach in cancel(), coroutine awaiting the discard. One survivor: no detach in the worker, which is redundant with cancel() and kept for a write stopped by another scope, as its comment says.
| assert requests[-1] == "abort_multipart_upload" | ||
|
|
||
| @pytest.mark.asyncio | ||
| async def test_transaction_write_cancelled_before_it_joins(self, tmp_path): |
There was a problem hiding this comment.
Independent follow-up 3, finding 3. P3 — The registration-gap repair still lacks a behavioral regression test (reviewed at test_s3_async.py:2133). The before_open case releases the worker and waits for pool shutdown before exiting the transaction, so the previous append-then-_opened() implementation finishes detaching before commit and passes this case.
There was a problem hiding this comment.
Verified and fixed in 6894974: test_transaction_write_cancelled_before_it_joins replaces transaction.files with a deque whose append() signals and blocks. After two cancellations it releases the write and exits the transaction as soon as the file joins (or after 1 s). The previous open-then-detach ordering joins and is committed at exit. The current gated registration never appends. Killed by the ungated-registration mutant.
The fourth follow-up review found that a transaction write stopped by an enclosing stop scope, rather than by a cancellation of its coroutine, discarded its file and still returned. It now raises concurrent.futures.CancelledError after the discard, which reaches the coroutine as asyncio.CancelledError. The test of the registration gap keeps the transaction open until the file joins it or the write has ended, instead of for at most one second. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
| if stopped: | ||
| self.discard() | ||
| if stopped: | ||
| raise CancelledError("Stopped before the transaction could commit the write.") |
There was a problem hiding this comment.
Independent follow-up review 4 (relayed): Codex CLI codex exec -s read-only (model gpt-6-astra, session 01a1cdb5-1a2f-7113-9bb1-f3b34be3cfba), static review only, repair 727a1dd612a34752ed12d6199122dc0807ad0c4e..6894974af56a995593b9a392daa65f7beebe44e9 (same base 6589eb6efeeddbd83e9faf07b879f6bd4270afa1), detached snapshot unchanged.
Covered: the repair diff; executor stop/return ordering; worker/coroutine discard ownership; first and repeated cancellation; enclosing stop scopes; deque-removal retries; uncancelled transaction writes; raise_if_stopped() callers; the changed tests; fsspec registration/commit; docs/aio.md. Result: FINDINGS (2).
1. P2 — Stopped writes can report success after discarding their data (reviewed at s3_async.py:162). An enclosing stop_retries_on(event) can stop without cancelling the coroutine. A small transaction _pipe_file() then buffers and closes normally, but the finalizer discards it and returns None; both the await and the transaction exit succeed although the object is omitted.
There was a problem hiding this comment.
Verified and fixed in 440b652: _S3TransactionWrite.run() raises concurrent.futures.CancelledError after discarding the file of a stopped write, also when the write itself succeeded. asyncio.wrap_future() delivers it to the coroutine as asyncio.CancelledError, as for any stopped writer call. test_transaction_write_stopped_by_enclosing_scope covers it (a small _pipe_file() under a set stop_retries_on() event raises, leaves no file in the transaction, sends nothing) and is killed when the raise is removed. Within PyAthena no stop scope encloses these coroutines; the case needs caller code that enters one.
| await await_bounded(task, timeout=2) | ||
| release.set() | ||
| # Exits once the file has joined, or once the write has ended. | ||
| await wait_until(lambda: joined.is_set() or write_ended.is_set()) |
There was a problem hiding this comment.
Independent follow-up 4, finding 2. P3 — The registration-gap test can pass with the regression restored (reviewed at test_s3_async.py:2201). If the released worker is delayed beyond one second, the test exits the transaction before the write opens its file; with fsspec open()-then-detach restored, the registration then bypasses HeldFiles, and the stopped small write makes no S3 request, so both assertions can pass. Keep the transaction active until registration is trapped or worker completion is confirmed.
There was a problem hiding this comment.
Verified and fixed in 440b652: the test keeps the transaction open until the file joins it or the write has ended. The end is signalled by a wrapper of _S3TransactionWrite.run that calls the original. wait_until fails after EVENT_TIMEOUT if neither happens. Killed by the ungated-registration mutant.
| self._cancelled = False | ||
| self._ended = False | ||
|
|
||
| def run(self) -> None: |
There was a problem hiding this comment.
Independent follow-up review 5 (relayed): Codex CLI codex exec -s read-only (model gpt-6-astra, session 01a1cdcb-bf92-7940-bc05-a7760341c5c5), static review only, repair 6894974af56a995593b9a392daa65f7beebe44e9..440b65267648754a4501d7b61395b16a1b56d935 (same base 6589eb6efeeddbd83e9faf07b879f6bd4270afa1), detached snapshot unchanged.
Covered: _S3TransactionWrite.run() stop detection, detachment, discard, and exception propagation; S3Executor.run() returned/stop ordering, future and waiter completion, exception retrieval, and a cancellation racing with the worker's raise; _write_in_transaction cleanup ownership, repeated cancellation, transaction exclusion, and the docs/aio.md claims; both changed tests.
CLEAN. The worker finishes its discard before raising; the executor's finally records completion before the error reaches the waiter, so its cancellation handler neither starts another wait nor sets the stop; the transaction handler sees that the cleanup was done; a concurrent cancellation preserves the cleanup, and a repeated one can end the wait without preventing the worker's completion. By inspection, removing the raise fails the enclosing-scope test, and restoring the premature registration makes the registration test observe the join; a delayed worker now times out instead of passing.
WHAT
Step 4b of the #1332 design:
AioS3FileSystembehaves asS3FileSystemfor the five defects of #1337.S3Executor.run()keeps the future of its call, tracked before it is scheduled. On a cancellation of the waiting task it stops the call's scope, cancels a call that has not started, and waits for a running one before it re-raises, asmap_async()does (B6). A second cancellation stops the wait._put_file(),_pipe_file(), or_get_file()now ends its transfer and raises after it, instead of committing in the background._cat(),_cat_ranges(),_pipe(),_put(),_get(), and_copy()as tasks, and a cancellation of the batch coroutine left them running (fsspec/asyn.py:380). These coroutines are now overridden to run fsspec's through_S3BatchTasks. Each per-object coroutine of the filesystem adds its task to the batch of its context, and on a failure or cancellation the batch cancels the tasks and waits for them. A task that is already being cancelled, asasyncio.run()cancels every task at shutdown, is not cancelled again, so its own wait is not stopped._mv()keepsasyncio.gather, whose cancellation already cancels the copies and waits for them._put_file()or_pipe_file()is kept out of the transaction and discarded once its write has ended. A write cancelled while its last part ran has nothing left to stop, and the transaction would otherwise commit it._S3TransactionWriteopens the file with_open(autocommit=False)and adds it totransaction.filesitself, under the lock that the cancellation takes, unless the write has been cancelled or stopped, so a cancelled write never joins the transaction, even if it opens its file later.run(); the coroutine never waits for a discard: it drops the file of a write that ended before it was stopped and discards it in a thread. A write stopped by an enclosing stop scope raisesCancelledErrorafter the discard instead of returning.pyathena.util.retries_stopped()reports whether a stop scope has stopped;raise_if_stopped()uses it._put()no longer calls_makedirs()for each local directory, which sent HeadBucket, or CreateBucket when allowed, or raisedPermissionErrorbefore any upload. The pairing ofS3FileSystem.put()moves toS3FileSystem._put_pairs(), and_put()hands only the file pairs to fsspec's_put()._expand_path(..., **kwargs)runsS3FileSystem.expand_path(), as_find()runsfind(), sorefresh=andprefix=reach the listing instead of raisingTypeError.glob()lists the whole directory. To keep those requests,S3FileSystem.glob()now passes the stem asprefix=tofind()as the asynchronous glob does, so the synchronous filesystem sends the narrower listing too. The matches are unchanged (test_glob_matches_fsspec)._rm(batch_size=)is the number of DeleteObjects requests sent at once,max_workersby default and with-1.max_workersis a deprecated alias that warns, and both together raiseTypeError, through the helper thatS3FileSystem.rm()now shares (_delete_batch_limit()). Before,batch_sizewas ignored andmax_workerswas sent to DeleteObjects._cp_file()takes the unusedrecursive,maxdepth, andon_errorofS3FileSystem.cp_file(), which were sent to CopyObject.cat(paths, batch_size=N)ofS3FileSystemwas fixed by Run the batch operations of S3FileSystem concurrently (GH-1322) #1344.S3FileSystem(max_workers=None)was fixed by Share one request pool per filesystem among its files and operations (GH-1332) #1359;test_request_pool_sizechecks thatNoneresolves to the client's connection pool.docs/aio.mddescribes the cancellation of the calls, of the transaction writes, and of the batch coroutines,_put(), and_rm(batch_size=).Changes that callers observe (release notes):
AioS3FileSystem, including one cancelled by the timeout ofasyncio.wait_for(), raises once the call that it runs in a thread has ended, instead of at once. A request already sent runs to its end first; a_cat_file()of a large object reads its whole body, which it gets in one GetObject request. A second cancellation raises at once._put_file()or_pipe_file()is left out of the commit and discarded; it was committed when it had been cancelled during its last part._rm(batch_size=)is validated: a value that is neither positive,-1, nor None raisesValueError(it was ignored)._rm(max_workers=)emits aDeprecationWarning._put()passes its keyword arguments to the localexpand_path(), asS3FileSystem.put()does since Run the batch operations of S3FileSystem concurrently (GH-1322) #1344. Both accept a single path-like destination again, which the pairing of Run the batch operations of S3FileSystem concurrently (GH-1322) #1344 iterated.S3FileSystem.glob(), and the expansion ofget(),copy(),cat(), andrm(), list only the keys that start with the stem before the first wildcard. The matches are unchanged.Not changed (follow-up candidates):
ValueError, apply toS3FileSystemonly.AsyncFileSystem._expand_path()takes no**kwargs, the synchronousglob()passes noprefix=, and_run_coros_in_chunks()leaves its tasks running when it is cancelled.WHY
Closes #1337. Part of #1332 (step 4b) and #1327 (defects 14 to 18).
TEST
Tested at 440b652; the local full filesystem suite was rerun after each review repair (1bf31e1: 1625 passed; 727a1dd: 1626; 6894974: 1627; 6 skipped each).
just lint,just docs lint.uv run --with 'botocore[crt]' --env-file .env pytest -n 4 tests/pyathena/filesystem/→ 1628 passed, 6 skipped.just benchmark test→ 340 passed, 1 skipped.test_run_cancelled_waits_for_call(one and two cancellations);test_put_file_cancelled_stops_upload(no further UploadPart, no completion, abort before the cancellation is raised);test_batch_cancelled_waits_for_calls(7 coroutines × cancelling the batch, the calls first and then the batch, or theasyncio.run()shutdown);test_batch_failure_waits_for_calls(_pipe,_put,_get);test_transaction_put_file_cancelled_is_discarded;test_transaction_put_file_cancelled_late(cancelled twice before the file is opened, during the last part, or just after the write ended; once after the write returned), each requiring the task to end while the thread is held;test_transaction_write_cancelled_before_it_joins(a heldappend()oftransaction.filescatches a file that joins after the cancellations);test_transaction_write_stopped_by_enclosing_scope;test_put_pipe_return_values(return values, path-like destination);test_put_recursive_creates_no_directories,test_expand_path_kwargs,test_rm_batch_size,test_cp_file_fsspec_arguments,test_glob_lists_stem_prefix.test_run_cancelled_while_retryingnow releases the call before awaiting the task, which waits for it.mvtests ofS3FileSystemcompare the copies by destination: they run in parallel since Run the batch operations of S3FileSystem concurrently (GH-1322) #1344, and a full run failed once on their order.run, also after a failure; cancellation of a queued call; drain, no re-cancellation, and tracking of batch tasks;_putdirectories;_expand_pathdelegation; glob prefix;_rmbatch size and alias;_cp_filearguments; the return values of_put/_pipe; the path-like destination; the gated registration, self-discard, coroutine discard without a wait, detach on cancellation, and the raise of_S3TransactionWrite). 21 are killed by their tests. The survivor removes the worker's own detach, whichcancel()repeats for a cancellation; it is kept for a write stopped by another scope. Reverting_mvtoasyncio.gatherwas not killed, so that change was dropped.🤖 Generated with Claude Code