From 4f90bd42aba97d41c455421bcb38c892592ee547 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 7 Nov 2026 02:39:44 +0900 Subject: [PATCH 1/8] Add a fan-out primitive to the S3 executors (GH-1332) S3Executor.map() and map_async() run a callable for each item in the executor and return the outcomes in input order. The calls check the stop of the batch before they start, and the first failure that counts is recorded in the thread of the failed call, so that no call starts after the failure that stops the batch. An interrupt or a cancellation of the waiting thread or task sets the stop, cancels the calls that have not started, waits for the running ones, and is re-raised; a second one stops that wait. submit_to() now cancels the submission when the caller's future is cancelled, so that a scope drops a waiting call at once. The range reads of both file classes, recursive chmod, rm, and clear_multipart_uploads of S3FileSystem, and the DeleteObjects requests and annotation copies of AioS3FileSystem use it. A failed or interrupted read, and a cancelled asyncio rm or multipart copy, now raise once their running requests have ended; an interrupt of rm or clear_multipart_uploads cancels the requests that have not started. Co-Authored-By: Claude Opus 5.5 --- docs/aio.md | 8 + docs/filesystem.md | 14 +- pyathena/filesystem/s3.py | 175 +++----- pyathena/filesystem/s3_async.py | 150 +++---- pyathena/filesystem/s3_executor.py | 303 +++++++++++++- tests/pyathena/filesystem/test_s3.py | 104 ++++- tests/pyathena/filesystem/test_s3_async.py | 185 ++++++--- tests/pyathena/filesystem/test_s3_executor.py | 389 +++++++++++++++++- 8 files changed, 1055 insertions(+), 273 deletions(-) diff --git a/docs/aio.md b/docs/aio.md index 5bf72b50b..f64e7d677 100644 --- a/docs/aio.md +++ b/docs/aio.md @@ -334,6 +334,10 @@ ends. A request that has already been sent still runs to its end in its thread. requests that run to their end after a cancellation keep their retries: the commit of a file, the CompleteMultipartUpload request of a multipart copy, and the aborts of multipart uploads. +The range reads of a file, the DeleteObjects requests of `_rm()`, and the annotation +copies of a multipart copy are waited for before the cancellation is raised, also when +the shutdown of the event loop at the end of `asyncio.run()` cancels the coroutine; a +second cancellation stops that wait. ### Executor strategy @@ -343,6 +347,10 @@ and runs them in a pool of `max_workers` threads; in a forked child process, the creates new threads. `S3ThreadPoolExecutor` and `S3AioThreadPool` are aliases of `S3ThreadPool`. `S3AioExecutor` is kept for compatibility: it is an `S3ThreadPool` that ignores its `loop` argument, with a `DeprecationWarning` if one is given. +`S3Executor.map()` and `S3Executor.map_async()` run a callable for each item in the +executor and return the outcomes in the order of the items. With `stop_on_failure`, +the default, the first failure cancels the calls that have not started, stops the +retries of the running ones, and is raised once they have ended. The file handles of `AioS3FileSystem.open()` run their multipart uploads and parallel range reads in the request pool of the filesystem, at most the `max_workers` of the diff --git a/docs/filesystem.md b/docs/filesystem.md index 15f58f551..2cb3663a4 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -116,10 +116,14 @@ In a forked child process, create a new filesystem. A file that the parent opene `RuntimeError` there from its reads, writes, and flushes, `commit()`, and `discard()`, and closing it only marks it closed, without uploading or aborting anything. -An interrupt, such as Ctrl+C, while a file reads ranges in parallel cancels the range -reads that have not started. Those that are running finish, and their data is released -once they end. Closing the file waits for its running requests; an interrupt of that -wait stops it, and the requests still finish. An interrupt while a +An interrupt, such as Ctrl+C, or a failure while a file reads ranges in parallel +cancels the range reads that have not started, and the running ones make no further +attempt. The interrupt or the error is raised once they have ended, and their data is +released. A second interrupt stops that wait, and the reads still finish. `rm` and +`clear_multipart_uploads` send every request also after one fails; an interrupt cancels +those that have not started and is raised once the running ones have ended. Closing +the file waits for its running requests; an interrupt of that wait stops it, and the +requests still finish. An interrupt while a request is being submitted can let one more request than `max_workers` run at once until the request pool is idle. These guarantees hold for one interrupt at a time. If another interrupt arrives while the cleanup of an interrupt or an error runs, the call still does not wait forever, but the cleanup may stop early: @@ -299,7 +303,7 @@ fs.chmod("s3://YOUR_S3_BUCKET/path/to/", "private", recursive=True) A recursive `chmod` updates the keys in parallel. Once an update fails, the updates that have not started are cancelled, the running ones make no further attempt, and the -error is raised. +error is raised once they have ended. Note that `setxattr` rewrites the object by copying it onto itself (S3 does not allow updating the metadata of an existing object in place), which updates its last-modified diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 480f377b9..f8c309e60 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -3,15 +3,13 @@ from __future__ import annotations import contextlib +import functools import hashlib import logging import mimetypes import os.path -import threading import time -import traceback from collections.abc import Callable, Mapping -from concurrent.futures import FIRST_EXCEPTION, Future, as_completed, wait from copy import deepcopy from dataclasses import dataclass, field from datetime import datetime @@ -48,7 +46,7 @@ from pyathena.filesystem.s3_path import S3Path from pyathena.filesystem.s3_path_pairing import S3PathPairing from pyathena.filesystem.s3_writer import _S3WriteSession -from pyathena.util import RetryConfig, override, stop_retries_on +from pyathena.util import RetryConfig, override _logger = logging.getLogger(__name__) @@ -1173,10 +1171,16 @@ def _delete_objects(self, paths: list[str], max_workers: int | None = None, **kw return max_workers = max_workers if max_workers else self.max_workers - with self._create_executor(max_workers=max_workers) as executor: - fs = [executor.submit(self._delete_batch, batch, **params) for batch in batches] - # The executor has waited for every request, also after a failure. - self._raise_delete_errors([f.exception() or f.result() for f in fs]) + executor = self._create_executor(max_workers=max_workers) + try: + # Every request runs also after another failed, and an interrupt + # cancels those that have not started. + results = executor.map( + functools.partial(self._delete_batch, **params), batches, stop_on_failure=False + ) + finally: + executor.shutdown(wait=False) + self._raise_delete_errors(results.outcomes) @staticmethod def _delete_batches(paths: list[str], **kwargs) -> tuple[list[S3DeleteBatch], dict[str, Any]]: @@ -2257,31 +2261,16 @@ def chmod(self, path: str, acl: str, recursive: bool = False, **kwargs) -> None: raise ValueError(f"ACL not in {self.core.OBJECT_ACLS}.") if recursive: paths = self.find(path, withdirs=False) - # Set on failure, so that the running requests make no further attempt. - stop = threading.Event() - with self._create_executor(max_workers=self.max_workers) as executor: - futures: list[Future[None]] = [] - try: - with stop_retries_on(stop): - for p in paths: - # Tracked before it is scheduled, so that an - # interrupt while it is being scheduled still lets - # it be cancelled. - future: Future[None] = Future() - futures.append(future) - executor.submit_to( - future, self.chmod, p, acl, recursive=False, **kwargs - ) - for future in as_completed(futures): - future.result() - except BaseException: - # No more ACL is applied after one failed. The updates that - # have not started are cancelled before the stop, which - # would let a thread that it frees take one of them. - for future in futures: - future.cancel() - stop.set() - raise + executor = self._create_executor(max_workers=self.max_workers) + try: + executor.map( + lambda p: self.chmod(p, acl, recursive=False, **kwargs), + paths, + ) + finally: + # The map has waited for the updates it started, unless a + # second interrupt stopped it. + executor.shutdown(wait=False) if s3_path.key: # A key prefix is not an object itself; only the objects # below it have ACLs. @@ -2392,30 +2381,37 @@ def clear_multipart_uploads(self, path: str) -> None: uploads = self.list_multipart_uploads(path) if not uploads: return - error: Exception | None = None - with self._create_executor(max_workers=self.max_workers) as executor: - futures = [ - executor.submit( - self.core.abort_multipart_upload, - upload, - ) - for upload in uploads - ] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - cause = e.__cause__ - if ( - isinstance(e, FileNotFoundError) - and isinstance(cause, botocore.exceptions.ClientError) - and S3ClientError(cause).code == "NoSuchUpload" - ): - continue - if error is None: - error = e - if error is not None: - raise error + executor = self._create_executor(max_workers=self.max_workers) + try: + # Every abort runs also after another failed, and an interrupt + # cancels those that have not started. + results = executor.map( + self.core.abort_multipart_upload, + uploads, + stop_on_failure=False, + counts=self._is_not_no_such_upload, + ) + finally: + executor.shutdown(wait=False) + if results.failure is not None: + raise results.failure + + @staticmethod + def _is_not_no_such_upload(error: BaseException) -> bool: + """Return whether an error of an abort is other than an upload that no longer exists. + + Args: + error: The error of ``abort_multipart_upload``. + + Returns: + False for the ``FileNotFoundError`` of a ``NoSuchUpload`` error. + """ + cause = error.__cause__ + return not ( + isinstance(error, FileNotFoundError) + and isinstance(cause, botocore.exceptions.ClientError) + and S3ClientError(cause).code == "NoSuchUpload" + ) def created(self, path: str) -> datetime: """Return the creation time of the path. @@ -3408,9 +3404,11 @@ def _fetch_range(self, start: int, end: int) -> bytes: may request a range that is empty or reaches past the end of the object. S3 would answer the former with the whole object and a range starting past the end with an ``InvalidRange`` error. A range larger - than the block size is read with parallel requests; when one fails, - those that have not started are cancelled, the running ones make no - further attempt, and its error is raised. + than the block size is read with parallel requests (see + :meth:`S3Executor.map() `); + when one fails, those that have not started are cancelled, the running + ones make no further attempt, and its error is raised once they have + ended. Args: start: The offset of the first byte to read. @@ -3428,54 +3426,17 @@ def _fetch_range(self, start: int, end: int) -> bytes: path = S3Path(self.bucket, self.key, self.version_id) request_kwargs = self._get_request_kwargs("get_object") if len(ranges) > 1: - # Set on failure, so that the running reads make no further attempt. - stop = threading.Event() - futures: list[Future[bytes]] = [] - done: set[Future[bytes]] = set() - try: - with stop_retries_on(stop): - for r in ranges: - # Tracked before it is scheduled, so that an interrupt - # while it is being scheduled still lets it be - # cancelled; not by a variable, which would keep it in - # this frame, which the traceback of an error keeps. - futures.append(Future()) - self._executor.submit_to( - futures[-1], self.fs.core.get_object, path, r, **request_kwargs - ) - done = wait(futures, return_when=FIRST_EXCEPTION).done - failed = next( - (f for f in done if not f.cancelled() and f.exception() is not None), None + # Not kept by a variable of this frame, which the traceback of an + # interrupt keeps (F5). + return b"".join( + cast( + list[bytes], + self._executor.map( + functools.partial(self.fs.core.get_object, path, **request_kwargs), + ranges, + ).outcomes, ) - if failed is not None: - failed.result() - return b"".join(f.result() for f in futures) - except BaseException as e: - # A failed read, or an interrupt at any point, including after - # the wait: the reads that have not started are cancelled, and - # the futures of the others are dropped, so that this frame, - # which the traceback of the error keeps, does not keep their - # data, also if another interrupt stops this cleanup. The - # frames of the wait, which keep the futures, are cleared, and - # clear_frames() skips this frame, which is executing. - try: - # Popped, not iterated, so that no variable of this frame - # keeps a future. - while futures: - futures.pop().cancel() - finally: - # Also if an interrupt stops the cancellation, whose error - # keeps this one as its context; the frames first, as - # they keep the futures of the wait. - try: - traceback.clear_frames(e.__traceback__) - finally: - done.clear() - futures.clear() - # After the cancellation, which a thread that it frees - # could otherwise precede by taking a queued read. - stop.set() - raise + ) return self.fs.core.get_object(path, ranges[0], **request_kwargs) @staticmethod diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 8a20ce546..e0c1ea3e6 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -15,8 +15,6 @@ import logging import mimetypes import os -import threading -import traceback import warnings from concurrent.futures import Future, wait from io import BytesIO @@ -43,7 +41,6 @@ from pyathena.filesystem.s3_path import S3Path from pyathena.filesystem.s3_path_pairing import S3PathPairing from pyathena.filesystem.s3_writer import _S3WriteSession, _S3WriteState -from pyathena.util import stop_retries_on if TYPE_CHECKING: from datetime import datetime @@ -350,8 +347,9 @@ async def _rm( """Delete objects with DeleteObjects requests. See :meth:`S3FileSystem.rm`. The DeleteObjects requests run in - parallel with ``asyncio.gather``, at most ``max_workers`` at once, in - the request pool of the internal ``S3FileSystem``. + parallel, at most ``max_workers`` at once, in the request pool of the + internal ``S3FileSystem``. A cancellation cancels the requests that + have not started and waits for the others. Args: path: S3 path (s3://bucket/key) or list of paths to delete. @@ -371,7 +369,7 @@ async def _rm( await self._delete_objects(paths, **kwargs) async def _delete_objects(self, paths: list[str], **kwargs) -> None: - """Delete objects with DeleteObjects requests run with ``asyncio.gather``. + """Delete objects with parallel DeleteObjects requests. See :meth:`S3FileSystem._delete_objects`. @@ -389,11 +387,15 @@ async def _delete_objects(self, paths: list[str], **kwargs) -> None: # so that a large deletion does not queue ahead of the requests of # the files that share the request pool. requests = self._sync_fs._create_executor(max_workers=self._sync_fs.max_workers) - results = await asyncio.gather( - *[requests.run(self._sync_fs._delete_batch, batch, **params) for batch in batches], - return_exceptions=True, - ) - self._sync_fs._raise_delete_errors(results) + try: + results = await requests.map_async( + functools.partial(self._sync_fs._delete_batch, **params), + batches, + stop_on_failure=False, + ) + finally: + requests.shutdown(wait=False) + self._sync_fs._raise_delete_errors(results.outcomes) async def _mv(self, path1, path2, recursive=False, maxdepth=None, **kwargs) -> None: """Move files from one S3 location to another. @@ -612,11 +614,12 @@ async def _copy_object_with_multipart_upload( bounds the requests of all its copies and files together, with at most ``max_workers`` of them in flight; each one is sent once an earlier one has finished, and the upload is completed once all of them have. - The annotation copies then run in parallel as asyncio tasks, at most - ``max_workers`` at once, in the same pool. On a cancellation, the - requests that have not started, including the completion, are - cancelled, the running ones are waited for in a thread of the - filesystem, the upload is aborted if it was created and has not + The annotation copies then run in parallel, at most ``max_workers`` + at once, in the same pool, and a failure or a cancellation stops them + as :meth:`~pyathena.filesystem.s3_executor.S3Executor.map_async` does. + On a cancellation, the requests that have not started, including the + completion, are cancelled, the running ones are waited for in a thread + of the filesystem, the upload is aborted if it was created and has not completed, and the cancellation is re-raised. On a failure or cancellation, the running part copies make no further attempt. This also holds when the shutdown of the event loop, such as at the end of @@ -722,62 +725,20 @@ def _abort() -> None: # At most max_workers copies at once in the request pool. requests = self._sync_fs._create_executor(max_workers=max_workers) - # Set when an annotation copy fails, so that the running ones make no - # further attempt, and the waiting ones do not start. - stop = threading.Event() - # The index of the copy that failed first, whose error is raised - # instead of those of the copies that its failure stopped. Recorded - # with the stop, in the thread of the failed copy. - first_failure: int | None = None - failure_lock = threading.Lock() - - def _copy_annotation_in_thread(i: int, name: str) -> None: - """Copy an annotation in a scope that a failed copy stops. - - Args: - i: The index of the annotation. - name: The name of the annotation. - """ - nonlocal first_failure - if stop.is_set(): - # Waited for a thread of the request pool until a copy failed. - return - try: - with stop_retries_on(stop): - self.core.copy_object_annotation( - name, - plan.source, - plan.destination, - completed.version_id, - completed.etag, - **kwargs, - ) - except Exception: - with failure_lock: - if not stop.is_set(): - first_failure = i - # Set in this thread, before it takes a waiting copy. - stop.set() - raise - - async def _copy_annotation(i: int, name: str) -> None: - try: - await requests.run(_copy_annotation_in_thread, i, name) - except Exception: - # The running copies make no further attempt, and the waiting - # ones do not start, as in S3FileSystem. - stop.set() - raise - - results = await asyncio.gather( - *[_copy_annotation(i, name) for i, name in enumerate(plan.annotations)], - return_exceptions=True, - ) - if first_failure is not None: - raise cast(BaseException, results[first_failure]) - for result in results: - if isinstance(result, BaseException): - raise result + try: + await requests.map_async( + lambda name: self.core.copy_object_annotation( + name, + plan.source, + plan.destination, + completed.version_id, + completed.etag, + **kwargs, + ), + plan.annotations, + ) + finally: + requests.shutdown(wait=False) async def _find( self, @@ -1770,8 +1731,9 @@ async def _fetch_range(self, start: int, end: int) -> bytes: The range is clamped to the size of the object, and split into ranges of the block size, read at most ``max_workers`` at once in the request pool of the filesystem. When a read fails, the reads that have - not finished are cancelled, which stops the retries of the running - ones, and its error is raised, as :meth:`S3File._fetch_range` does. + not started are cancelled, the running ones make no further attempt, + and its error is raised once they have ended, as + :meth:`S3File._fetch_range` does; a cancellation does the same. Args: start: The offset of the first byte to read. @@ -1793,32 +1755,16 @@ async def _fetch_range(self, start: int, end: int) -> bytes: ) path = S3Path(self.bucket, self.key, self.version_id) request_kwargs = self._get_request_kwargs("get_object") - reads = [ - asyncio.ensure_future( - self._requests.run(self.fs.core.get_object, path, r, **request_kwargs) + # Not kept by a variable of this frame, which the traceback of a + # cancellation keeps (F5). + return b"".join( + cast( + list[bytes], + ( + await self._requests.map_async( + functools.partial(self.fs.core.get_object, path, **request_kwargs), + ranges, + ) + ).outcomes, ) - for r in ranges - ] - try: - done = (await asyncio.wait(reads, return_when=asyncio.FIRST_EXCEPTION))[0] - except BaseException as e: - # The reads are cancelled with the read of the range, and the - # frames of the wait, which keep their tasks, are cleared; - # clear_frames() skips this frame, which is executing. Cancelling a - # failed read also keeps asyncio from logging its error. - traceback.clear_frames(e.__traceback__) - while reads: - reads.pop().cancel() - raise - # A comprehension, whose variable does not keep a read in this frame. - failures = [read for read in done if not read.cancelled() and read.exception() is not None] - if failures: - # The reads that have not finished are cancelled, and the tasks of - # the others are dropped, so that this frame, which the traceback - # of the error keeps, does not keep their data. Cancelling a failed - # read also keeps asyncio from logging its error. - while reads: - reads.pop().cancel() - done.clear() - failures[0].result() - return b"".join(read.result() for read in reads) + ) diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index a5f43b9ef..e43a6ea65 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -18,15 +18,17 @@ import weakref from abc import ABCMeta, abstractmethod from collections import deque -from collections.abc import Callable -from concurrent.futures import Future +from collections.abc import Callable, Iterable +from concurrent.futures import FIRST_EXCEPTION, CancelledError, Future, wait from concurrent.futures.thread import ThreadPoolExecutor +from dataclasses import dataclass from multiprocessing import cpu_count from typing import Any, ClassVar, Generic, TypeVar, cast from pyathena.util import override, stop_retries_on T = TypeVar("T") +ItemT = TypeVar("ItemT") class _S3WorkItem(Generic[T]): @@ -115,6 +117,195 @@ def _take(self) -> Future[T]: return future +@dataclass(frozen=True) +class S3MapResult(Generic[T]): + """The outcomes of :meth:`S3Executor.map` and :meth:`S3Executor.map_async`. + + Attributes: + outcomes: The result of each item, or the exception that its call + raised, in the order of the items. + failure: The first failure that counted, in the order the calls + ended, or None. + """ + + outcomes: list[T | BaseException] + failure: BaseException | None + + +class _S3Batch(Generic[ItemT, T]): + """The calls of one map of :class:`S3Executor`, and the state that they share. + + A call checks the stop of the batch before it starts, and records the + first failure that counts in its own thread, before that thread can take + another call, so that no call starts after the failure that stops the + batch, and the failure recorded is the first in the order the calls + ended. + + Args: + fn: The callable to run for each item. + stop_on_failure: Whether a failure that counts stops the batch. + counts: Whether a failure counts; every failure if None. + """ + + def __init__( + self, + fn: Callable[[ItemT], T], + stop_on_failure: bool, + counts: Callable[[BaseException], bool] | None, + ) -> None: + """Initialize the batch. + + Args: + fn: The callable to run for each item. + stop_on_failure: Whether a failure that counts stops the batch. + counts: Whether a failure counts; every failure if None. + """ + self._fn = fn + self._stop_on_failure = stop_on_failure + self._counts = counts + # Set when the batch stops, so that the calls that have not started + # make no request and the running ones make no further attempt. + self.stop = threading.Event() + self._failure: BaseException | None = None + self._lock = threading.Lock() + # The futures of the calls, each tracked before it is scheduled, so + # that an interrupt while it is being scheduled still lets it be + # cancelled. + self.futures: list[Future[T]] = [] + + def submit(self, executor: S3Executor, items: Iterable[ItemT]) -> None: + """Submit a call for each item. + + Args: + executor: The executor that runs the calls. + items: The items. + """ + with stop_retries_on(self.stop): + for item in items: + self.futures.append(Future()) + executor.submit_to(self.futures[-1], self._call, item) + + def wait(self) -> None: + """Wait until every call has ended or a failure has stopped the batch.""" + pending = set(self.futures) + while pending and not self._stopped_by_failure(): + pending = wait(pending, return_when=FIRST_EXCEPTION).not_done + + async def wait_async(self) -> None: + """Await until every call has ended or a failure has stopped the batch.""" + pending = {self._wrap(future) for future in self.futures} + while pending and not self._stopped_by_failure(): + pending = (await asyncio.wait(pending, return_when=asyncio.FIRST_EXCEPTION))[1] + + def result(self) -> S3MapResult[T]: + """Return the outcomes of the calls, which have all ended. + + Returns: + The outcomes. + + Raises: + BaseException: The failure that stopped the batch. + """ + if self._stopped_by_failure(): + raise cast(BaseException, self._failure) + return S3MapResult( + [ + f.result() if f.exception() is None else cast(Any, f.exception()) + for f in self.futures + ], + self._failure, + ) + + def drain(self) -> None: + """Stop the batch, cancel the calls that have not started, and wait for the others. + + The states are also checked every second, since an interrupt of a + cancellation can skip the notification that ``wait()`` relies on. + """ + started = self._stop() + try: + while not all(future.done() for future in started): + wait(started, timeout=1.0) + finally: + started.clear() + + async def drain_async(self) -> None: + """Stop the batch, cancel the calls that have not started, and await the others.""" + started = self._stop() + try: + waiters = [self._wrap(future) for future in started] + while not all(future.done() for future in started): + await asyncio.wait(waiters, timeout=1.0) + finally: + started.clear() + + def _stop(self) -> list[Future[T]]: + """Stop the batch and cancel the calls that have not started. + + The stop is set first, so that a further interrupt during the + cancellation still stops the retries of the running calls; a thread + that the cancellation frees finds the stop before it starts a call. + The futures are popped, so that no frame keeps them. + + Returns: + The futures of the calls that have started. + """ + self.stop.set() + started = [] + while self.futures: + future = self.futures.pop() + if not future.cancel(): + started.append(future) + return started + + def _stopped_by_failure(self) -> bool: + """Return whether a failure has stopped the batch. + + Returns: + Whether a failure that counts has stopped the batch. + """ + return self._stop_on_failure and self._failure is not None + + def _call(self, item: ItemT) -> T: + """Run the callable for an item unless the batch has stopped. + + Args: + item: The item. + + Returns: + The return value of the callable. + + Raises: + CancelledError: If the batch has stopped before the call started. + """ + if self.stop.is_set(): + raise CancelledError + try: + return self._fn(item) + except BaseException as e: + if not isinstance(e, Exception) or self._counts is None or self._counts(e): + with self._lock: + if self._failure is None: + self._failure = e + if self._stop_on_failure: + self.stop.set() + raise + + @staticmethod + def _wrap(future: Future[T]) -> asyncio.Future[T]: + """Wrap a future for asyncio, retrieving its outcome so that asyncio does not log it. + + Args: + future: The future. + + Returns: + The asyncio future. + """ + waiter = asyncio.wrap_future(future) + waiter.add_done_callback(lambda f: f.cancelled() or f.exception()) + return waiter + + class S3Executor(metaclass=ABCMeta): """Abstract executor for parallel S3 operations. @@ -255,7 +446,113 @@ def submit_to(self, future: Future[T], fn: Callable[..., T], *args: Any, **kwarg **kwargs: Keyword arguments passed to the callable. """ item = _S3WorkItem(future, fn, args, kwargs) - self.submit(item.run).add_done_callback(item.settle) + submitted = self.submit(item.run) + submitted.add_done_callback(item.settle) + # A cancellation of the future also cancels the submission, so that a + # scope drops a waiting call at once instead of letting it pass + # through the pool without running. + future.add_done_callback(lambda f: f.cancelled() and submitted.cancel()) + + def map( + self, + fn: Callable[[ItemT], T], + items: Iterable[ItemT], + *, + stop_on_failure: bool = True, + counts: Callable[[BaseException], bool] | None = None, + ) -> S3MapResult[T]: + """Run a callable for each item in the executor and wait for the calls. + + Every item is submitted at once, and the executor applies its own + limit, such as the ``max_workers`` of a scope. Each call runs in a + scope of :func:`~pyathena.util.stop_retries_on` that the stop of the + batch ends. + + With ``stop_on_failure``, the first failure that counts stops the + batch: the calls that have not started are cancelled, the running + ones make no further attempt, and the failure is raised once they + have ended. An interrupt of the waiting thread, such as + ``KeyboardInterrupt``, does the same whatever ``stop_on_failure`` is, + and is raised once the running calls have ended. A further interrupt + stops that wait; the running calls then end in the executor. + + Args: + fn: The callable to run for each item. + items: The items. + stop_on_failure: Whether a failure that counts stops the batch. + counts: Whether a failure counts, by default every one. A failure + that does not count neither stops the batch nor becomes + :attr:`S3MapResult.failure`. + + Returns: + The outcomes of the calls. + + Raises: + BaseException: The failure that stopped the batch. + RuntimeError: If the executor has been shut down. + """ + batch = _S3Batch(fn, stop_on_failure, counts) + try: + batch.submit(self, items) + batch.wait() + return batch.result() + except BaseException as e: + try: + batch.drain() + finally: + # The frames below this one keep the futures, whose outcomes, + # such as the bodies of reads, the traceback of the error would + # otherwise keep; clear_frames() skips this frame, which is + # executing. + try: + traceback.clear_frames(e.__traceback__) + finally: + batch.futures.clear() + raise + + async def map_async( + self, + fn: Callable[[ItemT], T], + items: Iterable[ItemT], + *, + stop_on_failure: bool = True, + counts: Callable[[BaseException], bool] | None = None, + ) -> S3MapResult[T]: + """Run a callable for each item in the executor and await the calls. + + As :meth:`map`, with a cancellation of the waiting task in place of + an interrupt. The task waits for the running calls also when the + event loop shuts down with ``asyncio.run()``, which cancels it and + runs the loop until it ends; a loop that is closed without + cancelling it leaves the running calls to end in the executor. + + Args: + fn: The callable to run for each item. + items: The items. + stop_on_failure: Whether a failure that counts stops the batch. + counts: Whether a failure counts, by default every one. + + Returns: + The outcomes of the calls. + + Raises: + BaseException: The failure that stopped the batch. + RuntimeError: If the executor has been shut down. + """ + batch = _S3Batch(fn, stop_on_failure, counts) + try: + batch.submit(self, items) + await batch.wait_async() + return batch.result() + except BaseException as e: + try: + await batch.drain_async() + finally: + try: + traceback.clear_frames(e.__traceback__) + finally: + batch.futures.clear() + raise def __enter__(self) -> S3Executor: return self diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 3808fcc11..98b87169f 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -41,7 +41,7 @@ from pyathena.filesystem.s3 import CompressedBuffer, S3File, S3FileSystem from pyathena.filesystem.s3_core import S3Core, S3DeleteBatch, S3MultipartCopyPlan from pyathena.filesystem.s3_errors import S3ClientError -from pyathena.filesystem.s3_executor import S3ThreadPool, _S3PoolScope +from pyathena.filesystem.s3_executor import S3Executor, S3ThreadPool, _S3PoolScope from pyathena.filesystem.s3_object import S3MultipartUpload, S3Object, S3ObjectType, S3StorageClass from pyathena.filesystem.s3_path import S3Path from pyathena.filesystem.s3_writer import S3MultipartStateError, _S3WriteSession @@ -61,6 +61,7 @@ poll_until, retried, run_in_forked_child, + shutdown_bounded, sigint_when, stub_multipart_copy, ) @@ -565,6 +566,77 @@ def call(method, **request): "Failed to delete objects: b2/b (AccessDenied: Access Denied)" ] + @pytest.mark.skipif( + threading.current_thread() is not threading.main_thread(), + reason="SIGINT interrupts the main thread.", + ) + @pytest.mark.parametrize("operation", ["rm", "clear_multipart_uploads"]) + def test_interrupted_operation_cancels_queued_requests(self, operation): + # GH-1332 (B3): an interrupt of rm() or clear_multipart_uploads(), + # which run every request also after one failed, cancels the requests + # that have not started. The queued requests used to be sent after + # the interrupt. + fs = self._make_fs() + fs.max_workers = 1 + started = threading.Event() + requested = [] + + def call(method, **request): + requested.append(request["Bucket"] if operation == "rm" else request["UploadId"]) + started.set() + time.sleep(0.2) + return {} + + fs._call.side_effect = call + if operation == "rm": + run = functools.partial(fs.rm, ["s3://b1/a", "s3://b2/b"]) + else: + fs.list_multipart_uploads = mock.MagicMock( + return_value=[ + S3MultipartUpload({"Bucket": "bucket", "Key": "key", "UploadId": u}) + for u in ("u1", "u2") + ] + ) + run = functools.partial(fs.clear_multipart_uploads, "s3://bucket/") + try: + with sigint_when(started.is_set), pytest.raises(KeyboardInterrupt): + run() + finally: + shutdown_bounded(fs._request_pool) + assert len(requested) == 1 + + def test_clear_multipart_uploads_raises_first_failure(self): + # GH-1332: the error raised is the first abort failure, other than an + # upload that no longer exists, in the order the aborts ended. + fs = self._make_fs() + fs.max_workers = 3 + ended = {u: threading.Event() for u in ("gone", "late", "early")} + + def abort_multipart_upload(upload): + # "gone" ends first, then "early", then "late". + name = upload.upload_id + if name == "early": + assert ended["gone"].wait(EVENT_TIMEOUT) + if name == "late": + assert ended["early"].wait(EVENT_TIMEOUT) + ended[name].set() + if name == "gone": + error = botocore.exceptions.ClientError( + {"Error": {"Code": "NoSuchUpload", "Message": name}}, "AbortMultipartUpload" + ) + raise FileNotFoundError(name) from error + raise PermissionError(name) + + fs.core.abort_multipart_upload = abort_multipart_upload + fs.list_multipart_uploads = mock.MagicMock( + return_value=[ + S3MultipartUpload({"Bucket": "bucket", "Key": "key", "UploadId": u}) + for u in ("gone", "late", "early") + ] + ) + with pytest.raises(PermissionError, match="early"): + fs.clear_multipart_uploads("s3://bucket/") + def test_rm_requests_invalidate_shared_parent(self): # The request threads invalidate the shared parent "bucket/dir" at # once. DirCache.pop() reads before it deletes, so the second delete @@ -3586,6 +3658,30 @@ def get_object(path, range_, **kwargs): assert time.monotonic() - started < EVENT_TIMEOUT assert operation.calls == 1 + def test_failed_read_waits_for_running_reads(self): + # GH-1332 (B6): a failed parallel read raises once its running reads + # have ended, which used to run on after it raised. + fs = self._make_fs() + fs.info = mock.MagicMock(return_value=self._file_object("key")) + fs.info.return_value.size = 32 + ended = [] + both_started = threading.Barrier(2, timeout=EVENT_TIMEOUT) + + def get_object(path, range_, **kwargs): + both_started.wait() + if range_[0] == 0: + raise PermissionError("Access Denied") + time.sleep(0.2) + ended.append(range_) + return bytes(range_[1] - range_[0]) + + fs.core.get_object = get_object + f = fs.open("s3://bucket/key", "rb", cache_type="none", block_size=16, max_workers=2) + with pytest.raises(PermissionError, match="Access Denied"): + f.read() + assert ended == [(16, 32)] + f.close() + def test_read_interrupted_submission_stops_running_reads(self): # GH-1203: an interrupt while the reads are being submitted stops the # running ones, whose backoff wait would take at least EVENT_TIMEOUT @@ -3698,7 +3794,7 @@ def get_object(path, range_, **kwargs): fs.core.get_object = get_object f = fs.open("s3://bucket/key", "rb", cache_type="none", block_size=16, max_workers=4) try: - with mock.patch("pyathena.filesystem.s3.wait", interrupted_wait): + with mock.patch("pyathena.filesystem.s3_executor.wait", interrupted_wait): try: interrupt(f._fetch_range, 0, 64) error = interrupt.error @@ -3714,7 +3810,7 @@ def get_object(path, range_, **kwargs): f.close() fs._request_pool.shutdown() - assert interrupt_each_opcode(S3File._fetch_range, case) > 0 + assert interrupt_each_opcode(S3Executor.map, case) > 0 @pytest.mark.skipif( threading.current_thread() is not threading.main_thread(), @@ -3791,7 +3887,7 @@ def interrupted_wait(futures, **kwargs): f = fs.open("s3://bucket/key", "rb", cache_type="none", block_size=16, max_workers=4) with ( - mock.patch("pyathena.filesystem.s3.wait", side_effect=interrupted_wait), + mock.patch("pyathena.filesystem.s3_executor.wait", interrupted_wait), pytest.raises(KeyboardInterrupt) as info, ): f.read() diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index f6f5eaee8..8c6d9565d 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -45,6 +45,7 @@ await_bounded, cancel_when_started, forking, + interrupt_each_opcode, mock_connection, overlapping_calls, retried, @@ -1247,6 +1248,54 @@ def copy_object_annotation(name, *args, **kwargs): ) assert operation.calls == 1 + @pytest.mark.asyncio + async def test_copy_object_with_multipart_upload_cancelled_annotation_copy(self): + # GH-1332 (B6): a cancelled copy raises once its running annotation + # copy has ended, which used to run on after the cancellation. + fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) + sync_fs = fs._sync_fs + plan = S3MultipartCopyPlan( + source=S3Path("bucket", "src"), + destination=S3Path("bucket", "dst"), + size=1, + ranges=((0, 1),), + annotations=("a1",), + ) + sync_fs.core.plan_multipart_copy = mock.MagicMock(return_value=plan) + sync_fs.core.create_multipart_upload = mock.MagicMock( + return_value=S3MultipartUpload({"Bucket": "bucket", "Key": "dst", "UploadId": "id"}) + ) + sync_fs.core.upload_part_copy = mock.MagicMock( + side_effect=lambda **kw: SimpleNamespace(etag='"e"', part_number=kw["part_number"]) + ) + sync_fs.core.complete_multipart_upload = mock.MagicMock( + return_value=SimpleNamespace(version_id=None, etag='"e"') + ) + started = threading.Event() + release = threading.Event() + ended = [] + + def copy_object_annotation(name, *args, **kwargs): + started.set() + release.wait(EVENT_TIMEOUT) + ended.append(name) + + sync_fs.core.copy_object_annotation = mock.MagicMock(side_effect=copy_object_annotation) + task = asyncio.ensure_future( + fs._copy_object_with_multipart_upload(plan.source, plan.destination) + ) + try: + assert await asyncio.to_thread(started.wait, EVENT_TIMEOUT) + task.cancel() + await asyncio.sleep(0.1) + assert not task.done() + release.set() + with pytest.raises(asyncio.CancelledError): + await await_bounded(task) + assert ended == ["a1"] + finally: + release.set() + @pytest.mark.asyncio async def test_copy_object_with_multipart_upload_failed_annotation_skips_queued_copies(self): # GH-1203: the annotation copies that wait for a thread of the filesystem @@ -1642,47 +1691,37 @@ def call(method, **request): ] @pytest.mark.asyncio - async def test_rm_cancel_invalidates_cache_after_each_request(self): - # The request threads keep running after _rm() is cancelled, so each - # request invalidates the cache of its objects when it finishes. + async def test_rm_cancel_waits_for_running_requests(self): + # GH-1332 (B6): a cancelled _rm() raises once its running requests + # have ended, each of which invalidates the cache of its objects. The + # requests used to run on after the cancellation was raised. fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) started = {"b1": threading.Event(), "b2": threading.Event()} - release = {"b1": threading.Event(), "b2": threading.Event()} + release = threading.Event() def call(method, **request): started[request["Bucket"]].set() - release[request["Bucket"]].wait(10) + release.wait(EVENT_TIMEOUT) return {} - async def wait_invalidated(path): - for _ in range(100): - if path not in fs.dircache: - return - await asyncio.sleep(0.01) - fs._sync_fs._call = fs._sync_fs._core.call = call task = asyncio.create_task(fs._rm(["s3://b1/a", "s3://b2/b"])) - for event in started.values(): - await asyncio.to_thread(event.wait, 10) - # Cancel every other task, as asyncio.run() does at shutdown, so the - # request tasks are cancelled directly, not only through _rm(). - for other in asyncio.all_tasks(): - if other is not asyncio.current_task(): - other.cancel() - with pytest.raises(asyncio.CancelledError): - await task - # Cached while the requests still run, e.g. by a concurrent info(). - fs.dircache["b1/a"] = [] - fs.dircache["b2/b"] = [] - - release["b1"].set() - await wait_invalidated("b1/a") - assert "b1/a" not in fs.dircache - assert "b2/b" in fs.dircache - - release["b2"].set() - await wait_invalidated("b2/b") - assert "b2/b" not in fs.dircache + try: + for event in started.values(): + assert await asyncio.to_thread(event.wait, EVENT_TIMEOUT) + # Cached while the requests run, e.g. by a concurrent info(). + fs.dircache["b1/a"] = [] + fs.dircache["b2/b"] = [] + task.cancel() + await asyncio.sleep(0.1) + assert not task.done() + release.set() + with pytest.raises(asyncio.CancelledError): + await await_bounded(task) + assert "b1/a" not in fs.dircache + assert "b2/b" not in fs.dircache + finally: + release.set() @pytest.mark.asyncio async def test_rm_maxdepth(self): @@ -2737,9 +2776,62 @@ async def flush_and_exit(): fs.core.complete_multipart_upload.assert_called_once() fs.core.abort_multipart_upload.assert_not_called() + def test_close_after_interrupted_cancel(self): + # GH-1332 (F2, I1): an interrupt at any point of the cancellation of a + # request that a close waits for, including one that skips the + # callbacks of its future, does not keep the close waiting. + def case(interrupt): + fs = AioS3FileSystem( + connection=mock_connection(), max_workers=1, skip_instance_cache=True + ) + info = S3Object( + init={"Key": "key", "ETag": '"etag"'}, + type=S3ObjectType.S3_OBJECT_TYPE_FILE, + bucket="bucket", + key="key", + ) + info.size = 1 + fs._sync_fs.info = mock.MagicMock(return_value=info) + pool = fs._sync_fs._request_pool + release = threading.Event() + loop = asyncio.new_event_loop() + thread = threading.Thread(target=loop.run_forever, daemon=True) + thread.start() + try: + # Holds the only thread of the pool, longer than the waits + # below, so that the pool cannot reach the cancelled request. + _S3PoolScope(pool, 1).submit(release.wait, 3 * EVENT_TIMEOUT) + f = asyncio.run_coroutine_threadsafe( + fs.open_async("s3://bucket/key", "rb"), loop + ).result(EVENT_TIMEOUT) + future = f._requests.submit(sum, [1, 2]) + waiting = threading.Event() + add_done_callback = future.add_done_callback + + def watch(fn): + add_done_callback(fn) + waiting.set() + + future.add_done_callback = watch # type: ignore[method-assign] + closed = asyncio.run_coroutine_threadsafe(f.close(), loop) + assert waiting.wait(EVENT_TIMEOUT) + interrupt(future.cancel) + if not future.cancelled(): + # Interrupted before it changed the state. + assert future.cancel() + closed.result(EVENT_TIMEOUT) + finally: + release.set() + loop.call_soon_threadsafe(loop.stop) + thread.join(EVENT_TIMEOUT) + loop.close() + pool.shutdown() + + assert interrupt_each_opcode(Future.cancel, case) > 0 + @pytest.mark.asyncio async def test_close_waits_for_cancelled_read(self, fs): - # GH-1332 (F2): closing a file whose read was cancelled while its + # GH-1332 (F2): closing a file whose read was cancelled twice while its # request ran waits for the request, which used to run on after the # close returned. block_size = 2**10 @@ -2767,6 +2859,11 @@ def get_object(path, range_, **kwargs): close = None try: assert await asyncio.to_thread(started.wait, EVENT_TIMEOUT) + # The read waits for its running request (B6) until a second + # cancellation stops that wait. + read.cancel() + await asyncio.sleep(0.1) + assert not read.done() read.cancel() with pytest.raises(asyncio.CancelledError): await read @@ -3178,11 +3275,10 @@ async def test_cancelled_creation_stops_retries(self, fs): @pytest.mark.asyncio async def test_cancelled_read_frees_ranges(self, fs): - # GH-1186: a cancelled parallel read cancels the reads that have not - # finished, which stops their retries, and frees the data of the + # GH-1186: a cancelled parallel read stops the retries of the running + # reads, waits for them (GH-1332, B6), and frees the data of the # others, also while the cancellation is held, as a failed read does. block_size = 2**10 - loop = asyncio.get_running_loop() read = [] all_read = threading.Barrier(4, timeout=EVENT_TIMEOUT) held = threading.Event() @@ -3190,24 +3286,11 @@ async def test_cancelled_read_frees_ranges(self, fs): operation = ThrottledOperation(failures=1) throttled = retried(operation, RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT)) - async def others_done(): - def running_reads(): - return [ - t - for t in asyncio.all_tasks() - if t.get_coro().__qualname__ == "S3Executor.run" and not t.done() - ] - - # At least the held read is found. - assert running_reads() - await wait_until(lambda: len(running_reads()) <= 1) - def get_object(path, range_, **kwargs): if range_[0] == 0: + # Retrying when the read is cancelled, once the other ranges + # have been read. all_read.wait() - # Retrying when the read is cancelled, once the tasks of the - # other ranges have taken their results. - asyncio.run_coroutine_threadsafe(others_done(), loop).result(EVENT_TIMEOUT) held.set() try: throttled() diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index e774599c9..660a61ceb 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -30,6 +30,7 @@ forking, interrupt_each_opcode, overlapping_calls, + retried, run_in_forked_child, shutdown_bounded, sigint_when, @@ -1070,11 +1071,397 @@ def test_submit_to_cancelled_executor_failure(self): executor = _DeferringExecutor() future: Future[int] = Future() executor.submit_to(future, sum, [1, 2]) - assert future.cancel() (submitted,) = executor.submitted + # Started, so that the cancellation of the future does not reach it. submitted.set_running_or_notify_cancel() + assert future.cancel() submitted.set_exception(RuntimeError("cannot schedule new futures after shutdown")) _, not_done = wait([future], timeout=EVENT_TIMEOUT) assert not not_done assert future.cancelled() + + def test_submit_to_cancel_frees_waiting_call(self): + # GH-1332 (S8): cancelling the future of a call that submit_to() put + # in a scope's queue drops the call at once, so that its arguments are + # not kept until a place frees. + class Body: + pass + + pool = S3ThreadPool(1) + scope = _S3PoolScope(pool, 1) + release = threading.Event() + body = Body() + body_ref = weakref.ref(body) + try: + scope.submit(release.wait, EVENT_TIMEOUT) + future: Future[None] = Future() + scope.submit_to(future, lambda b: None, body) + del body + assert future.cancel() + gc.collect() + assert body_ref() is None + finally: + release.set() + shutdown_bounded(scope) + pool.shutdown() + + def test_map(self): + # GH-1332 (B1): the outcomes are in the order of the items, at most + # the executor's limit at once. + call, state = overlapping_calls(2) + pool = S3ThreadPool(4) + scope = _S3PoolScope(pool, 2) + try: + + def square(i): + call() + return i * i + + result = scope.map(square, range(6)) + assert result.outcomes == [0, 1, 4, 9, 16, 25] + assert result.failure is None + assert state["peak"] == 2 + finally: + shutdown_bounded(scope) + pool.shutdown() + + def test_map_executor_without_limit(self): + # GH-1332: an executor that has only submit() and shutdown(), as one + # given to S3File(executor=), maps its items. + class Executor(S3Executor): + def __init__(self): + self._pool = ThreadPoolExecutor(2) + + @override + def submit(self, fn, *args, **kwargs): + return self._pool.submit(contextvars.copy_context().run, fn, *args, **kwargs) + + @override + def shutdown(self, wait=True): + self._pool.shutdown(wait=wait) + + with Executor() as executor: + assert executor.map(str, [1, 2]).outcomes == ["1", "2"] + + def test_map_failure_stops_batch(self): + # GH-1332 (B2): a failure starts no further call, stops the retries + # of the running one, whose backoff wait would take at least + # EVENT_TIMEOUT seconds, and is raised once that call has ended. The + # call that the failure frees the place for makes no request, also + # when its thread takes it before the waiting thread resumes. + operation = ThrottledOperation() + throttled = retried(operation, RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT)) + ended = [] + pool = S3ThreadPool(2) + scope = _S3PoolScope(pool, 2) + + def call(i): + try: + if i == 0: + return throttled() + if i == 1: + assert operation.started.wait(EVENT_TIMEOUT) + assert waiting.wait(EVENT_TIMEOUT) + raise ValueError("failed") + return i + finally: + ended.append(i) + + waiting = threading.Event() + + def late_wait(futures, **kwargs): + # The waiting thread resumes after the thread of the failed call + # has taken the next one. + waiting.set() + done = wait(futures, **kwargs) + time.sleep(0.2) + return done + + try: + started = time.monotonic() + with ( + mock.patch("pyathena.filesystem.s3_executor.wait", late_wait), + pytest.raises(ValueError, match="failed"), + ): + scope.map(call, range(4)) + assert time.monotonic() - started < EVENT_TIMEOUT + assert sorted(ended) == [0, 1] + assert operation.calls == 1 + finally: + shutdown_bounded(scope) + pool.shutdown() + + def test_map_raises_first_failure_to_end(self): + # GH-1332 (B2): the failure raised is the one that stopped the batch, + # the first to end, not the first item that failed. + pool = S3ThreadPool(2) + scope = _S3PoolScope(pool, 2) + second_failed = threading.Event() + + def call(i): + if i == 0: + assert second_failed.wait(EVENT_TIMEOUT) + raise ValueError("first item") + second_failed.set() + raise KeyError("second item") + + try: + with pytest.raises(KeyError, match="second item"): + scope.map(call, range(2)) + finally: + shutdown_bounded(scope) + pool.shutdown() + + def test_map_counts_failures(self): + # GH-1332 (B5): without stop_on_failure, every call runs, and the + # failure reported is the first that counts in the order they ended. + pool = S3ThreadPool(3) + scope = _S3PoolScope(pool, 3) + ended = {i: threading.Event() for i in range(3)} + + def call(i): + # Item 0 ends first and does not count; item 2 ends before item 1. + if i == 2: + assert ended[0].wait(EVENT_TIMEOUT) + if i == 1: + assert ended[2].wait(EVENT_TIMEOUT) + ended[i].set() + raise (LookupError if i == 0 else ValueError)(i) + + try: + result = scope.map( + call, + range(3), + stop_on_failure=False, + counts=lambda e: not isinstance(e, LookupError), + ) + assert [type(e) for e in result.outcomes] == [LookupError, ValueError, ValueError] + assert result.failure is result.outcomes[2] + finally: + shutdown_bounded(scope) + pool.shutdown() + + def test_map_interrupted_without_stop_on_failure(self): + # GH-1332 (B3): an interrupt of the wait cancels the calls that have + # not started and stops the retries of the running one, whose backoff + # wait would take at least EVENT_TIMEOUT seconds, also without + # stop_on_failure, and is raised once that call has ended. + operation = ThrottledOperation() + throttled = retried(operation, RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT)) + ran = [] + waits = [] + pool = S3ThreadPool(1) + scope = _S3PoolScope(pool, 1) + + def call(i): + ran.append(i) + return throttled() + + def interrupted_wait(futures, **kwargs): + waits.append(futures) + if len(waits) == 1: + assert operation.started.wait(EVENT_TIMEOUT) + raise KeyboardInterrupt + return wait(futures, **kwargs) + + try: + started = time.monotonic() + with ( + mock.patch("pyathena.filesystem.s3_executor.wait", interrupted_wait), + pytest.raises(KeyboardInterrupt), + ): + scope.map(call, range(2), stop_on_failure=False) + assert time.monotonic() - started < EVENT_TIMEOUT + assert ran == [0] + assert operation.calls == 1 + assert scope.pending() == [] + finally: + waits.clear() + shutdown_bounded(scope) + pool.shutdown() + + def test_map_interrupted_twice(self): + # GH-1332 (B3, I1 amended): a second interrupt stops the wait for the + # running calls, which then end in the executor. + release = threading.Event() + started = threading.Event() + pool = S3ThreadPool(1) + scope = _S3PoolScope(pool, 1) + + def call(i): + started.set() + release.wait(EVENT_TIMEOUT) + return i + + def interrupted_wait(futures, **kwargs): + assert started.wait(EVENT_TIMEOUT) + raise KeyboardInterrupt + + try: + with ( + mock.patch("pyathena.filesystem.s3_executor.wait", interrupted_wait), + pytest.raises(KeyboardInterrupt) as info, + ): + scope.map(call, range(1)) + assert isinstance(info.value.__context__, KeyboardInterrupt) + assert not release.is_set() + assert len(scope.pending()) == 1 + finally: + release.set() + shutdown_bounded(scope) + pool.shutdown() + + def test_map_failure_frees_outcomes(self): + # GH-1332 (F5): the outcome of a call that ends while a failed map + # waits for it is not kept by the traceback of the failure. + class Body: + pass + + refs = [] + failed = threading.Event() + both_started = threading.Barrier(2, timeout=EVENT_TIMEOUT) + pool = S3ThreadPool(2) + scope = _S3PoolScope(pool, 2) + + def call(i): + both_started.wait() + if i == 0: + failed.set() + raise ValueError("failed") + assert failed.wait(EVENT_TIMEOUT) + # Ends after the failure, while the map waits for it. + time.sleep(0.1) + body = Body() + refs.append(weakref.ref(body)) + return body + + try: + with pytest.raises(ValueError, match="failed") as info: + scope.map(call, range(2)) + gc.collect() + assert len(refs) == 1 + assert refs[0]() is None + assert info.value.__traceback__ is not None + finally: + shutdown_bounded(scope) + pool.shutdown() + + def test_map_interrupted_at_each_opcode(self): + # GH-1332 (B6, S1, S5): an interrupt at any opcode of map() leaves no + # call in the scope, and every call that it started has ended when it + # raises. + def case(interrupt): + pool = S3ThreadPool(2) + scope = _S3PoolScope(pool, 2) + running = [] + lock = threading.Lock() + + def call(i): + with lock: + running.append(i) + time.sleep(0.01) + with lock: + running.remove(i) + return i + + try: + result = interrupt(scope.map, call, range(4)) + if not interrupt.fired: + assert result.outcomes == [0, 1, 2, 3] + return + assert running == [] + assert scope.pending() == [] + finally: + shutdown_bounded(scope) + pool.shutdown() + + assert interrupt_each_opcode(S3Executor.map, case) > 0 + + @pytest.mark.asyncio + async def test_map_async_cancelled(self): + # GH-1332 (B3, B6): a cancellation of the waiting task cancels the + # calls that have not started, stops the retries of the running one, + # and is raised once that call has ended. + operation = ThrottledOperation() + throttled = retried(operation, RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT)) + ran = [] + ended = [] + pool = S3ThreadPool(1) + scope = _S3PoolScope(pool, 1) + + def call(i): + ran.append(i) + try: + return throttled() + finally: + ended.append(i) + + try: + task = asyncio.ensure_future(scope.map_async(call, range(2))) + assert await asyncio.to_thread(operation.started.wait, EVENT_TIMEOUT) + started = time.monotonic() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, EVENT_TIMEOUT) + assert time.monotonic() - started < EVENT_TIMEOUT + assert ran == ended == [0] + assert operation.calls == 1 + finally: + await asyncio.to_thread(shutdown_bounded, scope) + pool.shutdown() + + @pytest.mark.asyncio + async def test_map_async_failures_are_retrieved(self): + # GH-1332: the failures of the calls that a failure stopped are not + # logged as never retrieved. + loop = asyncio.get_running_loop() + logged = [] + previous = loop.get_exception_handler() + loop.set_exception_handler(lambda loop, context: logged.append(context)) + pool = S3ThreadPool(2) + scope = _S3PoolScope(pool, 2) + barrier = threading.Barrier(2, timeout=EVENT_TIMEOUT) + + def call(i): + barrier.wait() + raise ValueError(i) + + try: + with pytest.raises(ValueError, match=r"^[01]$"): + await scope.map_async(call, range(2)) + gc.collect() + await asyncio.sleep(0) + assert logged == [] + finally: + loop.set_exception_handler(previous) + await asyncio.to_thread(shutdown_bounded, scope) + pool.shutdown() + + def test_map_async_at_loop_shutdown(self): + # GH-1332 (B6): the end of asyncio.run() cancels a waiting map, which + # waits for its running call before the loop closes. + release = threading.Event() + started = threading.Event() + ended = [] + pool = S3ThreadPool(1) + scope = _S3PoolScope(pool, 1) + + def call(i): + started.set() + release.wait(0.5) + ended.append(i) + + tasks = [] + + async def main(): + tasks.append(asyncio.ensure_future(scope.map_async(call, range(1)))) + await asyncio.to_thread(started.wait, EVENT_TIMEOUT) + + try: + asyncio.run(main()) + assert ended == [0] + finally: + release.set() + shutdown_bounded(scope) + pool.shutdown() From 6590afd939ca00c7ebcae40c2a451b1eee0aeac0 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 7 Nov 2026 03:05:00 +0900 Subject: [PATCH 2/8] Share the scope, waiter, and error-code helpers of the fan-out (GH-1332) From the /simplify review of the fan-out primitive: - S3FileSystem._map() and _map_async() own the one-shot scope of an operation and its shutdown without another wait, instead of five copies at the sites. - S3Executor.waiter() wraps a future for asyncio and retrieves its outcome, for wait_shielded(), the map, and the streamed drain; the asyncio map reuses its wrappers in the drain. - _S3WorkItem.link() owns both links between the caller's future and the submission, and drops the submission once the future has ended. - A map without stop_on_failure waits for all calls at once. - S3ClientError.code_of() returns the code of the ClientError that an error was raised from, for the InvalidRange and NoSuchUpload checks. Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3.py | 128 ++++++++++++++++++--------- pyathena/filesystem/s3_async.py | 51 +++++------ pyathena/filesystem/s3_errors.py | 21 ++++- pyathena/filesystem/s3_executor.py | 134 ++++++++++++++++++++--------- 4 files changed, 215 insertions(+), 119 deletions(-) diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index f8c309e60..2527d44df 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -9,7 +9,7 @@ import mimetypes import os.path import time -from collections.abc import Callable, Mapping +from collections.abc import Callable, Iterable, Mapping from copy import deepcopy from dataclasses import dataclass, field from datetime import datetime @@ -33,7 +33,14 @@ from pyathena.connection import Connection from pyathena.filesystem.s3_core import S3Core, S3DeleteBatch, S3DeleteResult from pyathena.filesystem.s3_errors import S3ClientError -from pyathena.filesystem.s3_executor import S3Executor, S3ThreadPool, _S3PoolScope +from pyathena.filesystem.s3_executor import ( + ItemT, + S3Executor, + S3MapResult, + S3ThreadPool, + T, + _S3PoolScope, +) from pyathena.filesystem.s3_object import ( S3Metadata, S3MultipartUpload, @@ -1152,6 +1159,63 @@ def _create_executor(self, max_workers: int) -> _S3PoolScope: """ return _S3PoolScope(self._request_pool, max_workers) + def _map( + self, + fn: Callable[[ItemT], T], + items: Iterable[ItemT], + max_workers: int | None = None, + **kwargs: Any, + ) -> S3MapResult[T]: + """Run the parallel requests of an operation in a scope of their own. + + See :meth:`S3Executor.map() `. + The scope is shut down without another wait, since the map has waited + for the calls that it started, unless a second interrupt stopped it. + + Args: + fn: The request to run for each item. + items: The items. + max_workers: Maximum number of parallel requests. Defaults to + ``self.max_workers``. + **kwargs: The options of ``S3Executor.map()``. + + Returns: + The outcomes of the requests. + """ + executor = self._create_executor(max_workers=max_workers or self.max_workers) + try: + return executor.map(fn, items, **kwargs) + finally: + executor.shutdown(wait=False) + + async def _map_async( + self, + fn: Callable[[ItemT], T], + items: Iterable[ItemT], + max_workers: int | None = None, + **kwargs: Any, + ) -> S3MapResult[T]: + """Run the parallel requests of an asyncio operation in a scope of their own. + + See :meth:`_map` and + :meth:`S3Executor.map_async() `. + + Args: + fn: The request to run for each item. + items: The items. + max_workers: Maximum number of parallel requests. Defaults to + ``self.max_workers``. + **kwargs: The options of ``S3Executor.map_async()``. + + Returns: + The outcomes of the requests. + """ + executor = self._create_executor(max_workers=max_workers or self.max_workers) + try: + return await executor.map_async(fn, items, **kwargs) + finally: + executor.shutdown(wait=False) + def _delete_objects(self, paths: list[str], max_workers: int | None = None, **kwargs) -> None: """Delete objects with DeleteObjects requests grouped by bucket. @@ -1170,16 +1234,14 @@ def _delete_objects(self, paths: list[str], max_workers: int | None = None, **kw if not batches: return - max_workers = max_workers if max_workers else self.max_workers - executor = self._create_executor(max_workers=max_workers) - try: - # Every request runs also after another failed, and an interrupt - # cancels those that have not started. - results = executor.map( - functools.partial(self._delete_batch, **params), batches, stop_on_failure=False - ) - finally: - executor.shutdown(wait=False) + # Every request runs also after another failed, and an interrupt + # cancels those that have not started. + results = self._map( + functools.partial(self._delete_batch, **params), + batches, + max_workers=max_workers, + stop_on_failure=False, + ) self._raise_delete_errors(results.outcomes) @staticmethod @@ -1927,11 +1989,7 @@ def cat_file( try: return self.core.get_object(s3_path.with_version_id(version_id), ranges, **kwargs) except OSError as e: - if ( - ranges - and isinstance(e.__cause__, botocore.exceptions.ClientError) - and S3ClientError(e.__cause__).code == "InvalidRange" - ): + if ranges and S3ClientError.code_of(e) == "InvalidRange": # The range starts at or past the end of the object. return b"" raise @@ -2261,16 +2319,7 @@ def chmod(self, path: str, acl: str, recursive: bool = False, **kwargs) -> None: raise ValueError(f"ACL not in {self.core.OBJECT_ACLS}.") if recursive: paths = self.find(path, withdirs=False) - executor = self._create_executor(max_workers=self.max_workers) - try: - executor.map( - lambda p: self.chmod(p, acl, recursive=False, **kwargs), - paths, - ) - finally: - # The map has waited for the updates it started, unless a - # second interrupt stopped it. - executor.shutdown(wait=False) + self._map(lambda p: self.chmod(p, acl, recursive=False, **kwargs), paths) if s3_path.key: # A key prefix is not an object itself; only the objects # below it have ACLs. @@ -2381,18 +2430,14 @@ def clear_multipart_uploads(self, path: str) -> None: uploads = self.list_multipart_uploads(path) if not uploads: return - executor = self._create_executor(max_workers=self.max_workers) - try: - # Every abort runs also after another failed, and an interrupt - # cancels those that have not started. - results = executor.map( - self.core.abort_multipart_upload, - uploads, - stop_on_failure=False, - counts=self._is_not_no_such_upload, - ) - finally: - executor.shutdown(wait=False) + # Every abort runs also after another failed, and an interrupt cancels + # those that have not started. + results = self._map( + self.core.abort_multipart_upload, + uploads, + stop_on_failure=False, + counts=self._is_not_no_such_upload, + ) if results.failure is not None: raise results.failure @@ -2406,11 +2451,8 @@ def _is_not_no_such_upload(error: BaseException) -> bool: Returns: False for the ``FileNotFoundError`` of a ``NoSuchUpload`` error. """ - cause = error.__cause__ return not ( - isinstance(error, FileNotFoundError) - and isinstance(cause, botocore.exceptions.ClientError) - and S3ClientError(cause).code == "NoSuchUpload" + isinstance(error, FileNotFoundError) and S3ClientError.code_of(error) == "NoSuchUpload" ) def created(self, path: str) -> datetime: diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index e0c1ea3e6..827474834 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -29,7 +29,7 @@ from pyathena.aio.util import create_in_thread from pyathena.filesystem.s3 import CompressedBuffer, S3File, S3FileSystem, S3FileTarget from pyathena.filesystem.s3_core import S3Core -from pyathena.filesystem.s3_executor import S3ThreadPool +from pyathena.filesystem.s3_executor import S3Executor, S3ThreadPool from pyathena.filesystem.s3_object import ( S3CompleteMultipartUpload, S3Metadata, @@ -386,15 +386,11 @@ async def _delete_objects(self, paths: list[str], **kwargs) -> None: # At most max_workers requests at once, as S3FileSystem.rm() sends, # so that a large deletion does not queue ahead of the requests of # the files that share the request pool. - requests = self._sync_fs._create_executor(max_workers=self._sync_fs.max_workers) - try: - results = await requests.map_async( - functools.partial(self._sync_fs._delete_batch, **params), - batches, - stop_on_failure=False, - ) - finally: - requests.shutdown(wait=False) + results = await self._sync_fs._map_async( + functools.partial(self._sync_fs._delete_batch, **params), + batches, + stop_on_failure=False, + ) self._sync_fs._raise_delete_errors(results.outcomes) async def _mv(self, path1, path2, recursive=False, maxdepth=None, **kwargs) -> None: @@ -724,21 +720,18 @@ def _abort() -> None: raise # At most max_workers copies at once in the request pool. - requests = self._sync_fs._create_executor(max_workers=max_workers) - try: - await requests.map_async( - lambda name: self.core.copy_object_annotation( - name, - plan.source, - plan.destination, - completed.version_id, - completed.etag, - **kwargs, - ), - plan.annotations, - ) - finally: - requests.shutdown(wait=False) + await self._sync_fs._map_async( + lambda name: self.core.copy_object_annotation( + name, + plan.source, + plan.destination, + completed.version_id, + completed.etag, + **kwargs, + ), + plan.annotations, + max_workers=max_workers, + ) async def _find( self, @@ -1483,11 +1476,9 @@ async def _drain(self) -> None: pending.append(future) if not pending: return - waiters = [asyncio.wrap_future(future) for future in pending] - for waiter in waiters: - # The outcomes belong to the reads, the commit, and the abort that - # awaited them; retrieved here, so that asyncio does not log them. - waiter.add_done_callback(lambda f: f.cancelled() or f.exception()) + # The outcomes belong to the reads, the commit, and the abort that + # awaited them; retrieved, so that asyncio does not log them. + waiters = [S3Executor.waiter(future) for future in pending] cancelled = False # The states are also checked every second, since an interrupt of a # cancellation can skip the callbacks that the waiters rely on. diff --git a/pyathena/filesystem/s3_errors.py b/pyathena/filesystem/s3_errors.py index d7f167fdd..7e3050028 100644 --- a/pyathena/filesystem/s3_errors.py +++ b/pyathena/filesystem/s3_errors.py @@ -20,10 +20,9 @@ from __future__ import annotations import errno -from typing import TYPE_CHECKING, ClassVar +from typing import ClassVar -if TYPE_CHECKING: - import botocore.exceptions +import botocore.exceptions class S3ClientError: @@ -136,6 +135,22 @@ def _translate(self) -> OSError: errno_ = errno.EIO return OSError(errno_, self._message) + @classmethod + def code_of(cls, error: BaseException) -> str | None: + """Return the S3 error code of an error raised from a ``ClientError``. + + Args: + error: An error, such as the ``OSError`` of a request. + + Returns: + The code of the ``ClientError`` that the error was raised from, + or None if it was not raised from one. + """ + cause = error.__cause__ + if isinstance(cause, botocore.exceptions.ClientError): + return cls(cause).code + return None + @property def code(self) -> str: """The S3 error code that uniquely identifies the error condition.""" diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index e43a6ea65..8cbdf0e51 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -19,7 +19,7 @@ from abc import ABCMeta, abstractmethod from collections import deque from collections.abc import Callable, Iterable -from concurrent.futures import FIRST_EXCEPTION, CancelledError, Future, wait +from concurrent.futures import ALL_COMPLETED, FIRST_EXCEPTION, CancelledError, Future, wait from concurrent.futures.thread import ThreadPoolExecutor from dataclasses import dataclass from multiprocessing import cpu_count @@ -65,6 +65,36 @@ def __init__( self._args = args self._kwargs = kwargs self._claim = threading.Lock() + self._submitted: Future[Any] | None = None + + def link(self, future: Future[T], submitted: Future[Any]) -> None: + """Link the future of the caller with the submission of :meth:`run`. + + The submission settles the future if the executor drops the callable + (:meth:`settle`), and a cancellation of the future cancels the + submission, so that a scope drops a waiting call at once instead of + letting it pass through the pool without running. + + Args: + future: The future of the caller. + submitted: The future of the submitted :meth:`run`. + """ + submitted.add_done_callback(self.settle) + self._submitted = submitted + future.add_done_callback(self._cancel_submission) + + def _cancel_submission(self, future: Future[T]) -> None: + """Cancel the submission if the future of the caller was cancelled. + + The future keeps its callbacks after it has ended, so the reference + to the submission is dropped here. + + Args: + future: The future of the caller, which has ended. + """ + submitted, self._submitted = self._submitted, None + if submitted is not None and future.cancelled(): + submitted.cancel() def run(self) -> None: """Run the callable and resolve the future unless it was cancelled.""" @@ -172,6 +202,8 @@ def __init__( # that an interrupt while it is being scheduled still lets it be # cancelled. self.futures: list[Future[T]] = [] + # The asyncio wrappers of the futures, by future. + self._waiters: dict[Future[T], asyncio.Future[T]] = {} def submit(self, executor: S3Executor, items: Iterable[ItemT]) -> None: """Submit a call for each item. @@ -187,15 +219,19 @@ def submit(self, executor: S3Executor, items: Iterable[ItemT]) -> None: def wait(self) -> None: """Wait until every call has ended or a failure has stopped the batch.""" + return_when = FIRST_EXCEPTION if self._stop_on_failure else ALL_COMPLETED pending = set(self.futures) while pending and not self._stopped_by_failure(): - pending = wait(pending, return_when=FIRST_EXCEPTION).not_done + pending = wait(pending, return_when=return_when).not_done async def wait_async(self) -> None: """Await until every call has ended or a failure has stopped the batch.""" - pending = {self._wrap(future) for future in self.futures} + return_when = asyncio.FIRST_EXCEPTION if self._stop_on_failure else asyncio.ALL_COMPLETED + for future in self.futures: + self._waiters[future] = S3Executor.waiter(future) + pending = set(self._waiters.values()) while pending and not self._stopped_by_failure(): - pending = (await asyncio.wait(pending, return_when=asyncio.FIRST_EXCEPTION))[1] + pending = (await asyncio.wait(pending, return_when=return_when))[1] def result(self) -> S3MapResult[T]: """Return the outcomes of the calls, which have all ended. @@ -209,10 +245,7 @@ def result(self) -> S3MapResult[T]: if self._stopped_by_failure(): raise cast(BaseException, self._failure) return S3MapResult( - [ - f.result() if f.exception() is None else cast(Any, f.exception()) - for f in self.futures - ], + [e if (e := f.exception()) is not None else f.result() for f in self.futures], self._failure, ) @@ -233,12 +266,29 @@ async def drain_async(self) -> None: """Stop the batch, cancel the calls that have not started, and await the others.""" started = self._stop() try: - waiters = [self._wrap(future) for future in started] + waiters = [self._waiters.get(f) or S3Executor.waiter(f) for f in started] + self._waiters.clear() while not all(future.done() for future in started): await asyncio.wait(waiters, timeout=1.0) finally: started.clear() + def release(self, error: BaseException) -> None: + """Drop the futures of a batch that raises, and their outcomes, such as the bodies of reads. + + The frames below the map, which the traceback of the error keeps, are + cleared; :func:`traceback.clear_frames` skips the frames that are + executing. + + Args: + error: The error that the map raises. + """ + try: + traceback.clear_frames(error.__traceback__) + finally: + self.futures.clear() + self._waiters.clear() + def _stop(self) -> list[Future[T]]: """Stop the batch and cancel the calls that have not started. @@ -291,20 +341,6 @@ def _call(self, item: ItemT) -> T: self.stop.set() raise - @staticmethod - def _wrap(future: Future[T]) -> asyncio.Future[T]: - """Wrap a future for asyncio, retrieving its outcome so that asyncio does not log it. - - Args: - future: The future. - - Returns: - The asyncio future. - """ - waiter = asyncio.wrap_future(future) - waiter.add_done_callback(lambda f: f.cancelled() or f.exception()) - return waiter - class S3Executor(metaclass=ABCMeta): """Abstract executor for parallel S3 operations. @@ -412,14 +448,41 @@ async def wait_shielded(self, future: Future[T]) -> T: Returns: The return value of the call. """ - waiter = asyncio.wrap_future(future) - waiter.add_done_callback(lambda f: f.cancelled() or f.exception()) + waiter = self.waiter(future) # Unlike shield(), which on Python 3.14 and later also logs the failure # of a future that ends after a cancellation, wait() leaves the future # alone when the waiting task is cancelled. await asyncio.wait([waiter]) return waiter.result() + @staticmethod + def waiter(future: Future[T]) -> asyncio.Future[T]: + """Wrap a future of an executor for asyncio, retrieving its outcome. + + The outcome belongs to the caller that awaits the wrapper; it is + retrieved, so that asyncio does not log a failure as never retrieved + when no one awaits it, such as after a cancellation. + + Args: + future: The future. + + Returns: + The asyncio future. + """ + waiter = asyncio.wrap_future(future) + waiter.add_done_callback(S3Executor._retrieve) + return waiter + + @staticmethod + def _retrieve(waiter: asyncio.Future[Any]) -> None: + """Retrieve the outcome of a waiter, so that asyncio does not log it. + + Args: + waiter: The asyncio future, which has ended. + """ + if not waiter.cancelled(): + waiter.exception() + def submit_to(self, future: Future[T], fn: Callable[..., T], *args: Any, **kwargs: Any) -> None: """Submit a callable whose outcome resolves a future of the caller. @@ -446,12 +509,7 @@ def submit_to(self, future: Future[T], fn: Callable[..., T], *args: Any, **kwarg **kwargs: Keyword arguments passed to the callable. """ item = _S3WorkItem(future, fn, args, kwargs) - submitted = self.submit(item.run) - submitted.add_done_callback(item.settle) - # A cancellation of the future also cancels the submission, so that a - # scope drops a waiting call at once instead of letting it pass - # through the pool without running. - future.add_done_callback(lambda f: f.cancelled() and submitted.cancel()) + item.link(future, self.submit(item.run)) def map( self, @@ -500,14 +558,7 @@ def map( try: batch.drain() finally: - # The frames below this one keep the futures, whose outcomes, - # such as the bodies of reads, the traceback of the error would - # otherwise keep; clear_frames() skips this frame, which is - # executing. - try: - traceback.clear_frames(e.__traceback__) - finally: - batch.futures.clear() + batch.release(e) raise async def map_async( @@ -548,10 +599,7 @@ async def map_async( try: await batch.drain_async() finally: - try: - traceback.clear_frames(e.__traceback__) - finally: - batch.futures.clear() + batch.release(e) raise def __enter__(self) -> S3Executor: From 5ff74f6f0dc2f4ac0148ddfceefcf733cb7c2507 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 7 Nov 2026 03:28:10 +0900 Subject: [PATCH 3/8] Count failures that the executor reports in a map (GH-1332) A call that the executor failed or cancelled without running it, such as one whose thread could not start, settled its future without passing through the batch, so it neither stopped the batch nor became its failure: a parallel read then failed joining an exception, and chmod or clear_multipart_uploads ignored it. The waiting thread now records such outcomes, and a cancelled call's outcome is a CancelledError instead of raising from result(). Also: - release() sets the stop, also if a further interrupt stopped the drain before it did. - _map_async() lives on AioS3FileSystem, its only user. - The asyncio _delete_objects() returns at once without batches, as the synchronous one does. Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3.py | 28 ---------- pyathena/filesystem/s3_async.py | 39 ++++++++++++-- pyathena/filesystem/s3_executor.py | 52 +++++++++++++++---- tests/pyathena/filesystem/test_s3_executor.py | 37 ++++++++++++- 4 files changed, 115 insertions(+), 41 deletions(-) diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 2527d44df..e799c39ec 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -1188,34 +1188,6 @@ def _map( finally: executor.shutdown(wait=False) - async def _map_async( - self, - fn: Callable[[ItemT], T], - items: Iterable[ItemT], - max_workers: int | None = None, - **kwargs: Any, - ) -> S3MapResult[T]: - """Run the parallel requests of an asyncio operation in a scope of their own. - - See :meth:`_map` and - :meth:`S3Executor.map_async() `. - - Args: - fn: The request to run for each item. - items: The items. - max_workers: Maximum number of parallel requests. Defaults to - ``self.max_workers``. - **kwargs: The options of ``S3Executor.map_async()``. - - Returns: - The outcomes of the requests. - """ - executor = self._create_executor(max_workers=max_workers or self.max_workers) - try: - return await executor.map_async(fn, items, **kwargs) - finally: - executor.shutdown(wait=False) - def _delete_objects(self, paths: list[str], max_workers: int | None = None, **kwargs) -> None: """Delete objects with DeleteObjects requests grouped by bucket. diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 827474834..8eea54c2e 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -16,6 +16,7 @@ import mimetypes import os import warnings +from collections.abc import Callable, Iterable from concurrent.futures import Future, wait from io import BytesIO from typing import TYPE_CHECKING, Any, cast @@ -29,7 +30,7 @@ from pyathena.aio.util import create_in_thread from pyathena.filesystem.s3 import CompressedBuffer, S3File, S3FileSystem, S3FileTarget from pyathena.filesystem.s3_core import S3Core -from pyathena.filesystem.s3_executor import S3Executor, S3ThreadPool +from pyathena.filesystem.s3_executor import ItemT, S3Executor, S3MapResult, S3ThreadPool, T from pyathena.filesystem.s3_object import ( S3CompleteMultipartUpload, S3Metadata, @@ -383,16 +384,48 @@ async def _delete_objects(self, paths: list[str], **kwargs) -> None: OSError: If S3 could not delete some of the objects. """ batches, params = self._sync_fs._delete_batches(paths, **kwargs) + if not batches: + return # At most max_workers requests at once, as S3FileSystem.rm() sends, # so that a large deletion does not queue ahead of the requests of # the files that share the request pool. - results = await self._sync_fs._map_async( + results = await self._map_async( functools.partial(self._sync_fs._delete_batch, **params), batches, stop_on_failure=False, ) self._sync_fs._raise_delete_errors(results.outcomes) + async def _map_async( + self, + fn: Callable[[ItemT], T], + items: Iterable[ItemT], + max_workers: int | None = None, + **kwargs: Any, + ) -> S3MapResult[T]: + """Run the parallel requests of an operation in a scope of their own. + + As :meth:`S3FileSystem._map`, with + :meth:`S3Executor.map_async() `. + + Args: + fn: The request to run for each item. + items: The items. + max_workers: Maximum number of parallel requests. Defaults to the + ``max_workers`` of the internal ``S3FileSystem``. + **kwargs: The options of ``S3Executor.map_async()``. + + Returns: + The outcomes of the requests. + """ + executor = self._sync_fs._create_executor( + max_workers=max_workers or self._sync_fs.max_workers + ) + try: + return await executor.map_async(fn, items, **kwargs) + finally: + executor.shutdown(wait=False) + async def _mv(self, path1, path2, recursive=False, maxdepth=None, **kwargs) -> None: """Move files from one S3 location to another. @@ -720,7 +753,7 @@ def _abort() -> None: raise # At most max_workers copies at once in the request pool. - await self._sync_fs._map_async( + await self._map_async( lambda name: self.core.copy_object_annotation( name, plan.source, diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index 8cbdf0e51..b078e55f3 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -222,7 +222,9 @@ def wait(self) -> None: return_when = FIRST_EXCEPTION if self._stop_on_failure else ALL_COMPLETED pending = set(self.futures) while pending and not self._stopped_by_failure(): - pending = wait(pending, return_when=return_when).not_done + done, pending = wait(pending, return_when=return_when) + for future in done: + self._record(self._outcome(future)) async def wait_async(self) -> None: """Await until every call has ended or a failure has stopped the batch.""" @@ -231,7 +233,9 @@ async def wait_async(self) -> None: self._waiters[future] = S3Executor.waiter(future) pending = set(self._waiters.values()) while pending and not self._stopped_by_failure(): - pending = (await asyncio.wait(pending, return_when=return_when))[1] + done, pending = await asyncio.wait(pending, return_when=return_when) + for waiter in done: + self._record(CancelledError() if waiter.cancelled() else waiter.exception()) def result(self) -> S3MapResult[T]: """Return the outcomes of the calls, which have all ended. @@ -245,10 +249,23 @@ def result(self) -> S3MapResult[T]: if self._stopped_by_failure(): raise cast(BaseException, self._failure) return S3MapResult( - [e if (e := f.exception()) is not None else f.result() for f in self.futures], + [e if (e := self._outcome(f)) is not None else f.result() for f in self.futures], self._failure, ) + @staticmethod + def _outcome(future: Future[Any]) -> BaseException | None: + """Return the failure of a call that has ended, or None if it succeeded. + + Args: + future: The future of the call. + + Returns: + Its exception, a ``CancelledError`` if the executor cancelled it, + or None. + """ + return CancelledError() if future.cancelled() else future.exception() + def drain(self) -> None: """Stop the batch, cancel the calls that have not started, and wait for the others. @@ -286,6 +303,8 @@ def release(self, error: BaseException) -> None: try: traceback.clear_frames(error.__traceback__) finally: + # Also if a further interrupt stopped the drain before its stop. + self.stop.set() self.futures.clear() self._waiters.clear() @@ -333,14 +352,29 @@ def _call(self, item: ItemT) -> T: try: return self._fn(item) except BaseException as e: - if not isinstance(e, Exception) or self._counts is None or self._counts(e): - with self._lock: - if self._failure is None: - self._failure = e - if self._stop_on_failure: - self.stop.set() + self._record(e) raise + def _record(self, error: BaseException | None) -> None: + """Record the first failure that counts, and stop the batch with it if requested. + + A call records its own failure in its thread; the waiting thread also + records those that the executor reports without running the call, + such as a failure to start a thread. + + Args: + error: The failure of a call, or None if it succeeded. + """ + if error is None or ( + isinstance(error, Exception) and self._counts is not None and not self._counts(error) + ): + return + with self._lock: + if self._failure is None: + self._failure = error + if self._stop_on_failure: + self.stop.set() + class S3Executor(metaclass=ABCMeta): """Abstract executor for parallel S3 operations. diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index 660a61ceb..a162490de 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -14,7 +14,7 @@ import time import traceback import weakref -from concurrent.futures import Future, ThreadPoolExecutor, wait +from concurrent.futures import CancelledError, Future, ThreadPoolExecutor, wait from typing import Any, cast from unittest import mock @@ -1144,6 +1144,41 @@ def shutdown(self, wait=True): with Executor() as executor: assert executor.map(str, [1, 2]).outcomes == ["1", "2"] + @pytest.mark.parametrize("rejection", ["failed", "cancelled"]) + def test_map_executor_rejection(self, rejection): + # GH-1332: a call that the executor fails or cancels without running + # it counts as a failure: it stops the batch, or is the failure of a + # batch without stop_on_failure, instead of being returned as a value. + class Rejecting(S3Executor): + def __init__(self): + self._pool = ThreadPoolExecutor(2) + self.calls = 0 + + @override + def submit(self, fn, *args, **kwargs): + self.calls += 1 + if self.calls != 2: + return self._pool.submit(contextvars.copy_context().run, fn, *args, **kwargs) + future: Future[Any] = Future() + if rejection == "cancelled": + future.cancel() + else: + future.set_running_or_notify_cancel() + future.set_exception(RuntimeError("cannot start a thread")) + return future + + @override + def shutdown(self, wait=True): + self._pool.shutdown(wait=wait) + + expected = CancelledError if rejection == "cancelled" else RuntimeError + with Rejecting() as executor, pytest.raises(expected): + executor.map(str, range(3)) + with Rejecting() as executor: + result = executor.map(str, range(3), stop_on_failure=False) + assert isinstance(result.failure, expected) + assert isinstance(result.outcomes[1], expected) + def test_map_failure_stops_batch(self): # GH-1332 (B2): a failure starts no further call, stops the retries # of the running one, whose backoff wait would take at least From 58ef05d2c57c2d6de311895a95ce309b1abb612d Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 7 Nov 2026 03:29:22 +0900 Subject: [PATCH 4/8] Correct the asyncio docs on how the fan-outs run (GH-1332) The bulk delete and the annotation copies of AioS3FileSystem no longer use asyncio.gather, and the synchronous filesystem runs its parallel requests in its request pool since #1359. The cancellation note covers the files of open_async(), whose reads are coroutines. Co-Authored-By: Claude Opus 5.5 --- docs/aio.md | 8 ++++---- pyathena/filesystem/s3_async.py | 12 +++++++----- 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/docs/aio.md b/docs/aio.md index f64e7d677..79f786bbf 100644 --- a/docs/aio.md +++ b/docs/aio.md @@ -281,10 +281,10 @@ its coroutines run the blocking boto3 calls in threads that PyAthena owns, sized | | S3FileSystem | AioS3FileSystem | |---|---|---| | **Blocking calls** | Caller's thread | Thread pool of the filesystem | -| **Parallelism** | `ThreadPoolExecutor` | `asyncio.gather` on a thread pool | +| **Parallelism** | Request pool of the filesystem | `asyncio.gather` on a thread pool; requests in the request pool | | **File handles** | `S3File` with thread pool | `AioS3File` with thread pool; `AioS3StreamedFile` from `open_async()` | -| **Bulk delete** | Thread pool per batch | `asyncio.gather` per batch | -| **Multipart copy** | Thread pool per copy | Thread pool of the filesystem | +| **Bulk delete** | Request pool, `max_workers` requests at once | Request pool, `max_workers` requests at once | +| **Multipart copy** | Request pool, `max_workers` requests per copy | Request pool, `max_workers` requests per copy | | **Best for** | Synchronous applications | Async frameworks (FastAPI, aiohttp, etc.) | ### Threads @@ -334,7 +334,7 @@ ends. A request that has already been sent still runs to its end in its thread. requests that run to their end after a cancellation keep their retries: the commit of a file, the CompleteMultipartUpload request of a multipart copy, and the aborts of multipart uploads. -The range reads of a file, the DeleteObjects requests of `_rm()`, and the annotation +The range reads of a file from `open_async()`, the DeleteObjects requests of `_rm()`, and the annotation copies of a multipart copy are waited for before the cancellation is raised, also when the shutdown of the event loop at the end of `asyncio.run()` cancels the coroutine; a second cancellation stops that wait. diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 8eea54c2e..e3ef8f16f 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -57,11 +57,13 @@ class AioS3FileSystem(AsyncFileSystem): This class wraps ``S3FileSystem`` to provide native asyncio support. Its coroutines run the blocking boto3 calls in an ``S3ThreadPool`` with ``max_workers`` threads that the filesystem owns, instead of the default - executor of the event loop, and run parallel requests with - ``asyncio.gather``. The single requests that these calls run in - parallel, such as the parts of a file or a copy and the DeleteObjects - batches of ``_rm()``, run in the request pool of the wrapped - ``S3FileSystem``, which its files share. + executor of the event loop, and run parallel per-object calls, such as + the copies of ``_mv()``, with ``asyncio.gather``. The single requests + that these calls run in parallel, such as the parts of a file or a copy, + and the DeleteObjects batches of ``_rm()``, run in the request pool of + the wrapped ``S3FileSystem``, which its files share, through + :meth:`~pyathena.filesystem.s3_executor.S3Executor.map_async` or a write + session. The implementation uses composition: an internal ``S3FileSystem`` instance handles all boto3 calls, while this class delegates to it in the threads of its pool. From 81a0c9c86507cd09698023b5e6e76f0b39500514 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 7 Nov 2026 04:25:42 +0900 Subject: [PATCH 5/8] Record executor failures when their futures settle (GH-1332) From the independent review of the fan-out primitive: - A failure that the executor reports without running a call is now recorded by a done callback of its future, when it is set and before the setting thread takes another call, instead of by the waiting thread, which let later calls start and later failures win the latch. The batch drops its futures on every path, since the callbacks would otherwise form a cycle that keeps the outcomes until the cyclic garbage collector runs. - S3Executor.map() and docs/aio.md state that a pool must not map on itself; only scopes run such calls in the calling thread. - Tests: the completion-order tests wait for the recording of the preceding failure; the asyncio rm test checks the invalidation of each request while the other runs; the retrieval test fails the second call during the drain; the synchronous map tests run under call_bounded(); test_map_result_frees_outcomes checks the cycle. Co-Authored-By: Claude Opus 5.5 --- docs/aio.md | 4 +- pyathena/filesystem/s3_executor.py | 48 +++++-- tests/pyathena/filesystem/test_s3.py | 20 ++- tests/pyathena/filesystem/test_s3_async.py | 19 +-- tests/pyathena/filesystem/test_s3_executor.py | 136 +++++++++++++----- tests/pyathena/util.py | 34 +++++ 6 files changed, 200 insertions(+), 61 deletions(-) diff --git a/docs/aio.md b/docs/aio.md index 79f786bbf..cb71fc6e7 100644 --- a/docs/aio.md +++ b/docs/aio.md @@ -350,7 +350,9 @@ creates new threads. `S3ThreadPoolExecutor` and `S3AioThreadPool` are aliases of `S3Executor.map()` and `S3Executor.map_async()` run a callable for each item in the executor and return the outcomes in the order of the items. With `stop_on_failure`, the default, the first failure cancels the calls that have not started, stops the -retries of the running ones, and is raised once they have ended. +retries of the running ones, and is raised once they have ended. A call that runs in an +`S3ThreadPool` must not map on that pool, which would wait for the threads that the +call holds; a scope of the filesystem runs such calls in the calling thread. The file handles of `AioS3FileSystem.open()` run their multipart uploads and parallel range reads in the request pool of the filesystem, at most the `max_workers` of the diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index b078e55f3..9c61a9665 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -215,6 +215,10 @@ def submit(self, executor: S3Executor, items: Iterable[ItemT]) -> None: with stop_retries_on(self.stop): for item in items: self.futures.append(Future()) + # Records a failure that the executor reports without running + # the call, such as a failure to start a thread, when it is + # set, before the thread that sets it takes another call. + self.futures[-1].add_done_callback(self._settled) executor.submit_to(self.futures[-1], self._call, item) def wait(self) -> None: @@ -222,9 +226,7 @@ def wait(self) -> None: return_when = FIRST_EXCEPTION if self._stop_on_failure else ALL_COMPLETED pending = set(self.futures) while pending and not self._stopped_by_failure(): - done, pending = wait(pending, return_when=return_when) - for future in done: - self._record(self._outcome(future)) + pending = wait(pending, return_when=return_when).not_done async def wait_async(self) -> None: """Await until every call has ended or a failure has stopped the batch.""" @@ -233,9 +235,7 @@ async def wait_async(self) -> None: self._waiters[future] = S3Executor.waiter(future) pending = set(self._waiters.values()) while pending and not self._stopped_by_failure(): - done, pending = await asyncio.wait(pending, return_when=return_when) - for waiter in done: - self._record(CancelledError() if waiter.cancelled() else waiter.exception()) + pending = (await asyncio.wait(pending, return_when=return_when))[1] def result(self) -> S3MapResult[T]: """Return the outcomes of the calls, which have all ended. @@ -305,8 +305,17 @@ def release(self, error: BaseException) -> None: finally: # Also if a further interrupt stopped the drain before its stop. self.stop.set() - self.futures.clear() - self._waiters.clear() + self.clear() + + def clear(self) -> None: + """Drop the futures of the calls. + + The done callback of each future refers to the batch, so a batch + that kept its futures would form a reference cycle with them, and + keep their outcomes until the cyclic garbage collector frees them. + """ + self.futures.clear() + self._waiters.clear() def _stop(self) -> list[Future[T]]: """Stop the batch and cancel the calls that have not started. @@ -355,12 +364,20 @@ def _call(self, item: ItemT) -> T: self._record(e) raise + def _settled(self, future: Future[T]) -> None: + """Record the failure of a call whose future has been set. + + Args: + future: The future of the call. + """ + self._record(self._outcome(future)) + def _record(self, error: BaseException | None) -> None: """Record the first failure that counts, and stop the batch with it if requested. - A call records its own failure in its thread; the waiting thread also - records those that the executor reports without running the call, - such as a failure to start a thread. + A call records its own failure in its thread before its future is + set; the done callback of the future records those that the executor + reports without running the call. Args: error: The failure of a call, or None if it succeeded. @@ -568,6 +585,11 @@ def map( and is raised once the running calls have ended. A further interrupt stops that wait; the running calls then end in the executor. + A thread of a scope's pool that maps on that scope runs the calls + itself. A thread of a bare :class:`S3ThreadPool` that maps on that + pool waits for threads that it may hold, so a pool must not map on + itself. + Args: fn: The callable to run for each item. items: The items. @@ -594,6 +616,8 @@ def map( finally: batch.release(e) raise + finally: + batch.clear() async def map_async( self, @@ -635,6 +659,8 @@ async def map_async( finally: batch.release(e) raise + finally: + batch.clear() def __enter__(self) -> S3Executor: return self diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 98b87169f..dd9a5306a 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -41,7 +41,7 @@ from pyathena.filesystem.s3 import CompressedBuffer, S3File, S3FileSystem from pyathena.filesystem.s3_core import S3Core, S3DeleteBatch, S3MultipartCopyPlan from pyathena.filesystem.s3_errors import S3ClientError -from pyathena.filesystem.s3_executor import S3Executor, S3ThreadPool, _S3PoolScope +from pyathena.filesystem.s3_executor import S3Executor, S3ThreadPool, _S3Batch, _S3PoolScope from pyathena.filesystem.s3_object import S3MultipartUpload, S3Object, S3ObjectType, S3StorageClass from pyathena.filesystem.s3_path import S3Path from pyathena.filesystem.s3_writer import S3MultipartStateError, _S3WriteSession @@ -610,16 +610,21 @@ def test_clear_multipart_uploads_raises_first_failure(self): # upload that no longer exists, in the order the aborts ended. fs = self._make_fs() fs.max_workers = 3 - ended = {u: threading.Event() for u in ("gone", "late", "early")} + recorded = {u: threading.Event() for u in ("gone", "late", "early")} + record = _S3Batch._record + + def recording(batch, error): + record(batch, error) + if error is not None: + recorded[str(error)].set() def abort_multipart_upload(upload): # "gone" ends first, then "early", then "late". name = upload.upload_id if name == "early": - assert ended["gone"].wait(EVENT_TIMEOUT) + assert recorded["gone"].wait(EVENT_TIMEOUT) if name == "late": - assert ended["early"].wait(EVENT_TIMEOUT) - ended[name].set() + assert recorded["early"].wait(EVENT_TIMEOUT) if name == "gone": error = botocore.exceptions.ClientError( {"Error": {"Code": "NoSuchUpload", "Message": name}}, "AbortMultipartUpload" @@ -634,7 +639,10 @@ def abort_multipart_upload(upload): for u in ("gone", "late", "early") ] ) - with pytest.raises(PermissionError, match="early"): + with ( + mock.patch.object(_S3Batch, "_record", recording), + pytest.raises(PermissionError, match="early"), + ): fs.clear_multipart_uploads("s3://bucket/") def test_rm_requests_invalidate_shared_parent(self): diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index 8c6d9565d..1ce42ca62 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -1693,15 +1693,16 @@ def call(method, **request): @pytest.mark.asyncio async def test_rm_cancel_waits_for_running_requests(self): # GH-1332 (B6): a cancelled _rm() raises once its running requests - # have ended, each of which invalidates the cache of its objects. The - # requests used to run on after the cancellation was raised. + # have ended, and each request invalidates the cache of its objects + # when it ends. The requests used to run on after the cancellation was + # raised. fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) started = {"b1": threading.Event(), "b2": threading.Event()} - release = threading.Event() + release = {"b1": threading.Event(), "b2": threading.Event()} def call(method, **request): started[request["Bucket"]].set() - release.wait(EVENT_TIMEOUT) + release[request["Bucket"]].wait(EVENT_TIMEOUT) return {} fs._sync_fs._call = fs._sync_fs._core.call = call @@ -1713,15 +1714,17 @@ def call(method, **request): fs.dircache["b1/a"] = [] fs.dircache["b2/b"] = [] task.cancel() - await asyncio.sleep(0.1) + release["b1"].set() + await wait_until(lambda: "b1/a" not in fs.dircache) + assert "b2/b" in fs.dircache assert not task.done() - release.set() + release["b2"].set() with pytest.raises(asyncio.CancelledError): await await_bounded(task) - assert "b1/a" not in fs.dircache assert "b2/b" not in fs.dircache finally: - release.set() + for event in release.values(): + event.set() @pytest.mark.asyncio async def test_rm_maxdepth(self): diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index a162490de..0c8c3db52 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -21,15 +21,23 @@ import pytest from botocore.exceptions import ClientError -from pyathena.filesystem.s3_executor import S3AioExecutor, S3Executor, S3ThreadPool, _S3PoolScope +from pyathena.filesystem.s3_executor import ( + S3AioExecutor, + S3Executor, + S3ThreadPool, + _S3Batch, + _S3PoolScope, +) from pyathena.util import RetryConfig, override, retry_api_call, stop_retries_on from tests.pyathena.util import ( EVENT_TIMEOUT, OpcodeInterrupter, ThrottledOperation, + call_bounded, forking, interrupt_each_opcode, overlapping_calls, + poll_until, retried, run_in_forked_child, shutdown_bounded, @@ -1118,7 +1126,7 @@ def square(i): call() return i * i - result = scope.map(square, range(6)) + result = call_bounded(scope.map, square, range(6)) assert result.outcomes == [0, 1, 4, 9, 16, 25] assert result.failure is None assert state["peak"] == 2 @@ -1142,13 +1150,14 @@ def shutdown(self, wait=True): self._pool.shutdown(wait=wait) with Executor() as executor: - assert executor.map(str, [1, 2]).outcomes == ["1", "2"] + assert call_bounded(executor.map, str, [1, 2]).outcomes == ["1", "2"] @pytest.mark.parametrize("rejection", ["failed", "cancelled"]) def test_map_executor_rejection(self, rejection): # GH-1332: a call that the executor fails or cancels without running - # it counts as a failure: it stops the batch, or is the failure of a - # batch without stop_on_failure, instead of being returned as a value. + # it counts as a failure when its future is set: it stops the batch + # before the next call starts, or is the failure of a batch without + # stop_on_failure, instead of being returned as a value. class Rejecting(S3Executor): def __init__(self): self._pool = ThreadPoolExecutor(2) @@ -1172,10 +1181,18 @@ def shutdown(self, wait=True): self._pool.shutdown(wait=wait) expected = CancelledError if rejection == "cancelled" else RuntimeError + ran = [] + + def call(i): + ran.append(i) + return str(i) + with Rejecting() as executor, pytest.raises(expected): - executor.map(str, range(3)) + call_bounded(executor.map, call, range(3)) + # The item submitted after the rejection does not run. + assert 2 not in ran with Rejecting() as executor: - result = executor.map(str, range(3), stop_on_failure=False) + result = call_bounded(executor.map, str, range(3), stop_on_failure=False) assert isinstance(result.failure, expected) assert isinstance(result.outcomes[1], expected) @@ -1219,7 +1236,7 @@ def late_wait(futures, **kwargs): mock.patch("pyathena.filesystem.s3_executor.wait", late_wait), pytest.raises(ValueError, match="failed"), ): - scope.map(call, range(4)) + call_bounded(scope.map, call, range(4)) assert time.monotonic() - started < EVENT_TIMEOUT assert sorted(ended) == [0, 1] assert operation.calls == 1 @@ -1232,18 +1249,27 @@ def test_map_raises_first_failure_to_end(self): # the first to end, not the first item that failed. pool = S3ThreadPool(2) scope = _S3PoolScope(pool, 2) - second_failed = threading.Event() + recorded = threading.Event() + record = _S3Batch._record + + def recording(batch, error): + record(batch, error) + if error is not None: + recorded.set() def call(i): if i == 0: - assert second_failed.wait(EVENT_TIMEOUT) + # Fails once the failure of the second item is recorded. + assert recorded.wait(EVENT_TIMEOUT) raise ValueError("first item") - second_failed.set() raise KeyError("second item") try: - with pytest.raises(KeyError, match="second item"): - scope.map(call, range(2)) + with ( + mock.patch.object(_S3Batch, "_record", recording), + pytest.raises(KeyError, match="second item"), + ): + call_bounded(scope.map, call, range(2)) finally: shutdown_bounded(scope) pool.shutdown() @@ -1253,24 +1279,31 @@ def test_map_counts_failures(self): # failure reported is the first that counts in the order they ended. pool = S3ThreadPool(3) scope = _S3PoolScope(pool, 3) - ended = {i: threading.Event() for i in range(3)} + recorded = {str(i): threading.Event() for i in range(3)} + record = _S3Batch._record + + def recording(batch, error): + record(batch, error) + if error is not None: + recorded[str(error)].set() def call(i): - # Item 0 ends first and does not count; item 2 ends before item 1. + # Item 0 fails first and does not count; item 2 fails before item 1. if i == 2: - assert ended[0].wait(EVENT_TIMEOUT) + assert recorded["0"].wait(EVENT_TIMEOUT) if i == 1: - assert ended[2].wait(EVENT_TIMEOUT) - ended[i].set() + assert recorded["2"].wait(EVENT_TIMEOUT) raise (LookupError if i == 0 else ValueError)(i) try: - result = scope.map( - call, - range(3), - stop_on_failure=False, - counts=lambda e: not isinstance(e, LookupError), - ) + with mock.patch.object(_S3Batch, "_record", recording): + result = call_bounded( + scope.map, + call, + range(3), + stop_on_failure=False, + counts=lambda e: not isinstance(e, LookupError), + ) assert [type(e) for e in result.outcomes] == [LookupError, ValueError, ValueError] assert result.failure is result.outcomes[2] finally: @@ -1306,7 +1339,7 @@ def interrupted_wait(futures, **kwargs): mock.patch("pyathena.filesystem.s3_executor.wait", interrupted_wait), pytest.raises(KeyboardInterrupt), ): - scope.map(call, range(2), stop_on_failure=False) + call_bounded(scope.map, call, range(2), stop_on_failure=False) assert time.monotonic() - started < EVENT_TIMEOUT assert ran == [0] assert operation.calls == 1 @@ -1338,7 +1371,7 @@ def interrupted_wait(futures, **kwargs): mock.patch("pyathena.filesystem.s3_executor.wait", interrupted_wait), pytest.raises(KeyboardInterrupt) as info, ): - scope.map(call, range(1)) + call_bounded(scope.map, call, range(1)) assert isinstance(info.value.__context__, KeyboardInterrupt) assert not release.is_set() assert len(scope.pending()) == 1 @@ -1347,6 +1380,34 @@ def interrupted_wait(futures, **kwargs): shutdown_bounded(scope) pool.shutdown() + def test_map_result_frees_outcomes(self): + # GH-1332 (F5): the outcomes of a map are freed once its result is + # dropped, without the cyclic garbage collector, so that the bodies of + # parallel reads are not kept until it runs. + class Body: + pass + + refs = [] + pool = S3ThreadPool(2) + scope = _S3PoolScope(pool, 2) + + def call(i): + body = Body() + refs.append(weakref.ref(body)) + return body + + gc.disable() + try: + result = call_bounded(scope.map, call, range(2)) + assert len(result.outcomes) == 2 + del result + # The threads leave the frames of the calls right after they end. + poll_until(lambda: all(ref() is None for ref in refs), "an outcome was kept") + finally: + gc.enable() + shutdown_bounded(scope) + pool.shutdown() + def test_map_failure_frees_outcomes(self): # GH-1332 (F5): the outcome of a call that ends while a failed map # waits for it is not kept by the traceback of the failure. @@ -1373,7 +1434,7 @@ def call(i): try: with pytest.raises(ValueError, match="failed") as info: - scope.map(call, range(2)) + call_bounded(scope.map, call, range(2)) gc.collect() assert len(refs) == 1 assert refs[0]() is None @@ -1448,25 +1509,30 @@ def call(i): @pytest.mark.asyncio async def test_map_async_failures_are_retrieved(self): - # GH-1332: the failures of the calls that a failure stopped are not - # logged as never retrieved. + # GH-1332: the failures of the calls that a failure stopped, which end + # while the map drains, are not logged as never retrieved. loop = asyncio.get_running_loop() logged = [] previous = loop.get_exception_handler() loop.set_exception_handler(lambda loop, context: logged.append(context)) pool = S3ThreadPool(2) scope = _S3PoolScope(pool, 2) - barrier = threading.Barrier(2, timeout=EVENT_TIMEOUT) + both_started = threading.Barrier(2, timeout=EVENT_TIMEOUT) + # The second call fails once the first has stopped its retries. + throttled = retried(ThrottledOperation(), RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT)) def call(i): - barrier.wait() - raise ValueError(i) + both_started.wait() + if i == 0: + raise ValueError(i) + return throttled() try: - with pytest.raises(ValueError, match=r"^[01]$"): + with pytest.raises(ValueError, match=r"^0$"): await scope.map_async(call, range(2)) - gc.collect() - await asyncio.sleep(0) + for _ in range(3): + gc.collect() + await asyncio.sleep(0) assert logged == [] finally: loop.set_exception_handler(previous) diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index 5871aa124..a74e9656a 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -517,6 +517,40 @@ def mock_connection(max_pool_connections=10): return connection +def call_bounded(fn, *args, timeout=EVENT_TIMEOUT, **kwargs): + """Call a function in a daemon thread, failing if it does not return within the timeout. + + A regression that blocks the call fails the test instead of hanging it. + + Args: + fn: The function. + *args: Positional arguments passed to the function. + timeout: The longest call, in seconds. + **kwargs: Keyword arguments passed to the function. + + Returns: + The return value of the function. + + Raises: + BaseException: The error that the function raised. + """ + outcome = {} + + def run(): + try: + outcome["result"] = fn(*args, **kwargs) + except BaseException as e: + outcome["error"] = e + + runner = threading.Thread(target=run, daemon=True) + runner.start() + runner.join(timeout) + assert not runner.is_alive(), f"{fn!r} did not return within {timeout} s" + if "error" in outcome: + raise outcome.pop("error") + return outcome.get("result") + + def run_bounded(main, timeout=EVENT_TIMEOUT): """Run a coroutine with ``asyncio.run()``, failing if it does not end within the timeout. From e0f0bfa03807c022c2c728a45e600ecc6efe6ae9 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 7 Nov 2026 05:12:46 +0900 Subject: [PATCH 6/8] Close the remaining gaps in failure recording of a map (GH-1332) From the follow-up of the independent review: - An error of the counts predicate in a done callback is the failure of the map instead of escaping the callback, which skipped the callbacks after it, including the one that wakes an asyncio map. - A future wakes its waiters before it runs its callbacks, so the waiting thread also records the futures that it sees end; result() records only finished futures, since exception() waits for one that has not ended, which would have kept a stopped batch waiting for its queued calls. - result() drops the futures itself, so that an interrupt after it returns finds none kept by the batch, instead of a finally after it. - Tests: the rejection test runs the other calls in the submitting thread, so the call after the rejection starts before the map waits; test_map_failure_cancels_queued_calls and test_map_async_counts_error_in_callback are new; the counts test orders its failures so that the first to end is not the last item; the opcode injection test is bounded. Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3_executor.py | 42 +++++++-- tests/pyathena/filesystem/test_s3_executor.py | 91 +++++++++++++++++-- 2 files changed, 113 insertions(+), 20 deletions(-) diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index 9c61a9665..398cd8613 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -226,7 +226,11 @@ def wait(self) -> None: return_when = FIRST_EXCEPTION if self._stop_on_failure else ALL_COMPLETED pending = set(self.futures) while pending and not self._stopped_by_failure(): - pending = wait(pending, return_when=return_when).not_done + done, pending = wait(pending, return_when=return_when) + # Also here, since a future wakes its waiters before it runs the + # callback that records it. + for future in done: + self._settled(future) async def wait_async(self) -> None: """Await until every call has ended or a failure has stopped the batch.""" @@ -236,9 +240,18 @@ async def wait_async(self) -> None: pending = set(self._waiters.values()) while pending and not self._stopped_by_failure(): pending = (await asyncio.wait(pending, return_when=return_when))[1] + for future in self.futures: + if future.done(): + self._settled(future) def result(self) -> S3MapResult[T]: - """Return the outcomes of the calls, which have all ended. + """Return the outcomes of the calls, which have all ended, and drop their futures. + + The futures are dropped here, not by the map, so that an interrupt + after this method returns finds no future kept by the batch; their + callbacks refer to the batch, which would otherwise form a reference + cycle with them that keeps the outcomes until the cyclic garbage + collector frees them. Returns: The outcomes. @@ -246,12 +259,18 @@ def result(self) -> S3MapResult[T]: Raises: BaseException: The failure that stopped the batch. """ + for future in self.futures: + if future.done(): + # exception() would wait for a call that has not ended. + self._settled(future) if self._stopped_by_failure(): raise cast(BaseException, self._failure) - return S3MapResult( + result = S3MapResult( [e if (e := self._outcome(f)) is not None else f.result() for f in self.futures], self._failure, ) + self.clear() + return result @staticmethod def _outcome(future: Future[Any]) -> BaseException | None: @@ -382,10 +401,17 @@ def _record(self, error: BaseException | None) -> None: Args: error: The failure of a call, or None if it succeeded. """ - if error is None or ( - isinstance(error, Exception) and self._counts is not None and not self._counts(error) - ): + if error is None: return + if isinstance(error, Exception) and self._counts is not None: + try: + if not self._counts(error): + return + except BaseException as e: + # A done callback must not raise, which would skip the + # callbacks after it; the error of the predicate is the + # failure. + error = e with self._lock: if self._failure is None: self._failure = error @@ -616,8 +642,6 @@ def map( finally: batch.release(e) raise - finally: - batch.clear() async def map_async( self, @@ -659,8 +683,6 @@ async def map_async( finally: batch.release(e) raise - finally: - batch.clear() def __enter__(self) -> S3Executor: return self diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index 0c8c3db52..eb5178d7d 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -33,6 +33,7 @@ EVENT_TIMEOUT, OpcodeInterrupter, ThrottledOperation, + await_bounded, call_bounded, forking, interrupt_each_opcode, @@ -1159,16 +1160,20 @@ def test_map_executor_rejection(self, rejection): # before the next call starts, or is the failure of a batch without # stop_on_failure, instead of being returned as a value. class Rejecting(S3Executor): + # Runs the other calls in the submitting thread, before submit() + # returns, so that the call after the rejection starts before the + # map waits. def __init__(self): - self._pool = ThreadPoolExecutor(2) self.calls = 0 @override def submit(self, fn, *args, **kwargs): self.calls += 1 - if self.calls != 2: - return self._pool.submit(contextvars.copy_context().run, fn, *args, **kwargs) future: Future[Any] = Future() + if self.calls != 2: + future.set_running_or_notify_cancel() + future.set_result(contextvars.copy_context().run(fn, *args, **kwargs)) + return future if rejection == "cancelled": future.cancel() else: @@ -1178,7 +1183,7 @@ def submit(self, fn, *args, **kwargs): @override def shutdown(self, wait=True): - self._pool.shutdown(wait=wait) + pass expected = CancelledError if rejection == "cancelled" else RuntimeError ran = [] @@ -1196,6 +1201,38 @@ def call(i): assert isinstance(result.failure, expected) assert isinstance(result.outcomes[1], expected) + def test_map_failure_cancels_queued_calls(self): + # GH-1332 (B2): a failure cancels the calls that have not started + # without waiting for them to start, also when the executor would + # never start them. + class Holding(S3Executor): + # Runs the first call, and holds the others until they are + # cancelled. + def __init__(self): + self._pool = ThreadPoolExecutor(1) + self.held: list[Future[Any]] = [] + self.calls = 0 + + @override + def submit(self, fn, *args, **kwargs): + self.calls += 1 + if self.calls > 1: + self.held.append(Future()) + return self.held[-1] + return self._pool.submit(contextvars.copy_context().run, fn, *args, **kwargs) + + @override + def shutdown(self, wait=True): + self._pool.shutdown(wait=wait) + + def call(i): + raise ValueError(i) + + with Holding() as executor: + with pytest.raises(ValueError, match=r"^0$"): + call_bounded(executor.map, call, range(2)) + assert all(future.cancelled() for future in executor.held) + def test_map_failure_stops_batch(self): # GH-1332 (B2): a failure starts no further call, stops the retries # of the running one, whose backoff wait would take at least @@ -1288,11 +1325,11 @@ def recording(batch, error): recorded[str(error)].set() def call(i): - # Item 0 fails first and does not count; item 2 fails before item 1. - if i == 2: - assert recorded["0"].wait(EVENT_TIMEOUT) + # Item 0 fails first and does not count; item 1 fails before item 2. if i == 1: - assert recorded["2"].wait(EVENT_TIMEOUT) + assert recorded["0"].wait(EVENT_TIMEOUT) + if i == 2: + assert recorded["1"].wait(EVENT_TIMEOUT) raise (LookupError if i == 0 else ValueError)(i) try: @@ -1305,7 +1342,7 @@ def call(i): counts=lambda e: not isinstance(e, LookupError), ) assert [type(e) for e in result.outcomes] == [LookupError, ValueError, ValueError] - assert result.failure is result.outcomes[2] + assert result.failure is result.outcomes[1] finally: shutdown_bounded(scope) pool.shutdown() @@ -1472,7 +1509,8 @@ def call(i): shutdown_bounded(scope) pool.shutdown() - assert interrupt_each_opcode(S3Executor.map, case) > 0 + # The tracer is installed in the thread that runs the case. + assert call_bounded(interrupt_each_opcode, S3Executor.map, case, timeout=120) > 0 @pytest.mark.asyncio async def test_map_async_cancelled(self): @@ -1507,6 +1545,39 @@ def call(i): await asyncio.to_thread(shutdown_bounded, scope) pool.shutdown() + @pytest.mark.asyncio + async def test_map_async_counts_error_in_callback(self): + # GH-1332: an error of the counts predicate in the done callback that + # records a failure the executor reports later, in another thread, is + # the failure of the map. It used to skip the callbacks after it, + # including the one that wakes the map, which then waited forever. + class Deferring(S3Executor): + @override + def submit(self, fn, *args, **kwargs): + future: Future[Any] = Future() + + def reject(): + future.set_running_or_notify_cancel() + future.set_exception(RuntimeError("rejected")) + + threading.Timer(0.05, reject).start() + return future + + @override + def shutdown(self, wait=True): + pass + + class Abort(BaseException): + # Not an Exception, which concurrent.futures would only log; not + # SystemExit, which asyncio raises out of the event loop. + pass + + def counts(error): + raise Abort("predicate") + + with pytest.raises(Abort, match="predicate"): + await await_bounded(Deferring().map_async(str, range(1), counts=counts)) + @pytest.mark.asyncio async def test_map_async_failures_are_retrieved(self): # GH-1332: the failures of the calls that a failure stopped, which end From 13860dfbdd9fb3b001a8562fc3a4ce8829394f09 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 7 Nov 2026 06:13:34 +0900 Subject: [PATCH 7/8] Wait for one future that ends a map (GH-1332) From the second follow-up of the independent review: - The waiting thread waits for one future, which the done callbacks of the calls set when every call has settled or when a failure stops the batch, with the states checked every second. A cancellation that the executor reported after the wait began did not wake a wait for the first exception, so a call held until it was cancelled kept the map waiting forever. It also replaces the repeated waits on every pending future. - Only the done callback catches an error of the counts predicate; one in the waiting thread, such as an interrupt, is raised. - release() also clears the frames of an interrupt that stopped the drain, whose traceback keeps the futures of the drain's wait. - Tests: deferred cancellation (sync and asyncio), an interrupt in the predicate, a failure seen before its callback runs, the predicate's error not escaping into the thread that sets the future, and retention in the opcode injection over map(). The interrupted-read tests make their reads wait until the wait is entered, since the wait is skipped when every call has ended. Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3_executor.py | 121 +++++++++----- tests/pyathena/filesystem/test_s3.py | 11 ++ tests/pyathena/filesystem/test_s3_executor.py | 158 ++++++++++++++++-- 3 files changed, 236 insertions(+), 54 deletions(-) diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index 398cd8613..172e50c9a 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -12,6 +12,7 @@ import asyncio import contextvars import os +import sys import threading import traceback import warnings @@ -19,7 +20,7 @@ from abc import ABCMeta, abstractmethod from collections import deque from collections.abc import Callable, Iterable -from concurrent.futures import ALL_COMPLETED, FIRST_EXCEPTION, CancelledError, Future, wait +from concurrent.futures import CancelledError, Future, wait from concurrent.futures.thread import ThreadPoolExecutor from dataclasses import dataclass from multiprocessing import cpu_count @@ -198,12 +199,16 @@ def __init__( self.stop = threading.Event() self._failure: BaseException | None = None self._lock = threading.Lock() + # Set when every call has settled, or when a failure stops the + # batch, so that the waiting thread waits for one future, which also + # wakes it for a failure that is a cancellation. + self._ended: Future[None] = Future() + self._remaining = 0 + self._submitting = True # The futures of the calls, each tracked before it is scheduled, so # that an interrupt while it is being scheduled still lets it be # cancelled. self.futures: list[Future[T]] = [] - # The asyncio wrappers of the futures, by future. - self._waiters: dict[Future[T], asyncio.Future[T]] = {} def submit(self, executor: S3Executor, items: Iterable[ItemT]) -> None: """Submit a call for each item. @@ -214,35 +219,34 @@ def submit(self, executor: S3Executor, items: Iterable[ItemT]) -> None: """ with stop_retries_on(self.stop): for item in items: + with self._lock: + self._remaining += 1 self.futures.append(Future()) - # Records a failure that the executor reports without running - # the call, such as a failure to start a thread, when it is - # set, before the thread that sets it takes another call. + # Records the outcome when the future is set, before the + # thread that sets it takes another call, also of a call that + # the executor fails or cancels without running it. self.futures[-1].add_done_callback(self._settled) executor.submit_to(self.futures[-1], self._call, item) + with self._lock: + self._submitting = False + ended = self._remaining == 0 + if ended: + self._end() def wait(self) -> None: - """Wait until every call has ended or a failure has stopped the batch.""" - return_when = FIRST_EXCEPTION if self._stop_on_failure else ALL_COMPLETED - pending = set(self.futures) - while pending and not self._stopped_by_failure(): - done, pending = wait(pending, return_when=return_when) - # Also here, since a future wakes its waiters before it runs the - # callback that records it. - for future in done: - self._settled(future) + """Wait until every call has ended or a failure has stopped the batch. + + The states are also checked every second, since an interrupt of a + cancellation can skip the callbacks that end the wait. + """ + while not self._ended.done() and not all(f.done() for f in self.futures): + wait([self._ended], timeout=1.0) async def wait_async(self) -> None: """Await until every call has ended or a failure has stopped the batch.""" - return_when = asyncio.FIRST_EXCEPTION if self._stop_on_failure else asyncio.ALL_COMPLETED - for future in self.futures: - self._waiters[future] = S3Executor.waiter(future) - pending = set(self._waiters.values()) - while pending and not self._stopped_by_failure(): - pending = (await asyncio.wait(pending, return_when=return_when))[1] - for future in self.futures: - if future.done(): - self._settled(future) + ended = S3Executor.waiter(self._ended) + while not self._ended.done() and not all(f.done() for f in self.futures): + await asyncio.wait([ended], timeout=1.0) def result(self) -> S3MapResult[T]: """Return the outcomes of the calls, which have all ended, and drop their futures. @@ -261,8 +265,10 @@ def result(self) -> S3MapResult[T]: """ for future in self.futures: if future.done(): - # exception() would wait for a call that has not ended. - self._settled(future) + # Also here, since a future wakes its waiters before it runs + # the callback that records it; exception() would wait for a + # call that has not ended. + self._record(self._outcome(future)) if self._stopped_by_failure(): raise cast(BaseException, self._failure) result = S3MapResult( @@ -302,8 +308,7 @@ async def drain_async(self) -> None: """Stop the batch, cancel the calls that have not started, and await the others.""" started = self._stop() try: - waiters = [self._waiters.get(f) or S3Executor.waiter(f) for f in started] - self._waiters.clear() + waiters = [S3Executor.waiter(future) for future in started] while not all(future.done() for future in started): await asyncio.wait(waiters, timeout=1.0) finally: @@ -321,6 +326,11 @@ def release(self, error: BaseException) -> None: """ try: traceback.clear_frames(error.__traceback__) + # A further interrupt that stopped the drain propagates instead; + # the frames of the drain, which keep the futures, are in its + # traceback. + if (propagating := sys.exception()) is not None and propagating is not error: + traceback.clear_frames(propagating.__traceback__) finally: # Also if a further interrupt stopped the drain before its stop. self.stop.set() @@ -334,7 +344,6 @@ def clear(self) -> None: keep their outcomes until the cyclic garbage collector frees them. """ self.futures.clear() - self._waiters.clear() def _stop(self) -> list[Future[T]]: """Stop the batch and cancel the calls that have not started. @@ -384,12 +393,25 @@ def _call(self, item: ItemT) -> T: raise def _settled(self, future: Future[T]) -> None: - """Record the failure of a call whose future has been set. + """Record the failure of a call whose future has been set, as its done callback. + + A done callback must not raise, which would skip the callbacks after + it, such as the one that wakes an asyncio map, so an error of the + ``counts`` predicate is the failure. Args: future: The future of the call. """ - self._record(self._outcome(future)) + try: + self._record(self._outcome(future)) + except BaseException as e: + self._latch(e) + finally: + with self._lock: + self._remaining -= 1 + ended = not self._submitting and self._remaining == 0 + if ended: + self._end() def _record(self, error: BaseException | None) -> None: """Record the first failure that counts, and stop the batch with it if requested. @@ -403,20 +425,31 @@ def _record(self, error: BaseException | None) -> None: """ if error is None: return - if isinstance(error, Exception) and self._counts is not None: - try: - if not self._counts(error): - return - except BaseException as e: - # A done callback must not raise, which would skip the - # callbacks after it; the error of the predicate is the - # failure. - error = e + if isinstance(error, Exception) and self._counts is not None and not self._counts(error): + return + self._latch(error) + + def _latch(self, error: BaseException) -> None: + """Record a failure unless one was recorded, and stop the batch with it if requested. + + Args: + error: The failure. + """ + with self._lock: + if self._failure is not None: + return + self._failure = error + if self._stop_on_failure: + self.stop.set() + if self._stop_on_failure: + self._end() + + def _end(self) -> None: + """End the wait of the waiting thread, once.""" with self._lock: - if self._failure is None: - self._failure = error - if self._stop_on_failure: - self.stop.set() + if self._ended.done(): + return + self._ended.set_result(None) class S3Executor(metaclass=ABCMeta): diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index dd9a5306a..9cdd102fe 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -3781,8 +3781,11 @@ def test_read_interrupted_twice_frees_ranges(self): # context, also frees the data of the ranges that were read. The # frame of the wait keeps the futures, as concurrent.futures.wait() # does. + entered = threading.Event() + def interrupted_wait(futures, **kwargs): # A set of its own, as concurrent.futures.wait() makes. + entered.set() held = set(futures) wait(held) raise KeyboardInterrupt @@ -3794,11 +3797,14 @@ def case(interrupt): read = [] def get_object(path, range_, **kwargs): + # Ends once the read waits, so that its wait is interrupted. + assert entered.wait(EVENT_TIMEOUT) # Unlike bytes, an array can be weakly referenced. data = array.array("B", bytes(range_[1] - range_[0])) read.append(weakref.ref(data)) return data + entered.clear() fs.core.get_object = get_object f = fs.open("s3://bucket/key", "rb", cache_type="none", block_size=16, max_workers=4) try: @@ -3878,7 +3884,11 @@ def test_interrupted_read_frees_ranges(self): fs.info.return_value.size = 64 read = [] + entered = threading.Event() + def get_object(path, range_, **kwargs): + # Ends once the read waits, so that its wait is interrupted. + assert entered.wait(EVENT_TIMEOUT) # Unlike bytes, an array can be weakly referenced. data = array.array("B", bytes(range_[1] - range_[0])) read.append(weakref.ref(data)) @@ -3889,6 +3899,7 @@ def get_object(path, range_, **kwargs): def interrupted_wait(futures, **kwargs): # As the wait of concurrent.futures, whose frame keeps the futures # in a set, interrupted once the reads have finished. + entered.set() futures = set(futures) assert not wait(futures, timeout=EVENT_TIMEOUT).not_done raise KeyboardInterrupt diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index eb5178d7d..82143937e 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -1233,6 +1233,115 @@ def call(i): call_bounded(executor.map, call, range(2)) assert all(future.cancelled() for future in executor.held) + def test_map_deferred_cancellation_stops_batch(self): + # GH-1332 (B2): a call that the executor cancels after the map began + # waiting stops the batch, also while another call is held until it + # is cancelled; a cancellation does not wake a wait for the first + # exception. + class Holding(S3Executor): + def __init__(self): + self.held: list[Future[Any]] = [] + + @override + def submit(self, fn, *args, **kwargs): + self.held.append(Future()) + return self.held[-1] + + @override + def shutdown(self, wait=True): + pass + + executor = Holding() + threading.Timer(0.1, lambda: executor.held[0].cancel()).start() + with pytest.raises(CancelledError): + call_bounded(executor.map, str, range(2)) + assert executor.held[1].cancelled() + + def test_map_records_failure_before_callbacks_run(self): + # GH-1332: a future wakes its waiters before it runs its callbacks, so + # the map records a failure that it sees, also before the callback + # that would record it has run. + release = threading.Event() + + class LateCallbacks(Future): + def _invoke_callbacks(self): + def invoke(): + release.wait(EVENT_TIMEOUT) + Future._invoke_callbacks(self) + + threading.Thread(target=invoke, daemon=True).start() + + class Rejecting(S3Executor): + @override + def submit(self, fn, *args, **kwargs): + future: Future[Any] = Future() + future.set_running_or_notify_cancel() + future.set_exception(RuntimeError("cannot start a thread")) + return future + + @override + def shutdown(self, wait=True): + pass + + try: + with ( + mock.patch("pyathena.filesystem.s3_executor.Future", LateCallbacks), + pytest.raises(RuntimeError, match="cannot start"), + ): + call_bounded(Rejecting().map, str, range(1)) + finally: + release.set() + + def test_map_interrupted_in_counts(self): + # GH-1332 (B3): an interrupt of the waiting thread while it evaluates + # the counts predicate is raised, also after a failure was recorded. + pool = S3ThreadPool(1) + scope = _S3PoolScope(pool, 1) + waiting = [] + + def counts(error): + if threading.get_ident() in waiting: + raise KeyboardInterrupt + return True + + def call(i): + raise ValueError(i) + + def run(): + waiting.append(threading.get_ident()) + return scope.map(call, range(1), stop_on_failure=False, counts=counts) + + try: + with pytest.raises(KeyboardInterrupt): + call_bounded(run) + finally: + shutdown_bounded(scope) + pool.shutdown() + + @pytest.mark.asyncio + async def test_map_async_deferred_cancellation_stops_batch(self): + # GH-1332 (B2): as test_map_deferred_cancellation_stops_batch, for + # map_async(). + class Holding(S3Executor): + def __init__(self): + self.held: list[Future[Any]] = [] + + @override + def submit(self, fn, *args, **kwargs): + self.held.append(Future()) + return self.held[-1] + + @override + def shutdown(self, wait=True): + pass + + executor = Holding() + threading.Timer(0.1, lambda: executor.held[0].cancel()).start() + # The failure of the call, not a cancellation of the waiting task. + with pytest.raises(CancelledError): + await await_bounded(executor.map_async(str, range(2))) + assert executor.held[1].cancelled() + def test_map_failure_stops_batch(self): # GH-1332 (B2): a failure starts no further call, stops the retries # of the running one, whose backoff wait would take at least @@ -1481,30 +1590,40 @@ def call(i): pool.shutdown() def test_map_interrupted_at_each_opcode(self): - # GH-1332 (B6, S1, S5): an interrupt at any opcode of map() leaves no - # call in the scope, and every call that it started has ended when it - # raises. + # GH-1332 (B6, S1, S5, F5): an interrupt at any opcode of map() leaves + # no call in the scope, every call that it started has ended when it + # raises, and the outcomes are freed while the caller holds the + # interrupt. + class Body: + pass + def case(interrupt): pool = S3ThreadPool(2) scope = _S3PoolScope(pool, 2) running = [] + refs = [] lock = threading.Lock() def call(i): with lock: running.append(i) time.sleep(0.01) + body = Body() with lock: + refs.append(weakref.ref(body)) running.remove(i) - return i + return body try: result = interrupt(scope.map, call, range(4)) if not interrupt.fired: - assert result.outcomes == [0, 1, 2, 3] + assert len(result.outcomes) == 4 return + assert interrupt.error.__traceback__ is not None assert running == [] assert scope.pending() == [] + # The threads leave the frames of the calls right after they end. + poll_until(lambda: all(ref() is None for ref in refs), "an outcome was kept") finally: shutdown_bounded(scope) pool.shutdown() @@ -1549,18 +1668,32 @@ def call(i): async def test_map_async_counts_error_in_callback(self): # GH-1332: an error of the counts predicate in the done callback that # records a failure the executor reports later, in another thread, is - # the failure of the map. It used to skip the callbacks after it, - # including the one that wakes the map, which then waited forever. + # the failure of the map, and does not escape the callback into that + # thread, which would skip the callbacks after it. + # Rejects once the map has wrapped the future of the call for asyncio. + wrapped = threading.Event() + escaped = [] + waiter = S3Executor.waiter + + def watched_waiter(future): + result = waiter(future) + wrapped.set() + return result + class Deferring(S3Executor): @override def submit(self, fn, *args, **kwargs): future: Future[Any] = Future() def reject(): + assert wrapped.wait(EVENT_TIMEOUT) future.set_running_or_notify_cancel() - future.set_exception(RuntimeError("rejected")) + try: + future.set_exception(RuntimeError("rejected")) + except BaseException as e: + escaped.append(e) - threading.Timer(0.05, reject).start() + threading.Thread(target=reject, daemon=True).start() return future @override @@ -1575,8 +1708,13 @@ class Abort(BaseException): def counts(error): raise Abort("predicate") - with pytest.raises(Abort, match="predicate"): + with ( + mock.patch.object(S3Executor, "waiter", staticmethod(watched_waiter)), + pytest.raises(Abort, match="predicate"), + ): await await_bounded(Deferring().map_async(str, range(1), counts=counts)) + # Not raised in the thread that set the future either. + assert escaped == [] @pytest.mark.asyncio async def test_map_async_failures_are_retrieved(self): From 6b8f5f17e879323a683ac318cb66ea4c07701c90 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 7 Nov 2026 06:41:45 +0900 Subject: [PATCH 8/8] Resume the end of a map's wait after an interrupt in its signal (GH-1332) From the third follow-up of the independent review: - A later failure, or the same one recorded again by the callback that caught an interrupt of the signal, sets the stop and ends the wait again, and the wait also ends once a failure has stopped the batch, so that an interrupt in the signal does not keep the map waiting for a call held until it is cancelled. - Tests: the deferred cancellations happen once the map waits, not after a timer; the callback test checks the rejecting thread once its callbacks have returned; test_map_interrupted_while_ending_wait injects an interrupt at each opcode of the signal. Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3_executor.py | 30 ++++++-- tests/pyathena/filesystem/test_s3_executor.py | 77 +++++++++++++++++-- 2 files changed, 95 insertions(+), 12 deletions(-) diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index 172e50c9a..832d5ffa0 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -239,13 +239,13 @@ def wait(self) -> None: The states are also checked every second, since an interrupt of a cancellation can skip the callbacks that end the wait. """ - while not self._ended.done() and not all(f.done() for f in self.futures): + while not self._ended_or_settled(): wait([self._ended], timeout=1.0) async def wait_async(self) -> None: """Await until every call has ended or a failure has stopped the batch.""" ended = S3Executor.waiter(self._ended) - while not self._ended.done() and not all(f.done() for f in self.futures): + while not self._ended_or_settled(): await asyncio.wait([ended], timeout=1.0) def result(self) -> S3MapResult[T]: @@ -364,6 +364,22 @@ def _stop(self) -> list[Future[T]]: started.append(future) return started + def _ended_or_settled(self) -> bool: + """Return whether the wait of the waiting thread can end. + + Also checks the states that end it, in case an interrupt in a + callback stopped the signal. + + Returns: + Whether the wait has been ended, a failure has stopped the batch, + or every call has settled. + """ + return ( + self._ended.done() + or self._stopped_by_failure() + or all(future.done() for future in self.futures) + ) + def _stopped_by_failure(self) -> bool: """Return whether a failure has stopped the batch. @@ -436,12 +452,12 @@ def _latch(self, error: BaseException) -> None: error: The failure. """ with self._lock: - if self._failure is not None: - return - self._failure = error - if self._stop_on_failure: - self.stop.set() + if self._failure is None: + self._failure = error if self._stop_on_failure: + # Also for a later failure, which resumes a signal that an + # interrupt stopped. + self.stop.set() self._end() def _end(self) -> None: diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index 82143937e..6925f3767 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -1252,11 +1252,57 @@ def shutdown(self, wait=True): pass executor = Holding() - threading.Timer(0.1, lambda: executor.held[0].cancel()).start() - with pytest.raises(CancelledError): + waiting = threading.Event() + + def entered_wait(futures, **kwargs): + waiting.set() + return wait(futures, **kwargs) + + def cancel(): + # Once the map waits. + assert waiting.wait(EVENT_TIMEOUT) + executor.held[0].cancel() + + threading.Thread(target=cancel, daemon=True).start() + with ( + mock.patch("pyathena.filesystem.s3_executor.wait", entered_wait), + pytest.raises(CancelledError), + ): call_bounded(executor.map, str, range(2)) assert executor.held[1].cancelled() + def test_map_interrupted_while_ending_wait(self): + # GH-1332 (I1): an interrupt at any opcode of the signal that ends the + # wait, in a callback that records a failure, does not keep the map + # waiting forever, also while another call is held until it is + # cancelled. + class Holding(S3Executor): + # Rejects the first call in the submitting thread, and holds the + # others until they are cancelled. + def __init__(self): + self.held: list[Future[Any]] = [] + + @override + def submit(self, fn, *args, **kwargs): + future: Future[Any] = Future() + if not self.held: + future.set_running_or_notify_cancel() + future.set_exception(RuntimeError("cannot start a thread")) + self.held.append(future) + return future + + @override + def shutdown(self, wait=True): + pass + + def case(interrupt): + # The map raises the rejection, or the interrupt, which the + # callback records as the failure or which propagates from it. + with contextlib.suppress(RuntimeError): + interrupt(Holding().map, str, range(2)) + + assert call_bounded(interrupt_each_opcode, _S3Batch._end, case, timeout=120) > 0 + def test_map_records_failure_before_callbacks_run(self): # GH-1332: a future wakes its waiters before it runs its callbacks, so # the map records a failure that it sees, also before the callback @@ -1336,9 +1382,25 @@ def shutdown(self, wait=True): pass executor = Holding() - threading.Timer(0.1, lambda: executor.held[0].cancel()).start() + waiting = threading.Event() + waiter = S3Executor.waiter + + def watched_waiter(future): + result = waiter(future) + waiting.set() + return result + + def cancel(): + # Once the map has wrapped the futures that it waits for. + assert waiting.wait(EVENT_TIMEOUT) + executor.held[0].cancel() + + threading.Thread(target=cancel, daemon=True).start() # The failure of the call, not a cancellation of the waiting task. - with pytest.raises(CancelledError): + with ( + mock.patch.object(S3Executor, "waiter", staticmethod(watched_waiter)), + pytest.raises(CancelledError), + ): await await_bounded(executor.map_async(str, range(2))) assert executor.held[1].cancelled() @@ -1672,6 +1734,7 @@ async def test_map_async_counts_error_in_callback(self): # thread, which would skip the callbacks after it. # Rejects once the map has wrapped the future of the call for asyncio. wrapped = threading.Event() + rejected = threading.Event() escaped = [] waiter = S3Executor.waiter @@ -1692,6 +1755,8 @@ def reject(): future.set_exception(RuntimeError("rejected")) except BaseException as e: escaped.append(e) + finally: + rejected.set() threading.Thread(target=reject, daemon=True).start() return future @@ -1713,7 +1778,9 @@ def counts(error): pytest.raises(Abort, match="predicate"), ): await await_bounded(Deferring().map_async(str, range(1), counts=counts)) - # Not raised in the thread that set the future either. + # Not raised in the thread that set the future either, once its + # callbacks have returned. + assert await asyncio.to_thread(rejected.wait, EVENT_TIMEOUT) assert escaped == [] @pytest.mark.asyncio