diff --git a/docs/aio.md b/docs/aio.md index 5bf72b50b..cb71fc6e7 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,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 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. ### Executor strategy @@ -343,6 +347,12 @@ 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. 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/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..e799c39ec 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 collections.abc import Callable, Iterable, Mapping from copy import deepcopy from dataclasses import dataclass, field from datetime import datetime @@ -35,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, @@ -48,7 +53,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__) @@ -1154,6 +1159,35 @@ 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) + def _delete_objects(self, paths: list[str], max_workers: int | None = None, **kwargs) -> None: """Delete objects with DeleteObjects requests grouped by bucket. @@ -1172,11 +1206,15 @@ 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 - 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]) + # 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 def _delete_batches(paths: list[str], **kwargs) -> tuple[list[S3DeleteBatch], dict[str, Any]]: @@ -1923,11 +1961,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 @@ -2257,31 +2291,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) - # 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 + 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. @@ -2392,30 +2402,30 @@ 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 + # 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 + + @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. + """ + return not ( + isinstance(error, FileNotFoundError) and S3ClientError.code_of(error) == "NoSuchUpload" + ) def created(self, path: str) -> datetime: """Return the creation time of the path. @@ -3408,9 +3418,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 +3440,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..e3ef8f16f 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -15,9 +15,8 @@ import logging import mimetypes import os -import threading -import traceback 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 @@ -31,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 S3ThreadPool +from pyathena.filesystem.s3_executor import ItemT, S3Executor, S3MapResult, S3ThreadPool, T from pyathena.filesystem.s3_object import ( S3CompleteMultipartUpload, S3Metadata, @@ -43,7 +42,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 @@ -59,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. @@ -350,8 +350,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 +372,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`. @@ -385,15 +386,47 @@ 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. - 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, + 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 ) - self._sync_fs._raise_delete_errors(results) + 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. @@ -612,11 +645,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 @@ -721,63 +755,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) - # 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, + await self._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, ) - if first_failure is not None: - raise cast(BaseException, results[first_failure]) - for result in results: - if isinstance(result, BaseException): - raise result async def _find( self, @@ -1522,11 +1511,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. @@ -1770,8 +1757,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 +1781,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_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 a5f43b9ef..832d5ffa0 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -12,21 +12,24 @@ import asyncio import contextvars import os +import sys import threading import traceback import warnings 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 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]): @@ -63,6 +66,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.""" @@ -115,6 +148,326 @@ 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() + # 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]] = [] + + 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: + with self._lock: + self._remaining += 1 + self.futures.append(Future()) + # 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. + + 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_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_or_settled(): + 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. + + 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. + + Raises: + BaseException: The failure that stopped the batch. + """ + for future in self.futures: + if future.done(): + # 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( + [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: + """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. + + 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 = [S3Executor.waiter(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 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__) + # 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() + 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() + + 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 _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. + + 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: + self._record(e) + raise + + def _settled(self, future: Future[T]) -> None: + """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. + """ + 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. + + 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. + """ + if error is None: + return + 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 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: + """End the wait of the waiting thread, once.""" + with self._lock: + if self._ended.done(): + return + self._ended.set_result(None) + + class S3Executor(metaclass=ABCMeta): """Abstract executor for parallel S3 operations. @@ -221,14 +574,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. @@ -255,7 +635,103 @@ 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) + item.link(future, self.submit(item.run)) + + 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. + + 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. + 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: + batch.release(e) + 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: + batch.release(e) + 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..9cdd102fe 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, _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 @@ -61,6 +61,7 @@ poll_until, retried, run_in_forked_child, + shutdown_bounded, sigint_when, stub_multipart_copy, ) @@ -565,6 +566,85 @@ 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 + 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 recorded["gone"].wait(EVENT_TIMEOUT) + if name == "late": + assert recorded["early"].wait(EVENT_TIMEOUT) + 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 ( + 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): # The request threads invalidate the shared parent "bucket/dir" at # once. DirCache.pop() reads before it deletes, so the second delete @@ -3586,6 +3666,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 @@ -3677,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 @@ -3690,15 +3797,18 @@ 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: - 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 +3824,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(), @@ -3774,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)) @@ -3785,13 +3899,14 @@ 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 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..1ce42ca62 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,40 @@ 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, 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 = {"b1": threading.Event(), "b2": threading.Event()} def call(method, **request): started[request["Bucket"]].set() - release[request["Bucket"]].wait(10) + release[request["Bucket"]].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() + release["b1"].set() + await wait_until(lambda: "b1/a" not in fs.dircache) + assert "b2/b" in fs.dircache + assert not task.done() + release["b2"].set() + with pytest.raises(asyncio.CancelledError): + await await_bounded(task) + assert "b2/b" not in fs.dircache + finally: + for event in release.values(): + event.set() @pytest.mark.asyncio async def test_rm_maxdepth(self): @@ -2737,9 +2779,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 +2862,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 +3278,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 +3289,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..6925f3767 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -14,22 +14,32 @@ 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 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, + await_bounded, + call_bounded, forking, interrupt_each_opcode, overlapping_calls, + poll_until, + retried, run_in_forked_child, shutdown_bounded, sigint_when, @@ -1070,11 +1080,765 @@ 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 = 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 + 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 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 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): + # 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.calls = 0 + + @override + def submit(self, fn, *args, **kwargs): + self.calls += 1 + 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: + future.set_running_or_notify_cancel() + future.set_exception(RuntimeError("cannot start a thread")) + return future + + @override + def shutdown(self, wait=True): + pass + + expected = CancelledError if rejection == "cancelled" else RuntimeError + ran = [] + + def call(i): + ran.append(i) + return str(i) + + with Rejecting() as executor, pytest.raises(expected): + 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 = call_bounded(executor.map, str, range(3), stop_on_failure=False) + 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_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() + 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 + # 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() + 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 ( + 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() + + 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"), + ): + call_bounded(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) + 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: + # Fails once the failure of the second item is recorded. + assert recorded.wait(EVENT_TIMEOUT) + raise ValueError("first item") + raise KeyError("second item") + + try: + 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() + + 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) + 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 fails first and does not count; item 1 fails before item 2. + if i == 1: + assert recorded["0"].wait(EVENT_TIMEOUT) + if i == 2: + assert recorded["1"].wait(EVENT_TIMEOUT) + raise (LookupError if i == 0 else ValueError)(i) + + try: + 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[1] + 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), + ): + call_bounded(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, + ): + call_bounded(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_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. + 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: + call_bounded(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, 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 body + + try: + result = interrupt(scope.map, call, range(4)) + if not interrupt.fired: + 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() + + # 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): + # 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_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, 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() + rejected = 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() + try: + future.set_exception(RuntimeError("rejected")) + except BaseException as e: + escaped.append(e) + finally: + rejected.set() + + threading.Thread(target=reject, daemon=True).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 ( + 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, once its + # callbacks have returned. + assert await asyncio.to_thread(rejected.wait, EVENT_TIMEOUT) + assert escaped == [] + + @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 + # 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) + 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): + both_started.wait() + if i == 0: + raise ValueError(i) + return throttled() + + try: + with pytest.raises(ValueError, match=r"^0$"): + await scope.map_async(call, range(2)) + for _ in range(3): + 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() 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.