Free the data of failed S3 requests and clean up async copies at loop shutdown

docs/filesystem.md

@@ -136,9 +136,11 @@ from that version, which needs `s3:GetObjectVersion` and, to copy the tags,
 upload completes, so the destination exists without them until the last one is
 written. If an annotation fails to copy, the error is raised and the destination is
 kept. A failed part copy aborts the multipart upload. So does an interrupt, or the
-cancellation of an `AioS3FileSystem` copy, unless the upload has already completed. A
+cancellation of an `AioS3FileSystem` copy, including one by the shutdown of the event
+loop at the end of `asyncio.run()`, unless the upload has already completed. A
 CreateMultipartUpload request and the part copies in flight finish first, and so does
-the CompleteMultipartUpload request of an `AioS3FileSystem` copy. A further interrupt
+the CompleteMultipartUpload request of an `AioS3FileSystem` copy; the requests of an
+`AioS3FileSystem` copy that have not started are cancelled. A further interrupt
 while the CreateMultipartUpload request is waited for is ignored. One while the part
 copies are waited for or the upload is aborted stops the cleanup, and the upload ID is
 logged.
@@ -540,6 +542,7 @@ It drops the buffer and closes the file before waiting, so a later close or defe
 It suppresses Ctrl-C during the blocking recovery wait and re-raises the original error when cleanup finishes.
 An additional interruption outside that wait can stop cleanup; once the upload identity is known, `discard()` can retry the abort.
 An abort failure retains the upload identity for `discard()` to retry.
+A failed multipart `commit()` also closes the file without its buffered data, so a later `commit()` neither completes that upload nor uploads the last part alone with PutObject.
 Cancelling an async buffered write through `asyncio.to_thread` still allows its writer thread to finish.
 
 ### Multipart requests

pyathena/filesystem/s3.py

@@ -8,6 +8,7 @@
 import mimetypes
 import os.path
 import time
+import traceback
 from collections.abc import Callable, Mapping
 from concurrent.futures import FIRST_COMPLETED, FIRST_EXCEPTION, Future, as_completed, wait
 from copy import deepcopy
@@ -3280,7 +3281,8 @@ def commit(self) -> None:
         as soon as it fails, without waiting for the parts before it. If the
         abort also fails, the upload is kept so that :meth:`discard` can
         abort it. Invalidates the cache of the path afterwards. Does nothing
-        for a file whose failed write dropped the written data.
+        for a file whose failed write or multipart commit dropped the written
+        data.
 
         Raises:
             FileExistsError: If an object was created at the path after the
