From 4ac521df86ff361b3fafa8fc9ae8e6d4737523f9 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 15 Oct 2026 23:50:29 +0900 Subject: [PATCH 01/18] Free failed parts without the GC and abort copies at loop shutdown S3Executor.submit_to() clears the frames below the callable and the references of its own frame before it sets the error, so the body of a failed part is freed when the part fails. raise_failure() raises the error of a finished future after the caller's set is emptied, and S3AioThreadPool.run_shielded() and AioS3StreamedFile._discard() drop their references to failed futures, so that an S3File, an AioS3File or an AioS3StreamedFile whose part failed is freed with its data by reference counting once the error is dropped. The async multipart copy of AioS3FileSystem keeps the creation and the completion as futures instead of tasks, and keeps the request of each part, so that when the shutdown of the event loop at the end of asyncio.run() cancels the copy, the cleanup still waits for them and aborts an upload created during the cancellation, after its running parts have ended, and not during a completion. Closes #1186 Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3.py | 24 ++-- pyathena/filesystem/s3_async.py | 67 ++++++---- pyathena/filesystem/s3_executor.py | 61 +++++++++- tests/pyathena/filesystem/test_s3.py | 28 +++++ tests/pyathena/filesystem/test_s3_async.py | 114 ++++++++++++++++++ tests/pyathena/filesystem/test_s3_executor.py | 98 ++++++++++++++- 6 files changed, 346 insertions(+), 46 deletions(-) diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 2ad7d1b1a..8d5e41686 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -34,7 +34,7 @@ from pyathena.connection import Connection from pyathena.filesystem.s3_core import S3Core, S3DeleteBatch, S3DeleteResult from pyathena.filesystem.s3_errors import S3ClientError -from pyathena.filesystem.s3_executor import S3Executor, S3ThreadPoolExecutor +from pyathena.filesystem.s3_executor import S3Executor, S3ThreadPoolExecutor, raise_failure from pyathena.filesystem.s3_object import ( S3CompleteMultipartUpload, S3Metadata, @@ -3236,19 +3236,17 @@ def _upload_chunk(self, final: bool = False) -> bool: if self._copies_in_flight: copied = {copy for copy in self._copies_in_flight if copy.done()} self._copies_in_flight -= copied - for copy in copied: - # Otherwise the rest of the data is uploaded after a - # failed part copy. - copy.result() + # Otherwise the rest of the data is uploaded after a + # failed part copy. + raise_failure(copied) if len(self._uploads_in_flight) >= self.max_workers: # Otherwise the parts waiting for a worker pile up in memory. done, self._uploads_in_flight = wait( self._uploads_in_flight, return_when=FIRST_COMPLETED ) - for upload in done: - # A failed part would otherwise keep its data until - # commit() while the rest of the data is uploaded. - upload.result() + # A failed part would otherwise keep its data until + # commit() while the rest of the data is uploaded. + raise_failure(done) # Tracked before it is scheduled, so that an interrupt while # it is being scheduled still lets discard() cancel it or wait # for it before the abort. @@ -3310,11 +3308,9 @@ def commit(self) -> None: upload_id = cast(str, self.multipart_upload.upload_id) try: done, _ = wait(self.multipart_upload_parts, return_when=FIRST_EXCEPTION) - for future in done: - if not future.cancelled() and future.exception() is not None: - # Raised without waiting for the parts before it, so - # that the parts queued after it are cancelled. - future.result() + # Raised without waiting for the parts before it, so that the + # parts queued after it are cancelled. + raise_failure(done) parts = [future.result() for future in self.multipart_upload_parts] self._get_multipart_writer().complete( parts, **self._get_request_kwargs("complete_multipart_upload") diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 1bdee1556..763b9e8f3 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -33,6 +33,7 @@ S3AioThreadPool, S3Executor, S3ThreadPoolExecutor, + raise_failure, ) from pyathena.filesystem.s3_object import ( S3CompleteMultipartUpload, @@ -579,11 +580,13 @@ 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 + 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. This also holds when + the shutdown of the event loop, such as at the end of + ``asyncio.run()``, cancels the tasks of the copy. A repeated + cancellation returns without stopping this cleanup. Args: source: Source S3 path, with the version ID to copy, if any. @@ -612,6 +615,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 +626,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,6 +634,9 @@ 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 @@ -642,8 +651,13 @@ async def _abort() -> 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. + # stored after the abort, so wait for the running parts first, + # including those whose tasks were cancelled. await asyncio.gather(*tasks, return_exceptions=True) + await asyncio.gather( + *[asyncio.wrap_future(c) for c in part_copies if not c.cancel()], + return_exceptions=True, + ) if completion is not None: await asyncio.wait([completion]) if not completion.cancelled() and completion.exception() is None: @@ -656,11 +670,12 @@ async def _abort() -> None: 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 - # nothing can fail between the two. - creation = asyncio.ensure_future( - self._thread_pool.run( + # Not a task, which the shutdown of the event loop would cancel while + # its request keeps running, so that _abort() can wait for an upload + # that is created after a cancellation; scheduled right before the + # try, so that nothing can fail between the two. + creation = asyncio.wrap_future( + self._thread_pool.submit( self.core.create_multipart_upload, plan.destination, **plan.create_params ) ) @@ -678,8 +693,9 @@ 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( + # Not a task either; see the creation. + completion = asyncio.wrap_future( + self._thread_pool.submit( self.core.complete_multipart_upload, multipart_upload, cast(list[S3MultipartUploadPart], parts), @@ -1501,19 +1517,17 @@ async def _upload_chunk(self, final: bool = False) -> bool: if self._copies_in_flight: copied = {copy for copy in self._copies_in_flight if copy.done()} self._copies_in_flight -= copied - for copy in copied: - # Otherwise the rest of the data is uploaded after a - # failed part copy, as S3File does. - copy.result() + # Otherwise the rest of the data is uploaded after a + # failed part copy, as S3File does. + raise_failure(copied) if len(self._uploads_in_flight) >= self.max_workers: # Otherwise the parts waiting for a thread pile up in memory. done, self._uploads_in_flight = await asyncio.wait( self._uploads_in_flight, return_when=asyncio.FIRST_COMPLETED ) - for upload in done: - # A failed part would otherwise keep its data until - # the commit while the rest of the data is uploaded. - upload.result() + # A failed part would otherwise keep its data until the + # commit while the rest of the data is uploaded. + raise_failure(done) # Tracked before it is scheduled; see _initiate_upload(). part: Future[S3MultipartUploadPart] = Future() self._parts.append(part) @@ -1578,9 +1592,7 @@ def _complete_upload(self) -> None: # all started. A failed part is raised without waiting for # the parts before it, as S3File does. done, _ = wait(self._parts, return_when=FIRST_EXCEPTION) - for part in done: - if not part.cancelled() and part.exception() is not None: - part.result() + raise_failure(done) parts = [part.result() for part in self._parts] writer.complete(parts, **self._get_request_kwargs("complete_multipart_upload")) except BaseException: @@ -1607,6 +1619,11 @@ def _discard(self) -> None: _logger.exception( f"Failed to abort the multipart upload to s3://{self.bucket}/{self.key}." ) + # Dropped as S3File.discard() drops them, so that the file does not + # keep the errors of failed parts, whose tracebacks may keep the file. + self._parts = [] + self._uploads_in_flight = set() + self._copies_in_flight = set() async def _fetch_range(self, start: int, end: int) -> bytes: """Read a byte range of the object with parallel range requests. diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index ce9b57112..8d4714630 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,7 +22,8 @@ from concurrent.futures import Future from concurrent.futures.thread import ThreadPoolExecutor from multiprocessing import cpu_count -from typing import Any, ClassVar, TypeVar +from types import TracebackType +from typing import Any, ClassVar, TypeVar, cast from pyathena.util import override @@ -59,6 +61,12 @@ 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 error of a failed callable keeps the frames that it was raised + from in its traceback, without their local variables, so that the + arguments, such as the body of a part, are freed when the callable + fails, and the future and its error are freed without the garbage + collector once they are no longer referenced. + Args: future: A pending future that no executor has started. fn: The callable to run. @@ -68,14 +76,24 @@ def submit_to(self, future: Future[T], fn: Callable[..., T], *args: Any, **kwarg def run() -> None: """Run the callable and resolve the future unless it was cancelled.""" + nonlocal future, fn, args, kwargs 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) + # The frames below this one hold the arguments. + traceback.clear_frames(cast(TracebackType, e.__traceback__).tb_next) + # The traceback also keeps this frame, which would keep the + # arguments and form a reference cycle with the future. Once + # run() has resolved the future, settle() does not use it. + resolved = future + del future, fn, args, kwargs + resolved.set_exception(e) + del resolved + return + # Deleted only on the failure path, which returns. + future.set_result(result) # noqa: F821 def settle(submitted: Future[None]) -> None: """Resolve the future if the executor dropped the callable before it ran. @@ -108,6 +126,34 @@ def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: self.shutdown(wait=True) +def raise_failure(done: set[Future[Any]] | set[asyncio.Future[Any]]) -> None: + """Raise the error of a failed future among finished ones, if any. + + Otherwise, the cancellation of a cancelled future is raised, if any. The + set is emptied first. The traceback of the error keeps the frames that it + is raised through, so a frame that still referenced the future, which + keeps the error, would form a reference cycle that only the garbage + collector frees, with the local variables of every frame of the + traceback, such as the data of a write. + + Args: + done: Finished futures that the caller no longer needs. + + Raises: + BaseException: The error of a failed future, or the cancellation + error of a cancelled one. + """ + failed = next((f for f in done if not f.cancelled() and f.exception() is not None), None) + if failed is None: + failed = next((f for f in done if f.cancelled()), None) + done.clear() + try: + if failed is not None: + failed.result() + finally: + del failed + + class S3ThreadPoolExecutor(S3Executor): """Executor that delegates to a ``ThreadPoolExecutor``. @@ -245,7 +291,12 @@ async def run_shielded(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> # 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) + try: + return await asyncio.shield(future) + finally: + # The traceback of an error keeps this frame, which would form a + # reference cycle with the future of the error. + del future @override def shutdown(self, wait: bool = True) -> None: diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 7eaa1608d..054887ef9 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -7244,6 +7244,34 @@ def upload_part(**kwargs): fs.core.complete_multipart_upload.assert_not_called() assert f.multipart_upload is None + @pytest.mark.parametrize("stage", ["write", "commit"]) + def test_failed_part_freed_by_reference_counting(self, stage): + # GH-1186: the file and its data are freed as soon as the error of a + # failed part is no longer referenced, whether a write or the commit + # raised it. The traceback of the error used to keep them in a + # reference cycle with the future of the part until the garbage + # collector ran. + fs = self._make_append_fs(b"") + + def upload_part(**kwargs): + raise PermissionError("Access Denied") + + fs.core.upload_part.side_effect = upload_part + f = S3File(fs, "s3://bucket/key.txt", mode="wb", block_size=4, max_workers=1) + f.write(b"a" * 4) + file_ref = weakref.ref(f) + buffer_ref = weakref.ref(f.buffer) + fail = functools.partial(f.write, b"a" * 4) if stage == "write" else f.close + gc.disable() + try: + with pytest.raises(PermissionError, match="Access Denied"): + fail() + del f, fail + assert file_ref() is None + assert buffer_ref() is None + finally: + gc.enable() + @pytest.mark.parametrize("max_workers", [1, 4]) def test_append_failed_part_copy(self, max_workers): # GH-1178: a failed part copy of an append stops the write before it diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index 80de1537d..0056f8d6b 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -1,5 +1,7 @@ import asyncio import contextlib +import functools +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 @@ -576,6 +579,88 @@ def create_multipart_upload(*args, **kw): assert events == (["create"] if creation_fails else ["create", ("abort", "uploadid")]) sync_fs.core.upload_part_copy.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"]), + ], + ) + def test_copy_object_with_multipart_upload_loop_shutdown(self, stage, expected): + # GH-1186: when the shutdown of the event loop at the end of + # asyncio.run() cancels a copy, the cleanup still waits for the + # running requests and aborts the upload unless it completed. The + # shutdown also cancelled the tasks of the requests, 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. + assert copy is not None + deadline = time.monotonic() + 5 + while not copy.cancelling(): + 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) + + asyncio.run(main()) + + assert events == expected + @pytest.mark.parametrize( "block_size", [ @@ -2796,6 +2881,35 @@ async def test_write_failed_part_while_waiting(self, fs): fs.core.abort_multipart_upload.assert_called_once() fs.core.complete_multipart_upload.assert_not_called() + @pytest.mark.parametrize("stage", ["write", "commit"]) + @pytest.mark.asyncio + async def test_failed_part_freed_by_reference_counting(self, fs, stage): + # GH-1186: the file and its data are freed as soon as the error of a + # failed part is no longer referenced, whether a write or the commit + # raised it, as with S3File. + block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE + + def upload_part(**kwargs): + raise PermissionError("Access Denied") + + fs.core.upload_part = mock.MagicMock(side_effect=upload_part) + f = await fs.open_async("s3://bucket/key", "wb", block_size=block_size, max_workers=1) + await f.write(b"a" * block_size) + file_ref = weakref.ref(f) + buffer_ref = weakref.ref(f.buffer) + fail = functools.partial(f.write, b"a" * block_size) if stage == "write" else f.close + gc.disable() + try: + with pytest.raises(PermissionError, match="Access Denied"): + await fail() + del f, fail + # Lets the event loop run the callbacks that hold the futures. + await asyncio.sleep(0) + assert file_ref() is None + assert buffer_ref() is None + finally: + gc.enable() + @pytest.mark.parametrize("blocks", [0, 2]) @pytest.mark.asyncio async def test_exit_with_exception(self, fs, blocks): diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index b4aa849f9..4ba03cae5 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -7,15 +7,23 @@ import asyncio import contextvars +import gc import os import threading import time -from concurrent.futures import Future, ThreadPoolExecutor, wait +import traceback +import weakref +from concurrent.futures import CancelledError, Future, ThreadPoolExecutor, wait from unittest import mock import pytest -from pyathena.filesystem.s3_executor import S3AioExecutor, S3AioThreadPool, S3ThreadPoolExecutor +from pyathena.filesystem.s3_executor import ( + S3AioExecutor, + S3AioThreadPool, + S3ThreadPoolExecutor, + raise_failure, +) from tests.pyathena.util import EVENT_TIMEOUT, overlapping_calls @@ -395,6 +403,41 @@ 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, and the error is freed by reference counting. The + # traceback used to keep them in a reference cycle with the future + # until the garbage collector ran. + class Body: + pass + + class DeniedError(PermissionError): + # Unlike PermissionError, it can be weakly referenced. + pass + + def upload(body): + raise DeniedError("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 body_ref() is None + assert isinstance(error, DeniedError) + assert traceback.extract_tb(error.__traceback__)[-1].name == "upload" + error_ref = weakref.ref(error) + del error, future + assert error_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 = [] @@ -482,3 +525,54 @@ async def main(): assert not not_done assert future.cancelled() + + +@pytest.mark.parametrize( + ("outcomes", "expected"), + [ + ([], None), + (["result"], None), + (["result", "error"], PermissionError), + (["cancelled", "error"], PermissionError), + (["cancelled", "result"], CancelledError), + ], +) +def test_raise_failure(outcomes, expected): + # The error of a failed future is raised before the cancellation of a + # cancelled one, and the set is emptied. + done: set[Future[int]] = set() + for outcome in outcomes: + future: Future[int] = Future() + if outcome == "result": + future.set_result(1) + elif outcome == "error": + future.set_exception(PermissionError("Access Denied")) + else: + future.cancel() + done.add(future) + + if expected is None: + raise_failure(done) + else: + with pytest.raises(expected): + raise_failure(done) + assert done == set() + + +def test_raise_failure_frees_future(): + # GH-1186: the raised error does not keep the failed future, so that + # neither forms a reference cycle with the frames of its traceback. + future: Future[int] = Future() + future.set_exception(PermissionError("Access Denied")) + future_ref = weakref.ref(future) + done = {future} + del future + gc.disable() + try: + with pytest.raises(PermissionError, match="Access Denied") as info: + raise_failure(done) + # Neither the frames of its traceback nor the set keep the future. + assert info.value.__traceback__ is not None + assert future_ref() is None + finally: + gc.enable() From 779e213ee7101854e9fd9440ea892a0cd1d8298f Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Fri, 16 Oct 2026 00:00:03 +0900 Subject: [PATCH 02/18] Document the loop shutdown of async copies and raise_failure Co-Authored-By: Claude Opus 5.5 --- docs/api/filesystem.rst | 2 ++ docs/filesystem.md | 3 ++- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/docs/api/filesystem.rst b/docs/api/filesystem.rst index 722efcd4e..298c5c5c2 100644 --- a/docs/api/filesystem.rst +++ b/docs/api/filesystem.rst @@ -58,6 +58,8 @@ S3 Executor .. autoclass:: pyathena.filesystem.s3_executor.S3AioThreadPool :members: +.. autofunction:: pyathena.filesystem.s3_executor.raise_failure + S3 Paths -------- diff --git a/docs/filesystem.md b/docs/filesystem.md index 1855bb2f9..b67d19aad 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -136,7 +136,8 @@ 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 while the CreateMultipartUpload request is waited for is ignored. One while the part From 3c94685421b353bdfeb82ace305fa8d84d8c2510 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Fri, 16 Oct 2026 00:23:18 +0900 Subject: [PATCH 03/18] Free chained errors, multipart writes, and failed creations too - submit_to() also clears the frames of the errors chained to a failed callable's error, such as the ClientError that S3Core.call() raises an OSError from, whose frames held the body of the part. - _upload_chunk() of S3File and AioS3StreamedFile no longer keeps the last submitted part in its frame, which kept a write of several parts in a cycle with the error of a failed part. - A failed creation no longer keeps the file in a cycle: S3File drops it from the frame of _initiate_upload(), and AioS3StreamedFile drops it in _discard(). - The cleanup of an async multipart copy cancels the part copies that wait for a thread before it waits for the parts and aborts. Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3.py | 5 ++ pyathena/filesystem/s3_async.py | 17 ++-- pyathena/filesystem/s3_executor.py | 34 +++++--- tests/pyathena/filesystem/test_s3.py | 29 ++++--- tests/pyathena/filesystem/test_s3_async.py | 86 ++++++++++++++++--- tests/pyathena/filesystem/test_s3_executor.py | 23 +++-- 6 files changed, 151 insertions(+), 43 deletions(-) diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 8d5e41686..6ee2e15a5 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -3178,6 +3178,8 @@ def _initiate_upload(self) -> None: # A submitted creation may still return an ID after the wait is # interrupted. The writer records it before resolving the future. self._close_without_commit(creation=creation) + # Not kept by this frame, which the traceback of its error keeps. + del creation raise def _upload_chunk(self, final: bool = False) -> bool: @@ -3260,6 +3262,9 @@ def _upload_chunk(self, final: bool = False) -> bool: body=data, **self._get_request_kwargs("upload_part"), ) + # Not kept by this frame, which the traceback of the part's + # error keeps if a later wait of this call raises it. + del part except BaseException: self._close_without_commit() raise diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 763b9e8f3..00012c383 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -651,13 +651,15 @@ async def _abort() -> 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, - # including those whose tasks were cancelled. - await asyncio.gather(*tasks, return_exceptions=True) + # stored after the abort, so the part copies that have not started + # are cancelled, and the running ones are waited for first, + # including those whose tasks were cancelled. The parts that wait + # for the semaphore do not start any more. await asyncio.gather( *[asyncio.wrap_future(c) for c in part_copies if not c.cancel()], return_exceptions=True, ) + await asyncio.gather(*tasks, return_exceptions=True) if completion is not None: await asyncio.wait([completion]) if not completion.cancelled() and completion.exception() is None: @@ -1543,6 +1545,9 @@ async def _upload_chunk(self, final: bool = False) -> bool: body=data, **upload_kwargs, ) + # Not kept by this frame, which the traceback of the part's + # error keeps if a later wait of this call raises it. + del part, waiter except BaseException: await self._close_without_commit() raise @@ -1619,8 +1624,10 @@ def _discard(self) -> None: _logger.exception( f"Failed to abort the multipart upload to s3://{self.bucket}/{self.key}." ) - # Dropped as S3File.discard() drops them, so that the file does not - # keep the errors of failed parts, whose tracebacks may keep the file. + # Dropped as S3File.discard() drops the parts, so that the file does + # not keep the errors of failed requests, whose tracebacks may keep + # the file. + self._creation = None self._parts = [] self._uploads_in_flight = set() self._copies_in_flight = set() diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index 8d4714630..e74bf25c1 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -62,10 +62,11 @@ def submit_to(self, future: Future[T], fn: Callable[..., T], *args: Any, **kwarg the executor reported. The error of a failed callable keeps the frames that it was raised - from in its traceback, without their local variables, so that the - arguments, such as the body of a part, are freed when the callable - fails, and the future and its error are freed without the garbage - collector once they are no longer referenced. + from in its traceback, and so do the errors that it was raised from or + while handling, without their local variables, so that the arguments, + such as the body of a part, are freed when the callable fails, and the + future and its error are freed without the garbage collector once + they are no longer referenced. Args: future: A pending future that no executor has started. @@ -82,18 +83,29 @@ def run() -> None: try: result = fn(*args, **kwargs) except BaseException as e: - # The frames below this one hold the arguments. + # The frames below this one hold the arguments, and so may + # those of the errors chained to it, such as the ClientError + # that S3Core.call() raises an OSError from. traceback.clear_frames(cast(TracebackType, e.__traceback__).tb_next) + chained, seen = [e.__cause__, e.__context__], {id(e)} + while chained: + if (error := chained.pop()) is not None and id(error) not in seen: + seen.add(id(error)) + traceback.clear_frames(error.__traceback__) + chained += [error.__cause__, error.__context__] # The traceback also keeps this frame, which would keep the - # arguments and form a reference cycle with the future. Once - # run() has resolved the future, settle() does not use it. + # arguments and form a reference cycle with the future, and so + # do the frames of the executor that called run(), and the + # done callback of the submission. The variables are cleared + # in the closure that they share; settle() reads the future + # only when run() did not run. resolved = future - del future, fn, args, kwargs + del future, fn, args, kwargs, chained, error resolved.set_exception(e) del resolved - return - # Deleted only on the failure path, which returns. - future.set_result(result) # noqa: F821 + else: + # Deleted only on the failure path. + future.set_result(result) # noqa: F821 def settle(submitted: Future[None]) -> None: """Resolve the future if the executor dropped the callable before it ran. diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 054887ef9..380a5ea43 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -7244,29 +7244,34 @@ def upload_part(**kwargs): fs.core.complete_multipart_upload.assert_not_called() assert f.multipart_upload is None - @pytest.mark.parametrize("stage", ["write", "commit"]) - def test_failed_part_freed_by_reference_counting(self, stage): + @pytest.mark.parametrize("stage", ["creation", "write", "commit"]) + def test_failed_request_freed_by_reference_counting(self, stage): # GH-1186: the file and its data are freed as soon as the error of a - # failed part is no longer referenced, whether a write or the commit - # raised it. The traceback of the error used to keep them in a - # reference cycle with the future of the part until the garbage - # collector ran. + # failed request is no longer referenced, whether the creation of the + # upload, a write of several parts, or the commit raised it. The + # traceback of the error used to keep them in a reference cycle with + # the future of the request until the garbage collector ran. fs = self._make_append_fs(b"") - def upload_part(**kwargs): + def fail(*args, **kwargs): raise PermissionError("Access Denied") - fs.core.upload_part.side_effect = upload_part + if stage == "creation": + fs.core.create_multipart_upload.side_effect = fail + else: + fs.core.upload_part.side_effect = fail f = S3File(fs, "s3://bucket/key.txt", mode="wb", block_size=4, max_workers=1) - f.write(b"a" * 4) + if stage == "commit": + f.write(b"a" * 4) file_ref = weakref.ref(f) buffer_ref = weakref.ref(f.buffer) - fail = functools.partial(f.write, b"a" * 4) if stage == "write" else f.close + # The write submits a part before it waits for the failed one. + operation = f.close if stage == "commit" else functools.partial(f.write, b"a" * 12) gc.disable() try: with pytest.raises(PermissionError, match="Access Denied"): - fail() - del f, fail + operation() + del f, operation assert file_ref() is None assert buffer_ref() is None finally: diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index 0056f8d6b..ece72f2a0 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -579,6 +579,66 @@ def create_multipart_upload(*args, **kw): assert events == (["create"] if creation_fails else ["create", ("abort", "uploadid")]) sync_fs.core.upload_part_copy.assert_not_called() + @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"), [ @@ -2881,28 +2941,34 @@ async def test_write_failed_part_while_waiting(self, fs): fs.core.abort_multipart_upload.assert_called_once() fs.core.complete_multipart_upload.assert_not_called() - @pytest.mark.parametrize("stage", ["write", "commit"]) + @pytest.mark.parametrize("stage", ["creation", "write", "commit"]) @pytest.mark.asyncio - async def test_failed_part_freed_by_reference_counting(self, fs, stage): + async def test_failed_request_freed_by_reference_counting(self, fs, stage): # GH-1186: the file and its data are freed as soon as the error of a - # failed part is no longer referenced, whether a write or the commit - # raised it, as with S3File. + # failed request is no longer referenced, as with S3File. block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE - def upload_part(**kwargs): + def fail(*args, **kwargs): raise PermissionError("Access Denied") - fs.core.upload_part = mock.MagicMock(side_effect=upload_part) + if stage == "creation": + fs.core.create_multipart_upload = mock.MagicMock(side_effect=fail) + else: + fs.core.upload_part = mock.MagicMock(side_effect=fail) f = await fs.open_async("s3://bucket/key", "wb", block_size=block_size, max_workers=1) - await f.write(b"a" * block_size) + if stage == "commit": + await f.write(b"a" * block_size) file_ref = weakref.ref(f) buffer_ref = weakref.ref(f.buffer) - fail = functools.partial(f.write, b"a" * block_size) if stage == "write" else f.close + # The write submits a part before it waits for the failed one. + operation = ( + f.close if stage == "commit" else functools.partial(f.write, b"a" * block_size * 3) + ) gc.disable() try: with pytest.raises(PermissionError, match="Access Denied"): - await fail() - del f, fail + await operation() + del f, operation # Lets the event loop run the callbacks that hold the futures. await asyncio.sleep(0) assert file_ref() is None diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index 4ba03cae5..2cb81765d 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -403,13 +403,16 @@ def test_submit_to(self): with pytest.raises(ValueError, match="invalid literal"): failed.result() + @pytest.mark.parametrize("chained", [False, True]) @pytest.mark.parametrize("executor_class", [S3ThreadPoolExecutor, S3AioThreadPool]) - def test_submit_to_failure_frees_arguments(self, executor_class): + def test_submit_to_failure_frees_arguments(self, executor_class, chained): # 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, and the error is freed by reference counting. The - # traceback used to keep them in a reference cycle with the future - # until the garbage collector ran. + # its traceback, also when the frames of the error that it was raised + # from hold them, as S3Core.call() raises an OSError from a + # ClientError. The error is freed by reference counting. The traceback + # used to keep them in a reference cycle with the future until the + # garbage collector ran. class Body: pass @@ -417,8 +420,16 @@ class DeniedError(PermissionError): # Unlike PermissionError, it can be weakly referenced. pass + def send(body): + raise RuntimeError("AccessDenied") + def upload(body): - raise DeniedError("Access Denied") + if not chained: + raise DeniedError("Access Denied") + try: + send(body) + except RuntimeError as e: + raise DeniedError("Access Denied") from e body = Body() body_ref = weakref.ref(body) @@ -432,6 +443,8 @@ def upload(body): assert body_ref() is None assert isinstance(error, DeniedError) assert traceback.extract_tb(error.__traceback__)[-1].name == "upload" + if chained: + assert traceback.extract_tb(error.__cause__.__traceback__)[-1].name == "send" error_ref = weakref.ref(error) del error, future assert error_ref() is None From 1cce639d38cacf96f870cdec81526d061126ad4e Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Fri, 16 Oct 2026 00:33:21 +0900 Subject: [PATCH 04/18] Submit parts from a method and use raise_failure for copied parts - S3File and AioS3StreamedFile submit an uploaded part from _submit_upload(), so that no frame of _upload_chunk() keeps the future of a part whose error a later wait raises, instead of deleting it. - S3FileSystem._finish_multipart_upload() raises a failed part copy with raise_failure(), as the commit of a file does. - submit_to() clears the frames of the error and its chained errors in one walk; clear_frames() skips the frame of run(), which is executing. - The cleanup of an async multipart copy waits for the part copies and their tasks together. - The reference counting test of S3File also covers a write after a failed part copy of an append. Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3.py | 50 ++++++++++++++----------- pyathena/filesystem/s3_async.py | 55 ++++++++++++++++++---------- pyathena/filesystem/s3_executor.py | 14 +++---- tests/pyathena/filesystem/test_s3.py | 21 +++++++---- 4 files changed, 86 insertions(+), 54 deletions(-) diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 6ee2e15a5..3112dcf1f 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -1850,11 +1850,9 @@ def _finish_multipart_upload( request_kwargs = request_kwargs or {} try: done, _ = wait(futures, return_when=FIRST_EXCEPTION) - for future in done: - if not future.cancelled() and future.exception() is not None: - # Raised without waiting for the parts before it, so that - # the parts queued after it are cancelled. - future.result() + # Raised without waiting for the parts before it, so that the + # parts queued after it are cancelled. + raise_failure(done) # The futures are in part-number order. parts = [future.result() for future in futures] return self.core.complete_multipart_upload( @@ -3249,22 +3247,7 @@ def _upload_chunk(self, final: bool = False) -> bool: # A failed part would otherwise keep its data until # commit() while the rest of the data is uploaded. raise_failure(done) - # Tracked before it is scheduled, so that an interrupt while - # it is being scheduled still lets discard() cancel it or wait - # for it before the abort. - 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"), - ) - # Not kept by this frame, which the traceback of the part's - # error keeps if a later wait of this call raises it. - del part + self._submit_upload(writer, part_number, data) except BaseException: self._close_without_commit() raise @@ -3273,6 +3256,31 @@ def _upload_chunk(self, final: bool = False) -> bool: self.commit() return not final + def _submit_upload(self, writer: S3MultipartWriter, part_number: int, data: bytes) -> None: + """Submit the upload of a part and track it as in flight. + + The part is tracked before it is scheduled, so that an interrupt while + it is being scheduled still lets :meth:`discard` cancel it or wait for + it before the abort. Its future is not left in the frame of + :meth:`_upload_chunk`, which the traceback of the part's error keeps + if a later wait of the same call raises it. + + Args: + writer: The writer of the multipart upload. + part_number: The number of the part, from 1. + data: The data of the part. + """ + 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"), + ) + def commit(self) -> None: """Complete the upload of the written data. diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 00012c383..e98ec88ef 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -657,9 +657,9 @@ async def _abort() -> None: # for the semaphore do not start any more. await asyncio.gather( *[asyncio.wrap_future(c) for c in part_copies if not c.cancel()], + *tasks, return_exceptions=True, ) - await asyncio.gather(*tasks, return_exceptions=True) if completion is not None: await asyncio.wait([completion]) if not completion.cancelled() and completion.exception() is None: @@ -1530,24 +1530,7 @@ async def _upload_chunk(self, final: bool = False) -> bool: # A failed part would otherwise keep its data until the # commit while the rest of the data is uploaded. raise_failure(done) - # Tracked before it is scheduled; see _initiate_upload(). - part: Future[S3MultipartUploadPart] = Future() - self._parts.append(part) - waiter = asyncio.wrap_future(part) - # The commit reports a failure that no wait found; retrieved - # 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, - ) - # Not kept by this frame, which the traceback of the part's - # error keeps if a later wait of this call raises it. - del part, waiter + self._submit_upload(writer, part_number, data, upload_kwargs) except BaseException: await self._close_without_commit() raise @@ -1555,6 +1538,40 @@ async def _upload_chunk(self, final: bool = False) -> bool: await self._commit() return not final + def _submit_upload( + self, + writer: S3MultipartWriter, + part_number: int, + data: bytes, + request_kwargs: dict[str, Any], + ) -> None: + """Submit the upload of a part and track it as in flight. + + See :meth:`S3File._submit_upload`; the file waits for the part + through an asyncio future. + + Args: + writer: The writer of the multipart upload. + part_number: The number of the part, from 1. + data: The data of the part. + request_kwargs: The UploadPart parameters of the file. + """ + # Tracked before it is scheduled; see _initiate_upload(). + part: Future[S3MultipartUploadPart] = Future() + self._parts.append(part) + waiter = asyncio.wrap_future(part) + # The commit reports a failure that no wait found; retrieved 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, + **request_kwargs, + ) + async def _commit(self) -> None: """Commit the written data in a thread of the file. diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index e74bf25c1..dfc68f784 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -22,8 +22,7 @@ from concurrent.futures import Future from concurrent.futures.thread import ThreadPoolExecutor from multiprocessing import cpu_count -from types import TracebackType -from typing import Any, ClassVar, TypeVar, cast +from typing import Any, ClassVar, TypeVar from pyathena.util import override @@ -85,9 +84,10 @@ def run() -> None: except BaseException as e: # The frames below this one hold the arguments, and so may # those of the errors chained to it, such as the ClientError - # that S3Core.call() raises an OSError from. - traceback.clear_frames(cast(TracebackType, e.__traceback__).tb_next) - chained, seen = [e.__cause__, e.__context__], {id(e)} + # that S3Core.call() raises an OSError from. clear_frames() + # skips this frame, which is still executing. + chained: list[BaseException | None] = [e] + seen: set[int] = set() while chained: if (error := chained.pop()) is not None and id(error) not in seen: seen.add(id(error)) @@ -139,10 +139,10 @@ def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: def raise_failure(done: set[Future[Any]] | set[asyncio.Future[Any]]) -> None: - """Raise the error of a failed future among finished ones, if any. + """Empty a set of finished futures and raise the error of a failed one, if any. Otherwise, the cancellation of a cancelled future is raised, if any. The - set is emptied first. The traceback of the error keeps the frames that it + set is emptied before the error is raised. The traceback of the error keeps the frames that it is raised through, so a frame that still referenced the future, which keeps the error, would form a reference cycle that only the garbage collector frees, with the local variables of every frame of the diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 380a5ea43..3dccb0e36 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -7244,25 +7244,32 @@ def upload_part(**kwargs): fs.core.complete_multipart_upload.assert_not_called() assert f.multipart_upload is None - @pytest.mark.parametrize("stage", ["creation", "write", "commit"]) + @pytest.mark.parametrize("stage", ["creation", "write", "copy", "commit"]) def test_failed_request_freed_by_reference_counting(self, stage): # GH-1186: the file and its data are freed as soon as the error of a # failed request is no longer referenced, whether the creation of the - # upload, a write of several parts, or the commit raised it. The - # traceback of the error used to keep them in a reference cycle with - # the future of the request until the garbage collector ran. - fs = self._make_append_fs(b"") + # upload, a write of several parts, a write after a failed part copy + # of an append, or the commit raised it. The traceback of the error + # used to keep them in a reference cycle with the future of the + # request until the garbage collector ran. + fs = self._make_append_fs(b"a" * 6 if stage == "copy" else b"") def fail(*args, **kwargs): raise PermissionError("Access Denied") if stage == "creation": fs.core.create_multipart_upload.side_effect = fail + elif stage == "copy": + fs.core.upload_part_copy.side_effect = fail else: fs.core.upload_part.side_effect = fail - f = S3File(fs, "s3://bucket/key.txt", mode="wb", block_size=4, max_workers=1) - if stage == "commit": + mode = "ab" if stage == "copy" else "wb" + f = S3File(fs, "s3://bucket/key.txt", mode=mode, block_size=4, max_workers=1) + if stage in ("copy", "commit"): f.write(b"a" * 4) + if stage == "copy": + # The part copy has failed before the next write. + assert wait(f.multipart_upload_parts[:1], timeout=5).done file_ref = weakref.ref(f) buffer_ref = weakref.ref(f.buffer) # The write submits a part before it waits for the failed one. From 8dc6ae9a8fa2658e6d58e28d0f221538d26e3832 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Fri, 16 Oct 2026 08:17:22 +0900 Subject: [PATCH 05/18] Clear grouped errors and retrieve every error in raise_failure - submit_to() also clears the frames of the errors that an exception group of the callable holds. - raise_failure() retrieves the errors of all the failed futures in one pass, so that asyncio does not log those that it does not raise, and no longer fails on an empty set. - The append case of the reference counting test holds the part copy until the first write has returned, as the other append tests do. Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3_executor.py | 40 +++++++++------ tests/pyathena/filesystem/test_s3.py | 23 ++++++--- tests/pyathena/filesystem/test_s3_executor.py | 51 +++++++++++++++---- 3 files changed, 82 insertions(+), 32 deletions(-) diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index dfc68f784..d31cea86f 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -61,11 +61,11 @@ def submit_to(self, future: Future[T], fn: Callable[..., T], *args: Any, **kwarg the executor reported. The error of a failed callable keeps the frames that it was raised - from in its traceback, and so do the errors that it was raised from or - while handling, without their local variables, so that the arguments, - such as the body of a part, are freed when the callable fails, and the - future and its error are freed without the garbage collector once - they are no longer referenced. + from in its traceback, and so do the errors that it was raised from, + raised while handling, or groups, without their local variables, so + that the arguments, such as the body of a part, are freed when the + callable fails, and the future and its error are freed without the + garbage collector once they are no longer referenced. Args: future: A pending future that no executor has started. @@ -84,8 +84,8 @@ def run() -> None: except BaseException as e: # The frames below this one hold the arguments, and so may # those of the errors chained to it, such as the ClientError - # that S3Core.call() raises an OSError from. clear_frames() - # skips this frame, which is still executing. + # that S3Core.call() raises an OSError from, or grouped in it. + # clear_frames() skips this frame, which is still executing. chained: list[BaseException | None] = [e] seen: set[int] = set() while chained: @@ -93,6 +93,8 @@ def run() -> None: seen.add(id(error)) traceback.clear_frames(error.__traceback__) chained += [error.__cause__, error.__context__] + if isinstance(error, BaseExceptionGroup): + chained += error.exceptions # The traceback also keeps this frame, which would keep the # arguments and form a reference cycle with the future, and so # do the frames of the executor that called run(), and the @@ -142,9 +144,11 @@ def raise_failure(done: set[Future[Any]] | set[asyncio.Future[Any]]) -> None: """Empty a set of finished futures and raise the error of a failed one, if any. Otherwise, the cancellation of a cancelled future is raised, if any. The - set is emptied before the error is raised. The traceback of the error keeps the frames that it - is raised through, so a frame that still referenced the future, which - keeps the error, would form a reference cycle that only the garbage + errors of all the failed futures are retrieved, so that asyncio does not + log those that are not raised as never retrieved. The set is emptied + before the error is raised: the traceback of the error keeps the frames + that it is raised through, so a frame that still referenced the future, + which keeps the error, would form a reference cycle that only the garbage collector frees, with the local variables of every frame of the traceback, such as the data of a write. @@ -155,15 +159,19 @@ def raise_failure(done: set[Future[Any]] | set[asyncio.Future[Any]]) -> None: BaseException: The error of a failed future, or the cancellation error of a cancelled one. """ - failed = next((f for f in done if not f.cancelled() and f.exception() is not None), None) - if failed is None: - failed = next((f for f in done if f.cancelled()), None) + future: Future[Any] | asyncio.Future[Any] | None = None + failed = cancelled = None + for future in done: + if future.cancelled(): + cancelled = cancelled or future + elif future.exception() is not None: + failed = failed or future done.clear() try: - if failed is not None: - failed.result() + if (raised := failed or cancelled) is not None: + raised.result() finally: - del failed + del future, failed, cancelled, raised class S3ThreadPoolExecutor(S3Executor): diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 3dccb0e36..e89e10046 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -7253,23 +7253,34 @@ def test_failed_request_freed_by_reference_counting(self, stage): # used to keep them in a reference cycle with the future of the # request until the garbage collector ran. fs = self._make_append_fs(b"a" * 6 if stage == "copy" else b"") + release = threading.Event() def fail(*args, **kwargs): raise PermissionError("Access Denied") + def fail_copy(**kwargs): + # Copying until the first write has returned. The test always + # sets the event, so the wait needs no timeout. + release.wait() + fail() + if stage == "creation": fs.core.create_multipart_upload.side_effect = fail elif stage == "copy": - fs.core.upload_part_copy.side_effect = fail + fs.core.upload_part_copy.side_effect = fail_copy else: fs.core.upload_part.side_effect = fail mode = "ab" if stage == "copy" else "wb" f = S3File(fs, "s3://bucket/key.txt", mode=mode, block_size=4, max_workers=1) - if stage in ("copy", "commit"): - f.write(b"a" * 4) - if stage == "copy": - # The part copy has failed before the next write. - assert wait(f.multipart_upload_parts[:1], timeout=5).done + try: + if stage in ("copy", "commit"): + f.write(b"a" * 4) + release.set() + if stage == "copy": + # The part copy has failed before the next write. + assert wait(f.multipart_upload_parts[:1], timeout=5).done + finally: + release.set() file_ref = weakref.ref(f) buffer_ref = weakref.ref(f.buffer) # The write submits a part before it waits for the failed one. diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index 2cb81765d..c4d8f0380 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -403,16 +403,16 @@ def test_submit_to(self): with pytest.raises(ValueError, match="invalid literal"): failed.result() - @pytest.mark.parametrize("chained", [False, True]) + @pytest.mark.parametrize("chained", [None, "cause", "group"]) @pytest.mark.parametrize("executor_class", [S3ThreadPoolExecutor, S3AioThreadPool]) def test_submit_to_failure_frees_arguments(self, executor_class, chained): # 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, also when the frames of the error that it was raised - # from hold them, as S3Core.call() raises an OSError from a - # ClientError. The error is freed by reference counting. The traceback - # used to keep them in a reference cycle with the future until the - # garbage collector ran. + # its traceback, also when the frames of an error that it was raised + # from, as S3Core.call() raises an OSError from a ClientError, or that + # it groups hold them. The error is freed by reference counting. The + # traceback used to keep them in a reference cycle with the future + # until the garbage collector ran. class Body: pass @@ -424,12 +424,16 @@ def send(body): raise RuntimeError("AccessDenied") def upload(body): - if not chained: + if chained is None: raise DeniedError("Access Denied") try: send(body) except RuntimeError as e: - raise DeniedError("Access Denied") from e + if chained == "cause": + raise DeniedError("Access Denied") from e + error = e + # Raised outside the handler, so that the group has no context. + raise ExceptionGroup("Access Denied", [error]) body = Body() body_ref = weakref.ref(body) @@ -441,10 +445,16 @@ def upload(body): del body error = future.exception() assert body_ref() is None - assert isinstance(error, DeniedError) assert traceback.extract_tb(error.__traceback__)[-1].name == "upload" - if chained: + if chained == "cause": + assert isinstance(error.__cause__, RuntimeError) assert traceback.extract_tb(error.__cause__.__traceback__)[-1].name == "send" + elif chained == "group": + assert isinstance(error, ExceptionGroup) + assert error.__context__ is None + assert traceback.extract_tb(error.exceptions[0].__traceback__)[-1].name == "send" + else: + assert isinstance(error, DeniedError) error_ref = weakref.ref(error) del error, future assert error_ref() is None @@ -589,3 +599,24 @@ def test_raise_failure_frees_future(): assert future_ref() is None finally: gc.enable() + + +def test_raise_failure_retrieves_errors(): + # The errors of the failed asyncio futures that are not raised are + # retrieved, so that asyncio does not log them as never retrieved. + loop = asyncio.new_event_loop() + try: + logged = [] + loop.set_exception_handler(lambda _, context: logged.append(context)) + done: set[asyncio.Future[int]] = set() + for _ in range(2): + future = loop.create_future() + future.set_exception(PermissionError("Access Denied")) + done.add(future) + del future + with pytest.raises(PermissionError, match="Access Denied"): + raise_failure(done) + gc.collect() + assert logged == [] + finally: + loop.close() From 1fc62be2cdb04548ce993239abbf0e99e2e0672d Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Fri, 16 Oct 2026 08:28:45 +0900 Subject: [PATCH 06/18] Run the cleanup of an async multipart copy in a filesystem thread The shutdown of the event loop at the end of asyncio.run() cancels every remaining task, including the cleanup task that a cancellation of the copy has started, which then returned without aborting the upload. The cleanup now runs in a thread of the filesystem through run_shielded(), waiting for the requests as concurrent futures, so the event loop can no longer stop it, and _cleanup_tasks is no longer needed. The creation and the part copies that have not started are cancelled before the cleanup is scheduled, as S3FileSystem does, and the part tasks retrieve their errors themselves. Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3_async.py | 104 ++++++++++----------- tests/pyathena/filesystem/test_s3_async.py | 25 +++-- 2 files changed, 66 insertions(+), 63 deletions(-) diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index e98ec88ef..c8c88c0ea 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -54,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): @@ -581,12 +579,14 @@ async def _copy_object_with_multipart_upload( 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. This also holds when - the shutdown of the event loop, such as at the end of - ``asyncio.run()``, cancels the tasks of the copy. A repeated - cancellation returns without stopping this cleanup. + creation of the upload and the part copies that have not started are + cancelled, the running ones and the completion 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. @@ -642,50 +642,45 @@ async def _upload_part(i: int, range_: tuple[int, int]) -> S3MultipartUploadPart 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 the part copies that have not started - # are cancelled, and the running ones are waited for first, - # including those whose tasks were cancelled. The parts that wait - # for the semaphore do not start any more. - await asyncio.gather( - *[asyncio.wrap_future(c) for c in part_copies if not c.cancel()], - *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, - ) - - # Not a task, which the shutdown of the event loop would cancel while - # its request keeps running, so that _abort() can wait for an upload - # that is created after a cancellation; scheduled right before the - # try, so that nothing can fail between the two. - creation = asyncio.wrap_future( - self._thread_pool.submit( - self.core.create_multipart_upload, plan.destination, **plan.create_params - ) + self._sync_fs._abort_multipart_upload(creation.result(), plan.abort_params) + + # 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 = 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) + # shield keeps a cancellation from cancelling the creation, so + # that _abort() can wait for an upload that it creates. + multipart_upload = await asyncio.shield(asyncio.wrap_future(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(). @@ -695,28 +690,27 @@ async def _abort() -> None: raise error # The tasks are in part-number order. parts = [task.result() for task in tasks] - # Not a task either; see the creation. - completion = asyncio.wrap_future( - self._thread_pool.submit( - 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) + # shield keeps a cancellation from cancelling the completion, so + # that _abort() can wait for it. + completed = await asyncio.shield(asyncio.wrap_future(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() # 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 diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index ece72f2a0..e783f8f9f 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -647,13 +647,15 @@ def upload_part_copy(**kw): ("complete", ["start complete", "end complete"]), ], ) - def test_copy_object_with_multipart_upload_loop_shutdown(self, stage, expected): + @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, the cleanup still waits for the - # running requests and aborts the upload unless it completed. The - # shutdown also cancelled the tasks of the requests, so the cleanup - # left an upload created during the cancellation, aborted before - # the running parts ended, or aborted during the completion. + # 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 = [] @@ -663,10 +665,11 @@ def test_copy_object_with_multipart_upload_loop_shutdown(self, stage, expected): def request(name): with lock: events.append(f"start {name}") - # Held until asyncio.run() has cancelled the copy. + # 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 not copy.cancelling(): + while copy.cancelling() < (2 if cancelled else 1): assert time.monotonic() < deadline time.sleep(0.01) with lock: @@ -716,8 +719,14 @@ async def main(): 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 From 30d45e6886541e4f5ca2eb45dd84d4cd111ab35c Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Fri, 16 Oct 2026 09:15:09 +0900 Subject: [PATCH 07/18] Wait for shielded requests without logging their late failures After a cancellation, shield() leaves the asyncio wrapper of a creation or completion that fails later unretrieved, and Python 3.15 also logs such a failure itself. The new S3AioThreadPool.wait_shielded() waits for a submitted call with asyncio.wait(), which leaves the call alone when the waiting task is cancelled, and retrieves its outcome. The async multipart copy and run_shielded() use it. Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3_async.py | 12 +++---- pyathena/filesystem/s3_executor.py | 36 ++++++++++++++----- tests/pyathena/filesystem/test_s3_async.py | 18 ++++++++++ tests/pyathena/filesystem/test_s3_executor.py | 36 +++++++++++++++++++ 4 files changed, 88 insertions(+), 14 deletions(-) diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index c8c88c0ea..41c17e22b 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -673,9 +673,9 @@ def _abort() -> None: self.core.create_multipart_upload, plan.destination, **plan.create_params ) try: - # shield keeps a cancellation from cancelling the creation, so - # that _abort() can wait for an upload that it creates. - multipart_upload = await asyncio.shield(asyncio.wrap_future(creation)) + # A cancellation does not cancel the creation, so that _abort() + # can wait for an upload that it 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 @@ -696,9 +696,9 @@ def _abort() -> None: cast(list[S3MultipartUploadPart], parts), **plan.complete_params, ) - # shield keeps a cancellation from cancelling the completion, so - # that _abort() can wait for it. - completed = await asyncio.shield(asyncio.wrap_future(completion)) + # A cancellation does not cancel the completion, so that _abort() + # can wait for it. + completed = await self._thread_pool.wait_shielded(completion) except BaseException: # Also on cancellation, as S3FileSystem._finish_multipart_upload # does on an interrupt. diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index d31cea86f..aaf56d335 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -294,7 +294,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. @@ -307,16 +308,35 @@ 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 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()) try: - return await asyncio.shield(future) + # Unlike shield(), which also logs the failure of a future that + # ends after a cancellation since Python 3.15, wait() leaves the + # future alone when the waiting task is cancelled. + await asyncio.wait([waiter]) + return waiter.result() finally: # The traceback of an error keeps this frame, which would form a - # reference cycle with the future of the error. - del future + # reference cycle with the futures of the error. + del future, waiter @override def shutdown(self, wait: bool = True) -> None: diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index e783f8f9f..92985f141 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -494,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"), @@ -518,8 +522,13 @@ def complete_multipart_upload(*args, **kw): release.set() with pytest.raises(asyncio.CancelledError): await task + gc.collect() + 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 @@ -551,6 +560,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"), @@ -575,8 +588,13 @@ def create_multipart_upload(*args, **kw): release.set() with pytest.raises(asyncio.CancelledError): await task + gc.collect() + 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 diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index c4d8f0380..11489f8e9 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -95,6 +95,42 @@ 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() + release.wait(EVENT_TIMEOUT) + 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) + waiting = asyncio.ensure_future(pool.wait_shielded(running)) + await asyncio.to_thread(started.wait, EVENT_TIMEOUT) + waiting.cancel() + with pytest.raises(asyncio.CancelledError): + await waiting + release.set() + with pytest.raises(PermissionError, match="Access Denied"): + await asyncio.wrap_future(running) + del waiting + gc.collect() + # Lets the event loop run the callbacks of the failure. + await asyncio.sleep(0) + return logged + + assert asyncio.run(main()) == [] + pool.shutdown() + def test_shutdown(self): # The calls submitted before the shutdown still run. events = [] From 6b9253aa7fdae8c335e4e55950bf090d32bbfcd0 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Fri, 16 Oct 2026 09:16:46 +0900 Subject: [PATCH 08/18] Correct the Python version in the comment of wait_shielded Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3_executor.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index aaf56d335..c8e9a88f7 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -328,9 +328,9 @@ async def wait_shielded(self, future: Future[T]) -> T: waiter = asyncio.wrap_future(future) waiter.add_done_callback(lambda f: f.cancelled() or f.exception()) try: - # Unlike shield(), which also logs the failure of a future that - # ends after a cancellation since Python 3.15, wait() leaves the - # future alone when the waiting task is cancelled. + # 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() finally: From 395cb07283249549e4e78268173cca4c7f27c7a8 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Fri, 16 Oct 2026 09:38:33 +0900 Subject: [PATCH 09/18] Check the late failure logs after the threads have run their callbacks The tests of a late failure after a cancellation shut the pool down before they check the event loop's log, so that the worker has run the callbacks of the failed request, and drop the futures that keep its asyncio wrapper. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/filesystem/test_s3_async.py | 12 ++++++++++-- tests/pyathena/filesystem/test_s3_executor.py | 13 ++++++++----- 2 files changed, 18 insertions(+), 7 deletions(-) diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index 92985f141..ceae08a90 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -522,8 +522,12 @@ 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() - await asyncio.sleep(0) + 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 @@ -588,8 +592,12 @@ 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() - await asyncio.sleep(0) + 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 diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index 11489f8e9..c438ba6b6 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -120,16 +120,19 @@ async def main(): with pytest.raises(asyncio.CancelledError): await waiting 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"): - await asyncio.wrap_future(running) - del waiting + running.result() + # The callbacks of the future of the call keep its waiter. + del waiting, running gc.collect() - # Lets the event loop run the callbacks of the failure. - await asyncio.sleep(0) + for _ in range(3): + await asyncio.sleep(0) return logged assert asyncio.run(main()) == [] - pool.shutdown() def test_shutdown(self): # The calls submitted before the shutdown still run. From 0975505831910fc2b83d271e0266b0438735aa9f Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Fri, 16 Oct 2026 10:00:31 +0900 Subject: [PATCH 10/18] Wait until the task waits for the call before cancelling it test_wait_shielded_cancelled cancels the waiting task once it has added its callback to the future of the call, so that the cancellation always reaches wait_shielded(); otherwise a regression to shield() could pass on Python 3.15. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/filesystem/test_s3_executor.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index c438ba6b6..eaab600f9 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -116,6 +116,9 @@ async def main(): running = pool.submit(work) 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. + while not running._done_callbacks: + await asyncio.sleep(0) waiting.cancel() with pytest.raises(asyncio.CancelledError): await waiting From d9a9a8ccceec416802d27bbae7a8f67740b12c21 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Fri, 16 Oct 2026 10:11:30 +0900 Subject: [PATCH 11/18] Bound the wait of the shielded test and always release its call The call waits for the test without a timeout, so it cannot finish before the task waits for it, and the test releases the call and shuts the pool down in a finally block. The wait for the task to add its callback fails after EVENT_TIMEOUT instead of looping forever. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/filesystem/test_s3_executor.py | 32 +++++++++++-------- 1 file changed, 19 insertions(+), 13 deletions(-) diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index eaab600f9..aee2a2148 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -105,7 +105,8 @@ def test_wait_shielded_cancelled(self): def work(): started.set() - release.wait(EVENT_TIMEOUT) + # The finally block of the test always releases it. + release.wait() raise PermissionError("Access Denied") async def main(): @@ -114,18 +115,23 @@ async def main(): lambda _, context: logged.append(context["message"]) ) running = pool.submit(work) - 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. - while not running._done_callbacks: - await asyncio.sleep(0) - waiting.cancel() - with pytest.raises(asyncio.CancelledError): - await waiting - 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) + 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. From 0c2f5ea35764b2c8905c8e2907d74eb71a410b8f Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Fri, 16 Oct 2026 23:57:51 +0900 Subject: [PATCH 12/18] Free failed reads and retried API errors without the GC - S3File._fetch_range() waits for the parallel reads with FIRST_EXCEPTION, cancels the reads that have not started when one fails, and drops the futures of the others before raising, instead of keeping them in its frame in a cycle with the error. - AioS3StreamedFile._fetch_range() does the same with its read tasks, which gather() left running after a failure, and close() drops the cache of a read file, as fsspec does, which kept the file in a cycle. - S3AioExecutor.submit() clears the variables of its run() closure when the function fails, as ThreadPoolExecutor's _WorkItem does, and returns the future through a variable that run() does not clear. - retry_api_call() clears the frames of tenacity in the traceback of a raised error; their retry state kept the error, and so every frame of its traceback, in a reference cycle. - raise_failure() accepts a set of any kind of future, such as tasks. Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3.py | 15 +++-- pyathena/filesystem/s3_async.py | 33 +++++++-- pyathena/filesystem/s3_executor.py | 22 ++++-- pyathena/util.py | 17 ++++- tests/pyathena/filesystem/test_s3.py | 28 ++++++++ tests/pyathena/filesystem/test_s3_async.py | 78 ++++++++++++++++++++++ tests/pyathena/test_util.py | 26 ++++++++ 7 files changed, 204 insertions(+), 15 deletions(-) diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 3112dcf1f..de487f95a 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -3423,7 +3423,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. @@ -3445,9 +3447,14 @@ 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() + done = wait(futures, return_when=FIRST_EXCEPTION).done + if any(not f.cancelled() and f.exception() is not None for f in done): + # The reads that have not started are cancelled, and the + # futures of the others are dropped, so that the traceback of + # the error, which keeps this frame, does not keep their data. + while futures: + futures.pop().cancel() + raise_failure(done) 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 41c17e22b..afd6f571c 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -1408,6 +1408,10 @@ async def close(self) -> None: await self.flush(force=True) finally: self.closed = True + if self.mode == "rb": + # Dropped as fsspec's close() drops it, since the cache keeps + # a method of the file and so a reference cycle with it. + self.cache = None # A commit that is still running keeps its thread. self._pool.shutdown(wait=False) @@ -1647,7 +1651,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. @@ -1664,7 +1670,24 @@ 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: + # The reads are cancelled with the read of the range. + while reads: + reads.pop().cancel() + raise + if any(not read.cancelled() and read.exception() is not None for read in done): + # The reads that have not finished are cancelled, and the tasks of + # the others are dropped, so that the traceback of the error, which + # keeps this frame, does not keep their data. + while reads: + reads.pop().cancel() + raise_failure(done) + return b"".join(read.result() for read in reads) diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index c8e9a88f7..58c7ef73a 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -27,6 +27,8 @@ from pyathena.util import override T = TypeVar("T") +# A future of an executor or of asyncio. +F = TypeVar("F", bound="Future[Any] | asyncio.Future[Any]") class S3Executor(metaclass=ABCMeta): @@ -140,7 +142,7 @@ def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: self.shutdown(wait=True) -def raise_failure(done: set[Future[Any]] | set[asyncio.Future[Any]]) -> None: +def raise_failure(done: set[F]) -> None: """Empty a set of finished futures and raise the error of a failed one, if any. Otherwise, the cancellation of a cancelled future is raised, if any. The @@ -159,7 +161,7 @@ def raise_failure(done: set[Future[Any]] | set[asyncio.Future[Any]]) -> None: BaseException: The error of a failed future, or the cancellation error of a cancelled one. """ - future: Future[Any] | asyncio.Future[Any] | None = None + future: F | None = None failed = cancelled = None for future in done: if future.cancelled(): @@ -444,12 +446,21 @@ def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]: def run() -> None: """Run the function and resolve the future unless it was cancelled.""" + nonlocal future, fn, args, kwargs 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) + # The traceback keeps this frame, which would form a + # reference cycle with the future, as with the _WorkItem + # of ThreadPoolExecutor. The variables are cleared in the + # closure that settle() shares, which reads the future + # only when run() has not claimed it. + resolved = future + del future, fn, args, kwargs + resolved.set_exception(e) + del resolved else: future.set_result(result) @@ -472,12 +483,15 @@ def settle(task: Future[None]) -> None: elif future.set_running_or_notify_cancel(): future.set_exception(task.exception()) + # Returned through a variable that the closures do not share, + # since run() may clear theirs before this method returns. + submitted = future 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) - return future + return submitted raise RuntimeError( "S3AioExecutor requires a running event loop. " "Use S3ThreadPoolExecutor for synchronous usage." diff --git a/pyathena/util.py b/pyathena/util.py index a88dc6bf5..037ad076a 100644 --- a/pyathena/util.py +++ b/pyathena/util.py @@ -307,7 +307,9 @@ def retry_api_call( Note: Only retries on AWS exceptions listed in the RetryConfig.exceptions. This includes recognized Glue error codes wrapped in MetadataException. - Other errors are propagated without retrying. + Other errors are propagated without retrying. The frames of tenacity + in the traceback of a raised error have no local variables, so that + the error is freed by reference counting. """ retry = tenacity.Retrying( retry=retry_if_exception( @@ -325,4 +327,15 @@ def retry_api_call( after=after_log(logger, logger.getEffectiveLevel()) if logger else None, # type: ignore[arg-type] reraise=True, ) - return retry(func, *args, **kwargs) + try: + return retry(func, *args, **kwargs) + except BaseException as e: + # The frames of tenacity in the traceback keep its retry state, which + # keeps the error, so they would form a reference cycle that keeps + # every frame of the traceback until the garbage collector runs. + tb = e.__traceback__ + while tb is not None: + if tb.tb_frame.f_globals.get("__name__", "").startswith("tenacity"): + tb.tb_frame.clear() + tb = tb.tb_next + raise diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index e89e10046..9bc8c5672 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -3428,6 +3428,34 @@ 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_freed_by_reference_counting(self): + # GH-1186: a parallel read raises the error of a failed range, and + # the file is freed by reference counting once it is closed and the + # error is no longer referenced. The traceback of the error used to + # keep the file, and the data of the other ranges, in a reference + # cycle with the futures of the reads. + data = bytes(range(64)) + fs, _ = self._make_object_fs(data) + get_object = fs._client.get_object.side_effect + + def fail_first_range(**request): + if request["Range"].startswith("bytes=0-"): + raise PermissionError("Access Denied") + return get_object(**request) + + fs._client.get_object.side_effect = fail_first_range + f = fs.open("s3://bucket/key", "rb", cache_type="none", block_size=16, max_workers=4) + file_ref = weakref.ref(f) + gc.disable() + try: + with pytest.raises(PermissionError, match="Access Denied"): + f.read() + f.close() + del f + assert file_ref() is None + finally: + gc.enable() + @pytest.mark.parametrize("cache_type", ["bytes", "all", "first"]) def test_open_directory(self, cache_type): fs = self._make_fs() diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index ceae08a90..acb41bcc2 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -2630,6 +2630,47 @@ 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_freed_by_reference_counting(self): + # GH-1186: a parallel read raises the error of a failed range, and + # the file is freed by reference counting once it is closed and the + # error is no longer referenced. The executor of the file kept the + # future of the failed read in a reference cycle with its error. + fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + block_size = 2**10 + + def get_object(path, range_, **kwargs): + if range_[0] == 0: + raise PermissionError("Access Denied") + return b"a" * (range_[1] - range_[0]) + + 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" + ) + file_ref = weakref.ref(f) + gc.disable() + try: + with pytest.raises(PermissionError, match="Access Denied"): + await asyncio.to_thread(f.read) + await asyncio.to_thread(f.close) + del f + # Lets the event loop run the callbacks that hold the futures. + await asyncio.sleep(0) + assert file_ref() is None + finally: + gc.enable() + @pytest.mark.parametrize("asynchronous", [False, True]) @pytest.mark.asyncio async def test_open_parallel_requests(self, asynchronous): @@ -3287,6 +3328,43 @@ 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_freed_by_reference_counting(self, fs): + # GH-1186: a parallel read raises the error of a failed range, and + # the file is freed by reference counting once it is closed and the + # error is no longer referenced. The reads of the other ranges kept + # their data while the error was held, and a closed file kept itself + # in a reference cycle with its cache. + block_size = 2**10 + + def get_object(path, range_, **kwargs): + if range_[0] == 0: + raise PermissionError("Access Denied") + return b"a" * (range_[1] - range_[0]) + + 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) + file_ref = weakref.ref(f) + gc.disable() + try: + with pytest.raises(PermissionError, match="Access Denied"): + await f.read() + await f.close() + del f + # Lets the event loop run the callbacks that hold the futures. + await asyncio.sleep(0) + assert file_ref() is None + finally: + gc.enable() + @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/test_util.py b/tests/pyathena/test_util.py index 1c5eed6fa..d00a4d141 100644 --- a/tests/pyathena/test_util.py +++ b/tests/pyathena/test_util.py @@ -1,3 +1,5 @@ +import gc +import weakref from typing import Any import pytest @@ -100,6 +102,30 @@ def test_retry_api_call(): _test_retry(_WithCodeError(500)) +def test_retry_api_call_error_freed_by_reference_counting(): + # GH-1186: the error of a failed call is freed by reference counting, as + # are the frames of its traceback, such as those of a file that read with + # the call. The frames of tenacity kept its retry state, and so the error, + # in a reference cycle until the garbage collector ran. + class DeniedError(Exception): + # Unlike Exception, it can be weakly referenced. + pass + + def call(): + raise DeniedError("Access Denied") + + gc.disable() + try: + with pytest.raises(DeniedError) as info: + retry_api_call(call, config=RetryConfig(attempt=2)) + assert [entry.name for entry in info.traceback][-1] == "call" + error_ref = weakref.ref(info.value) + del info + assert error_ref() is None + finally: + gc.enable() + + def test_retry_api_call_with_none_error(): _test_retry(_NoResponseError()) From 658b2f9d0939f02c08ebf64b50ff534bae6e0538 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 00:08:16 +0900 Subject: [PATCH 13/18] Drop the parts before an abort and the buffer of a failed commit - S3File.discard() drops the parts once they have ended, before the abort, so that an abort failure, which keeps the upload for a retry, no longer keeps the errors of failed parts and so the file in a reference cycle. - commit() does nothing for a file whose upload discard() kept, instead of uploading the buffer with PutObject. - A failed multipart commit drops the buffer, which holds only the data of the last part; a later commit() uploaded it alone with PutObject. Co-Authored-By: Claude Opus 5.5 --- docs/filesystem.md | 1 + pyathena/filesystem/s3.py | 27 ++++++++++++++++---- tests/pyathena/filesystem/test_s3.py | 37 +++++++++++++++++++++------- 3 files changed, 51 insertions(+), 14 deletions(-) diff --git a/docs/filesystem.md b/docs/filesystem.md index b67d19aad..cb5e59d4b 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -541,6 +541,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 drops the 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 de487f95a..e2a88331c 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -3291,7 +3291,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, or whose upload :meth:`discard` kept after a failed abort. Raises: FileExistsError: If an object was created at the path after the @@ -3300,9 +3301,13 @@ 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.multipart_upload is not None and not self.multipart_upload_parts: + # discard() dropped the parts of an upload that it failed to + # abort, which is kept for discard(), not completed. return if self.tell() == 0: self.discard() @@ -3329,6 +3334,10 @@ def commit(self) -> None: parts, **self._get_request_kwargs("complete_multipart_upload") ) except BaseException: + # The buffer holds only the data of the last part, which a + # later commit() would otherwise upload alone with PutObject; + # dropped as _close_without_commit() drops it. + self.buffer = None # 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 @@ -3348,13 +3357,21 @@ 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 are waited for before the abort. The parts are dropped + then, so that :meth:`commit` does not complete an upload that this + method failed to abort, which is kept so that a later call retries + the abort. """ 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()]) + # Also dropped if the abort fails, so that the file does not keep + # the errors of failed parts, whose tracebacks may keep the file. + self.multipart_upload_parts = [] + self._uploads_in_flight = set() + self._copies_in_flight = set() self._get_multipart_writer().abort(**self._get_request_kwargs("abort_multipart_upload")) self.multipart_upload = None diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 9bc8c5672..817368075 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -7272,14 +7272,15 @@ def upload_part(**kwargs): fs.core.complete_multipart_upload.assert_not_called() assert f.multipart_upload is None - @pytest.mark.parametrize("stage", ["creation", "write", "copy", "commit"]) + @pytest.mark.parametrize("stage", ["creation", "write", "copy", "commit", "abort"]) def test_failed_request_freed_by_reference_counting(self, stage): # GH-1186: the file and its data are freed as soon as the error of a # failed request is no longer referenced, whether the creation of the # upload, a write of several parts, a write after a failed part copy - # of an append, or the commit raised it. The traceback of the error - # used to keep them in a reference cycle with the future of the - # request until the garbage collector ran. + # of an append, or the commit raised it, also when the abort of the + # upload fails. The traceback of the error used to keep them in a + # reference cycle with the future of the request until the garbage + # collector ran. fs = self._make_append_fs(b"a" * 6 if stage == "copy" else b"") release = threading.Event() @@ -7298,10 +7299,13 @@ def fail_copy(**kwargs): fs.core.upload_part_copy.side_effect = fail_copy else: fs.core.upload_part.side_effect = fail + if stage == "abort": + # A new error each time, which the mock does not keep. + fs._call.side_effect = fail mode = "ab" if stage == "copy" else "wb" f = S3File(fs, "s3://bucket/key.txt", mode=mode, block_size=4, max_workers=1) try: - if stage in ("copy", "commit"): + if stage in ("copy", "commit", "abort"): f.write(b"a" * 4) release.set() if stage == "copy": @@ -7312,10 +7316,17 @@ def fail_copy(**kwargs): file_ref = weakref.ref(f) buffer_ref = weakref.ref(f.buffer) # The write submits a part before it waits for the failed one. - operation = f.close if stage == "commit" else functools.partial(f.write, b"a" * 12) + operation = ( + f.close if stage in ("commit", "abort") else functools.partial(f.write, b"a" * 12) + ) gc.disable() try: - with pytest.raises(PermissionError, match="Access Denied"): + # The log capture of pytest would keep the logged abort failure, + # and so the frames of its traceback. + with ( + mock.patch("pyathena.filesystem.s3._logger"), + pytest.raises(PermissionError, match="Access Denied"), + ): operation() del f, operation assert file_ref() is None @@ -7585,10 +7596,17 @@ def test_commit_failure_and_discard(self, caplog, error, abort_fails): with pytest.raises(error, match="complete failed"): file.commit() assert (file.multipart_upload is not None) is abort_fails - assert bool(file.multipart_upload_parts) is abort_fails + # GH-1186: the parts are dropped also when the abort fails; they used + # to keep the errors of failed parts, and so the file, in a cycle. + assert file.multipart_upload_parts == [] assert ( "Failed to abort multipart upload uploadid to s3://bucket/key.txt." in caplog.text ) is abort_fails + # A commit does not complete the upload that is kept for discard(), + # nor upload the buffered data with PutObject. + 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 == [ @@ -7658,7 +7676,8 @@ def test_commit_failure_and_interrupted_abort(self): with pytest.raises(KeyboardInterrupt): file.commit() assert file.multipart_upload is not None - assert file.multipart_upload_parts + # GH-1186: the parts, which have ended, are dropped before the abort. + assert file.multipart_upload_parts == [] file.discard() assert file.fs._call.call_count == 2 From a8761d28fe6ef7b4463c95ad80d29f0cccbe5f01 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 00:15:04 +0900 Subject: [PATCH 14/18] Keep rows of a failed fetch, skip empty Arrow batches, move the clamp - fetchmany() and fetchall() of the result set and its aio variant return the rows that they took before a failed page request to the front of the buffer, instead of dropping them while counting them in rownumber, so that a later fetch returns them. - The Arrow result set skips record batches without rows, which ended the fetches early, such as the batch of an empty result file. - The size clamp to sys.maxsize moves from the base fetchmany() into the islice() of the pandas and Polars result sets, which need it. - The test helper _page is replaced with get_query_results_responses(). Co-Authored-By: Claude Opus 5.5 --- pyathena/aio/result_set.py | 16 +++++++--- pyathena/arrow/result_set.py | 35 +++++++++++---------- pyathena/pandas/result_set.py | 6 ++-- pyathena/polars/result_set.py | 6 ++-- pyathena/result_set.py | 36 ++++++++++++++++----- tests/pyathena/aio/test_result_set.py | 7 +++-- tests/pyathena/arrow/test_result_set.py | 42 +++++++++++++++++++++++++ tests/pyathena/test_result_set.py | 16 ++++------ 8 files changed, 119 insertions(+), 45 deletions(-) diff --git a/pyathena/aio/result_set.py b/pyathena/aio/result_set.py index 40bb2ad3c..765261d9c 100644 --- a/pyathena/aio/result_set.py +++ b/pyathena/aio/result_set.py @@ -241,6 +241,7 @@ async def _async_take_rows( The buffered rows are taken in bulk, and ``_async_fetch()`` refills the buffer when it is empty. + If a refill fails, the rows taken are returned to the buffer. Args: size: The maximum number of rows, or None for all remaining rows. @@ -249,12 +250,17 @@ async def _async_take_rows( The rows, fewer than ``size`` when the result is exhausted. """ rows: list[tuple[Any | None, ...] | dict[Any, Any | None]] = [] - while size is None or len(rows) < size: - if not self._rows: - await self._async_fetch() + try: + while size is None or len(rows) < size: if not self._rows: - break - self._take_buffered_rows(rows, size) + await self._async_fetch() + if not self._rows: + break + self._take_buffered_rows(rows, size) + except BaseException: + # Also on cancellation, which a later fetch can continue from. + self._return_rows(rows) + raise return rows @override diff --git a/pyathena/arrow/result_set.py b/pyathena/arrow/result_set.py index 508f644d7..cd471a103 100644 --- a/pyathena/arrow/result_set.py +++ b/pyathena/arrow/result_set.py @@ -197,24 +197,27 @@ def converters(self) -> dict[str, Callable[[str | None], Any | None]]: @override def _fetch(self) -> None: - try: - rows = next(self._batches) - except StopIteration: + """Buffer the rows of the next record batch that has any. + + A table can hold empty record batches between others, such as one of + an empty result file, which would otherwise end the fetches early. + """ + rows = next((batch for batch in self._batches if batch.num_rows), None) + if rows is None: return + # Read the columns and their converters by position; to_pydict() and the + # converters property keep one column per name. + columns = [column.to_pylist() for column in rows.columns] + description = self.description if self.description else [] + converters = [self._converter.get(d[1]) for d in description] + if any(convert is not _to_default for convert in converters): + processed_rows = [ + tuple(convert(v) for convert, v in zip(converters, row, strict=False)) + for row in zip(*columns, strict=False) + ] else: - # Read the columns and their converters by position; to_pydict() and the - # converters property keep one column per name. - columns = [column.to_pylist() for column in rows.columns] - description = self.description if self.description else [] - converters = [self._converter.get(d[1]) for d in description] - if any(convert is not _to_default for convert in converters): - processed_rows = [ - tuple(convert(v) for convert, v in zip(converters, row, strict=False)) - for row in zip(*columns, strict=False) - ] - else: - processed_rows = list(zip(*columns, strict=False)) - self._rows.extend(processed_rows) + processed_rows = list(zip(*columns, strict=False)) + self._rows.extend(processed_rows) def _open_csv(self, location: str) -> Any: """Open the CSV result file for ``pyarrow.csv.read_csv()``. diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index 2762b9b86..8592c7aa1 100644 --- a/pyathena/pandas/result_set.py +++ b/pyathena/pandas/result_set.py @@ -5,6 +5,7 @@ import csv import importlib import logging +import sys import threading from collections import abc from collections.abc import Callable, Iterable, Iterator @@ -992,8 +993,9 @@ def _take_rows(self, size: int | None) -> list[tuple[Any | None, ...] | dict[Any """ rows: list[tuple[Any | None, ...] | dict[Any, Any | None]] = [] try: - # extend() keeps the rows that it appended before an exception. - rows.extend(islice(self._row_iter, size)) + # extend() keeps the rows that it appended before an exception, and + # islice() takes no size above sys.maxsize. + rows.extend(islice(self._row_iter, None if size is None else min(size, sys.maxsize))) finally: if rows: self._count_rows(len(rows)) diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index 74a2e099f..b79802547 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -12,6 +12,7 @@ import csv import importlib.util import logging +import sys import threading from collections import abc from collections.abc import Callable, Iterable, Iterator @@ -647,8 +648,9 @@ def _take_rows(self, size: int | None) -> list[tuple[Any | None, ...] | dict[Any """ rows: list[tuple[Any | None, ...] | dict[Any, Any | None]] = [] try: - # extend() keeps the rows that it appended before an exception. - rows.extend(islice(self._row_iter, size)) + # extend() keeps the rows that it appended before an exception, and + # islice() takes no size above sys.maxsize. + rows.extend(islice(self._row_iter, None if size is None else min(size, sys.maxsize))) finally: if rows: self._count_rows(len(rows)) diff --git a/pyathena/result_set.py b/pyathena/result_set.py index c0a7cf9e3..7bf314dc7 100644 --- a/pyathena/result_set.py +++ b/pyathena/result_set.py @@ -5,7 +5,6 @@ import collections import logging import operator -import sys from datetime import datetime from itertools import islice from typing import ( @@ -529,8 +528,7 @@ def fetchmany( """ if not size or size <= 0: size = self._arraysize - # islice() in the subclasses' _take_rows() takes no size above sys.maxsize. - return self._take_rows(min(operator.index(size), sys.maxsize)) + return self._take_rows(operator.index(size)) @override def fetchall( @@ -549,6 +547,7 @@ def _take_rows(self, size: int | None) -> list[tuple[Any | None, ...] | dict[Any The buffered rows are taken in bulk, and ``_fetch()`` refills the buffer when it is empty. A subclass whose ``fetchone()`` does not use the buffer overrides this method too. + If a refill fails, the rows taken are returned to the buffer. Args: size: The maximum number of rows, or None for all remaining rows. @@ -557,12 +556,16 @@ def _take_rows(self, size: int | None) -> list[tuple[Any | None, ...] | dict[Any The rows, fewer than ``size`` when the result is exhausted. """ rows: list[tuple[Any | None, ...] | dict[Any, Any | None]] = [] - while size is None or len(rows) < size: - if not self._rows: - self._fetch() + try: + while size is None or len(rows) < size: if not self._rows: - break - self._take_buffered_rows(rows, size) + self._fetch() + if not self._rows: + break + self._take_buffered_rows(rows, size) + except BaseException: + self._return_rows(rows) + raise return rows def _take_buffered_rows( @@ -591,6 +594,23 @@ def _take_buffered_rows( rows.extend(buffered) self._count_rows(len(rows) - start) + def _return_rows(self, rows: list[tuple[Any | None, ...] | dict[Any, Any | None]]) -> None: + """Return the rows that a failed fetch took to the front of the buffer. + + The rows are no longer counted in ``rownumber``, so that a later fetch + returns them, as ``fetchone()`` leaves the rows after a failed page + request in the buffer. A result set that is closed keeps its empty + buffer. + + Args: + rows: The rows that the fetch took, in their order. + """ + rownumber = self._rownumber + if rownumber is None or not rows: + return + self._rows.extendleft(reversed(rows)) + self._rownumber = rownumber - len(rows) + def _count_rows(self, count: int) -> None: """Count the rows that a fetch took in ``rownumber``. diff --git a/tests/pyathena/aio/test_result_set.py b/tests/pyathena/aio/test_result_set.py index dc0292ba8..0312f6b33 100644 --- a/tests/pyathena/aio/test_result_set.py +++ b/tests/pyathena/aio/test_result_set.py @@ -100,13 +100,16 @@ async def test_fetchone_across_pages(self): assert result_set.connection.client.get_query_results.call_count == 2 @pytest.mark.parametrize(("method", "args"), [("fetchmany", (3,)), ("fetchall", ())]) - async def test_fetch_counts_rows_taken_before_error(self, method, args): - """The rows taken before a failed page request are counted in rownumber.""" + async def test_fetch_returns_rows_taken_before_error(self, method, args): + """The rows taken before a failed page request are returned to the buffer.""" + # GH-1186: they used to be dropped while counted in rownumber. result_set = await _create_result_set( AthenaAioResultSet, [("1", "2"), RuntimeError("failed")] ) with pytest.raises(OperationalError): await getattr(result_set, method)(*args) + assert result_set.rownumber == 0 + assert await result_set.fetchmany(2) == [(1,), (2,)] assert result_set.rownumber == 2 async def test_fetch_after_close(self): diff --git a/tests/pyathena/arrow/test_result_set.py b/tests/pyathena/arrow/test_result_set.py index 46be7dc41..c1eb15fa4 100644 --- a/tests/pyathena/arrow/test_result_set.py +++ b/tests/pyathena/arrow/test_result_set.py @@ -64,6 +64,48 @@ def test_fetch_across_batches(self): assert result_set.fetchone() is None assert result_set.rownumber == 3 + @pytest.mark.parametrize("method", ["fetchone", "fetchmany", "fetchall", "iter"]) + def test_fetch_across_empty_batches(self, method): + """An empty record batch between others does not end the fetches. + + GH-1186: a table read from result files that include an empty one has + such a batch, and the fetches used to stop at it. No AWS calls; the + first page request and the result table are mocked. + """ + table = pa.concat_tables( + [ + pa.table({"a": [1, 2, 3]}), + pa.table({"a": pa.array([], pa.int64())}), + pa.table({"a": [4]}), + ] + ) + with ( + patch.object(AthenaArrowResultSet, "_pre_fetch"), + patch.object(AthenaArrowResultSet, "_create_s3_file_system"), + patch.object(AthenaArrowResultSet, "_as_arrow", return_value=table), + ): + result_set = AthenaArrowResultSet( + connection=MagicMock(), + converter=DefaultArrowTypeConverter(), + query_execution=MagicMock( + state=AthenaQueryExecution.STATE_SUCCEEDED, + output_location="s3://bucket/query.csv", + ), + arraysize=2, + retry_config=RetryConfig(), + ) + assert 0 in [batch.num_rows for batch in table.to_batches(2)] + if method == "fetchone": + rows = list(iter(result_set.fetchone, None)) + elif method == "fetchmany": + rows = result_set.fetchmany(10) + elif method == "fetchall": + rows = result_set.fetchall() + else: + rows = list(result_set) + assert rows == [(1,), (2,), (3,), (4,)] + assert result_set.rownumber == 4 + def test_read_csv_timestamps_with_the_same_name(self): """Timestamp columns are converted by position, also with names that repeat. diff --git a/tests/pyathena/test_result_set.py b/tests/pyathena/test_result_set.py index 83df71187..3790237a8 100644 --- a/tests/pyathena/test_result_set.py +++ b/tests/pyathena/test_result_set.py @@ -20,13 +20,6 @@ from tests.pyathena.util import close_during_next_request, get_query_results_responses -def _page(values, next_token=None): - response = {"ResultSet": {"Rows": [{"Data": [{"VarCharValue": v}]} for v in values]}} - if next_token: - response["NextToken"] = next_token - return response - - def _row(*values): return {"Data": [{} if v is None else {"VarCharValue": v} for v in values]} @@ -112,11 +105,14 @@ def test_fetchmany_size_not_integer(self): assert result_set.fetchall() == [(1,), (2,)] @pytest.mark.parametrize(("method", "args"), [("fetchmany", (3,)), ("fetchall", ())]) - def test_fetch_counts_rows_taken_before_error(self, method, args): - """The rows taken before a failed page request are counted in rownumber.""" + def test_fetch_returns_rows_taken_before_error(self, method, args): + """The rows taken before a failed page request are returned to the buffer.""" + # GH-1186: they used to be dropped while counted in rownumber. result_set = _fetching_result_set(("1", "2"), RuntimeError("failed")) with pytest.raises(OperationalError): getattr(result_set, method)(*args) + assert result_set.rownumber == 0 + assert result_set.fetchmany(2) == [(1,), (2,)] assert result_set.rownumber == 2 def test_fetch_after_close(self): @@ -180,7 +176,7 @@ def test_fetch_all_rows_skips_column_labels_only_on_first_page(self): result_set._process_metadata( {"ResultSet": {"ResultSetMetadata": {"ColumnInfo": [{"Name": "a", "Type": "varchar"}]}}} ) - pages = [_page(["a", "1"], "token"), _page(["a", "2"])] + pages = get_query_results_responses(["a", "1"], ["a", "2"]) with patch.object(result_set, "_get_query_results", side_effect=pages): assert result_set._fetch_all_rows() == [("1",), ("a",), ("2",)] From e4662481af9abbc28482f892cfc10dc6429e00b2 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 00:59:39 +0900 Subject: [PATCH 15/18] Narrow the cleanup to the data of failed requests Keep the fixes of root causes and drop the per-site breaking of reference cycles, which only made files and caller buffers collectable without the garbage collector at the cost of an invariant that every caller had to keep: - A private _S3WorkItem owns the future, the callable, and its arguments of S3Executor.submit_to() and S3AioExecutor.submit(), and drops them when it runs or settles, instead of closures whose shared cells run() deleted. - S3Core.call() clears the frames below it in the traceback of a raised error, those of the retries and of botocore, which kept the request and, in tenacity's retry state, the error in a cycle; this replaces walking the chained errors in submit_to() and clearing tenacity frames in retry_api_call(). - raise_failure(), _submit_upload(), the deletions of locals, the resets of AioS3StreamedFile._discard(), the dropping of the parts in discard() with the guard of commit(), and the cache reset of close() are removed. - A failed multipart commit() also closes the file, and the cleanup of an AioS3FileSystem copy also cancels a completion that has not started. - The result set fixes move to a separate change. Co-Authored-By: Claude Opus 5.5 --- docs/api/filesystem.rst | 2 - docs/filesystem.md | 5 +- pyathena/aio/result_set.py | 16 +- pyathena/arrow/result_set.py | 35 ++- pyathena/filesystem/s3.py | 108 ++++---- pyathena/filesystem/s3_async.py | 101 +++---- pyathena/filesystem/s3_core.py | 15 +- pyathena/filesystem/s3_executor.py | 262 +++++++----------- pyathena/pandas/result_set.py | 6 +- pyathena/polars/result_set.py | 6 +- pyathena/result_set.py | 36 +-- pyathena/util.py | 17 +- tests/pyathena/aio/test_result_set.py | 7 +- tests/pyathena/arrow/test_result_set.py | 42 --- tests/pyathena/filesystem/test_s3.py | 126 +++------ tests/pyathena/filesystem/test_s3_async.py | 180 +++++++----- tests/pyathena/filesystem/test_s3_core.py | 69 +++++ tests/pyathena/filesystem/test_s3_executor.py | 124 +-------- tests/pyathena/test_result_set.py | 16 +- tests/pyathena/test_util.py | 26 -- 20 files changed, 467 insertions(+), 732 deletions(-) diff --git a/docs/api/filesystem.rst b/docs/api/filesystem.rst index 298c5c5c2..722efcd4e 100644 --- a/docs/api/filesystem.rst +++ b/docs/api/filesystem.rst @@ -58,8 +58,6 @@ S3 Executor .. autoclass:: pyathena.filesystem.s3_executor.S3AioThreadPool :members: -.. autofunction:: pyathena.filesystem.s3_executor.raise_failure - S3 Paths -------- diff --git a/docs/filesystem.md b/docs/filesystem.md index cb5e59d4b..892bce4d5 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -139,7 +139,8 @@ kept. A failed part copy aborts the multipart upload. So does an interrupt, or t 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. @@ -541,7 +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 drops the buffered data, so a later `commit()` neither completes that upload nor uploads the last part alone with PutObject. +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/aio/result_set.py b/pyathena/aio/result_set.py index 765261d9c..40bb2ad3c 100644 --- a/pyathena/aio/result_set.py +++ b/pyathena/aio/result_set.py @@ -241,7 +241,6 @@ async def _async_take_rows( The buffered rows are taken in bulk, and ``_async_fetch()`` refills the buffer when it is empty. - If a refill fails, the rows taken are returned to the buffer. Args: size: The maximum number of rows, or None for all remaining rows. @@ -250,17 +249,12 @@ async def _async_take_rows( The rows, fewer than ``size`` when the result is exhausted. """ rows: list[tuple[Any | None, ...] | dict[Any, Any | None]] = [] - try: - while size is None or len(rows) < size: + while size is None or len(rows) < size: + if not self._rows: + await self._async_fetch() if not self._rows: - await self._async_fetch() - if not self._rows: - break - self._take_buffered_rows(rows, size) - except BaseException: - # Also on cancellation, which a later fetch can continue from. - self._return_rows(rows) - raise + break + self._take_buffered_rows(rows, size) return rows @override diff --git a/pyathena/arrow/result_set.py b/pyathena/arrow/result_set.py index cd471a103..508f644d7 100644 --- a/pyathena/arrow/result_set.py +++ b/pyathena/arrow/result_set.py @@ -197,27 +197,24 @@ def converters(self) -> dict[str, Callable[[str | None], Any | None]]: @override def _fetch(self) -> None: - """Buffer the rows of the next record batch that has any. - - A table can hold empty record batches between others, such as one of - an empty result file, which would otherwise end the fetches early. - """ - rows = next((batch for batch in self._batches if batch.num_rows), None) - if rows is None: + try: + rows = next(self._batches) + except StopIteration: return - # Read the columns and their converters by position; to_pydict() and the - # converters property keep one column per name. - columns = [column.to_pylist() for column in rows.columns] - description = self.description if self.description else [] - converters = [self._converter.get(d[1]) for d in description] - if any(convert is not _to_default for convert in converters): - processed_rows = [ - tuple(convert(v) for convert, v in zip(converters, row, strict=False)) - for row in zip(*columns, strict=False) - ] else: - processed_rows = list(zip(*columns, strict=False)) - self._rows.extend(processed_rows) + # Read the columns and their converters by position; to_pydict() and the + # converters property keep one column per name. + columns = [column.to_pylist() for column in rows.columns] + description = self.description if self.description else [] + converters = [self._converter.get(d[1]) for d in description] + if any(convert is not _to_default for convert in converters): + processed_rows = [ + tuple(convert(v) for convert, v in zip(converters, row, strict=False)) + for row in zip(*columns, strict=False) + ] + else: + processed_rows = list(zip(*columns, strict=False)) + self._rows.extend(processed_rows) def _open_csv(self, location: str) -> Any: """Open the CSV result file for ``pyarrow.csv.read_csv()``. diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index e2a88331c..abe70dadb 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -34,7 +34,7 @@ from pyathena.connection import Connection from pyathena.filesystem.s3_core import S3Core, S3DeleteBatch, S3DeleteResult from pyathena.filesystem.s3_errors import S3ClientError -from pyathena.filesystem.s3_executor import S3Executor, S3ThreadPoolExecutor, raise_failure +from pyathena.filesystem.s3_executor import S3Executor, S3ThreadPoolExecutor from pyathena.filesystem.s3_object import ( S3CompleteMultipartUpload, S3Metadata, @@ -1850,9 +1850,11 @@ def _finish_multipart_upload( request_kwargs = request_kwargs or {} try: done, _ = wait(futures, return_when=FIRST_EXCEPTION) - # Raised without waiting for the parts before it, so that the - # parts queued after it are cancelled. - raise_failure(done) + for future in done: + if not future.cancelled() and future.exception() is not None: + # Raised without waiting for the parts before it, so that + # the parts queued after it are cancelled. + future.result() # The futures are in part-number order. parts = [future.result() for future in futures] return self.core.complete_multipart_upload( @@ -3176,8 +3178,6 @@ def _initiate_upload(self) -> None: # A submitted creation may still return an ID after the wait is # interrupted. The writer records it before resolving the future. self._close_without_commit(creation=creation) - # Not kept by this frame, which the traceback of its error keeps. - del creation raise def _upload_chunk(self, final: bool = False) -> bool: @@ -3236,18 +3236,32 @@ def _upload_chunk(self, final: bool = False) -> bool: if self._copies_in_flight: copied = {copy for copy in self._copies_in_flight if copy.done()} self._copies_in_flight -= copied - # Otherwise the rest of the data is uploaded after a - # failed part copy. - raise_failure(copied) + for copy in copied: + # Otherwise the rest of the data is uploaded after a + # failed part copy. + copy.result() if len(self._uploads_in_flight) >= self.max_workers: # Otherwise the parts waiting for a worker pile up in memory. done, self._uploads_in_flight = wait( self._uploads_in_flight, return_when=FIRST_COMPLETED ) - # A failed part would otherwise keep its data until - # commit() while the rest of the data is uploaded. - raise_failure(done) - self._submit_upload(writer, part_number, data) + for upload in done: + # A failed part would otherwise keep its data until + # commit() while the rest of the data is uploaded. + upload.result() + # Tracked before it is scheduled, so that an interrupt while + # it is being scheduled still lets discard() cancel it or wait + # for it before the abort. + 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"), + ) except BaseException: self._close_without_commit() raise @@ -3256,31 +3270,6 @@ def _upload_chunk(self, final: bool = False) -> bool: self.commit() return not final - def _submit_upload(self, writer: S3MultipartWriter, part_number: int, data: bytes) -> None: - """Submit the upload of a part and track it as in flight. - - The part is tracked before it is scheduled, so that an interrupt while - it is being scheduled still lets :meth:`discard` cancel it or wait for - it before the abort. Its future is not left in the frame of - :meth:`_upload_chunk`, which the traceback of the part's error keeps - if a later wait of the same call raises it. - - Args: - writer: The writer of the multipart upload. - part_number: The number of the part, from 1. - data: The data of the part. - """ - 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"), - ) - def commit(self) -> None: """Complete the upload of the written data. @@ -3292,7 +3281,7 @@ def commit(self) -> None: 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 or multipart commit dropped the written - data, or whose upload :meth:`discard` kept after a failed abort. + data. Raises: FileExistsError: If an object was created at the path after the @@ -3305,10 +3294,6 @@ def commit(self) -> None: # data. A multipart upload that they failed to abort is kept for # discard(), not completed. return - if self.multipart_upload is not None and not self.multipart_upload_parts: - # discard() dropped the parts of an upload that it failed to - # abort, which is kept for discard(), not completed. - return if self.tell() == 0: self.discard() self.fs.touch(self.path, **self._get_request_kwargs("put_object")) @@ -3326,18 +3311,23 @@ def commit(self) -> None: upload_id = cast(str, self.multipart_upload.upload_id) try: done, _ = wait(self.multipart_upload_parts, return_when=FIRST_EXCEPTION) - # Raised without waiting for the parts before it, so that the - # parts queued after it are cancelled. - raise_failure(done) + for future in done: + if not future.cancelled() and future.exception() is not None: + # Raised without waiting for the parts before it, so + # that the parts queued after it are cancelled. + future.result() parts = [future.result() for future in self.multipart_upload_parts] self._get_multipart_writer().complete( parts, **self._get_request_kwargs("complete_multipart_upload") ) except BaseException: # The buffer holds only the data of the last part, which a - # later commit() would otherwise upload alone with PutObject; - # dropped as _close_without_commit() drops it. + # later commit() would otherwise upload alone with PutObject, + # so the file is closed without it, as _close_without_commit() + # closes it. self.buffer = None + self.closed = True + self._executor.shutdown(wait=False) # 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 @@ -3357,21 +3347,13 @@ 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. The parts are dropped - then, so that :meth:`commit` does not complete an upload that this - method failed to abort, which is kept so that a later call retries - the abort. + running ones are waited for before the abort. """ 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()]) - # Also dropped if the abort fails, so that the file does not keep - # the errors of failed parts, whose tracebacks may keep the file. - self.multipart_upload_parts = [] - self._uploads_in_flight = set() - self._copies_in_flight = set() self._get_multipart_writer().abort(**self._get_request_kwargs("abort_multipart_upload")) self.multipart_upload = None @@ -3465,13 +3447,17 @@ def _fetch_range(self, start: int, end: int) -> bytes: for r in ranges ] done = wait(futures, return_when=FIRST_EXCEPTION).done - if any(not f.cancelled() and f.exception() is not None for f in done): + 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 the traceback of - # the error, which keeps this frame, does not keep their data. + # 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() - raise_failure(done) + 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 afd6f571c..d407980af 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -33,7 +33,6 @@ S3AioThreadPool, S3Executor, S3ThreadPoolExecutor, - raise_failure, ) from pyathena.filesystem.s3_object import ( S3CompleteMultipartUpload, @@ -579,10 +578,9 @@ async def _copy_object_with_multipart_upload( 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 and the part copies that have not started are - cancelled, the running ones and the completion 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 + 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 @@ -708,6 +706,8 @@ def _abort() -> None: 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. await self._thread_pool.run_shielded(_abort) @@ -1408,10 +1408,6 @@ async def close(self) -> None: await self.flush(force=True) finally: self.closed = True - if self.mode == "rb": - # Dropped as fsspec's close() drops it, since the cache keeps - # a method of the file and so a reference cycle with it. - self.cache = None # A commit that is still running keeps its thread. self._pool.shutdown(wait=False) @@ -1517,18 +1513,34 @@ async def _upload_chunk(self, final: bool = False) -> bool: if self._copies_in_flight: copied = {copy for copy in self._copies_in_flight if copy.done()} self._copies_in_flight -= copied - # Otherwise the rest of the data is uploaded after a - # failed part copy, as S3File does. - raise_failure(copied) + for copy in copied: + # Otherwise the rest of the data is uploaded after a + # failed part copy, as S3File does. + copy.result() if len(self._uploads_in_flight) >= self.max_workers: # Otherwise the parts waiting for a thread pile up in memory. done, self._uploads_in_flight = await asyncio.wait( self._uploads_in_flight, return_when=asyncio.FIRST_COMPLETED ) - # A failed part would otherwise keep its data until the - # commit while the rest of the data is uploaded. - raise_failure(done) - self._submit_upload(writer, part_number, data, upload_kwargs) + for upload in done: + # A failed part would otherwise keep its data until + # the commit while the rest of the data is uploaded. + upload.result() + # Tracked before it is scheduled; see _initiate_upload(). + part: Future[S3MultipartUploadPart] = Future() + self._parts.append(part) + waiter = asyncio.wrap_future(part) + # The commit reports a failure that no wait found; retrieved + # 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, + ) except BaseException: await self._close_without_commit() raise @@ -1536,40 +1548,6 @@ async def _upload_chunk(self, final: bool = False) -> bool: await self._commit() return not final - def _submit_upload( - self, - writer: S3MultipartWriter, - part_number: int, - data: bytes, - request_kwargs: dict[str, Any], - ) -> None: - """Submit the upload of a part and track it as in flight. - - See :meth:`S3File._submit_upload`; the file waits for the part - through an asyncio future. - - Args: - writer: The writer of the multipart upload. - part_number: The number of the part, from 1. - data: The data of the part. - request_kwargs: The UploadPart parameters of the file. - """ - # Tracked before it is scheduled; see _initiate_upload(). - part: Future[S3MultipartUploadPart] = Future() - self._parts.append(part) - waiter = asyncio.wrap_future(part) - # The commit reports a failure that no wait found; retrieved 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, - **request_kwargs, - ) - async def _commit(self) -> None: """Commit the written data in a thread of the file. @@ -1612,7 +1590,9 @@ def _complete_upload(self) -> None: # all started. A failed part is raised without waiting for # the parts before it, as S3File does. done, _ = wait(self._parts, return_when=FIRST_EXCEPTION) - raise_failure(done) + for part in done: + if not part.cancelled() and part.exception() is not None: + part.result() parts = [part.result() for part in self._parts] writer.complete(parts, **self._get_request_kwargs("complete_multipart_upload")) except BaseException: @@ -1639,13 +1619,6 @@ def _discard(self) -> None: _logger.exception( f"Failed to abort the multipart upload to s3://{self.bucket}/{self.key}." ) - # Dropped as S3File.discard() drops the parts, so that the file does - # not keep the errors of failed requests, whose tracebacks may keep - # the file. - self._creation = None - self._parts = [] - self._uploads_in_flight = set() - self._copies_in_flight = set() async def _fetch_range(self, start: int, end: int) -> bytes: """Read a byte range of the object with parallel range requests. @@ -1683,11 +1656,15 @@ async def _fetch_range(self, start: int, end: int) -> bytes: while reads: reads.pop().cancel() raise - if any(not read.cancelled() and read.exception() is not None for read in done): + # The errors are all retrieved, so that asyncio does not log those that + # are not raised as never retrieved. + 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 the traceback of the error, which - # keeps this frame, does not keep their data. + # the others are dropped, so that this frame, which the traceback + # of the error keeps, does not keep their data. while reads: reads.pop().cancel() - raise_failure(done) + 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 58c7ef73a..67e50823e 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -22,13 +22,97 @@ 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") -# A future of an executor or of asyncio. -F = TypeVar("F", bound="Future[Any] | asyncio.Future[Any]") + + +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): @@ -62,12 +146,9 @@ 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 error of a failed callable keeps the frames that it was raised - from in its traceback, and so do the errors that it was raised from, - raised while handling, or groups, without their local variables, so - that the arguments, such as the body of a part, are freed when the - callable fails, and the future and its error are freed without the - garbage collector once they are no longer referenced. + 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. Args: future: A pending future that no executor has started. @@ -75,65 +156,8 @@ def submit_to(self, future: Future[T], fn: Callable[..., T], *args: Any, **kwarg *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.""" - nonlocal future, fn, args, kwargs - if not future.set_running_or_notify_cancel(): - return - try: - result = fn(*args, **kwargs) - except BaseException as e: - # The frames below this one hold the arguments, and so may - # those of the errors chained to it, such as the ClientError - # that S3Core.call() raises an OSError from, or grouped in it. - # clear_frames() skips this frame, which is still executing. - chained: list[BaseException | None] = [e] - seen: set[int] = set() - while chained: - if (error := chained.pop()) is not None and id(error) not in seen: - seen.add(id(error)) - traceback.clear_frames(error.__traceback__) - chained += [error.__cause__, error.__context__] - if isinstance(error, BaseExceptionGroup): - chained += error.exceptions - # The traceback also keeps this frame, which would keep the - # arguments and form a reference cycle with the future, and so - # do the frames of the executor that called run(), and the - # done callback of the submission. The variables are cleared - # in the closure that they share; settle() reads the future - # only when run() did not run. - resolved = future - del future, fn, args, kwargs, chained, error - resolved.set_exception(e) - del resolved - else: - # Deleted only on the failure path. - future.set_result(result) # noqa: F821 - - 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 @@ -142,40 +166,6 @@ def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: self.shutdown(wait=True) -def raise_failure(done: set[F]) -> None: - """Empty a set of finished futures and raise the error of a failed one, if any. - - Otherwise, the cancellation of a cancelled future is raised, if any. The - errors of all the failed futures are retrieved, so that asyncio does not - log those that are not raised as never retrieved. The set is emptied - before the error is raised: the traceback of the error keeps the frames - that it is raised through, so a frame that still referenced the future, - which keeps the error, would form a reference cycle that only the garbage - collector frees, with the local variables of every frame of the - traceback, such as the data of a write. - - Args: - done: Finished futures that the caller no longer needs. - - Raises: - BaseException: The error of a failed future, or the cancellation - error of a cancelled one. - """ - future: F | None = None - failed = cancelled = None - for future in done: - if future.cancelled(): - cancelled = cancelled or future - elif future.exception() is not None: - failed = failed or future - done.clear() - try: - if (raised := failed or cancelled) is not None: - raised.result() - finally: - del future, failed, cancelled, raised - - class S3ThreadPoolExecutor(S3Executor): """Executor that delegates to a ``ThreadPoolExecutor``. @@ -329,16 +319,11 @@ async def wait_shielded(self, future: Future[T]) -> T: """ waiter = asyncio.wrap_future(future) waiter.add_done_callback(lambda f: f.cancelled() or f.exception()) - try: - # 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() - finally: - # The traceback of an error keeps this frame, which would form a - # reference cycle with the futures of the error. - del future, waiter + # 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: @@ -440,58 +425,13 @@ 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.""" - nonlocal future, fn, args, kwargs - if not claim.acquire(blocking=False) or not future.set_running_or_notify_cancel(): - return - try: - result = fn(*args, **kwargs) - except BaseException as e: - # The traceback keeps this frame, which would form a - # reference cycle with the future, as with the _WorkItem - # of ThreadPoolExecutor. The variables are cleared in the - # closure that settle() shares, which reads the future - # only when run() has not claimed it. - resolved = future - del future, fn, args, kwargs - resolved.set_exception(e) - del resolved - 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()) - - # Returned through a variable that the closures do not share, - # since run() may clear theirs before this method returns. - submitted = future + 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) - return submitted + 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. " "Use S3ThreadPoolExecutor for synchronous usage." diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index 8592c7aa1..2762b9b86 100644 --- a/pyathena/pandas/result_set.py +++ b/pyathena/pandas/result_set.py @@ -5,7 +5,6 @@ import csv import importlib import logging -import sys import threading from collections import abc from collections.abc import Callable, Iterable, Iterator @@ -993,9 +992,8 @@ def _take_rows(self, size: int | None) -> list[tuple[Any | None, ...] | dict[Any """ rows: list[tuple[Any | None, ...] | dict[Any, Any | None]] = [] try: - # extend() keeps the rows that it appended before an exception, and - # islice() takes no size above sys.maxsize. - rows.extend(islice(self._row_iter, None if size is None else min(size, sys.maxsize))) + # extend() keeps the rows that it appended before an exception. + rows.extend(islice(self._row_iter, size)) finally: if rows: self._count_rows(len(rows)) diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index b79802547..74a2e099f 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -12,7 +12,6 @@ import csv import importlib.util import logging -import sys import threading from collections import abc from collections.abc import Callable, Iterable, Iterator @@ -648,9 +647,8 @@ def _take_rows(self, size: int | None) -> list[tuple[Any | None, ...] | dict[Any """ rows: list[tuple[Any | None, ...] | dict[Any, Any | None]] = [] try: - # extend() keeps the rows that it appended before an exception, and - # islice() takes no size above sys.maxsize. - rows.extend(islice(self._row_iter, None if size is None else min(size, sys.maxsize))) + # extend() keeps the rows that it appended before an exception. + rows.extend(islice(self._row_iter, size)) finally: if rows: self._count_rows(len(rows)) diff --git a/pyathena/result_set.py b/pyathena/result_set.py index 7bf314dc7..c0a7cf9e3 100644 --- a/pyathena/result_set.py +++ b/pyathena/result_set.py @@ -5,6 +5,7 @@ import collections import logging import operator +import sys from datetime import datetime from itertools import islice from typing import ( @@ -528,7 +529,8 @@ def fetchmany( """ if not size or size <= 0: size = self._arraysize - return self._take_rows(operator.index(size)) + # islice() in the subclasses' _take_rows() takes no size above sys.maxsize. + return self._take_rows(min(operator.index(size), sys.maxsize)) @override def fetchall( @@ -547,7 +549,6 @@ def _take_rows(self, size: int | None) -> list[tuple[Any | None, ...] | dict[Any The buffered rows are taken in bulk, and ``_fetch()`` refills the buffer when it is empty. A subclass whose ``fetchone()`` does not use the buffer overrides this method too. - If a refill fails, the rows taken are returned to the buffer. Args: size: The maximum number of rows, or None for all remaining rows. @@ -556,16 +557,12 @@ def _take_rows(self, size: int | None) -> list[tuple[Any | None, ...] | dict[Any The rows, fewer than ``size`` when the result is exhausted. """ rows: list[tuple[Any | None, ...] | dict[Any, Any | None]] = [] - try: - while size is None or len(rows) < size: + while size is None or len(rows) < size: + if not self._rows: + self._fetch() if not self._rows: - self._fetch() - if not self._rows: - break - self._take_buffered_rows(rows, size) - except BaseException: - self._return_rows(rows) - raise + break + self._take_buffered_rows(rows, size) return rows def _take_buffered_rows( @@ -594,23 +591,6 @@ def _take_buffered_rows( rows.extend(buffered) self._count_rows(len(rows) - start) - def _return_rows(self, rows: list[tuple[Any | None, ...] | dict[Any, Any | None]]) -> None: - """Return the rows that a failed fetch took to the front of the buffer. - - The rows are no longer counted in ``rownumber``, so that a later fetch - returns them, as ``fetchone()`` leaves the rows after a failed page - request in the buffer. A result set that is closed keeps its empty - buffer. - - Args: - rows: The rows that the fetch took, in their order. - """ - rownumber = self._rownumber - if rownumber is None or not rows: - return - self._rows.extendleft(reversed(rows)) - self._rownumber = rownumber - len(rows) - def _count_rows(self, count: int) -> None: """Count the rows that a fetch took in ``rownumber``. diff --git a/pyathena/util.py b/pyathena/util.py index 037ad076a..a88dc6bf5 100644 --- a/pyathena/util.py +++ b/pyathena/util.py @@ -307,9 +307,7 @@ def retry_api_call( Note: Only retries on AWS exceptions listed in the RetryConfig.exceptions. This includes recognized Glue error codes wrapped in MetadataException. - Other errors are propagated without retrying. The frames of tenacity - in the traceback of a raised error have no local variables, so that - the error is freed by reference counting. + Other errors are propagated without retrying. """ retry = tenacity.Retrying( retry=retry_if_exception( @@ -327,15 +325,4 @@ def retry_api_call( after=after_log(logger, logger.getEffectiveLevel()) if logger else None, # type: ignore[arg-type] reraise=True, ) - try: - return retry(func, *args, **kwargs) - except BaseException as e: - # The frames of tenacity in the traceback keep its retry state, which - # keeps the error, so they would form a reference cycle that keeps - # every frame of the traceback until the garbage collector runs. - tb = e.__traceback__ - while tb is not None: - if tb.tb_frame.f_globals.get("__name__", "").startswith("tenacity"): - tb.tb_frame.clear() - tb = tb.tb_next - raise + return retry(func, *args, **kwargs) diff --git a/tests/pyathena/aio/test_result_set.py b/tests/pyathena/aio/test_result_set.py index 0312f6b33..dc0292ba8 100644 --- a/tests/pyathena/aio/test_result_set.py +++ b/tests/pyathena/aio/test_result_set.py @@ -100,16 +100,13 @@ async def test_fetchone_across_pages(self): assert result_set.connection.client.get_query_results.call_count == 2 @pytest.mark.parametrize(("method", "args"), [("fetchmany", (3,)), ("fetchall", ())]) - async def test_fetch_returns_rows_taken_before_error(self, method, args): - """The rows taken before a failed page request are returned to the buffer.""" - # GH-1186: they used to be dropped while counted in rownumber. + async def test_fetch_counts_rows_taken_before_error(self, method, args): + """The rows taken before a failed page request are counted in rownumber.""" result_set = await _create_result_set( AthenaAioResultSet, [("1", "2"), RuntimeError("failed")] ) with pytest.raises(OperationalError): await getattr(result_set, method)(*args) - assert result_set.rownumber == 0 - assert await result_set.fetchmany(2) == [(1,), (2,)] assert result_set.rownumber == 2 async def test_fetch_after_close(self): diff --git a/tests/pyathena/arrow/test_result_set.py b/tests/pyathena/arrow/test_result_set.py index c1eb15fa4..46be7dc41 100644 --- a/tests/pyathena/arrow/test_result_set.py +++ b/tests/pyathena/arrow/test_result_set.py @@ -64,48 +64,6 @@ def test_fetch_across_batches(self): assert result_set.fetchone() is None assert result_set.rownumber == 3 - @pytest.mark.parametrize("method", ["fetchone", "fetchmany", "fetchall", "iter"]) - def test_fetch_across_empty_batches(self, method): - """An empty record batch between others does not end the fetches. - - GH-1186: a table read from result files that include an empty one has - such a batch, and the fetches used to stop at it. No AWS calls; the - first page request and the result table are mocked. - """ - table = pa.concat_tables( - [ - pa.table({"a": [1, 2, 3]}), - pa.table({"a": pa.array([], pa.int64())}), - pa.table({"a": [4]}), - ] - ) - with ( - patch.object(AthenaArrowResultSet, "_pre_fetch"), - patch.object(AthenaArrowResultSet, "_create_s3_file_system"), - patch.object(AthenaArrowResultSet, "_as_arrow", return_value=table), - ): - result_set = AthenaArrowResultSet( - connection=MagicMock(), - converter=DefaultArrowTypeConverter(), - query_execution=MagicMock( - state=AthenaQueryExecution.STATE_SUCCEEDED, - output_location="s3://bucket/query.csv", - ), - arraysize=2, - retry_config=RetryConfig(), - ) - assert 0 in [batch.num_rows for batch in table.to_batches(2)] - if method == "fetchone": - rows = list(iter(result_set.fetchone, None)) - elif method == "fetchmany": - rows = result_set.fetchmany(10) - elif method == "fetchall": - rows = result_set.fetchall() - else: - rows = list(result_set) - assert rows == [(1,), (2,), (3,), (4,)] - assert result_set.rownumber == 4 - def test_read_csv_timestamps_with_the_same_name(self): """Timestamp columns are converted by position, also with names that repeat. diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 817368075..ec9df7924 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,33 +3429,36 @@ 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_freed_by_reference_counting(self): - # GH-1186: a parallel read raises the error of a failed range, and - # the file is freed by reference counting once it is closed and the - # error is no longer referenced. The traceback of the error used to - # keep the file, and the data of the other ranges, in a reference - # cycle with the futures of the reads. - data = bytes(range(64)) - fs, _ = self._make_object_fs(data) - get_object = fs._client.get_object.side_effect + 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 fail_first_range(**request): - if request["Range"].startswith("bytes=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") - return get_object(**request) + # 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._client.get_object.side_effect = fail_first_range + fs.core.get_object = get_object f = fs.open("s3://bucket/key", "rb", cache_type="none", block_size=16, max_workers=4) - file_ref = weakref.ref(f) - gc.disable() - try: - with pytest.raises(PermissionError, match="Access Denied"): - f.read() - f.close() - del f - assert file_ref() is None - finally: - gc.enable() + 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 @pytest.mark.parametrize("cache_type", ["bytes", "all", "first"]) def test_open_directory(self, cache_type): @@ -7272,68 +7276,6 @@ def upload_part(**kwargs): fs.core.complete_multipart_upload.assert_not_called() assert f.multipart_upload is None - @pytest.mark.parametrize("stage", ["creation", "write", "copy", "commit", "abort"]) - def test_failed_request_freed_by_reference_counting(self, stage): - # GH-1186: the file and its data are freed as soon as the error of a - # failed request is no longer referenced, whether the creation of the - # upload, a write of several parts, a write after a failed part copy - # of an append, or the commit raised it, also when the abort of the - # upload fails. The traceback of the error used to keep them in a - # reference cycle with the future of the request until the garbage - # collector ran. - fs = self._make_append_fs(b"a" * 6 if stage == "copy" else b"") - release = threading.Event() - - def fail(*args, **kwargs): - raise PermissionError("Access Denied") - - def fail_copy(**kwargs): - # Copying until the first write has returned. The test always - # sets the event, so the wait needs no timeout. - release.wait() - fail() - - if stage == "creation": - fs.core.create_multipart_upload.side_effect = fail - elif stage == "copy": - fs.core.upload_part_copy.side_effect = fail_copy - else: - fs.core.upload_part.side_effect = fail - if stage == "abort": - # A new error each time, which the mock does not keep. - fs._call.side_effect = fail - mode = "ab" if stage == "copy" else "wb" - f = S3File(fs, "s3://bucket/key.txt", mode=mode, block_size=4, max_workers=1) - try: - if stage in ("copy", "commit", "abort"): - f.write(b"a" * 4) - release.set() - if stage == "copy": - # The part copy has failed before the next write. - assert wait(f.multipart_upload_parts[:1], timeout=5).done - finally: - release.set() - file_ref = weakref.ref(f) - buffer_ref = weakref.ref(f.buffer) - # The write submits a part before it waits for the failed one. - operation = ( - f.close if stage in ("commit", "abort") else functools.partial(f.write, b"a" * 12) - ) - gc.disable() - try: - # The log capture of pytest would keep the logged abort failure, - # and so the frames of its traceback. - with ( - mock.patch("pyathena.filesystem.s3._logger"), - pytest.raises(PermissionError, match="Access Denied"), - ): - operation() - del f, operation - assert file_ref() is None - assert buffer_ref() is None - finally: - gc.enable() - @pytest.mark.parametrize("max_workers", [1, 4]) def test_append_failed_part_copy(self, max_workers): # GH-1178: a failed part copy of an append stops the write before it @@ -7596,14 +7538,15 @@ def test_commit_failure_and_discard(self, caplog, error, abort_fails): with pytest.raises(error, match="complete failed"): file.commit() assert (file.multipart_upload is not None) is abort_fails - # GH-1186: the parts are dropped also when the abort fails; they used - # to keep the errors of failed parts, and so the file, in a cycle. - assert file.multipart_upload_parts == [] + assert bool(file.multipart_upload_parts) is abort_fails assert ( "Failed to abort multipart upload uploadid to s3://bucket/key.txt." in caplog.text ) is abort_fails - # A commit does not complete the upload that is kept for discard(), - # nor upload the buffered data with PutObject. + # 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() @@ -7676,8 +7619,7 @@ def test_commit_failure_and_interrupted_abort(self): with pytest.raises(KeyboardInterrupt): file.commit() assert file.multipart_upload is not None - # GH-1186: the parts, which have ended, are dropped before the abort. - assert file.multipart_upload_parts == [] + assert file.multipart_upload_parts file.discard() assert file.fs._call.call_count == 2 diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index acb41bcc2..06390d18d 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -1,6 +1,6 @@ +import array import asyncio import contextlib -import functools import gc import gzip import os @@ -605,6 +605,72 @@ def create_multipart_upload(*args, **kw): 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") + ) + + 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: + # Until the completion waits for the thread. + deadline = time.monotonic() + 5 + while not fs._thread_pool._get_executor()._work_queue.qsize(): + assert time.monotonic() < deadline + await asyncio.sleep(0.01) + 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 @@ -2631,18 +2697,24 @@ def test_open_max_workers(self, max_workers): assert f.max_workers == (max_workers or fs._sync_fs.max_workers) @pytest.mark.asyncio - async def test_failed_read_freed_by_reference_counting(self): - # GH-1186: a parallel read raises the error of a failed range, and - # the file is freed by reference counting once it is closed and the - # error is no longer referenced. The executor of the file kept the - # future of the failed read in a reference cycle with its error. + 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") - return b"a" * (range_[1] - range_[0]) + # 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( @@ -2658,18 +2730,13 @@ def get_object(path, range_, **kwargs): f = await asyncio.to_thread( fs.open, "s3://bucket/key", "rb", block_size=block_size, cache_type="none" ) - file_ref = weakref.ref(f) - gc.disable() - try: - with pytest.raises(PermissionError, match="Access Denied"): - await asyncio.to_thread(f.read) - await asyncio.to_thread(f.close) - del f - # Lets the event loop run the callbacks that hold the futures. - await asyncio.sleep(0) - assert file_ref() is None - finally: - gc.enable() + 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 @@ -3017,41 +3084,6 @@ async def test_write_failed_part_while_waiting(self, fs): fs.core.abort_multipart_upload.assert_called_once() fs.core.complete_multipart_upload.assert_not_called() - @pytest.mark.parametrize("stage", ["creation", "write", "commit"]) - @pytest.mark.asyncio - async def test_failed_request_freed_by_reference_counting(self, fs, stage): - # GH-1186: the file and its data are freed as soon as the error of a - # failed request is no longer referenced, as with S3File. - block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE - - def fail(*args, **kwargs): - raise PermissionError("Access Denied") - - if stage == "creation": - fs.core.create_multipart_upload = mock.MagicMock(side_effect=fail) - else: - fs.core.upload_part = mock.MagicMock(side_effect=fail) - f = await fs.open_async("s3://bucket/key", "wb", block_size=block_size, max_workers=1) - if stage == "commit": - await f.write(b"a" * block_size) - file_ref = weakref.ref(f) - buffer_ref = weakref.ref(f.buffer) - # The write submits a part before it waits for the failed one. - operation = ( - f.close if stage == "commit" else functools.partial(f.write, b"a" * block_size * 3) - ) - gc.disable() - try: - with pytest.raises(PermissionError, match="Access Denied"): - await operation() - del f, operation - # Lets the event loop run the callbacks that hold the futures. - await asyncio.sleep(0) - assert file_ref() is None - assert buffer_ref() is None - finally: - gc.enable() - @pytest.mark.parametrize("blocks", [0, 2]) @pytest.mark.asyncio async def test_exit_with_exception(self, fs, blocks): @@ -3329,18 +3361,24 @@ async def test_close_small(self, fs, data): fs.core.create_multipart_upload.assert_not_called() @pytest.mark.asyncio - async def test_failed_read_freed_by_reference_counting(self, fs): - # GH-1186: a parallel read raises the error of a failed range, and - # the file is freed by reference counting once it is closed and the - # error is no longer referenced. The reads of the other ranges kept - # their data while the error was held, and a closed file kept itself - # in a reference cycle with its cache. + 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") - return b"a" * (range_[1] - range_[0]) + # 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( @@ -3351,19 +3389,13 @@ def get_object(path, range_, **kwargs): ) 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) - file_ref = weakref.ref(f) - gc.disable() - try: - with pytest.raises(PermissionError, match="Access Denied"): - await f.read() - await f.close() - del f - # Lets the event loop run the callbacks that hold the futures. - await asyncio.sleep(0) - assert file_ref() is None - finally: - gc.enable() + 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() + 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_read_parallel(self, fs): 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 aee2a2148..65bee13ef 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -13,17 +13,12 @@ import time import traceback import weakref -from concurrent.futures import CancelledError, Future, ThreadPoolExecutor, wait +from concurrent.futures import Future, ThreadPoolExecutor, wait from unittest import mock import pytest -from pyathena.filesystem.s3_executor import ( - S3AioExecutor, - S3AioThreadPool, - S3ThreadPoolExecutor, - raise_failure, -) +from pyathena.filesystem.s3_executor import S3AioExecutor, S3AioThreadPool, S3ThreadPoolExecutor from tests.pyathena.util import EVENT_TIMEOUT, overlapping_calls @@ -451,37 +446,17 @@ def test_submit_to(self): with pytest.raises(ValueError, match="invalid literal"): failed.result() - @pytest.mark.parametrize("chained", [None, "cause", "group"]) @pytest.mark.parametrize("executor_class", [S3ThreadPoolExecutor, S3AioThreadPool]) - def test_submit_to_failure_frees_arguments(self, executor_class, chained): + 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, also when the frames of an error that it was raised - # from, as S3Core.call() raises an OSError from a ClientError, or that - # it groups hold them. The error is freed by reference counting. The - # traceback used to keep them in a reference cycle with the future - # until the garbage collector ran. + # its traceback. The traceback used to keep them, in a reference + # cycle with the future until the garbage collector ran. class Body: pass - class DeniedError(PermissionError): - # Unlike PermissionError, it can be weakly referenced. - pass - - def send(body): - raise RuntimeError("AccessDenied") - def upload(body): - if chained is None: - raise DeniedError("Access Denied") - try: - send(body) - except RuntimeError as e: - if chained == "cause": - raise DeniedError("Access Denied") from e - error = e - # Raised outside the handler, so that the group has no context. - raise ExceptionGroup("Access Denied", [error]) + raise PermissionError("Access Denied") body = Body() body_ref = weakref.ref(body) @@ -492,20 +467,9 @@ def upload(body): executor.submit_to(future, upload, body) del body error = future.exception() - assert body_ref() is None + assert isinstance(error, PermissionError) assert traceback.extract_tb(error.__traceback__)[-1].name == "upload" - if chained == "cause": - assert isinstance(error.__cause__, RuntimeError) - assert traceback.extract_tb(error.__cause__.__traceback__)[-1].name == "send" - elif chained == "group": - assert isinstance(error, ExceptionGroup) - assert error.__context__ is None - assert traceback.extract_tb(error.exceptions[0].__traceback__)[-1].name == "send" - else: - assert isinstance(error, DeniedError) - error_ref = weakref.ref(error) - del error, future - assert error_ref() is None + assert body_ref() is None finally: gc.enable() @@ -596,75 +560,3 @@ async def main(): assert not not_done assert future.cancelled() - - -@pytest.mark.parametrize( - ("outcomes", "expected"), - [ - ([], None), - (["result"], None), - (["result", "error"], PermissionError), - (["cancelled", "error"], PermissionError), - (["cancelled", "result"], CancelledError), - ], -) -def test_raise_failure(outcomes, expected): - # The error of a failed future is raised before the cancellation of a - # cancelled one, and the set is emptied. - done: set[Future[int]] = set() - for outcome in outcomes: - future: Future[int] = Future() - if outcome == "result": - future.set_result(1) - elif outcome == "error": - future.set_exception(PermissionError("Access Denied")) - else: - future.cancel() - done.add(future) - - if expected is None: - raise_failure(done) - else: - with pytest.raises(expected): - raise_failure(done) - assert done == set() - - -def test_raise_failure_frees_future(): - # GH-1186: the raised error does not keep the failed future, so that - # neither forms a reference cycle with the frames of its traceback. - future: Future[int] = Future() - future.set_exception(PermissionError("Access Denied")) - future_ref = weakref.ref(future) - done = {future} - del future - gc.disable() - try: - with pytest.raises(PermissionError, match="Access Denied") as info: - raise_failure(done) - # Neither the frames of its traceback nor the set keep the future. - assert info.value.__traceback__ is not None - assert future_ref() is None - finally: - gc.enable() - - -def test_raise_failure_retrieves_errors(): - # The errors of the failed asyncio futures that are not raised are - # retrieved, so that asyncio does not log them as never retrieved. - loop = asyncio.new_event_loop() - try: - logged = [] - loop.set_exception_handler(lambda _, context: logged.append(context)) - done: set[asyncio.Future[int]] = set() - for _ in range(2): - future = loop.create_future() - future.set_exception(PermissionError("Access Denied")) - done.add(future) - del future - with pytest.raises(PermissionError, match="Access Denied"): - raise_failure(done) - gc.collect() - assert logged == [] - finally: - loop.close() diff --git a/tests/pyathena/test_result_set.py b/tests/pyathena/test_result_set.py index 3790237a8..83df71187 100644 --- a/tests/pyathena/test_result_set.py +++ b/tests/pyathena/test_result_set.py @@ -20,6 +20,13 @@ from tests.pyathena.util import close_during_next_request, get_query_results_responses +def _page(values, next_token=None): + response = {"ResultSet": {"Rows": [{"Data": [{"VarCharValue": v}]} for v in values]}} + if next_token: + response["NextToken"] = next_token + return response + + def _row(*values): return {"Data": [{} if v is None else {"VarCharValue": v} for v in values]} @@ -105,14 +112,11 @@ def test_fetchmany_size_not_integer(self): assert result_set.fetchall() == [(1,), (2,)] @pytest.mark.parametrize(("method", "args"), [("fetchmany", (3,)), ("fetchall", ())]) - def test_fetch_returns_rows_taken_before_error(self, method, args): - """The rows taken before a failed page request are returned to the buffer.""" - # GH-1186: they used to be dropped while counted in rownumber. + def test_fetch_counts_rows_taken_before_error(self, method, args): + """The rows taken before a failed page request are counted in rownumber.""" result_set = _fetching_result_set(("1", "2"), RuntimeError("failed")) with pytest.raises(OperationalError): getattr(result_set, method)(*args) - assert result_set.rownumber == 0 - assert result_set.fetchmany(2) == [(1,), (2,)] assert result_set.rownumber == 2 def test_fetch_after_close(self): @@ -176,7 +180,7 @@ def test_fetch_all_rows_skips_column_labels_only_on_first_page(self): result_set._process_metadata( {"ResultSet": {"ResultSetMetadata": {"ColumnInfo": [{"Name": "a", "Type": "varchar"}]}}} ) - pages = get_query_results_responses(["a", "1"], ["a", "2"]) + pages = [_page(["a", "1"], "token"), _page(["a", "2"])] with patch.object(result_set, "_get_query_results", side_effect=pages): assert result_set._fetch_all_rows() == [("1",), ("a",), ("2",)] diff --git a/tests/pyathena/test_util.py b/tests/pyathena/test_util.py index d00a4d141..1c5eed6fa 100644 --- a/tests/pyathena/test_util.py +++ b/tests/pyathena/test_util.py @@ -1,5 +1,3 @@ -import gc -import weakref from typing import Any import pytest @@ -102,30 +100,6 @@ def test_retry_api_call(): _test_retry(_WithCodeError(500)) -def test_retry_api_call_error_freed_by_reference_counting(): - # GH-1186: the error of a failed call is freed by reference counting, as - # are the frames of its traceback, such as those of a file that read with - # the call. The frames of tenacity kept its retry state, and so the error, - # in a reference cycle until the garbage collector ran. - class DeniedError(Exception): - # Unlike Exception, it can be weakly referenced. - pass - - def call(): - raise DeniedError("Access Denied") - - gc.disable() - try: - with pytest.raises(DeniedError) as info: - retry_api_call(call, config=RetryConfig(attempt=2)) - assert [entry.name for entry in info.traceback][-1] == "call" - error_ref = weakref.ref(info.value) - del info - assert error_ref() is None - finally: - gc.enable() - - def test_retry_api_call_with_none_error(): _test_retry(_NoResponseError()) From a333e17ffb8566a57939a8ac49abc127943adea0 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 01:01:40 +0900 Subject: [PATCH 16/18] Describe how a cancelled copy leaves its requests to the cleanup Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3_async.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index d407980af..520183f7c 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -580,9 +580,9 @@ async def _copy_object_with_multipart_upload( 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 + 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. @@ -671,8 +671,8 @@ def _abort() -> None: self.core.create_multipart_upload, plan.destination, **plan.create_params ) try: - # A cancellation does not cancel the creation, so that _abort() - # can wait for an upload that it creates. + # 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: @@ -694,8 +694,8 @@ def _abort() -> None: cast(list[S3MultipartUploadPart], parts), **plan.complete_params, ) - # A cancellation does not cancel the completion, so that _abort() - # can wait for it. + # 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 From 14b7ab38cecb88e53a6b8874dda97b83c9112d72 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 01:18:19 +0900 Subject: [PATCH 17/18] Cancel interrupted reads and reuse _close_without_commit() - S3File._fetch_range() also cancels the reads and drops their futures when the wait for them is interrupted, as AioS3StreamedFile does. - A failed multipart commit() closes the file with _close_without_commit() instead of repeating it. - The comments of AioS3StreamedFile._fetch_range() state that cancelling a failed read keeps asyncio from logging its error. Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3.py | 32 +++++++++++----------------- pyathena/filesystem/s3_async.py | 9 ++++---- tests/pyathena/filesystem/test_s3.py | 27 +++++++++++++++++++++++ 3 files changed, 45 insertions(+), 23 deletions(-) diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index abe70dadb..a93a2d41c 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -3308,7 +3308,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: @@ -3321,24 +3320,13 @@ def commit(self) -> None: parts, **self._get_request_kwargs("complete_multipart_upload") ) except BaseException: - # The buffer holds only the data of the last part, which a - # later commit() would otherwise upload alone with PutObject, - # so the file is closed without it, as _close_without_commit() - # closes it. - self.buffer = None - self.closed = True - self._executor.shutdown(wait=False) - # 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) @@ -3446,7 +3434,13 @@ 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 ] - done = wait(futures, return_when=FIRST_EXCEPTION).done + try: + done = wait(futures, return_when=FIRST_EXCEPTION).done + except BaseException: + # The reads are cancelled with the read of the range, as below. + while futures: + futures.pop().cancel() + raise failed = next( (f for f in done if not f.cancelled() and f.exception() is not None), None ) diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 520183f7c..25d0a5603 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -1652,17 +1652,18 @@ async def _fetch_range(self, start: int, end: int) -> bytes: try: done = (await asyncio.wait(reads, return_when=asyncio.FIRST_EXCEPTION))[0] except BaseException: - # The reads are cancelled with the read of the range. + # The reads are cancelled with the read of the range. Cancelling a + # failed one also keeps asyncio from logging its error. while reads: reads.pop().cancel() raise - # The errors are all retrieved, so that asyncio does not log those that - # are not raised as never retrieved. + # 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. + # 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() diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index ec9df7924..643699153 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -3460,6 +3460,33 @@ def get_object(path, range_, **kwargs): 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 + 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=KeyboardInterrupt), + pytest.raises(KeyboardInterrupt) as info, + ): + f.read() + # Waits for the reads that had started. + f.close() + assert [ref() for ref in read] == [None] * len(read) + 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() From 5fe9176a6dd59f64b3c1eb60c66f1ba54e07f50a Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 08:19:41 +0900 Subject: [PATCH 18/18] Clear the wait frames of interrupted reads and synchronize their tests - An interrupted or cancelled parallel read also clears the frames of the wait in the traceback, whose collections kept the finished reads. - The docstring of submit_to() states that the errors that a callable's error was raised from keep their frames, and that S3Core.call() clears those of a request's ClientError. - The interrupted-read tests wait as concurrent.futures.wait() and asyncio.wait() do, with the futures in their frames; the async read test waits for the threads of the file; the queued-completion test waits for the submission of the completion itself. Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3.py | 8 ++- pyathena/filesystem/s3_async.py | 10 +++- pyathena/filesystem/s3_executor.py | 5 +- tests/pyathena/filesystem/test_s3.py | 14 ++++- tests/pyathena/filesystem/test_s3_async.py | 68 ++++++++++++++++++++-- 5 files changed, 91 insertions(+), 14 deletions(-) diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index a93a2d41c..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 @@ -3436,8 +3437,11 @@ def _fetch_range(self, start: int, end: int) -> bytes: ] try: done = wait(futures, return_when=FIRST_EXCEPTION).done - except BaseException: - # The reads are cancelled with the read of the range, as below. + 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 diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 25d0a5603..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 @@ -1651,9 +1652,12 @@ async def _fetch_range(self, start: int, end: int) -> bytes: ] try: done = (await asyncio.wait(reads, return_when=asyncio.FIRST_EXCEPTION))[0] - except BaseException: - # The reads are cancelled with the read of the range. Cancelling a - # failed one also keeps asyncio from logging its error. + 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 diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index 67e50823e..425f87115 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -148,7 +148,10 @@ def submit_to(self, future: Future[T], fn: Callable[..., T], *args: Any, **kwarg 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. + 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. diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 643699153..4b0dffc35 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -3476,15 +3476,23 @@ def get_object(path, range_, **kwargs): 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=KeyboardInterrupt), + mock.patch("pyathena.filesystem.s3.wait", side_effect=interrupted_wait), pytest.raises(KeyboardInterrupt) as info, ): f.read() - # Waits for the reads that had started. f.close() - assert [ref() for ref in read] == [None] * len(read) + 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"]) diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index 06390d18d..a18f82524 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -642,6 +642,16 @@ def upload_part_copy(**kw): 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"), @@ -653,11 +663,8 @@ def upload_part_copy(**kw): ) ) try: - # Until the completion waits for the thread. - deadline = time.monotonic() + 5 - while not fs._thread_pool._get_executor()._work_queue.qsize(): - assert time.monotonic() < deadline - await asyncio.sleep(0.01) + # 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) @@ -3393,10 +3400,61 @@ def get_object(path, range_, **kwargs): 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