diff --git a/docs/filesystem.md b/docs/filesystem.md index 1855bb2f9..892bce4d5 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -136,9 +136,11 @@ from that version, which needs `s3:GetObjectVersion` and, to copy the tags, upload completes, so the destination exists without them until the last one is written. If an annotation fails to copy, the error is raised and the destination is kept. A failed part copy aborts the multipart upload. So does an interrupt, or the -cancellation of an `AioS3FileSystem` copy, unless the upload has already completed. A +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. A further interrupt +the CompleteMultipartUpload request of an `AioS3FileSystem` copy; 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. @@ -540,6 +542,7 @@ It drops the buffer and closes the file before waiting, so a later close or defe It suppresses Ctrl-C during the blocking recovery wait and re-raises the original error when cleanup finishes. An additional interruption outside that wait can stop cleanup; once the upload identity is known, `discard()` can retry the abort. An abort failure retains the upload identity for `discard()` to retry. +A failed multipart `commit()` also closes the file without its buffered data, so a later `commit()` neither completes that upload nor uploads the last part alone with PutObject. Cancelling an async buffered write through `asyncio.to_thread` still allows its writer thread to finish. ### Multipart requests diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 2ad7d1b1a..919bafad9 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -8,6 +8,7 @@ import mimetypes import os.path import time +import traceback from collections.abc import Callable, Mapping from concurrent.futures import FIRST_COMPLETED, FIRST_EXCEPTION, Future, as_completed, wait from copy import deepcopy @@ -3280,7 +3281,8 @@ def commit(self) -> None: as soon as it fails, without waiting for the parts before it. If the abort also fails, the upload is kept so that :meth:`discard` can abort it. Invalidates the cache of the path afterwards. Does nothing - for a file whose failed write dropped the written data. + for a file whose failed write or multipart commit dropped the written + data. Raises: FileExistsError: If an object was created at the path after the @@ -3289,9 +3291,9 @@ def commit(self) -> None: initialized. """ if self.buffer is None: - # _close_without_commit() dropped the written data. A multipart - # upload that it failed to abort is kept for discard(), not - # completed. + # _close_without_commit() or a failed commit dropped the written + # data. A multipart upload that they failed to abort is kept for + # discard(), not completed. return if self.tell() == 0: self.discard() @@ -3307,7 +3309,6 @@ def commit(self) -> None: if not self.multipart_upload: raise RuntimeError("Multipart upload is not initialized.") - upload_id = cast(str, self.multipart_upload.upload_id) try: done, _ = wait(self.multipart_upload_parts, return_when=FIRST_EXCEPTION) for future in done: @@ -3320,17 +3321,13 @@ def commit(self) -> None: parts, **self._get_request_kwargs("complete_multipart_upload") ) except BaseException: - # discard() keeps the upload if the abort fails or is - # interrupted, so that a later discard(), such as a - # transaction rollback, retries the abort. An abort failure + # Closed without the buffer, which holds only the data of the + # last part that a later commit() would otherwise upload alone + # with PutObject. discard() keeps the upload if the abort + # fails or is interrupted, so that a later discard(), such as + # a transaction rollback, retries the abort. An abort failure # is logged so that it does not mask the original error. - try: - self.discard() - except Exception: - _logger.exception( - f"Failed to abort multipart upload {upload_id} " - f"to s3://{self.bucket}/{self.key}." - ) + self._close_without_commit() raise self.fs.invalidate_cache(self.path) @@ -3414,7 +3411,9 @@ def _fetch_range(self, start: int, end: int) -> bytes: The range is clamped to the size of the object, since fsspec caches 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. + 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. Args: start: The offset of the first byte to read. @@ -3436,9 +3435,27 @@ def _fetch_range(self, start: int, end: int) -> bytes: self._executor.submit(self.fs.core.get_object, path, r, **request_kwargs) for r in ranges ] - for future in as_completed(futures): - # The failure of the first read to finish is raised. - future.result() + try: + done = wait(futures, return_when=FIRST_EXCEPTION).done + except BaseException as e: + # The reads are cancelled with the read of the range, as below, + # and the frames of the wait, which keep their futures, are + # cleared; clear_frames() skips this frame, which is executing. + traceback.clear_frames(e.__traceback__) + while futures: + futures.pop().cancel() + raise + failed = next( + (f for f in done if not f.cancelled() and f.exception() is not None), None + ) + if failed is not None: + # 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. + while futures: + futures.pop().cancel() + done.clear() + failed.result() return b"".join(f.result() for f in futures) return self.fs.core.get_object(path, ranges[0], **request_kwargs) diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 1bdee1556..8b1f26b2a 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -14,6 +14,7 @@ import logging import mimetypes import os +import traceback import warnings from concurrent.futures import FIRST_EXCEPTION, Future, wait from io import BytesIO @@ -53,8 +54,6 @@ from pyathena.connection import Connection _logger = logging.getLogger(__name__) -# The running cleanups of cancelled multipart copies. -_cleanup_tasks: set[asyncio.Future[None]] = set() class AioS3FileSystem(AsyncFileSystem): @@ -579,11 +578,14 @@ async def _copy_object_with_multipart_upload( See :meth:`S3FileSystem._copy_object_with_multipart_upload`. The part and annotation copies run in parallel as asyncio tasks, at most ``max_workers`` at once, in the threads of the filesystem, which bound - the requests of all its copies together. On a cancellation, the creation of the upload, - the running part copies and the completion are waited for, the - upload is aborted if it was created and has not completed, and the - cancellation is re-raised. A repeated cancellation returns without - stopping this cleanup. + the requests of all its copies together. On a cancellation, the + 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. Args: source: Source S3 path, with the version ID to copy, if any. @@ -612,6 +614,9 @@ async def _copy_object_with_multipart_upload( semaphore = asyncio.Semaphore(max_workers) failed = False + # 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]] = [] async def _upload_part(i: int, range_: tuple[int, int]) -> S3MultipartUploadPart | None: nonlocal failed @@ -620,7 +625,7 @@ async def _upload_part(i: int, range_: tuple[int, int]) -> S3MultipartUploadPart # The upload is being aborted; do not start more parts. return None try: - return await self._thread_pool.run( + part_copy = self._thread_pool.submit( self.core.upload_part_copy, upload=multipart_upload, part_number=i + 1, @@ -628,47 +633,53 @@ async def _upload_part(i: int, range_: tuple[int, int]) -> S3MultipartUploadPart range_=range_, **plan.part_params, ) + part_copies.append(part_copy) + copied = asyncio.wrap_future(part_copy) + return await copied except Exception: # Set before the semaphore lets a waiting part start. failed = True raise - tasks: list[asyncio.Future[S3MultipartUploadPart | None]] = [] - completion: asyncio.Future[S3CompleteMultipartUpload] | None = None + completion: Future[S3CompleteMultipartUpload] | None = None - async def _abort() -> None: - await asyncio.wait([creation]) + def _abort() -> None: + """Abort the upload, if any, once its requests have ended. + + Runs in a thread of the filesystem, so that neither a repeated + cancellation nor the shutdown of the event loop stops it. + """ + wait([creation]) if creation.cancelled() or creation.exception() is not None: # No upload was created. return # A part that is still copying when the upload is aborted may be - # stored after the abort, so wait for the running parts first. - await asyncio.gather(*tasks, return_exceptions=True) + # stored after the abort. + wait(part_copies) if completion is not None: - await asyncio.wait([completion]) + wait([completion]) if not completion.cancelled() and completion.exception() is None: # The upload completed despite the cancellation, so # there is nothing to abort. return - await self._thread_pool.run( - self._sync_fs._abort_multipart_upload, - creation.result(), - plan.abort_params, - ) + self._sync_fs._abort_multipart_upload(creation.result(), plan.abort_params) - # A task, so that _abort() can wait for an upload that is created - # after a cancellation; scheduled right before the try, so that + # The requests run in the threads of the filesystem, not as tasks, + # 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 = asyncio.ensure_future( - self._thread_pool.run( - self.core.create_multipart_upload, plan.destination, **plan.create_params - ) + creation = self._thread_pool.submit( + self.core.create_multipart_upload, plan.destination, **plan.create_params ) try: - # shield keeps a cancellation from cancelling the creation, whose - # thread would keep running, so that _abort() can wait for it. - multipart_upload = await asyncio.shield(creation) + # A cancellation of this task leaves the creation to the cleanup, + # which waits for an upload that a started creation creates. + multipart_upload = await self._thread_pool.wait_shielded(creation) tasks = [asyncio.ensure_future(_upload_part(i, r)) for i, r in enumerate(plan.ranges)] + for task in tasks: + # Retrieved, so that asyncio does not log the failure of a + # part that ends after this copy as never retrieved. + task.add_done_callback(lambda t: t.cancelled() or t.exception()) # Unlike gather, wait does not cancel the parts when this task is # cancelled; their threads would keep copying, so they are waited # for in _abort(). @@ -678,27 +689,29 @@ async def _abort() -> None: raise error # The tasks are in part-number order. parts = [task.result() for task in tasks] - completion = asyncio.ensure_future( - self._thread_pool.run( - self.core.complete_multipart_upload, - multipart_upload, - cast(list[S3MultipartUploadPart], parts), - **plan.complete_params, - ) + completion = self._thread_pool.submit( + self.core.complete_multipart_upload, + multipart_upload, + cast(list[S3MultipartUploadPart], parts), + **plan.complete_params, ) - # shield keeps a cancellation from cancelling the completion, whose - # thread would keep running, so that _abort() can wait for it. - completed = await asyncio.shield(completion) + # A cancellation of this task leaves the completion to the + # cleanup, which waits for a started completion. + completed = await self._thread_pool.wait_shielded(completion) except BaseException: # Also on cancellation, as S3FileSystem._finish_multipart_upload # does on an interrupt. failed = True - cleanup = asyncio.ensure_future(_abort()) + # The requests that have not started are cancelled here, before + # the cleanup waits for a thread after them. + creation.cancel() + for part_copy in part_copies: + part_copy.cancel() + if completion is not None: + completion.cancel() # A repeated cancellation of this task returns without stopping - # the cleanup; the event loop keeps only weak references to tasks. - _cleanup_tasks.add(cleanup) - cleanup.add_done_callback(_cleanup_tasks.discard) - await asyncio.shield(cleanup) + # the cleanup. + await self._thread_pool.run_shielded(_abort) raise failed = False @@ -1612,7 +1625,9 @@ async def _fetch_range(self, start: int, end: int) -> bytes: """Read a byte range of the object with parallel range requests. 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. + 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. Args: start: The offset of the first byte to read. @@ -1629,7 +1644,32 @@ 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") - results = await asyncio.gather( - *[self._pool.run(self.fs.core.get_object, path, r, **request_kwargs) for r in ranges] - ) - return b"".join(results) + reads = [ + asyncio.ensure_future( + self._pool.run(self.fs.core.get_object, path, r, **request_kwargs) + ) + 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_core.py b/pyathena/filesystem/s3_core.py index adb940eb2..73410d3fd 100644 --- a/pyathena/filesystem/s3_core.py +++ b/pyathena/filesystem/s3_core.py @@ -12,6 +12,7 @@ import contextlib import logging import math +import traceback from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass, field from datetime import datetime @@ -631,14 +632,24 @@ def call(self, method: str | Callable[..., Any], **request) -> dict[str, Any]: Raises: OSError: The translation of an S3 error response (see :class:`~pyathena.filesystem.s3_errors.S3ClientError`). + + Note: + The frames of a raised error below this method, those of the + retries and of botocore, have no local variables, so that they + keep neither the request, such as the body of a part, nor the + retry state, which keeps the error in a reference cycle. """ func = getattr(self._client, method) if isinstance(method, str) else method if self._request_kwargs: request = {**self.operation_params(func.__name__, self._request_kwargs), **request} try: response = retry_api_call(func, config=self._retry_config, logger=_logger, **request) - except botocore.exceptions.ClientError as e: - raise S3ClientError(e).os_error from e + except BaseException as e: + # clear_frames() skips this frame, which is still executing. + traceback.clear_frames(e.__traceback__) + if isinstance(e, botocore.exceptions.ClientError): + raise S3ClientError(e).os_error from e + raise return cast(dict[str, Any], response) def head_object(self, path: S3Path, **params) -> S3Metadata: diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index ce9b57112..425f87115 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -14,6 +14,7 @@ import os import sys import threading +import traceback import weakref from abc import ABCMeta, abstractmethod from collections.abc import Callable @@ -21,13 +22,99 @@ from concurrent.futures import Future from concurrent.futures.thread import ThreadPoolExecutor from multiprocessing import cpu_count -from typing import Any, ClassVar, TypeVar +from typing import Any, ClassVar, Generic, TypeVar, cast from pyathena.util import override T = TypeVar("T") +class _S3WorkItem(Generic[T]): + """A submitted callable and the future of the caller that its outcome resolves. + + The executor calls :meth:`run`, and :meth:`settle` is the done callback of + the submission, which resolves the future if the executor drops the + callable before it runs. Whichever comes first claims the future, so that + it is started or settled once, and drops the references of the item to + the future, the callable, and its arguments, which the frames and + callbacks of the executor may keep after the call. The arguments of a + failed callable, such as the body of a part, are therefore freed when it + fails, and the traceback of its error, which keeps the frame of + :meth:`run`, forms no reference cycle with the future. + """ + + def __init__( + self, + future: Future[T], + fn: Callable[..., T], + args: tuple[Any, ...], + kwargs: dict[str, Any], + ) -> None: + """Initialize the work item. + + Args: + future: A pending future that no executor has started. + fn: The callable to run. + args: Positional arguments passed to the callable. + kwargs: Keyword arguments passed to the callable. + """ + self._future: Future[T] | None = future + self._fn: Callable[..., T] | None = fn + self._args = args + self._kwargs = kwargs + self._claim = threading.Lock() + + def run(self) -> None: + """Run the callable and resolve the future unless it was cancelled.""" + if not self._claim.acquire(blocking=False): + return + if not cast(Future[T], self._future).set_running_or_notify_cancel(): + self._take() + return + try: + result = cast(Callable[..., T], self._fn)(*self._args, **self._kwargs) + except BaseException as e: + # The frames below this one hold the arguments. clear_frames() + # skips this frame, which is still executing. + traceback.clear_frames(e.__traceback__) + self._take().set_exception(e) + else: + self._take().set_result(result) + + def settle(self, submitted: Future[Any]) -> None: + """Resolve the future if the executor dropped the callable before it ran. + + This happens, for example, when the event loop of an + :class:`S3AioExecutor` shuts down. + + Args: + submitted: The finished future of the submitted :meth:`run`. + """ + if not self._claim.acquire(blocking=False): + # run() has claimed the future and resolves it. + return + future = self._take() + if submitted.cancelled(): + future.cancel() + # A cancelled future is only marked as notified, as an executor does + # when it drops a cancelled callable, so that wait() counts it as done. + if future.set_running_or_notify_cancel(): + future.set_exception( + submitted.exception() or RuntimeError("The executor did not run the call.") + ) + + def _take(self) -> Future[T]: + """Drop the references of the item and return its future. + + Returns: + The future of the caller. + """ + future = cast(Future[T], self._future) + self._future = self._fn = None + self._args, self._kwargs = (), {} + return future + + class S3Executor(metaclass=ABCMeta): """Abstract executor for parallel S3 operations. @@ -59,47 +146,21 @@ def submit_to(self, future: Future[T], fn: Callable[..., T], *args: Any, **kwarg event loop shuts down, the future is cancelled or gets the error that the executor reported. + The arguments of a failed callable, such as the body of a part, are + freed when it fails: its error keeps the frames that it was raised + from, without their local variables. The errors that it was raised + from keep theirs; :meth:`S3Core.call() + ` clears those of the + ``ClientError`` of a request. + Args: future: A pending future that no executor has started. fn: The callable to run. *args: Positional arguments passed to the callable. **kwargs: Keyword arguments passed to the callable. """ - - def run() -> None: - """Run the callable and resolve the future unless it was cancelled.""" - if not future.set_running_or_notify_cancel(): - return - try: - result = fn(*args, **kwargs) - except BaseException as e: - future.set_exception(e) - else: - future.set_result(result) - - def settle(submitted: Future[None]) -> None: - """Resolve the future if the executor dropped the callable before it ran. - - run() raises nothing for a future that no executor has started, - so a cancelled or failed submission means that it did not run, - and the future is still pending or cancelled by the caller. - - Args: - submitted: The finished future of the submitted run(). - """ - error: BaseException | None = None - if submitted.cancelled(): - future.cancel() - elif (error := submitted.exception()) is None: - # run() resolved the future. - return - # A future cancelled here or by the caller is only marked as - # notified, which an executor does when it drops a cancelled - # callable, so that wait() counts it as done. - if future.set_running_or_notify_cancel(): - future.set_exception(error) - - self.submit(run).add_done_callback(settle) + item = _S3WorkItem(future, fn, args, kwargs) + self.submit(item.run).add_done_callback(item.settle) def __enter__(self) -> S3Executor: return self @@ -228,7 +289,8 @@ async def run_shielded(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> """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. + task is cancelled; its outcome is then dropped. See + :meth:`wait_shielded`. Args: fn: The blocking callable. @@ -241,11 +303,30 @@ async def run_shielded(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Raises: RuntimeError: If the pool has been shut down. """ - future = asyncio.wrap_future(self.submit(fn, *args, **kwargs)) - # Retrieved, so that asyncio does not log the failure of a call whose - # waiting task was cancelled as never retrieved. - future.add_done_callback(lambda f: f.cancelled() or f.exception()) - return await asyncio.shield(future) + return await self.wait_shielded(self.submit(fn, *args, **kwargs)) + + async def wait_shielded(self, future: Future[T]) -> T: + """Wait for a submitted call that a cancellation of the waiting task does not cancel. + + Unlike awaiting ``asyncio.wrap_future(future)``, a cancellation of + the waiting task leaves the future to the caller, which can still + wait for it or cancel it, and its outcome is retrieved, so that + asyncio does not log a failure that ends after the cancellation as + never retrieved. + + Args: + future: The future of a call submitted to this pool. + + Returns: + The return value of the call. + """ + waiter = asyncio.wrap_future(future) + waiter.add_done_callback(lambda f: f.cancelled() or f.exception()) + # 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() @override def shutdown(self, wait: bool = True) -> None: @@ -347,45 +428,12 @@ def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]: # the function keeps running in its thread, so the returned future # is started and resolved by the function's thread instead. future: Future[T] = Future() - # Acquired once, by run() or by settle(), whichever comes first, - # so that the future is started or settled exactly once. - claim = threading.Lock() - - def run() -> None: - """Run the function and resolve the future unless it was cancelled.""" - if not claim.acquire(blocking=False) or not future.set_running_or_notify_cancel(): - return - try: - result = fn(*args, **kwargs) - except BaseException as e: - future.set_exception(e) - else: - future.set_result(result) - - def settle(task: Future[None]) -> None: - """Resolve the future if the task ended before the function started. - - This happens, for example, when the event loop shuts down. - - Args: - task: The finished future of the task that runs the function. - """ - if not claim.acquire(blocking=False): - # run() has claimed the future and resolves it. - return - if task.cancelled(): - future.cancel() - # Notify the waiters of the cancellation, as an executor - # does when it drops a cancelled function. - future.set_running_or_notify_cancel() - elif future.set_running_or_notify_cancel(): - future.set_exception(task.exception()) - + item = _S3WorkItem(future, fn, args, kwargs) with self._lock: if self._shutdown: raise RuntimeError("cannot schedule new futures after shutdown") - task = asyncio.run_coroutine_threadsafe(self._pool.run(run), self._loop) - task.add_done_callback(settle) + task = asyncio.run_coroutine_threadsafe(self._pool.run(item.run), self._loop) + task.add_done_callback(item.settle) return future raise RuntimeError( "S3AioExecutor requires a running event loop. " diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 7eaa1608d..4b0dffc35 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -1,3 +1,4 @@ +import array import asyncio import bz2 import contextlib @@ -3428,6 +3429,72 @@ def answer_first_range_last(**request): assert f.read() == data assert sorted(ranges) == ["bytes=0-15", "bytes=16-31", "bytes=32-47", "bytes=48-63"] + def test_failed_read_frees_other_ranges(self): + # GH-1186: a parallel read that fails frees the data of the other + # ranges, also while its error is held. Their futures stayed in the + # frame of the read, which the traceback of the error keeps. + fs = self._make_fs() + fs.info = mock.MagicMock(return_value=self._file_object("key")) + fs.info.return_value.size = 64 + read = [] + others_read = threading.Semaphore(0) + + def get_object(path, range_, **kwargs): + if range_[0] == 0: + for _ in range(3): + assert others_read.acquire(timeout=5) + raise PermissionError("Access Denied") + # Unlike bytes, an array can be weakly referenced. + data = array.array("B", bytes(range_[1] - range_[0])) + read.append(weakref.ref(data)) + others_read.release() + return data + + fs.core.get_object = get_object + f = fs.open("s3://bucket/key", "rb", cache_type="none", block_size=16, max_workers=4) + with pytest.raises(PermissionError, match="Access Denied") as info: + f.read() + # Waits for the reads, whose threads may still set the results. + f.close() + assert len(read) == 3 + assert [ref() for ref in read] == [None] * 3 + assert info.value.__traceback__ is not None + + 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 + # interrupt is held, as a failed read does. + fs = self._make_fs() + fs.info = mock.MagicMock(return_value=self._file_object("key")) + fs.info.return_value.size = 64 + read = [] + + def get_object(path, range_, **kwargs): + # Unlike bytes, an array can be weakly referenced. + data = array.array("B", bytes(range_[1] - range_[0])) + read.append(weakref.ref(data)) + return data + + fs.core.get_object = get_object + + 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. + futures = set(futures) + wait(futures) + 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), + pytest.raises(KeyboardInterrupt) as info, + ): + f.read() + f.close() + assert len(read) == 4 + assert [ref() for ref in read] == [None] * 4 + assert info.value.__traceback__ is not None + @pytest.mark.parametrize("cache_type", ["bytes", "all", "first"]) def test_open_directory(self, cache_type): fs = self._make_fs() @@ -7510,6 +7577,14 @@ def test_commit_failure_and_discard(self, caplog, error, abort_fails): assert ( "Failed to abort multipart upload uploadid to s3://bucket/key.txt." in caplog.text ) is abort_fails + # GH-1186: the file is closed without the buffer, which holds only the + # data of the last part, so a later commit() neither completes the + # upload nor uploads that data alone with PutObject. + assert file.closed + assert file.buffer is None + file.commit() + file.fs.core.complete_multipart_upload.assert_called_once() + file.fs.core.put_object.assert_not_called() file.discard() assert file.fs._call.call_args_list == [ diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index 80de1537d..a18f82524 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -1,5 +1,7 @@ +import array import asyncio import contextlib +import gc import gzip import os import tempfile @@ -8,6 +10,7 @@ import urllib.parse import urllib.request import uuid +import weakref from concurrent.futures import FIRST_COMPLETED, Future, ThreadPoolExecutor, wait from datetime import UTC, datetime from itertools import chain @@ -491,6 +494,10 @@ def complete_multipart_upload(*args, **kw): side_effect=lambda *args: events.append("abort") ) + logged = [] + asyncio.get_running_loop().set_exception_handler( + lambda _, context: logged.append(context["message"]) + ) task = asyncio.ensure_future( fs._copy_object_with_multipart_upload( S3Path("bucket", "src"), @@ -515,8 +522,17 @@ def complete_multipart_upload(*args, **kw): release.set() with pytest.raises(asyncio.CancelledError): await task + # The threads have run the callbacks of the requests once they are + # shut down, and the event loop runs theirs on the next iterations. + await asyncio.to_thread(fs._thread_pool.shutdown, True) + gc.collect() + for _ in range(3): + await asyncio.sleep(0) assert events == (["complete", "abort"] if completion_fails else ["complete"]) + # GH-1186: a completion that fails after the cancellation is not + # logged as never retrieved. + assert logged == [] @pytest.mark.parametrize("creation_fails", [False, True]) @pytest.mark.asyncio @@ -548,6 +564,10 @@ def create_multipart_upload(*args, **kw): side_effect=lambda upload, params: events.append(("abort", upload.upload_id)) ) + logged = [] + asyncio.get_running_loop().set_exception_handler( + lambda _, context: logged.append(context["message"]) + ) task = asyncio.ensure_future( fs._copy_object_with_multipart_upload( S3Path("bucket", "src"), @@ -572,10 +592,243 @@ def create_multipart_upload(*args, **kw): release.set() with pytest.raises(asyncio.CancelledError): await task + # The threads have run the callbacks of the requests once they are + # shut down, and the event loop runs theirs on the next iterations. + await asyncio.to_thread(fs._thread_pool.shutdown, True) + gc.collect() + for _ in range(3): + await asyncio.sleep(0) assert events == (["create"] if creation_fails else ["create", ("abort", "uploadid")]) + # GH-1186: a creation that fails after the cancellation is not logged + # as never retrieved. + assert logged == [] sync_fs.core.upload_part_copy.assert_not_called() + @pytest.mark.asyncio + async def test_copy_object_with_multipart_upload_cancelled_queued_completion(self): + # GH-1186: a cancellation cancels a CompleteMultipartUpload request + # that waits for a thread of the filesystem, and aborts the upload. The + # request used to run after the cancellation and complete the copy. + fs = AioS3FileSystem(connection=mock.MagicMock(), max_workers=1, skip_instance_cache=True) + sync_fs = fs._sync_fs + sync_fs.core.create_multipart_upload = mock.MagicMock( + return_value=S3MultipartUpload( + {"Bucket": "bucket", "Key": "dst", "UploadId": "uploadid"} + ) + ) + events = [] + release = threading.Event() + + def block(): + # The finally block of the test always releases it. + release.wait() + + def upload_part_copy(**kw): + if kw["part_number"] == 2: + # Holds the only thread of the filesystem, so that the + # completion waits for it. + fs._thread_pool.submit(block) + return SimpleNamespace(etag='"e"', part_number=kw["part_number"]) + + sync_fs.core.upload_part_copy = mock.MagicMock(side_effect=upload_part_copy) + sync_fs.core.complete_multipart_upload = mock.MagicMock( + side_effect=lambda *args, **kw: events.append("complete") + ) + # The HeadObject of the source, with 2 parts of the default block size. + size = 2 * S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE + sync_fs._call = sync_fs._core.call = mock.MagicMock(return_value={"ContentLength": size}) + sync_fs._abort_multipart_upload = mock.MagicMock( + side_effect=lambda upload, params: events.append("abort") + ) + + submit = fs._thread_pool.submit + completion_queued = threading.Event() + + def record_completion(fn, *args, **kwargs): + future = submit(fn, *args, **kwargs) + if fn is sync_fs.core.complete_multipart_upload: + completion_queued.set() + return future + + fs._thread_pool.submit = record_completion + task = asyncio.ensure_future( + fs._copy_object_with_multipart_upload( + S3Path("bucket", "src"), + S3Path("bucket", "dst"), + max_workers=1, + MetadataDirective="REPLACE", + TaggingDirective="REPLACE", + AnnotationDirective="EXCLUDE", + ) + ) + try: + # The completion waits for the thread, which block() holds. + assert await asyncio.to_thread(completion_queued.wait, 5) + task.cancel() + # Lets the copy cancel the completion and schedule its cleanup. + await asyncio.sleep(0.1) + except BaseException: + task.cancel() + raise + finally: + release.set() + with pytest.raises(asyncio.CancelledError): + await task + + assert events == ["abort"] + + @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 + # thread of the filesystem, whose threads bound all its copies, before + # it waits for the running ones and aborts the upload. The queued + # part copies used to run before the abort. + fs = AioS3FileSystem(connection=mock.MagicMock(), max_workers=1, skip_instance_cache=True) + sync_fs = fs._sync_fs + sync_fs.core.create_multipart_upload = mock.MagicMock( + return_value=S3MultipartUpload( + {"Bucket": "bucket", "Key": "dst", "UploadId": "uploadid"} + ) + ) + events = [] + started = threading.Event() + release = threading.Event() + + def upload_part_copy(**kw): + events.append(f"start {kw['part_number']}") + started.set() + # The finally block of the test always releases it. + release.wait() + events.append(f"end {kw['part_number']}") + return SimpleNamespace(etag='"e"', part_number=kw["part_number"]) + + sync_fs.core.upload_part_copy = mock.MagicMock(side_effect=upload_part_copy) + sync_fs.core.complete_multipart_upload = mock.MagicMock() + # The HeadObject of the source, with 2 parts of the default block size. + size = 2 * S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE + sync_fs._call = sync_fs._core.call = mock.MagicMock(return_value={"ContentLength": size}) + sync_fs._abort_multipart_upload = mock.MagicMock( + side_effect=lambda upload, params: events.append("abort") + ) + + task = asyncio.ensure_future( + fs._copy_object_with_multipart_upload( + S3Path("bucket", "src"), + S3Path("bucket", "dst"), + max_workers=2, + MetadataDirective="REPLACE", + TaggingDirective="REPLACE", + AnnotationDirective="EXCLUDE", + ) + ) + try: + assert await asyncio.to_thread(started.wait, 5) + task.cancel() + # Gives the cleanup time to cancel the queued part copy. + await asyncio.sleep(0.1) + except BaseException: + task.cancel() + raise + finally: + release.set() + with pytest.raises(asyncio.CancelledError): + await task + + assert events == ["start 1", "end 1", "abort"] + sync_fs.core.complete_multipart_upload.assert_not_called() + + @pytest.mark.parametrize( + ("stage", "expected"), + [ + ("create", ["start create", "end create", "abort"]), + ("part", ["start part", "start part", "end part", "end part", "abort"]), + ("complete", ["start complete", "end complete"]), + ], + ) + @pytest.mark.parametrize("cancelled", [False, True]) + def test_copy_object_with_multipart_upload_loop_shutdown(self, stage, expected, cancelled): + # GH-1186: when the shutdown of the event loop at the end of + # asyncio.run() cancels a copy, or a cleanup that a cancellation of + # the copy has started, the cleanup still waits for the running + # requests and aborts the upload unless it completed. The shutdown + # also cancelled the tasks of the requests and of the cleanup, so the + # cleanup left an upload created during the cancellation, aborted + # before the running parts ended, or aborted during the completion. + fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + sync_fs = fs._sync_fs + events = [] + lock = threading.Lock() + copy: asyncio.Task[None] | None = None + + def request(name): + with lock: + events.append(f"start {name}") + # Held until asyncio.run() has cancelled the copy, after the + # cancellation of the test, if any. + assert copy is not None + deadline = time.monotonic() + 5 + while copy.cancelling() < (2 if cancelled else 1): + assert time.monotonic() < deadline + time.sleep(0.01) + with lock: + events.append(f"end {name}") + + def create_multipart_upload(*args, **kw): + if stage == "create": + request("create") + return S3MultipartUpload({"Bucket": "bucket", "Key": "dst", "UploadId": "uploadid"}) + + def upload_part_copy(**kw): + if stage == "part": + request("part") + return SimpleNamespace(etag='"e"', part_number=kw["part_number"]) + + def complete_multipart_upload(*args, **kw): + request("complete") + return SimpleNamespace() + + sync_fs.core.create_multipart_upload = mock.MagicMock(side_effect=create_multipart_upload) + sync_fs.core.upload_part_copy = mock.MagicMock(side_effect=upload_part_copy) + sync_fs.core.complete_multipart_upload = mock.MagicMock( + side_effect=complete_multipart_upload + ) + # The HeadObject of the source, with 2 parts of the default block size. + size = 2 * S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE + sync_fs._call = sync_fs._core.call = mock.MagicMock(return_value={"ContentLength": size}) + sync_fs._abort_multipart_upload = mock.MagicMock( + side_effect=lambda upload, params: events.append("abort") + ) + + async def main(): + nonlocal copy + copy = asyncio.ensure_future( + fs._copy_object_with_multipart_upload( + S3Path("bucket", "src"), + S3Path("bucket", "dst"), + max_workers=2, + MetadataDirective="REPLACE", + TaggingDirective="REPLACE", + AnnotationDirective="EXCLUDE", + ) + ) + # Returns once the held requests have started. + started = expected.count(f"start {stage}") + deadline = time.monotonic() + 5 + while len(events) < started: + assert time.monotonic() < deadline + await asyncio.sleep(0.01) + if cancelled: + copy.cancel() + # Lets the copy start its cleanup. + await asyncio.sleep(0.1) + + asyncio.run(main()) + # The cleanup may end after asyncio.run() returns. + fs._thread_pool.shutdown(wait=True) + + assert events == expected + @pytest.mark.parametrize( "block_size", [ @@ -2450,6 +2703,48 @@ def test_open_max_workers(self, max_workers): assert isinstance(f, AioS3File) assert f.max_workers == (max_workers or fs._sync_fs.max_workers) + @pytest.mark.asyncio + async def test_failed_read_frees_other_ranges(self): + # GH-1186: a parallel read that fails frees the data of the other + # ranges, also while its error is held, as with S3File. + fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + block_size = 2**10 + read = [] + others_read = threading.Semaphore(0) + + def get_object(path, range_, **kwargs): + if range_[0] == 0: + for _ in range(3): + assert others_read.acquire(timeout=5) + raise PermissionError("Access Denied") + # Unlike bytes, an array can be weakly referenced. + data = array.array("B", bytes(range_[1] - range_[0])) + read.append(weakref.ref(data)) + others_read.release() + return data + + 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 * 4 + fs._sync_fs.info = mock.MagicMock(return_value=info) + # The executor of the file dispatches through the event loop, which + # the blocking calls must not hold. + f = await asyncio.to_thread( + fs.open, "s3://bucket/key", "rb", block_size=block_size, cache_type="none" + ) + with pytest.raises(PermissionError, match="Access Denied") as info_: + await asyncio.to_thread(f.read) + # Waits for the reads, whose threads may still set the results. + await asyncio.to_thread(f.close) + assert len(read) == 3 + assert [ref() for ref in read] == [None] * 3 + assert info_.value.__traceback__ is not None + @pytest.mark.parametrize("asynchronous", [False, True]) @pytest.mark.asyncio async def test_open_parallel_requests(self, asynchronous): @@ -3072,6 +3367,94 @@ async def test_close_small(self, fs, data): fs.core.put_object.assert_not_called() fs.core.create_multipart_upload.assert_not_called() + @pytest.mark.asyncio + async def test_failed_read_frees_other_ranges(self, fs): + # GH-1186: a parallel read that fails cancels the reads that have not + # finished and frees the data of the others, also while its error is + # held, as with S3File. + block_size = 2**10 + read = [] + others_read = threading.Semaphore(0) + + def get_object(path, range_, **kwargs): + if range_[0] == 0: + for _ in range(3): + assert others_read.acquire(timeout=5) + raise PermissionError("Access Denied") + # Unlike bytes, an array can be weakly referenced. + data = array.array("B", bytes(range_[1] - range_[0])) + read.append(weakref.ref(data)) + others_read.release() + return data + + 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 * 4 + fs._sync_fs.info = mock.MagicMock(return_value=info) + f = await fs.open_async("s3://bucket/key", "rb", block_size=block_size, max_workers=4) + with pytest.raises(PermissionError, match="Access Denied") as info_: + await f.read() + await f.close() + # The threads have set the results once the pool of the file is shut + # down, and the event loop runs their callbacks on the next iterations. + await asyncio.to_thread(f._pool.shutdown, True) + for _ in range(3): + await asyncio.sleep(0) + assert len(read) == 3 + assert [ref() for ref in read] == [None] * 3 + assert info_.value.__traceback__ is not None + + @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 and frees the data of the others, also while the + # cancellation is held, as a failed read does. + block_size = 2**10 + read = [] + + def get_object(path, range_, **kwargs): + # Unlike bytes, an array can be weakly referenced. + data = array.array("B", bytes(range_[1] - range_[0])) + read.append(weakref.ref(data)) + return data + + 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 * 4 + fs._sync_fs.info = mock.MagicMock(return_value=info) + wait_reads = asyncio.wait + + async def cancelled_wait(reads, **kwargs): + # As asyncio.wait(), whose frame keeps the tasks in a set, + # cancelled once the reads have finished. + reads = set(reads) + await wait_reads(reads) + raise asyncio.CancelledError + + f = await fs.open_async("s3://bucket/key", "rb", block_size=block_size, max_workers=4) + with ( + mock.patch.object(asyncio, "wait", side_effect=cancelled_wait), + pytest.raises(asyncio.CancelledError) as info_, + ): + await f.read() + await f.close() + await asyncio.to_thread(f._pool.shutdown, True) + for _ in range(3): + await asyncio.sleep(0) + assert len(read) == 4 + assert [ref() for ref in read] == [None] * 4 + assert info_.value.__traceback__ is not None + @pytest.mark.asyncio async def test_read_parallel(self, fs): # A read runs at most max_workers range requests at once, conditional diff --git a/tests/pyathena/filesystem/test_s3_core.py b/tests/pyathena/filesystem/test_s3_core.py index ef992573f..2a0695b30 100644 --- a/tests/pyathena/filesystem/test_s3_core.py +++ b/tests/pyathena/filesystem/test_s3_core.py @@ -6,10 +6,14 @@ # SPDX-License-Identifier: MIT import copy +import gc import io import re +import weakref +from concurrent.futures import Future from datetime import UTC, datetime from itertools import pairwise +from typing import Any from unittest import mock from urllib.parse import parse_qs, urlsplit @@ -33,6 +37,7 @@ S3MultipartCopyPlan, S3ObjectSummary, ) +from pyathena.filesystem.s3_executor import S3ThreadPoolExecutor from pyathena.filesystem.s3_object import S3Metadata, S3MultipartUpload, S3MultipartUploadPart from pyathena.filesystem.s3_path import S3Path from pyathena.util import RetryConfig @@ -123,6 +128,70 @@ def test_call_translates_errors(self, status, code, error): with stubber, pytest.raises(error): core.head_object(S3Path("bucket", "key")) + @pytest.mark.parametrize("translated", [True, False]) + def test_call_failure_frees_request(self, translated): + # GH-1186: the frames of the retries and of botocore below call() in + # the traceback of a raised error, also of the ClientError that the + # error is translated from, have no local variables. tenacity's retry + # state in them used to keep the error, and so the request, such as + # the body of a part, in a reference cycle until the garbage collector + # ran. + class Body: + pass + + def send(**params): + if translated: + raise botocore.exceptions.ClientError( + { + "Error": {"Code": "AccessDenied"}, + "ResponseMetadata": {"HTTPStatusCode": 403}, + }, + "PutObject", + ) + raise ConnectionError("Connection lost") + + core, _ = _make_core() + body = Body() + body_ref = weakref.ref(body) + gc.disable() + try: + with pytest.raises(PermissionError if translated else ConnectionError) as info: + core.call(send, Body=body) + del body + if translated: + assert isinstance(info.value.__cause__, botocore.exceptions.ClientError) + del info + assert body_ref() is None + finally: + gc.enable() + + def test_call_failure_through_executor_frees_request(self): + # GH-1186: through S3Executor.submit_to(), which clears the frames + # below it, the request of a failed call is freed while its error is + # still held. + class Body: + pass + + def send(**params): + raise botocore.exceptions.ClientError( + {"Error": {"Code": "AccessDenied"}, "ResponseMetadata": {"HTTPStatusCode": 403}}, + "UploadPart", + ) + + core, _ = _make_core() + body = Body() + body_ref = weakref.ref(body) + future: Future[dict[str, Any]] = Future() + gc.disable() + try: + with S3ThreadPoolExecutor(max_workers=1) as executor: + executor.submit_to(future, core.call, send, Body=body) + del body + assert isinstance(future.exception(), PermissionError) + assert body_ref() is None + finally: + gc.enable() + def test_head_object(self): core, stubber = _make_core() stubber.add_response( diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index b4aa849f9..65bee13ef 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -7,9 +7,12 @@ import asyncio import contextvars +import gc import os import threading import time +import traceback +import weakref from concurrent.futures import Future, ThreadPoolExecutor, wait from unittest import mock @@ -87,6 +90,54 @@ async def main(): assert events == (["pending finished"] if shielded else []) + def test_wait_shielded_cancelled(self): + # GH-1186: cancelling the waiting task neither cancels the call nor + # lets asyncio log its failure, which ends after the cancellation; + # the caller can still wait for its future. + started = threading.Event() + release = threading.Event() + pool = S3AioThreadPool(1) + + def work(): + started.set() + # The finally block of the test always releases it. + release.wait() + raise PermissionError("Access Denied") + + async def main(): + logged = [] + asyncio.get_running_loop().set_exception_handler( + lambda _, context: logged.append(context["message"]) + ) + running = pool.submit(work) + try: + waiting = asyncio.ensure_future(pool.wait_shielded(running)) + await asyncio.to_thread(started.wait, EVENT_TIMEOUT) + # Lets the task start waiting before it is cancelled. + deadline = time.monotonic() + EVENT_TIMEOUT + while not running._done_callbacks: + assert time.monotonic() < deadline + await asyncio.sleep(0) + waiting.cancel() + with pytest.raises(asyncio.CancelledError): + await waiting + finally: + release.set() + # The worker has run the callbacks of the call once it is + # shut down, and the event loop runs theirs on the next + # iterations. + await asyncio.to_thread(pool.shutdown, True) + with pytest.raises(PermissionError, match="Access Denied"): + running.result() + # The callbacks of the future of the call keep its waiter. + del waiting, running + gc.collect() + for _ in range(3): + await asyncio.sleep(0) + return logged + + assert asyncio.run(main()) == [] + def test_shutdown(self): # The calls submitted before the shutdown still run. events = [] @@ -395,6 +446,33 @@ def test_submit_to(self): with pytest.raises(ValueError, match="invalid literal"): failed.result() + @pytest.mark.parametrize("executor_class", [S3ThreadPoolExecutor, S3AioThreadPool]) + def test_submit_to_failure_frees_arguments(self, executor_class): + # GH-1186: the arguments of a failed callable, such as the body of a + # part, are freed when it fails, while the error keeps the frames of + # its traceback. The traceback used to keep them, in a reference + # cycle with the future until the garbage collector ran. + class Body: + pass + + def upload(body): + raise PermissionError("Access Denied") + + body = Body() + body_ref = weakref.ref(body) + future: Future[None] = Future() + gc.disable() + try: + with executor_class(max_workers=1) as executor: + executor.submit_to(future, upload, body) + del body + error = future.exception() + assert isinstance(error, PermissionError) + assert traceback.extract_tb(error.__traceback__)[-1].name == "upload" + assert body_ref() is None + finally: + gc.enable() + def test_submit_to_cancelled_before_start(self): # Cancelling the future keeps a callable that has not started from running. events = []