@@ -3289,9 +3291,9 @@ def commit(self) -> None:
                 initialized.
         """
         if self.buffer is None:
-            # _close_without_commit() dropped the written data. A multipart
-            # upload that it failed to abort is kept for discard(), not
-            # completed.
+            # _close_without_commit() or a failed commit dropped the written
+            # data. A multipart upload that they failed to abort is kept for
+            # discard(), not completed.
             return
         if self.tell() == 0:
             self.discard()
@@ -3307,7 +3309,6 @@ def commit(self) -> None:
             if not self.multipart_upload:
                 raise RuntimeError("Multipart upload is not initialized.")
 
-            upload_id = cast(str, self.multipart_upload.upload_id)
             try:
                 done, _ = wait(self.multipart_upload_parts, return_when=FIRST_EXCEPTION)
                 for future in done:
@@ -3320,17 +3321,13 @@ def commit(self) -> None:
                     parts, **self._get_request_kwargs("complete_multipart_upload")
                 )
             except BaseException:
-                # discard() keeps the upload if the abort fails or is
-                # interrupted, so that a later discard(), such as a
-                # transaction rollback, retries the abort. An abort failure
+                # Closed without the buffer, which holds only the data of the
+                # last part that a later commit() would otherwise upload alone
+                # with PutObject. discard() keeps the upload if the abort
+                # fails or is interrupted, so that a later discard(), such as
+                # a transaction rollback, retries the abort. An abort failure
                 # is logged so that it does not mask the original error.
-                try:
-                    self.discard()
-                except Exception:
-                    _logger.exception(
-                        f"Failed to abort multipart upload {upload_id} "
-                        f"to s3://{self.bucket}/{self.key}."
-                    )
+                self._close_without_commit()
                 raise
 
         self.fs.invalidate_cache(self.path)
@@ -3414,7 +3411,9 @@ def _fetch_range(self, start: int, end: int) -> bytes:
         The range is clamped to the size of the object, since fsspec caches
         may request a range that is empty or reaches past the end of the
         object. S3 would answer the former with the whole object and a range
-        starting past the end with an ``InvalidRange`` error.
+        starting past the end with an ``InvalidRange`` error. A range larger
+        than the block size is read with parallel requests; when one fails,
+        those that have not started are cancelled and its error is raised.
 
         Args:
             start: The offset of the first byte to read.
@@ -3436,9 +3435,27 @@ def _fetch_range(self, start: int, end: int) -> bytes:
                 self._executor.submit(self.fs.core.get_object, path, r, **request_kwargs)
                 for r in ranges
             ]
-            for future in as_completed(futures):
-                # The failure of the first read to finish is raised.
-                future.result()
+            try:
+                done = wait(futures, return_when=FIRST_EXCEPTION).done
+            except BaseException as e:
+                # The reads are cancelled with the read of the range, as below,
+                # and the frames of the wait, which keep their futures, are
+                # cleared; clear_frames() skips this frame, which is executing.
+                traceback.clear_frames(e.__traceback__)
+                while futures:
+                    futures.pop().cancel()
+                raise
+            failed = next(
+                (f for f in done if not f.cancelled() and f.exception() is not None), None
+            )
+            if failed is not None:
+                # The reads that have not started are cancelled, and the
+                # futures of the others are dropped, so that this frame, which
+                # the traceback of the error keeps, does not keep their data.
+                while futures:
+                    futures.pop().cancel()
+                done.clear()
+                failed.result()
             return b"".join(f.result() for f in futures)
         return self.fs.core.get_object(path, ranges[0], **request_kwargs)
 

pyathena/filesystem/s3_async.py

@@ -14,6 +14,7 @@
 import logging
 import mimetypes
 import os
+import traceback
 import warnings
 from concurrent.futures import FIRST_EXCEPTION, Future, wait
 from io import BytesIO
@@ -53,8 +54,6 @@
     from pyathena.connection import Connection
 
 _logger = logging.getLogger(__name__)
-# The running cleanups of cancelled multipart copies.
-_cleanup_tasks: set[asyncio.Future[None]] = set()
 
 
 class AioS3FileSystem(AsyncFileSystem):
@@ -579,11 +578,14 @@ async def _copy_object_with_multipart_upload(
         See :meth:`S3FileSystem._copy_object_with_multipart_upload`. The part
         and annotation copies run in parallel as asyncio tasks, at most
         ``max_workers`` at once, in the threads of the filesystem, which bound
-        the requests of all its copies together. On a cancellation, the creation of the upload,
-        the running part copies and the completion are waited for, the
-        upload is aborted if it was created and has not completed, and the
-        cancellation is re-raised. A repeated cancellation returns without
-        stopping this cleanup.
+        the requests of all its copies together. On a cancellation, the
+        requests that have not started are cancelled, the running ones are
+        waited for in a thread of the filesystem, the upload is aborted if it
+        was created and has not completed, and the cancellation is re-raised.
+        This also holds when the shutdown of the event loop, such as at the
+        end of ``asyncio.run()``, cancels the copy, also after its cleanup has
+        started. A repeated cancellation returns without stopping this
+        cleanup.
 
         Args:
             source: Source S3 path, with the version ID to copy, if any.
@@ -612,6 +614,9 @@ async def _copy_object_with_multipart_upload(
 
         semaphore = asyncio.Semaphore(max_workers)
         failed = False
+        # The requests of the parts, which keep running when their tasks are
+        # cancelled, such as by the shutdown of the event loop.
+        part_copies: list[Future[S3MultipartUploadPart]] = []
 
         async def _upload_part(i: int, range_: tuple[int, int]) -> S3MultipartUploadPart | None:
             nonlocal failed
@@ -620,55 +625,61 @@ 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,
                         source=plan.source,
                         range_=range_,
                         **plan.part_params,
                     )
+                    part_copies.append(part_copy)
+                    copied = asyncio.wrap_future(part_copy)
+                    return await copied
                 except Exception:
                     # Set before the semaphore lets a waiting part start.
                     failed = True
                     raise
 
-        tasks: list[asyncio.Future[S3MultipartUploadPart | None]] = []
-        completion: asyncio.Future[S3CompleteMultipartUpload] | None = None
+        completion: Future[S3CompleteMultipartUpload] | None = None
 
-        async def _abort() -> None:
-            await asyncio.wait([creation])
+        def _abort() -> None:
+            """Abort the upload, if any, once its requests have ended.
+
+            Runs in a thread of the filesystem, so that neither a repeated
+            cancellation nor the shutdown of the event loop stops it.
+            """
+            wait([creation])
             if creation.cancelled() or creation.exception() is not None:
                 # No upload was created.
                 return
             # A part that is still copying when the upload is aborted may be
-            # stored after the abort, so wait for the running parts first.
-            await asyncio.gather(*tasks, return_exceptions=True)
+            # stored after the abort.
+            wait(part_copies)
             if completion is not None:
-                await asyncio.wait([completion])
+                wait([completion])
                 if not completion.cancelled() and completion.exception() is None:
                     # The upload completed despite the cancellation, so
                     # there is nothing to abort.
                     return
-            await self._thread_pool.run(
-                self._sync_fs._abort_multipart_upload,
-                creation.result(),
-                plan.abort_params,
-            )
+            self._sync_fs._abort_multipart_upload(creation.result(), plan.abort_params)
 
-        # A task, so that _abort() can wait for an upload that is created
-        # after a cancellation; scheduled right before the try, so that
+        # The requests run in the threads of the filesystem, not as tasks,
+        # which the shutdown of the event loop would cancel while the
+        # requests keep running; scheduled right before the try, so that
         # nothing can fail between the two.
-        creation = asyncio.ensure_future(
-            self._thread_pool.run(
-                self.core.create_multipart_upload, plan.destination, **plan.create_params
-            )
+        creation = self._thread_pool.submit(
+            self.core.create_multipart_upload, plan.destination, **plan.create_params
         )
         try:
-            # shield keeps a cancellation from cancelling the creation, whose
-            # thread would keep running, so that _abort() can wait for it.
-            multipart_upload = await asyncio.shield(creation)
+            # A cancellation of this task leaves the creation to the cleanup,
+            # which waits for an upload that a started creation creates.
+            multipart_upload = await self._thread_pool.wait_shielded(creation)
             tasks = [asyncio.ensure_future(_upload_part(i, r)) for i, r in enumerate(plan.ranges)]
+            for task in tasks:
+                # Retrieved, so that asyncio does not log the failure of a
+                # part that ends after this copy as never retrieved.
+                task.add_done_callback(lambda t: t.cancelled() or t.exception())
             # Unlike gather, wait does not cancel the parts when this task is
             # cancelled; their threads would keep copying, so they are waited
             # for in _abort().
@@ -678,27 +689,29 @@ async def _abort() -> None:
                     raise error
             # The tasks are in part-number order.
             parts = [task.result() for task in tasks]
-            completion = asyncio.ensure_future(
-                self._thread_pool.run(
-                    self.core.complete_multipart_upload,
-                    multipart_upload,
-                    cast(list[S3MultipartUploadPart], parts),
-                    **plan.complete_params,
-                )
+            completion = self._thread_pool.submit(
+                self.core.complete_multipart_upload,
+                multipart_upload,
+                cast(list[S3MultipartUploadPart], parts),
+                **plan.complete_params,
             )
-            # shield keeps a cancellation from cancelling the completion, whose
-            # thread would keep running, so that _abort() can wait for it.
-            completed = await asyncio.shield(completion)
+            # A cancellation of this task leaves the completion to the
+            # cleanup, which waits for a started completion.
+            completed = await self._thread_pool.wait_shielded(completion)
         except BaseException:
             # Also on cancellation, as S3FileSystem._finish_multipart_upload
             # does on an interrupt.
             failed = True
-            cleanup = asyncio.ensure_future(_abort())
+            # The requests that have not started are cancelled here, before
+            # the cleanup waits for a thread after them.
+            creation.cancel()
+            for part_copy in part_copies:
+                part_copy.cancel()
+            if completion is not None:
+                completion.cancel()
             # A repeated cancellation of this task returns without stopping
-            # the cleanup; the event loop keeps only weak references to tasks.
-            _cleanup_tasks.add(cleanup)
-            cleanup.add_done_callback(_cleanup_tasks.discard)
-            await asyncio.shield(cleanup)
+            # the cleanup.
+            await self._thread_pool.run_shielded(_abort)
             raise
 
         failed = False
@@ -1612,7 +1625,9 @@ async def _fetch_range(self, start: int, end: int) -> bytes:
         """Read a byte range of the object with parallel range requests.
 
         The range is clamped to the size of the object, and split into ranges
