Skip to content

Align AioS3FileSystem with S3FileSystem for cancellation and arguments (GH-1337) - #1395

Merged
laughingman7743 merged 7 commits into
masterfrom
fix/1337-aio-alignment
Nov 12, 2026
Merged

laughingman7743 merged 7 commits into
masterfrom
fix/1337-aio-alignment

Conversation

@laughingman7743

@laughingman7743 laughingman7743 commented Nov 12, 2026 •

Copy link
Copy Markdown
Member

WHAT

Step 4b of the #1332 design: AioS3FileSystem behaves as S3FileSystem for the five defects of #1337.

  1. Cancellation.
    • 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, as map_async() does (B6). A second cancellation stops the wait.
    • Since Run the batch operations of S3FileSystem concurrently (GH-1322) #1344, a write in a stopped scope sends no further part, does not commit, and is aborted, and a download reads and writes no further block. So a cancelled _put_file(), _pipe_file(), or _get_file() now ends its transfer and raises after it, instead of committing in the background.
    • fsspec runs the per-object coroutines of _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, as asyncio.run() cancels every task at shutdown, is not cancelled again, so its own wait is not stopped.
    • _mv() keeps asyncio.gather, whose cancellation already cancels the copies and waits for them.
    • In a transaction, a cancelled _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.
      • _S3TransactionWrite opens the file with _open(autocommit=False) and adds it to transaction.files itself, 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.
      • A write that ends stopped drops and discards its own file in its thread, so a single cancellation waits for the discard through the drain of 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 raises CancelledError after the discard instead of returning.
      • pyathena.util.retries_stopped() reports whether a stop scope has stopped; raise_if_stopped() uses it.
  2. Recursive _put() no longer calls _makedirs() for each local directory, which sent HeadBucket, or CreateBucket when allowed, or raised PermissionError before any upload. The pairing of S3FileSystem.put() moves to S3FileSystem._put_pairs(), and _put() hands only the file pairs to fsspec's _put().
  3. _expand_path(..., **kwargs) runs S3FileSystem.expand_path(), as _find() runs find(), so refresh= and prefix= reach the listing instead of raising TypeError.
    • fsspec's asynchronous glob, which the expansion used, lists only the keys that start with the stem before the first wildcard; its synchronous glob() lists the whole directory. To keep those requests, S3FileSystem.glob() now passes the stem as prefix= to find() as the asynchronous glob does, so the synchronous filesystem sends the narrower listing too. The matches are unchanged (test_glob_matches_fsspec).
  4. Arguments.
    • _rm(batch_size=) is the number of DeleteObjects requests sent at once, max_workers by default and with -1. max_workers is a deprecated alias that warns, and both together raise TypeError, through the helper that S3FileSystem.rm() now shares (_delete_batch_limit()). Before, batch_size was ignored and max_workers was sent to DeleteObjects.
    • _cp_file() takes the unused recursive, maxdepth, and on_error of S3FileSystem.cp_file(), which were sent to CopyObject.
    • cat(paths, batch_size=N) of S3FileSystem was fixed by Run the batch operations of S3FileSystem concurrently (GH-1322) #1344.
  5. S3FileSystem(max_workers=None) was fixed by Share one request pool per filesystem among its files and operations (GH-1332) #1359; test_request_pool_size checks that None resolves to the client's connection pool.

docs/aio.md describes 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):

  • A cancelled coroutine of AioS3FileSystem, including one cancelled by the timeout of asyncio.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.
  • In a transaction, a cancelled _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 raises ValueError (it was ignored). _rm(max_workers=) emits a DeprecationWarning.
  • _put() passes its keyword arguments to the local expand_path(), as S3FileSystem.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 of get(), copy(), cat(), and rm(), list only the keys that start with the stem before the first wildcard. The matches are unchanged.

Not changed (follow-up candidates):

  • The conflict checks of Run the batch operations of S3FileSystem concurrently (GH-1322) #1344, such as duplicate destinations raising ValueError, apply to S3FileSystem only.
  • The upstream reports for fsspec are not filed yet: AsyncFileSystem._expand_path() takes no **kwargs, the synchronous glob() passes no prefix=, 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.
  • Local AWS: 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.
  • New tests:
    • 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 the asyncio.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 held append() of transaction.files catches 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_retrying now releases the call before awaiting the task, which waits for it.
  • The mv tests of S3FileSystem compare 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.
  • Mutation check: 22 mutants, one per change (drain of run, also after a failure; cancellation of a queued call; drain, no re-cancellation, and tracking of batch tasks; _put directories; _expand_path delegation; glob prefix; _rm batch size and alias; _cp_file arguments; 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, which cancel() repeats for a cancellation; it is kept for a write stopped by another scope. Reverting _mv to asyncio.gather was not killed, so that change was dropped.
  • AWS CI: runs when the PR is marked Ready.

🤖 Generated with Claude Code

laughingman7743 and others added 2 commits November 12, 2026 10:30
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:

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

  1. 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, making stopped flaky. Fixed in df64393: _held_until_stopped holds 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 in run, no drain/re-cancel/no tracking in _S3BatchTasks) are still killed.
  2. The first draft replaced asyncio.gather in _mv with a tracked variant. The mutation check showed it unnecessary: Task.cancel() cancels the awaited gather synchronously, so at asyncio.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; _mv stays covered by the parametrized test.
  3. Behavior changes to state in the PR body: _rm(batch_size=0) now raises ValueError (it was ignored), and _put passes its keyword arguments to the local expand_path() as S3FileSystem.put() does (fsspec's _put did 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.

Comment thread pyathena/filesystem/s3.py
return {f.name: f for f in files}
return [f.name for f in files]

def glob(self, path, maxdepth=None, **kwargs):

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Self-review round two (claims, compatibility, operations): base 6589eb6efeeddbd83e9faf07b879f6bd4270afa1, head df64393031089f3bb95cee1f5758d24eed52632d.

Claims checked:

  • "fsspec's asynchronous glob passes the stem as prefix=, the synchronous glob() does not": read in fsspec 2026.9.0 (asyn.py _glob, spec.py glob); 2026.9.0 is the floor in pyproject.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, which run() 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:380 and :391-396; seen failing with run() or _S3BatchTasks not 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 others batch_size at a time.
  • Defect 5: test_request_pool_size[None-...] checks max_workers=None resolves 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>
Comment thread pyathena/filesystem/s3_async.py Outdated
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:

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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(

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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().

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread pyathena/filesystem/s3.py
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)

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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):

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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>
Comment thread pyathena/filesystem/s3_async.py Outdated
except Exception:
_logger.exception(f"Failed to discard the cancelled write to s3://{file.path}.")

def _opened(self, file: S3File) -> None:

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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().

Comment thread pyathena/filesystem/s3_async.py Outdated
return
await self._thread_pool.run(self._sync_fs.pipe_file, path, value, mode=mode, **kwargs)

async def _write_in_transaction(

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 = []

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread pyathena/filesystem/s3_async.py Outdated
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:

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread pyathena/filesystem/s3_async.py Outdated
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)

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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():

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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():

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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):

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.")

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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())

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Align AioS3FileSystem with S3FileSystem for cancellation, recursive put, expand_path, and arguments

1 participant