diff --git a/docs/aio.md b/docs/aio.md index 37a43752a..57d7e66dd 100644 --- a/docs/aio.md +++ b/docs/aio.md @@ -287,6 +287,14 @@ The pools create threads as the calls need them, and each call runs with the cal context variables, as with `asyncio.to_thread()`. In a forked child process, a pool creates new threads. +Cancelling a coroutine of the filesystem also stops the retries of the S3 requests that +its call makes in these threads, including the part uploads and range reads that the +call runs in parallel: no further attempt is made, and the wait before the next attempt +ends. A request that has already been sent still runs to its end in its thread. The +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. + ### Executor strategy `S3FileSystem` and `S3File` use a pluggable executor abstraction (`S3Executor`) for @@ -319,18 +327,20 @@ async with await fs.open_async("s3://my-bucket/data/file.bin", "rb") as f: ``` - A read splits the requested range into ranges of the block size, and reads them in - parallel. The file has no read cache. + parallel. When a range fails, the reads that have not finished are cancelled, and the + running ones make no further attempt. The file has no read cache. - A write uploads each full block as a part of a multipart upload, with at most `max_workers` parts in flight. A write that would upload another part awaits one of them, so the event loop keeps running. A part found failed by that wait, or a part copied from the existing object in an append that has failed before the write uploads another part, stops the write with its error, and the multipart upload is aborted. A part that fails later raises its error when the file is closed. The commit cancels - the parts that have not started as soon as one fails, and waits for the running ones - before it aborts the upload and raises the error. + the parts that have not started as soon as one fails, and waits for the running ones, + which make no further attempt, before it aborts the upload and raises the error. - Closing the file commits the written data, as `S3File` does. - Leaving an `async with` block with an exception, or cancelling a write, closes the file - without committing it, and aborts the multipart upload after the running parts finish. + without committing it, and aborts the multipart upload after the running parts, which + make no further attempt, finish. The abort completes even if the task is cancelled again. - Once closing or the final flush has begun the commit, the commit completes even if the task that closes or flushes the file is cancelled, and leaving an `async with` diff --git a/docs/filesystem.md b/docs/filesystem.md index 892bce4d5..698f0afa0 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -138,9 +138,10 @@ written. If an annotation fails to copy, the error is raised and the destination kept. A failed part copy aborts the multipart upload. So does an interrupt, or the cancellation of an `AioS3FileSystem` copy, including one by the shutdown of the event loop at the end of `asyncio.run()`, unless the upload has already completed. A -CreateMultipartUpload request and the part copies in flight finish first, and so does -the CompleteMultipartUpload request of an `AioS3FileSystem` copy; the requests of an -`AioS3FileSystem` copy that have not started are cancelled. A further interrupt +CreateMultipartUpload request and the part copies in flight finish first, without a +further attempt. The CompleteMultipartUpload request of an `AioS3FileSystem` copy +also finishes first, with its retries; the requests of an `AioS3FileSystem` copy that +have not started are cancelled. A further interrupt while the CreateMultipartUpload request is waited for is ignored. One while the part copies are waited for or the upload is aborted stops the cleanup, and the upload ID is logged. @@ -262,6 +263,10 @@ fs.chmod("s3://YOUR_S3_BUCKET/path/to/object", "bucket-owner-full-control") 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. + 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 time. The copy keeps the system-defined metadata (such as `ContentType` and diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 919bafad9..dea436383 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -7,6 +7,7 @@ import logging import mimetypes import os.path +import threading import time import traceback from collections.abc import Callable, Mapping @@ -49,7 +50,7 @@ from pyathena.filesystem.s3_path import S3Path from pyathena.filesystem.s3_path_pairing import S3PathPairing from pyathena.filesystem.s3_writer import S3MultipartWriter -from pyathena.util import RetryConfig, override +from pyathena.util import RetryConfig, override, stop_retries_on _logger = logging.getLogger(__name__) @@ -1640,12 +1641,13 @@ def _copy_object_with_multipart_upload( annotations are copied onto the destination after the upload completes. A failed part or completion aborts the upload, and so does an interrupt, after the creation of the upload and the running part - copies have finished. Further interrupts while the creation is waited - for are suppressed; one while the part copies are waited for or the - upload is aborted stops the cleanup, and the upload ID is logged. A - failed annotation copy is raised and leaves the destination in place. - If HeadObject reports a size that fits in a single CopyObject request, - the reported version is copied with CopyObject instead. + copies, which then make no further attempt, have finished. Further + interrupts while the creation is waited for are suppressed; one while + the part copies are waited for or the upload is aborted stops the + cleanup, and the upload ID is logged. A failed annotation copy is + raised and leaves the destination in place. If HeadObject reports a + size that fits in a single CopyObject request, the reported version is + copied with CopyObject instead. Args: source: Source S3 path, with the version ID to copy, if any. @@ -1676,26 +1678,31 @@ def _copy_object_with_multipart_upload( # for. creation: Future[S3MultipartUpload] = Future() futures: list[Future[S3MultipartUploadPart]] = [] + # Set on failure, so that the creation and the running part copies + # make no further attempt before the abort. + stop = threading.Event() try: - executor.submit_to( - creation, - self.core.create_multipart_upload, - plan.destination, - **plan.create_params, - ) + with stop_retries_on(stop): + executor.submit_to( + creation, + self.core.create_multipart_upload, + plan.destination, + **plan.create_params, + ) multipart_upload = creation.result() for i, range_ in enumerate(plan.ranges): future: Future[S3MultipartUploadPart] = Future() futures.append(future) - executor.submit_to( - future, - self.core.upload_part_copy, - upload=multipart_upload, - part_number=i + 1, - source=plan.source, - range_=range_, - **plan.part_params, - ) + with stop_retries_on(stop): + executor.submit_to( + future, + self.core.upload_part_copy, + upload=multipart_upload, + part_number=i + 1, + source=plan.source, + range_=range_, + **plan.part_params, + ) completed = self._finish_multipart_upload( upload=multipart_upload, futures=futures, @@ -1707,6 +1714,12 @@ def _copy_object_with_multipart_upload( abort=False, ) except BaseException: + # The requests that have not started are cancelled before the + # stop, which would let a thread that it frees take one of them. + creation.cancel() + for future in futures: + future.cancel() + stop.set() if not creation.cancel(): while not creation.done(): try: @@ -2284,7 +2297,9 @@ def chmod(self, path: str, acl: str, recursive: bool = False, **kwargs) -> None: path: S3 path (s3://bucket or s3://bucket/key) to set the ACL on. acl: The value of the canned ACL to apply. recursive: Whether to apply the ACL to all keys below the given - path too. + path too. The keys are updated in parallel; once one fails, the + updates that have not started are cancelled, and the running + ones make no further attempt. **kwargs: Additional parameters passed to the PutObjectAcl or PutBucketAcl API. """ @@ -2296,13 +2311,32 @@ def chmod(self, path: str, acl: str, recursive: bool = False, **kwargs) -> None: if s3_path.key and acl not in self.core.OBJECT_ACLS: 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 = [ - executor.submit(self.chmod, p, acl, recursive=False, **kwargs) - for p in self.find(path, withdirs=False) - ] - for future in as_completed(futures): - future.result() + 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 if s3_path.key: # A key prefix is not an object itself; only the objects # below it have ACLs. @@ -3005,6 +3039,10 @@ def __init__( self.append_block = False self._multipart_writer: S3MultipartWriter | None = None self.multipart_upload_parts: list[Future[S3MultipartUploadPart]] = [] + # Set when the upload is discarded, so that its creation and the + # running part uploads and copies make no further attempt before the + # abort. + self._parts_stop = threading.Event() # The parts uploaded from the buffer that may not have finished. The # part copies of an append hold no data and are not counted. self._uploads_in_flight: set[Future[S3MultipartUploadPart]] = set() @@ -3059,6 +3097,14 @@ def _close_without_commit(self, *, creation: Future[S3MultipartUpload] | None = return self.buffer = None self.closed = True + # The requests that have not started are cancelled before the stop, + # which would let a thread that it frees take one of them, and the + # creation and the running parts make no further attempt. + if creation is not None: + creation.cancel() + for part in self.multipart_upload_parts: + part.cancel() + self._parts_stop.set() try: if creation is not None and not creation.cancel(): while not creation.done(): @@ -3146,9 +3192,10 @@ def _initiate_upload(self) -> None: # scheduled still lets it be cancelled or waited for. creation: Future[S3MultipartUpload] = Future() try: - self._executor.submit_to( - creation, writer.initiate, **self._get_request_kwargs("create_multipart_upload") - ) + with stop_retries_on(self._parts_stop): + self._executor.submit_to( + creation, writer.initiate, **self._get_request_kwargs("create_multipart_upload") + ) creation.result() if self.append_block: whole_source = self.tell() <= self.fs.core.MULTIPART_UPLOAD_MAX_PART_SIZE @@ -3167,14 +3214,15 @@ def _initiate_upload(self) -> None: part: Future[S3MultipartUploadPart] = Future() self.multipart_upload_parts.append(part) self._copies_in_flight.add(part) - self._executor.submit_to( - part, - writer.upload_part_copy, - part_number=part_number, - source=S3Path(self.bucket, self.key), - range_=None if whole_source else range_, - **self._get_request_kwargs("upload_part_copy"), - ) + with stop_retries_on(self._parts_stop): + self._executor.submit_to( + part, + writer.upload_part_copy, + part_number=part_number, + source=S3Path(self.bucket, self.key), + range_=None if whole_source else range_, + **self._get_request_kwargs("upload_part_copy"), + ) except BaseException: # A submitted creation may still return an ID after the wait is # interrupted. The writer records it before resolving the future. @@ -3256,13 +3304,14 @@ def _upload_chunk(self, final: bool = False) -> bool: part: Future[S3MultipartUploadPart] = Future() self.multipart_upload_parts.append(part) self._uploads_in_flight.add(part) - self._executor.submit_to( - part, - writer.upload_part, - part_number=part_number, - body=data, - **self._get_request_kwargs("upload_part"), - ) + with stop_retries_on(self._parts_stop): + self._executor.submit_to( + part, + writer.upload_part, + part_number=part_number, + body=data, + **self._get_request_kwargs("upload_part"), + ) except BaseException: self._close_without_commit() raise @@ -3336,13 +3385,17 @@ def discard(self) -> None: """Abort the multipart upload, if any. The part uploads that have not started are cancelled, and the - running ones are waited for before the abort. + running ones make no further attempt and are waited for before the + abort, whose own retries are not stopped. """ if self.multipart_upload: # A part that is still uploading when the upload is aborted may # be stored after the abort, so wait for the parts that could not - # be cancelled first. - wait([f for f in self.multipart_upload_parts if not f.cancel()]) + # be cancelled first. They are cancelled before the stop, which + # would let a thread that it frees take one of them. + running = [f for f in self.multipart_upload_parts if not f.cancel()] + self._parts_stop.set() + wait(running) self._get_multipart_writer().abort(**self._get_request_kwargs("abort_multipart_upload")) self.multipart_upload = None @@ -3413,7 +3466,8 @@ def _fetch_range(self, start: int, end: int) -> bytes: 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 and its error is raised. + those that have not started are cancelled, the running ones make no + further attempt, and its error is raised. Args: start: The offset of the first byte to read. @@ -3431,11 +3485,20 @@ 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: - futures = [ - self._executor.submit(self.fs.core.get_object, path, r, **request_kwargs) - for r in ranges - ] + # Set on failure, so that the running reads make no further attempt. + stop = threading.Event() + futures: list[Future[bytes]] = [] 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 except BaseException as e: # The reads are cancelled with the read of the range, as below, @@ -3444,6 +3507,9 @@ def _fetch_range(self, start: int, end: int) -> bytes: traceback.clear_frames(e.__traceback__) while futures: futures.pop().cancel() + # After the cancellation, which a thread that it frees could + # otherwise precede by taking a queued read. + stop.set() raise failed = next( (f for f in done if not f.cancelled() and f.exception() is not None), None @@ -3454,6 +3520,7 @@ def _fetch_range(self, start: int, end: int) -> bytes: # the traceback of the error keeps, does not keep their data. while futures: futures.pop().cancel() + stop.set() done.clear() failed.result() return b"".join(f.result() for f in futures) diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 8b1f26b2a..20bc1d700 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -14,6 +14,7 @@ import logging import mimetypes import os +import threading import traceback import warnings from concurrent.futures import FIRST_EXCEPTION, Future, wait @@ -47,6 +48,7 @@ from pyathena.filesystem.s3_path import S3Path from pyathena.filesystem.s3_path_pairing import S3PathPairing from pyathena.filesystem.s3_writer import S3MultipartWriter +from pyathena.util import stop_retries_on if TYPE_CHECKING: from datetime import datetime @@ -582,10 +584,11 @@ async def _copy_object_with_multipart_upload( requests that have not started 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. - This also holds when the shutdown of the event loop, such as at the - end of ``asyncio.run()``, cancels the copy, also after its cleanup has - started. A repeated cancellation returns without stopping this - cleanup. + 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 ``asyncio.run()``, cancels the copy, also after its + cleanup has started. A repeated cancellation returns without stopping + this cleanup. Args: source: Source S3 path, with the version ID to copy, if any. @@ -617,6 +620,9 @@ async def _copy_object_with_multipart_upload( # The requests of the parts, which keep running when their tasks are # cancelled, such as by the shutdown of the event loop. part_copies: list[Future[S3MultipartUploadPart]] = [] + # Set on a failure or cancellation, so that the creation and the + # running part and annotation copies make no further attempt. + stop = threading.Event() async def _upload_part(i: int, range_: tuple[int, int]) -> S3MultipartUploadPart | None: nonlocal failed @@ -625,14 +631,15 @@ async def _upload_part(i: int, range_: tuple[int, int]) -> S3MultipartUploadPart # The upload is being aborted; do not start more parts. return None try: - part_copy = self._thread_pool.submit( - self.core.upload_part_copy, - upload=multipart_upload, - part_number=i + 1, - source=plan.source, - range_=range_, - **plan.part_params, - ) + with stop_retries_on(stop): + part_copy = self._thread_pool.submit( + self.core.upload_part_copy, + upload=multipart_upload, + part_number=i + 1, + source=plan.source, + range_=range_, + **plan.part_params, + ) part_copies.append(part_copy) copied = asyncio.wrap_future(part_copy) return await copied @@ -668,9 +675,10 @@ def _abort() -> None: # which the shutdown of the event loop would cancel while the # requests keep running; scheduled right before the try, so that # nothing can fail between the two. - creation = self._thread_pool.submit( - self.core.create_multipart_upload, plan.destination, **plan.create_params - ) + with stop_retries_on(stop): + creation = self._thread_pool.submit( + self.core.create_multipart_upload, plan.destination, **plan.create_params + ) try: # A cancellation of this task leaves the creation to the cleanup, # which waits for an upload that a started creation creates. @@ -703,12 +711,14 @@ def _abort() -> None: # does on an interrupt. failed = True # The requests that have not started are cancelled here, before - # the cleanup waits for a thread after them. + # the cleanup waits for a thread after them, and before the stop, + # which would let a thread that it frees take one of them. creation.cancel() for part_copy in part_copies: part_copy.cancel() if completion is not None: completion.cancel() + stop.set() # A repeated cancellation of this task returns without stopping # the cleanup. await self._thread_pool.run_shielded(_abort) @@ -716,16 +726,26 @@ def _abort() -> None: failed = False - async def _copy_annotation(name: str) -> None: - nonlocal failed - async with semaphore: - if failed: - # Do not start more copies after one failed, as - # S3FileSystem does. - return - try: - await self._thread_pool.run( - self.core.copy_object_annotation, + # 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 filesystem until a copy failed. + return + try: + with stop_retries_on(stop): + self.core.copy_object_annotation( name, plan.source, plan.destination, @@ -733,13 +753,35 @@ async def _copy_annotation(name: str) -> None: 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: + nonlocal failed + async with semaphore: + if failed: + # Do not start more copies after one failed, as + # S3FileSystem does. + return + try: + await self._thread_pool.run(_copy_annotation_in_thread, i, name) except Exception: + # The running copies make no further attempt. failed = True + stop.set() raise results = await asyncio.gather( - *[_copy_annotation(name) for name in plan.annotations], return_exceptions=True + *[_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 @@ -1288,6 +1330,9 @@ def __init__( self._committing = False # The futures of the part uploads and copies, in part-number order. self._parts: list[Future[S3MultipartUploadPart]] = [] + # Set by _discard(), so that the creation and the running part uploads + # and copies make no further attempt before the abort. + self._parts_stop = threading.Event() # The parts uploaded from the buffer that may not have finished. The # part copies of an append hold no data and are not counted. self._uploads_in_flight: set[asyncio.Future[S3MultipartUploadPart]] = set() @@ -1434,6 +1479,15 @@ async def _close_without_commit(self) -> None: self._pool.shutdown(wait=False) return self.buffer = None + # Here, so that the creation and the running parts stop even while + # they hold every thread of the file, ahead of _discard(). The + # requests that have not started are cancelled before the stop, which + # would let a thread that it frees take one of them. + if self._creation is not None: + self._creation.cancel() + for part in self._parts: + part.cancel() + self._parts_stop.set() try: await self._pool.run_shielded(self._discard) finally: @@ -1447,11 +1501,12 @@ async def _initiate_upload(self) -> None: # being scheduled still lets the cleanup cancel it or wait for it. self._creation = Future() try: - self._pool.submit_to( - self._creation, - writer.initiate, - **self._get_request_kwargs("create_multipart_upload"), - ) + with stop_retries_on(self._parts_stop): + self._pool.submit_to( + self._creation, + writer.initiate, + **self._get_request_kwargs("create_multipart_upload"), + ) await asyncio.wrap_future(self._creation) if self.append_block: size = self._append_size @@ -1463,14 +1518,15 @@ async def _initiate_upload(self) -> None: copy: Future[S3MultipartUploadPart] = Future() self._parts.append(copy) self._copies_in_flight.add(copy) - self._pool.submit_to( - copy, - writer.upload_part_copy, - part_number=part_number, - source=source, - range_=None if whole_source else range_, - **copy_kwargs, - ) + with stop_retries_on(self._parts_stop): + self._pool.submit_to( + copy, + writer.upload_part_copy, + part_number=part_number, + source=source, + range_=None if whole_source else range_, + **copy_kwargs, + ) except BaseException: await self._close_without_commit() raise @@ -1535,13 +1591,14 @@ async def _upload_chunk(self, final: bool = False) -> bool: # here so that asyncio does not log it as never retrieved. waiter.add_done_callback(lambda w: w.cancelled() or w.exception()) self._uploads_in_flight.add(waiter) - self._pool.submit_to( - part, - writer.upload_part, - part_number=part_number, - body=data, - **upload_kwargs, - ) + with stop_retries_on(self._parts_stop): + self._pool.submit_to( + part, + writer.upload_part, + part_number=part_number, + body=data, + **upload_kwargs, + ) except BaseException: await self._close_without_commit() raise @@ -1605,16 +1662,23 @@ def _discard(self) -> None: """Abort the multipart upload, if any; blocks, and runs in a thread. The requests that have not started are cancelled, and the running - ones are waited for first, since a part that is still uploading when - the upload is aborted may be stored after the abort. A failure is + ones make no further attempt and are waited for first, since a part + that is still uploading when the upload is aborted may be stored + after the abort. The abort's own retries are not stopped. A failure is logged instead of raised, so that it does not mask the error that the caller is handling, also when no task waits for the abort any more. """ + # The requests that have not started are cancelled before the stop, + # which would let a thread that it frees take one of them, and the + # creation and the running parts make no further attempt. + creating = self._creation is not None and not self._creation.cancel() + running = [part for part in self._parts if not part.cancel()] + self._parts_stop.set() try: - if self._creation is not None and not self._creation.cancel(): - wait([self._creation]) + if creating: + wait([cast(Future[S3MultipartUpload], self._creation)]) if self._writer is not None and self._writer.upload is not None: - wait([part for part in self._parts if not part.cancel()]) + wait(running) self._writer.abort(**self._get_request_kwargs("abort_multipart_upload")) except Exception: _logger.exception( @@ -1626,8 +1690,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 in the threads of the file. When a read fails, - the reads that have not finished are cancelled and its error is - raised, as :meth:`S3File._fetch_range` does. + 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. Args: start: The offset of the first byte to read. diff --git a/pyathena/filesystem/s3_core.py b/pyathena/filesystem/s3_core.py index 73410d3fd..737d2e066 100644 --- a/pyathena/filesystem/s3_core.py +++ b/pyathena/filesystem/s3_core.py @@ -32,7 +32,7 @@ S3PutObject, ) from pyathena.filesystem.s3_path import S3Path -from pyathena.util import RetryConfig, override, retry_api_call +from pyathena.util import RetryConfig, override, retry_api_call, unstoppable_retries _logger = logging.getLogger(__name__) @@ -1073,6 +1073,9 @@ def complete_multipart_upload( def abort_multipart_upload(self, upload: S3MultipartUpload, **params) -> None: """Abort a multipart upload with AbortMultipartUpload. + An abort cleans up, so the stop scopes of the caller do not stop its + retries (see :func:`~pyathena.util.unstoppable_retries`). + Args: upload: The multipart upload returned by creation or listing. **params: Additional request parameters. The fields that the @@ -1085,7 +1088,8 @@ def abort_multipart_upload(self, upload: S3MultipartUpload, **params) -> None: because it was completed or aborted. """ request = self._multipart_upload_request(upload) - self.call(self._client.abort_multipart_upload, **{**params, **request}) + with unstoppable_retries(): + self.call(self._client.abort_multipart_upload, **{**params, **request}) def part_ranges(self, size: int, block_size: int) -> list[tuple[int, int]]: """Split an object into the source ranges of the parts that copy it. diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index 425f87115..262bc3c1e 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -24,7 +24,7 @@ from multiprocessing import cpu_count from typing import Any, ClassVar, Generic, TypeVar, cast -from pyathena.util import override +from pyathena.util import override, stop_retries_on T = TypeVar("T") @@ -122,11 +122,19 @@ class S3Executor(metaclass=ABCMeta): work to run in parallel and for shutting down the executor when done. Both ``submit`` and ``shutdown`` mirror the ``concurrent.futures.Executor`` interface so that ``as_completed()`` and ``Future.cancel()`` work unchanged. + Each submitted callable runs with a copy of the context variables of its + submitter, so that the stop scopes of :func:`~pyathena.util.stop_retries_on` + reach the requests that it makes. """ @abstractmethod def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]: - """Submit a callable for execution and return a Future.""" + """Submit a callable for execution and return a Future. + + An implementation runs the callable in a copy of the submitter's + context, such as with ``contextvars.copy_context().run``, so that the + stop scopes of the submitter reach it. + """ ... @abstractmethod @@ -186,7 +194,8 @@ def __init__(self, max_workers: int) -> None: @override def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]: - return self._executor.submit(fn, *args, **kwargs) + context = contextvars.copy_context() + return self._executor.submit(context.run, fn, *args, **kwargs) @override def shutdown(self, wait: bool = True) -> None: @@ -269,8 +278,10 @@ def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]: async def run(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> T: """Run a blocking call in a thread of the pool and wait for it. - Cancelling the waiting task cancels a call that has not started; a - running call keeps running, as with ``asyncio.to_thread()``. + Cancelling the waiting task cancels a call that has not started, and + stops the retries of the AWS API calls that a running call makes, as + :func:`~pyathena.aio.util.run_in_thread` does; the call itself keeps + running until it returns or raises. Args: fn: The blocking callable. @@ -283,14 +294,37 @@ async def run(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> T: Raises: RuntimeError: If the pool has been shut down. """ - return await asyncio.wrap_future(self.submit(fn, *args, **kwargs)) + stop = threading.Event() + # Orders the end of the call with a cancellation, after which work + # that the call left running, such as the parts of a write in a + # transaction, is not stopped. + lock = threading.Lock() + returned = False + + def call() -> T: + """Run the call in a scope that the cancellation stops until it returns.""" + nonlocal returned + try: + with stop_retries_on(stop): + return fn(*args, **kwargs) + finally: + with lock: + returned = True + + try: + return await asyncio.wrap_future(self.submit(call)) + except asyncio.CancelledError: + with lock: + if not returned: + stop.set() + raise async def run_shielded(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> T: """Run a blocking call that a cancellation of the waiting task does not stop. The call runs to its end, even if it has not started when the waiting - task is cancelled; its outcome is then dropped. See - :meth:`wait_shielded`. + task is cancelled, and its retries are not stopped; its outcome is then + dropped. See :meth:`wait_shielded`. Args: fn: The blocking callable. @@ -432,7 +466,10 @@ def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]: with self._lock: if self._shutdown: raise RuntimeError("cannot schedule new futures after shutdown") - task = asyncio.run_coroutine_threadsafe(self._pool.run(item.run), self._loop) + # The task runs in a copy of the submitter's context, which + # call_soon_threadsafe() schedules it with, and so does the + # function in the thread of the pool. + task = asyncio.run_coroutine_threadsafe(self._dispatch(item.run), self._loop) task.add_done_callback(item.settle) return future raise RuntimeError( @@ -440,6 +477,18 @@ def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]: "Use S3ThreadPoolExecutor for synchronous usage." ) + async def _dispatch(self, run: Callable[[], None]) -> None: + """Hand a submitted function to the pool and wait for it. + + Unlike :meth:`S3AioThreadPool.run`, a cancellation of this task, such + as by the shutdown of the event loop, does not stop the retries of a + function that has started, which the caller still waits for. + + Args: + run: The function of the work item. + """ + await asyncio.wrap_future(self._pool.submit(run)) + async def _shutdown_pool(self) -> None: """Shut down the pool once the functions submitted before have reached it. diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 4b0dffc35..91b378bb8 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -40,19 +40,22 @@ import pyathena from pyathena.filesystem import register_s3_filesystem from pyathena.filesystem.s3 import CompressedBuffer, S3File, S3FileSystem -from pyathena.filesystem.s3_core import S3Core, S3DeleteBatch +from pyathena.filesystem.s3_core import S3Core, S3DeleteBatch, S3MultipartCopyPlan from pyathena.filesystem.s3_errors import S3ClientError from pyathena.filesystem.s3_executor import S3AioExecutor, S3Executor, S3ThreadPoolExecutor from pyathena.filesystem.s3_object import S3MultipartUpload, S3Object, S3ObjectType, S3StorageClass from pyathena.filesystem.s3_path import S3Path from pyathena.filesystem.s3_path_pairing import S3PathPairing -from pyathena.util import RetryConfig +from pyathena.util import RetryConfig, stop_retries_on from tests import ENV from tests.pyathena.conftest import connect from tests.pyathena.util import ( + EVENT_TIMEOUT, MULTIPART_COPY_BLOCK_SIZE, MULTIPART_COPY_KWARGS, MULTIPART_COPY_SIZE, + ThrottledOperation, + retried, stub_multipart_copy, ) @@ -2133,6 +2136,54 @@ def test_copy_object_with_multipart_upload_failed_part(self): self._multipart_copy(fs, **MULTIPART_COPY_KWARGS) stubber.assert_no_pending_responses() + def test_copy_object_with_multipart_upload_failed_part_stops_running_parts(self): + # GH-1203: a failed part copy stops the retries of the running ones + # before the upload is aborted; their backoff wait would take at least + # EVENT_TIMEOUT seconds. + fs = self._make_fs() + plan = S3MultipartCopyPlan( + source=S3Path("bucket", "src"), + destination=S3Path("bucket", "dst"), + size=3, + ranges=((0, 1), (1, 2), (2, 3)), + ) + fs.core.plan_multipart_copy = mock.MagicMock(return_value=plan) + fs.core.create_multipart_upload = mock.MagicMock( + return_value=S3MultipartUpload({"Bucket": "bucket", "Key": "dst", "UploadId": "id"}) + ) + fs.core.abort_multipart_upload = mock.MagicMock() + operation = ThrottledOperation(operation_name="UploadPartCopy") + throttled = retried(operation, RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT)) + + def upload_part_copy(part_number, **kwargs): + if part_number == 1: + assert operation.started.wait(EVENT_TIMEOUT) + raise PermissionError("Access Denied") + if part_number == 2: + return throttled() + return SimpleNamespace(etag='"e"', part_number=part_number) + + fs.core.upload_part_copy = upload_part_copy + started = time.monotonic() + with pytest.raises(PermissionError, match="Access Denied"): + fs._copy_object_with_multipart_upload(plan.source, plan.destination, max_workers=2) + assert time.monotonic() - started < EVENT_TIMEOUT + assert operation.calls == 1 + fs.core.abort_multipart_upload.assert_called_once() + + def test_abort_multipart_upload_in_stopped_scope(self): + # GH-1203: a stop scope of the caller, such as that of a cancelled + # AioS3FileSystem call, does not stop the retries of an abort. + fs = self._make_fs() + operation = ThrottledOperation(failures=1, operation_name="AbortMultipartUpload") + fs.core.call = retried(operation, RetryConfig(attempt=3, multiplier=0)) + stop = threading.Event() + stop.set() + upload = S3MultipartUpload({"Bucket": "bucket", "Key": "dst", "UploadId": "id"}) + with stop_retries_on(stop): + fs._abort_multipart_upload(upload, {}) + assert operation.calls == 2 + def test_copy_object_with_multipart_upload_failed_annotation(self): # GH-973: a failed annotation copy is raised; the completed # destination is neither aborted nor deleted. @@ -3460,6 +3511,78 @@ def get_object(path, range_, **kwargs): assert [ref() for ref in read] == [None] * 3 assert info.value.__traceback__ is not None + def test_failed_read_stops_running_reads(self): + # GH-1203: a failed parallel read stops the retries of the running + # reads, whose backoff wait would take at least EVENT_TIMEOUT seconds. + fs = self._make_fs() + fs.info = mock.MagicMock(return_value=self._file_object("key")) + fs.info.return_value.size = 32 + operation = ThrottledOperation() + throttled = retried(operation, RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT)) + + def get_object(path, range_, **kwargs): + if range_[0] == 0: + assert operation.started.wait(EVENT_TIMEOUT) + raise PermissionError("Access Denied") + return throttled() + + fs.core.get_object = get_object + f = fs.open("s3://bucket/key", "rb", cache_type="none", block_size=16, max_workers=2) + started = time.monotonic() + with pytest.raises(PermissionError, match="Access Denied"): + f.read() + # Waits for the running read. + f.close() + assert time.monotonic() - started < EVENT_TIMEOUT + assert operation.calls == 1 + + 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 + # seconds. + # The read that is queued when the interrupt comes is not sent. + fs = self._make_fs() + fs.info = mock.MagicMock(return_value=self._file_object("key")) + fs.info.return_value.size = 48 + operation = ThrottledOperation() + throttled = retried(operation, RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT)) + read = [] + + def get_object(path, range_, **kwargs): + read.append(range_[0]) + return throttled() + + fs.core.get_object = get_object + f = fs.open("s3://bucket/key", "rb", cache_type="none", block_size=16, max_workers=2) + executor = f._executor + submit = executor.submit + + submitted = [] + + def interrupted_submit(fn, *args, **kwargs): + submitted.append(fn) + if len(submitted) == 3: + # The third read waits for the two threads of the running reads. + deadline = time.monotonic() + EVENT_TIMEOUT + while len(read) < 2: + assert time.monotonic() < deadline + time.sleep(0.01) + future = submit(fn, *args, **kwargs) + if len(submitted) == 3: + # As an interrupt after the call is queued. + raise KeyboardInterrupt + return future + + executor.submit = interrupted_submit # type: ignore[method-assign] + started = time.monotonic() + with pytest.raises(KeyboardInterrupt): + f.read() + # Waits for the running reads. + f.close() + assert time.monotonic() - started < EVENT_TIMEOUT + assert sorted(read) == [0, 16] + assert operation.calls == 2 + def test_interrupted_read_frees_ranges(self): # GH-1186: an interrupted parallel read cancels the reads that have # not started and frees the data of the others, also while the @@ -4461,6 +4584,68 @@ def test_chmod_recursive(self): fs._client.put_object_acl, Bucket="bucket", Key="key2", ACL="private" ) + def test_chmod_recursive_failure_stops_running_requests(self): + # GH-1203: a failed ACL update stops the retries of the running ones, + # whose backoff wait would take at least EVENT_TIMEOUT seconds. + fs = self._make_fs() + fs.find = mock.MagicMock(return_value=["bucket/key1", "bucket/key2"]) + operation = ThrottledOperation(operation_name="PutObjectAcl") + throttled = retried(operation, RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT)) + + def put_object_acl(path, acl, **kwargs): + if path.key == "key1": + assert operation.started.wait(EVENT_TIMEOUT) + raise PermissionError("Access Denied") + return throttled() + + fs.core.put_object_acl = put_object_acl + started = time.monotonic() + with pytest.raises(PermissionError, match="Access Denied"): + fs.chmod("s3://bucket/path", "private", recursive=True) + assert time.monotonic() - started < EVENT_TIMEOUT + assert operation.calls == 1 + + def test_chmod_recursive_interrupted_submission_stops_running_requests(self): + # GH-1203: an interrupt while the updates are being submitted stops + # the running ones, whose backoff wait would take at least + # EVENT_TIMEOUT seconds. + # The update that is queued when the interrupt comes is not sent. + fs = self._make_fs() + fs.find = mock.MagicMock(return_value=["bucket/key1", "bucket/key2"]) + operation = ThrottledOperation(operation_name="PutObjectAcl") + throttled = retried(operation, RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT)) + updated = [] + + def put_object_acl(path, acl, **kwargs): + updated.append(path.key) + return throttled() + + fs.core.put_object_acl = put_object_acl + executor = S3ThreadPoolExecutor(max_workers=1) + submit = executor.submit + + submitted = [] + + def interrupted_submit(fn, *args, **kwargs): + submitted.append(fn) + if len(submitted) == 2: + # The second update waits for the only thread. + assert operation.started.wait(EVENT_TIMEOUT) + future = submit(fn, *args, **kwargs) + if len(submitted) == 2: + # As an interrupt after the call is queued. + raise KeyboardInterrupt + return future + + executor.submit = interrupted_submit # type: ignore[method-assign] + fs._create_executor = mock.MagicMock(return_value=executor) + started = time.monotonic() + with pytest.raises(KeyboardInterrupt): + fs.chmod("s3://bucket/path", "private", recursive=True) + assert time.monotonic() - started < EVENT_TIMEOUT + assert updated == ["key1"] + assert operation.calls == 1 + def test_put_tags_merge(self): fs = self._make_fs() fs._call.side_effect = [ @@ -6777,6 +6962,7 @@ def _make_write_file(data: bytes, autocommit: bool): file._multipart_writer = None file.multipart_upload = None file.multipart_upload_parts = [] + file._parts_stop = threading.Event() file.buffer = io.BytesIO(data) file.loc = len(data) # tell() returns self.loc # Not reached by these files, as the mocked executor never finishes @@ -7661,6 +7847,85 @@ def test_commit_failure_and_interrupted_abort(self): assert file.multipart_upload is None assert file.multipart_upload_parts == [] + def test_discard_stops_running_parts(self): + # GH-1203: a failed part stops the retries of the running parts before + # the upload is aborted; their backoff wait would take at least + # EVENT_TIMEOUT seconds. + fs = TestS3FileSystem._make_fs() + fs.core.create_multipart_upload = mock.MagicMock( + return_value=S3MultipartUpload({"Bucket": "bucket", "Key": "key", "UploadId": "id"}) + ) + fs.core.abort_multipart_upload = mock.MagicMock() + operation = ThrottledOperation(operation_name="UploadPart") + throttled = retried(operation, RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT)) + + def upload_part(part_number, **kwargs): + if part_number == 1: + assert operation.started.wait(EVENT_TIMEOUT) + raise PermissionError("Access Denied") + return throttled() + + fs.core.upload_part = upload_part + block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE + f = fs.open( + "s3://bucket/key", "wb", block_size=block_size, cache_type="none", max_workers=2 + ) + f.write(b"x" * (2 * block_size)) + started = time.monotonic() + with pytest.raises(PermissionError, match="Access Denied"): + f.close() + assert time.monotonic() - started < EVENT_TIMEOUT + assert operation.calls == 1 + fs.core.abort_multipart_upload.assert_called_once() + + def test_interrupted_creation_stops_retries(self): + # GH-1203: an interrupted write stops the retries of the creation of + # its multipart upload, which the cleanup waits for; their backoff + # wait would take at least EVENT_TIMEOUT seconds. + fs = TestS3FileSystem._make_fs() + operation = ThrottledOperation(operation_name="CreateMultipartUpload") + fs.core.create_multipart_upload = retried( + operation, RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT) + ) + executor = S3ThreadPoolExecutor(max_workers=1) + submit_to = executor.submit_to + + def submit_creation(future, fn, *args, **kwargs): + submit_to(future, fn, *args, **kwargs) + + def interrupted(timeout=None): + # As SIGINT interrupts the writer's wait. + assert operation.started.wait(EVENT_TIMEOUT) + raise KeyboardInterrupt + + future.result = interrupted # type: ignore[method-assign] + + executor.submit_to = submit_creation # type: ignore[method-assign] + block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE + file = S3File(fs, "s3://bucket/key", mode="wb", block_size=block_size, executor=executor) + started = time.monotonic() + with pytest.raises(KeyboardInterrupt): + file.write(b"x" * block_size) + assert time.monotonic() - started < EVENT_TIMEOUT + assert operation.calls == 1 + assert file.closed + + def test_discard_in_stopped_scope(self): + # GH-1203: a stop scope of the caller does not stop the retries of the + # abort of discard(). + file = self._make_write_file(b"", autocommit=False) + file.multipart_upload = S3MultipartUpload( + {"Bucket": "bucket", "Key": "key.txt", "UploadId": "uploadid"} + ) + operation = ThrottledOperation(failures=1, operation_name="AbortMultipartUpload") + file.fs.core.call = retried(operation, RetryConfig(attempt=3, multiplier=0)) + stop = threading.Event() + stop.set() + with stop_retries_on(stop): + file.discard() + assert operation.calls == 2 + assert file.multipart_upload is None + def test_discard_on_event_loop_thread(self): # GH-976: the parts that have not started are cancelled and not # waited for, so a rollback on the thread of the event loop that diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index a18f82524..7bb1f4a7b 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -26,7 +26,7 @@ from pyathena.filesystem.s3 import S3File, S3FileSystem from pyathena.filesystem.s3_async import AioS3File, AioS3FileSystem, AioS3StreamedFile -from pyathena.filesystem.s3_core import S3Core +from pyathena.filesystem.s3_core import S3Core, S3MultipartCopyPlan from pyathena.filesystem.s3_executor import S3AioThreadPool, S3Executor from pyathena.filesystem.s3_object import ( S3MultipartUpload, @@ -36,14 +36,16 @@ S3StorageClass, ) from pyathena.filesystem.s3_path import S3Path -from pyathena.util import RetryConfig +from pyathena.util import RetryConfig, stop_retries_on from tests import ENV from tests.pyathena.conftest import connect from tests.pyathena.util import ( EVENT_TIMEOUT, MULTIPART_COPY_BLOCK_SIZE, MULTIPART_COPY_KWARGS, + ThrottledOperation, overlapping_calls, + retried, stub_multipart_copy, ) @@ -678,6 +680,184 @@ def record_completion(fn, *args, **kwargs): assert events == ["abort"] + @pytest.mark.asyncio + async def test_cancelled_call_stops_retries(self): + # GH-1203: cancelling a call stops the retries of its request in the + # thread of the filesystem, whose backoff wait would take at least + # EVENT_TIMEOUT seconds. + fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + operation = ThrottledOperation(block=True, operation_name="HeadObject") + fs._sync_fs.info = retried(operation, RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT)) + task = asyncio.ensure_future(fs._info("s3://bucket/key")) + try: + assert await asyncio.to_thread(operation.started.wait, EVENT_TIMEOUT) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + finally: + operation.release.set() + started = time.monotonic() + # Waits for the thread of the cancelled call. + await asyncio.to_thread(fs._thread_pool.shutdown, True) + assert time.monotonic() - started < EVENT_TIMEOUT + assert operation.calls == 1 + + @pytest.mark.asyncio + async def test_copy_object_with_multipart_upload_failed_part_stops_running_parts(self): + # GH-1203: a failed part copy stops the retries of the running ones + # before the upload is aborted; their backoff wait would take at least + # EVENT_TIMEOUT seconds. + fs = AioS3FileSystem(connection=mock.MagicMock(), max_workers=4, skip_instance_cache=True) + sync_fs = fs._sync_fs + plan = S3MultipartCopyPlan( + source=S3Path("bucket", "src"), + destination=S3Path("bucket", "dst"), + size=3, + ranges=((0, 1), (1, 2), (2, 3)), + ) + 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._abort_multipart_upload = mock.MagicMock() + operation = ThrottledOperation(operation_name="UploadPartCopy") + throttled = retried(operation, RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT)) + + def upload_part_copy(part_number, **kwargs): + if part_number == 1: + assert operation.started.wait(EVENT_TIMEOUT) + raise PermissionError("Access Denied") + if part_number == 2: + return throttled() + return SimpleNamespace(etag='"e"', part_number=part_number) + + sync_fs.core.upload_part_copy = mock.MagicMock(side_effect=upload_part_copy) + started = time.monotonic() + with pytest.raises(PermissionError, match="Access Denied"): + await fs._copy_object_with_multipart_upload( + plan.source, plan.destination, max_workers=2 + ) + assert time.monotonic() - started < EVENT_TIMEOUT + assert operation.calls == 1 + sync_fs._abort_multipart_upload.assert_called_once() + + @pytest.mark.asyncio + async def test_copy_object_with_multipart_upload_cancelled_creation_stops_retries(self): + # GH-1203: cancelling a copy stops the retries of the creation of its + # upload, which the cleanup waits for; their backoff wait would take at + # least EVENT_TIMEOUT seconds. + fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + sync_fs = fs._sync_fs + plan = S3MultipartCopyPlan( + source=S3Path("bucket", "src"), + destination=S3Path("bucket", "dst"), + size=2, + ranges=((0, 1), (1, 2)), + ) + sync_fs.core.plan_multipart_copy = mock.MagicMock(return_value=plan) + operation = ThrottledOperation(block=True, operation_name="CreateMultipartUpload") + sync_fs.core.create_multipart_upload = retried( + operation, RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT) + ) + sync_fs._abort_multipart_upload = mock.MagicMock() + task = asyncio.ensure_future( + fs._copy_object_with_multipart_upload(plan.source, plan.destination) + ) + try: + assert await asyncio.to_thread(operation.started.wait, EVENT_TIMEOUT) + task.cancel() + finally: + operation.release.set() + started = time.monotonic() + with pytest.raises(asyncio.CancelledError): + await task + assert time.monotonic() - started < EVENT_TIMEOUT + assert operation.calls == 1 + # No upload was created. + sync_fs._abort_multipart_upload.assert_not_called() + + @pytest.mark.parametrize("failing", ["a1", "a2"]) + @pytest.mark.asyncio + async def test_copy_object_with_multipart_upload_failed_annotation_stops_running_copies( + self, failing + ): + # GH-1203: a failed annotation copy stops the retries of the running + # ones, whose backoff wait would take at least EVENT_TIMEOUT seconds, + # and its error is raised, whichever annotation comes first. + fs = AioS3FileSystem(connection=mock.MagicMock(), 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", "a2"), + ) + 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"') + ) + operation = ThrottledOperation(operation_name="PutObjectAnnotation") + throttled = retried(operation, RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT)) + + def copy_object_annotation(name, *args, **kwargs): + if name == failing: + assert operation.started.wait(EVENT_TIMEOUT) + raise PermissionError("Access Denied") + return throttled() + + sync_fs.core.copy_object_annotation = mock.MagicMock(side_effect=copy_object_annotation) + started = time.monotonic() + with pytest.raises(PermissionError, match="Access Denied"): + await fs._copy_object_with_multipart_upload(plan.source, plan.destination) + assert time.monotonic() - started < EVENT_TIMEOUT + assert operation.calls == 1 + + @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 + # when one fails are not sent. + fs = AioS3FileSystem(connection=mock.MagicMock(), max_workers=1, 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", "a2", "a3"), + ) + 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"') + ) + + def copy_object_annotation(name, *args, **kwargs): + # Fails once the other copies wait for the only thread. + deadline = time.monotonic() + EVENT_TIMEOUT + while fs._thread_pool._get_executor()._work_queue.qsize() < 2: + assert time.monotonic() < deadline + time.sleep(0.01) + raise PermissionError("Access Denied") + + sync_fs.core.copy_object_annotation = mock.MagicMock(side_effect=copy_object_annotation) + with pytest.raises(PermissionError, match="Access Denied"): + await fs._copy_object_with_multipart_upload( + plan.source, plan.destination, max_workers=3 + ) + sync_fs.core.copy_object_annotation.assert_called_once() + @pytest.mark.asyncio async def test_copy_object_with_multipart_upload_cancelled_queued_parts(self): # GH-1186: a cancellation cancels the part copies that wait for a @@ -3409,6 +3589,114 @@ def get_object(path, range_, **kwargs): assert [ref() for ref in read] == [None] * 3 assert info_.value.__traceback__ is not None + @pytest.mark.asyncio + async def test_failed_read_stops_running_reads(self, fs): + # GH-1203: a failed parallel read stops the retries of the running + # reads, whose backoff wait would take at least EVENT_TIMEOUT seconds. + block_size = 2**10 + operation = ThrottledOperation() + throttled = retried(operation, RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT)) + + def get_object(path, range_, **kwargs): + if range_[0] == 0: + assert operation.started.wait(EVENT_TIMEOUT) + raise PermissionError("Access Denied") + return throttled() + + fs.core.get_object = mock.MagicMock(side_effect=get_object) + info = S3Object( + init={"Key": "key", "ETag": '"etag"'}, + type=S3ObjectType.S3_OBJECT_TYPE_FILE, + bucket="bucket", + key="key", + ) + info.size = block_size * 2 + fs._sync_fs.info = mock.MagicMock(return_value=info) + f = await fs.open_async("s3://bucket/key", "rb", block_size=block_size, max_workers=2) + started = time.monotonic() + with pytest.raises(PermissionError, match="Access Denied"): + await f.read() + await f.close() + # Waits for the thread of the running read. + await asyncio.to_thread(f._pool.shutdown, True) + assert time.monotonic() - started < EVENT_TIMEOUT + assert operation.calls == 1 + + @pytest.mark.asyncio + async def test_failed_part_stops_running_parts(self, fs): + # GH-1203: a failed part stops the retries of the running parts before + # the upload is aborted; their backoff wait would take at least + # EVENT_TIMEOUT seconds. + operation = ThrottledOperation(operation_name="UploadPart") + throttled = retried(operation, RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT)) + + def upload_part(part_number, **kwargs): + if part_number == 1: + assert operation.started.wait(EVENT_TIMEOUT) + raise PermissionError("Access Denied") + return throttled() + + fs.core.upload_part = mock.MagicMock(side_effect=upload_part) + block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE + f = await fs.open_async("s3://bucket/key", "wb", block_size=block_size, max_workers=2) + await f.write(b"x" * (2 * block_size)) + started = time.monotonic() + with pytest.raises(PermissionError, match="Access Denied"): + await f.close() + assert time.monotonic() - started < EVENT_TIMEOUT + assert operation.calls == 1 + fs.core.abort_multipart_upload.assert_called_once() + + @pytest.mark.asyncio + async def test_cancelled_creation_stops_retries(self, fs): + # GH-1203: cancelling a write stops the retries of the creation of its + # multipart upload, which the cleanup waits for; their backoff wait + # would take at least EVENT_TIMEOUT seconds. + operation = ThrottledOperation(block=True, operation_name="CreateMultipartUpload") + fs.core.create_multipart_upload = retried( + operation, RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT) + ) + block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE + # The creation holds the only thread, after which the cleanup waits. + f = await fs.open_async("s3://bucket/key", "wb", block_size=block_size, max_workers=1) + task = asyncio.ensure_future(f.write(b"x" * block_size)) + try: + assert await asyncio.to_thread(operation.started.wait, EVENT_TIMEOUT) + task.cancel() + finally: + operation.release.set() + started = time.monotonic() + with pytest.raises(asyncio.CancelledError): + await task + assert time.monotonic() - started < EVENT_TIMEOUT + assert operation.calls == 1 + assert f.closed + + @pytest.mark.asyncio + async def test_discard_in_stopped_scope(self, fs): + # GH-1203: a stop scope does not stop the retries of the abort of + # _discard(). + block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE + f = await fs.open_async("s3://bucket/key", "wb", block_size=block_size) + await f.write(b"x" * block_size) + await f.flush() + # The abort of the core, whose request is stubbed. + del fs.core.abort_multipart_upload + operation = ThrottledOperation(failures=1, operation_name="AbortMultipartUpload") + fs.core.call = retried(operation, RetryConfig(attempt=3, multiplier=0)) + stop = threading.Event() + stop.set() + + def discard(): + with stop_retries_on(stop): + f._discard() + + await asyncio.to_thread(discard) + assert operation.calls == 2 + # Nothing is left to abort. + await f._close_without_commit() + assert operation.calls == 2 + @pytest.mark.asyncio async def test_cancelled_read_frees_ranges(self, fs): # GH-1186: a cancelled parallel read cancels the reads that have not diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index 65bee13ef..601d0d048 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -17,9 +17,11 @@ from unittest import mock import pytest +from botocore.exceptions import ClientError from pyathena.filesystem.s3_executor import S3AioExecutor, S3AioThreadPool, S3ThreadPoolExecutor -from tests.pyathena.util import EVENT_TIMEOUT, overlapping_calls +from pyathena.util import RetryConfig, retry_api_call, stop_retries_on +from tests.pyathena.util import EVENT_TIMEOUT, ThrottledOperation, overlapping_calls class TestS3AioThreadPool: @@ -58,6 +60,83 @@ async def main(): assert asyncio.run(main()) == "caller" pool.shutdown() + @pytest.mark.parametrize("shielded", [False, True]) + def test_run_cancelled_while_retrying(self, shielded): + # GH-1203: cancelling the waiting task stops the retries of a running + # call, unless it runs shielded. Without the stop, the backoff wait + # would take at least EVENT_TIMEOUT seconds. + operation = ThrottledOperation(block=True, failures=None if not shielded else 2) + finished = threading.Event() + pool = S3AioThreadPool(1) + + def work(): + try: + config = RetryConfig(attempt=3, multiplier=0 if shielded else EVENT_TIMEOUT) + return retry_api_call(operation, config) + finally: + finished.set() + + async def main(): + run = pool.run_shielded if shielded else pool.run + task = asyncio.ensure_future(run(work)) + await asyncio.to_thread(operation.started.wait, EVENT_TIMEOUT) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + operation.release.set() + started = time.monotonic() + assert await asyncio.to_thread(finished.wait, EVENT_TIMEOUT) + return time.monotonic() - started + + try: + assert asyncio.run(main()) < EVENT_TIMEOUT + finally: + operation.release.set() + pool.shutdown() + + assert operation.calls == (3 if shielded else 1) + + def test_run_cancelled_after_return(self): + # GH-1203: a cancellation that arrives after the call has returned does + # not stop the work that the call left running, such as the parts of a + # write in a transaction. + operation = ThrottledOperation(block=True, failures=1) + returned = threading.Event() + left = [] + pool = S3AioThreadPool(1) + + def work(): + context = contextvars.copy_context() + thread = threading.Thread( + target=context.run, + args=(retry_api_call, operation, RetryConfig(attempt=2, multiplier=0)), + ) + thread.start() + left.append(thread) + returned.set() + + async def main(): + task = asyncio.ensure_future(pool.run(work)) + await asyncio.sleep(0) + # Blocks the event loop, so that the task is cancelled before it + # resumes after the call has returned, which follows the event at + # once. + assert returned.wait(EVENT_TIMEOUT) + time.sleep(0.1) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + try: + asyncio.run(main()) + operation.release.set() + left[0].join(EVENT_TIMEOUT) + finally: + operation.release.set() + pool.shutdown() + + assert operation.calls == 2 + @pytest.mark.parametrize("shielded", [False, True]) def test_run_cancelled_before_start(self, shielded): # Cancelling the waiting task drops a call that has not started, @@ -193,6 +272,36 @@ async def main(): with pytest.raises(ValueError, match="invalid literal"): failed.result() + def test_submit_context(self): + # GH-1203: a function runs with the context of its submitter, as the + # task that dispatches it does, so that the stop scope of the + # submitter reaches its retries. + variable: contextvars.ContextVar[str] = contextvars.ContextVar("variable") + stop = threading.Event() + stop.set() + operation = ThrottledOperation() + + def work(): + with pytest.raises(ClientError): + retry_api_call(operation, RetryConfig(attempt=3, multiplier=0)) + return variable.get() + + async def main(): + executor = S3AioExecutor(loop=asyncio.get_running_loop()) + + def submit(): + variable.set("submitter") + with stop_retries_on(stop): + return executor.submit(work).result(EVENT_TIMEOUT) + + variable.set("loop") + result = await asyncio.to_thread(submit) + executor.shutdown() + return result + + assert asyncio.run(main()) == "submitter" + assert operation.calls == 1 + def test_submit_beyond_default_executor(self): # GH-1159: max_workers functions run at once, whatever the size of the # default executor of the event loop. They used to run in it. @@ -433,6 +542,27 @@ async def main(): class TestS3Executor: + @pytest.mark.parametrize("executor_class", [S3ThreadPoolExecutor, S3AioThreadPool]) + def test_submit_context(self, executor_class): + # GH-1203: a callable runs with a copy of the context of its submitter, + # so that the stop scope of the submitter reaches its retries. + variable: contextvars.ContextVar[str] = contextvars.ContextVar("variable") + stop = threading.Event() + stop.set() + operation = ThrottledOperation() + + def work(): + with pytest.raises(ClientError): + retry_api_call(operation, RetryConfig(attempt=3, multiplier=0)) + return variable.get() + + variable.set("submitter") + with executor_class(max_workers=1) as executor, stop_retries_on(stop): + future: Future[str] = Future() + executor.submit_to(future, work) + assert future.result(EVENT_TIMEOUT) == "submitter" + assert operation.calls == 1 + def test_submit_to(self): # The callable's outcome resolves the future of the caller. with S3ThreadPoolExecutor(max_workers=1) as executor: diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index f34636435..95d4a053c 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -26,7 +26,7 @@ from pyathena.filesystem.s3_core import S3Core from pyathena.glue import GlueMetadataClient from pyathena.model import AthenaCalculationExecutionStatus, AthenaQueryExecution -from pyathena.util import RetryConfig +from pyathena.util import RetryConfig, retry_api_call _queries = Environment( loader=FileSystemLoader(Path(__file__).parents[1].resolve() / "resources" / "queries") @@ -133,6 +133,28 @@ def __call__(self, *args, **kwargs): return "success" +def retried(operation, config): + """Return a stub of an ``S3Core`` request method that retries an operation. + + The stub sends the operation with :func:`~pyathena.util.retry_api_call` + in the thread that calls it, as ``S3Core.call()`` does, so that the stop + scopes of that thread reach its retries. + + Args: + operation: The stubbed AWS API operation, such as a + :class:`ThrottledOperation`. + config: The retry policy. + + Returns: + The stub, which accepts and ignores any arguments. + """ + + def request(*args, **kwargs): + return retry_api_call(operation, config) + + return request + + class StoppingOperation: """A stubbed AWS API operation that always throttles and stops its retries.