-        of the block size, read in the threads of the file.
+        of the block size, read in the threads of the file. When a read fails,
+        the reads that have not finished are cancelled and its error is
+        raised, as :meth:`S3File._fetch_range` does.
 
         Args:
             start: The offset of the first byte to read.
@@ -1629,7 +1644,32 @@ async def _fetch_range(self, start: int, end: int) -> bytes:
         )
         path = S3Path(self.bucket, self.key, self.version_id)
         request_kwargs = self._get_request_kwargs("get_object")
-        results = await asyncio.gather(
-            *[self._pool.run(self.fs.core.get_object, path, r, **request_kwargs) for r in ranges]
-        )
-        return b"".join(results)
+        reads = [
+            asyncio.ensure_future(
+                self._pool.run(self.fs.core.get_object, path, r, **request_kwargs)
+            )
+            for r in ranges
+        ]
+        try:
+            done = (await asyncio.wait(reads, return_when=asyncio.FIRST_EXCEPTION))[0]
+        except BaseException as e:
+            # The reads are cancelled with the read of the range, and the
+            # frames of the wait, which keep their tasks, are cleared;
+            # clear_frames() skips this frame, which is executing. Cancelling a
+            # failed read also keeps asyncio from logging its error.
+            traceback.clear_frames(e.__traceback__)
+            while reads:
+                reads.pop().cancel()
+            raise
+        # A comprehension, whose variable does not keep a read in this frame.
+        failures = [read for read in done if not read.cancelled() and read.exception() is not None]
+        if failures:
+            # The reads that have not finished are cancelled, and the tasks of
+            # the others are dropped, so that this frame, which the traceback
+            # of the error keeps, does not keep their data. Cancelling a failed
+            # read also keeps asyncio from logging its error.
+            while reads:
+                reads.pop().cancel()
+            done.clear()
+            failures[0].result()
+        return b"".join(read.result() for read in reads)

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: