Add a fan-out primitive to the S3 executors (GH-1332)

docs/aio.md

@@ -281,10 +281,10 @@ its coroutines run the blocking boto3 calls in threads that PyAthena owns, sized
 | | S3FileSystem | AioS3FileSystem |
 |---|---|---|
 | **Blocking calls** | Caller's thread | Thread pool of the filesystem |
-| **Parallelism** | `ThreadPoolExecutor` | `asyncio.gather` on a thread pool |
+| **Parallelism** | Request pool of the filesystem | `asyncio.gather` on a thread pool; requests in the request pool |
 | **File handles** | `S3File` with thread pool | `AioS3File` with thread pool; `AioS3StreamedFile` from `open_async()` |
-| **Bulk delete** | Thread pool per batch | `asyncio.gather` per batch |
-| **Multipart copy** | Thread pool per copy | Thread pool of the filesystem |
+| **Bulk delete** | Request pool, `max_workers` requests at once | Request pool, `max_workers` requests at once |
+| **Multipart copy** | Request pool, `max_workers` requests per copy | Request pool, `max_workers` requests per copy |
 | **Best for** | Synchronous applications | Async frameworks (FastAPI, aiohttp, etc.) |
 
 ### Threads
@@ -334,6 +334,10 @@ ends. A request that has already been sent still runs to its end in its thread.
 requests that run to their end after a cancellation keep their retries: the commit of a
 file, the CompleteMultipartUpload request of a multipart copy, and the aborts of
 multipart uploads.
+The range reads of a file from `open_async()`, the DeleteObjects requests of `_rm()`, and the annotation
+copies of a multipart copy are waited for before the cancellation is raised, also when
+the shutdown of the event loop at the end of `asyncio.run()` cancels the coroutine; a
+second cancellation stops that wait.
 
 ### Executor strategy
 
@@ -343,6 +347,12 @@ and runs them in a pool of `max_workers` threads; in a forked child process, the
 creates new threads. `S3ThreadPoolExecutor` and `S3AioThreadPool` are aliases of `S3ThreadPool`.
 `S3AioExecutor` is kept for compatibility: it is an `S3ThreadPool` that ignores its
 `loop` argument, with a `DeprecationWarning` if one is given.
+`S3Executor.map()` and `S3Executor.map_async()` run a callable for each item in the
+executor and return the outcomes in the order of the items. With `stop_on_failure`,
+the default, the first failure cancels the calls that have not started, stops the
+retries of the running ones, and is raised once they have ended. A call that runs in an
+`S3ThreadPool` must not map on that pool, which would wait for the threads that the
+call holds; a scope of the filesystem runs such calls in the calling thread.
 
 The file handles of `AioS3FileSystem.open()` run their multipart uploads and parallel
 range reads in the request pool of the filesystem, at most the `max_workers` of the

docs/filesystem.md

@@ -116,10 +116,14 @@ In a forked child process, create a new filesystem. A file that the parent opene
 `RuntimeError` there from its reads, writes, and flushes, `commit()`, and `discard()`,
 and closing it only marks it closed, without uploading or aborting anything.
 
-An interrupt, such as Ctrl+C, while a file reads ranges in parallel cancels the range
-reads that have not started. Those that are running finish, and their data is released
-once they end. Closing the file waits for its running requests; an interrupt of that
-wait stops it, and the requests still finish. An interrupt while a
+An interrupt, such as Ctrl+C, or a failure while a file reads ranges in parallel
+cancels the range reads that have not started, and the running ones make no further
+attempt. The interrupt or the error is raised once they have ended, and their data is
+released. A second interrupt stops that wait, and the reads still finish. `rm` and
+`clear_multipart_uploads` send every request also after one fails; an interrupt cancels
+those that have not started and is raised once the running ones have ended. Closing
+the file waits for its running requests; an interrupt of that wait stops it, and the
+requests still finish. An interrupt while a
 request is being submitted can let one more request than `max_workers` run at once
 until the request pool is idle. These guarantees hold for one interrupt at a time. If another interrupt arrives while the cleanup of an interrupt or
 an error runs, the call still does not wait forever, but the cleanup may stop early:
@@ -299,7 +303,7 @@ fs.chmod("s3://YOUR_S3_BUCKET/path/to/", "private", recursive=True)
 
 A recursive `chmod` updates the keys in parallel. Once an update fails, the updates
 that have not started are cancelled, the running ones make no further attempt, and the
-error is raised.
+error is raised once they have ended.
 
 Note that `setxattr` rewrites the object by copying it onto itself (S3 does not allow
 updating the metadata of an existing object in place), which updates its last-modified

pyathena/filesystem/s3.py

@@ -3,15 +3,13 @@
 from __future__ import annotations
 
 import contextlib
+import functools
 import hashlib
 import logging
 import mimetypes
 import os.path
-import threading
 import time
-import traceback
-from collections.abc import Callable, Mapping
-from concurrent.futures import FIRST_EXCEPTION, Future, as_completed, wait
+from collections.abc import Callable, Iterable, Mapping
 from copy import deepcopy
 from dataclasses import dataclass, field
 from datetime import datetime
@@ -35,7 +33,14 @@
 from pyathena.connection import Connection
 from pyathena.filesystem.s3_core import S3Core, S3DeleteBatch, S3DeleteResult
 from pyathena.filesystem.s3_errors import S3ClientError
-from pyathena.filesystem.s3_executor import S3Executor, S3ThreadPool, _S3PoolScope
+from pyathena.filesystem.s3_executor import (
+    ItemT,
+    S3Executor,
+    S3MapResult,
+    S3ThreadPool,
+    T,
+    _S3PoolScope,
+)
 from pyathena.filesystem.s3_object import (
     S3Metadata,
     S3MultipartUpload,
@@ -48,7 +53,7 @@
 from pyathena.filesystem.s3_path import S3Path
 from pyathena.filesystem.s3_path_pairing import S3PathPairing
 from pyathena.filesystem.s3_writer import _S3WriteSession
-from pyathena.util import RetryConfig, override, stop_retries_on
+from pyathena.util import RetryConfig, override
 
 _logger = logging.getLogger(__name__)
 
@@ -1154,6 +1159,35 @@ def _create_executor(self, max_workers: int) -> _S3PoolScope:
         """
         return _S3PoolScope(self._request_pool, max_workers)
 
+    def _map(
+        self,
+        fn: Callable[[ItemT], T],
+        items: Iterable[ItemT],
+        max_workers: int | None = None,
+        **kwargs: Any,
+    ) -> S3MapResult[T]:
+        """Run the parallel requests of an operation in a scope of their own.
+
+        See :meth:`S3Executor.map() <pyathena.filesystem.s3_executor.S3Executor.map>`.
+        The scope is shut down without another wait, since the map has waited
+        for the calls that it started, unless a second interrupt stopped it.
+
+        Args:
+            fn: The request to run for each item.
+            items: The items.
+            max_workers: Maximum number of parallel requests. Defaults to
+                ``self.max_workers``.
+            **kwargs: The options of ``S3Executor.map()``.
+
+        Returns:
+            The outcomes of the requests.
+        """
+        executor = self._create_executor(max_workers=max_workers or self.max_workers)
+        try:
+            return executor.map(fn, items, **kwargs)
+        finally:
+            executor.shutdown(wait=False)
+
     def _delete_objects(self, paths: list[str], max_workers: int | None = None, **kwargs) -> None:
         """Delete objects with DeleteObjects requests grouped by bucket.
 
@@ -1172,11 +1206,15 @@ def _delete_objects(self, paths: list[str], max_workers: int | None = None, **kw
         if not batches:
             return
 
-        max_workers = max_workers if max_workers else self.max_workers
-        with self._create_executor(max_workers=max_workers) as executor:
-            fs = [executor.submit(self._delete_batch, batch, **params) for batch in batches]
-        # The executor has waited for every request, also after a failure.
-        self._raise_delete_errors([f.exception() or f.result() for f in fs])
+        # Every request runs also after another failed, and an interrupt
+        # cancels those that have not started.
+        results = self._map(
+            functools.partial(self._delete_batch, **params),
+            batches,
+            max_workers=max_workers,
+            stop_on_failure=False,
+        )
+        self._raise_delete_errors(results.outcomes)
 
     @staticmethod
     def _delete_batches(paths: list[str], **kwargs) -> tuple[list[S3DeleteBatch], dict[str, Any]]:
@@ -1923,11 +1961,7 @@ def cat_file(
         try:
             return self.core.get_object(s3_path.with_version_id(version_id), ranges, **kwargs)
         except OSError as e:
-            if (
-                ranges
-                and isinstance(e.__cause__, botocore.exceptions.ClientError)
-                and S3ClientError(e.__cause__).code == "InvalidRange"
-            ):
+            if ranges and S3ClientError.code_of(e) == "InvalidRange":
                 # The range starts at or past the end of the object.
                 return b""
             raise
@@ -2257,31 +2291,7 @@ def chmod(self, path: str, acl: str, recursive: bool = False, **kwargs) -> None:
             raise ValueError(f"ACL not in {self.core.OBJECT_ACLS}.")
         if recursive:
             paths = self.find(path, withdirs=False)
-            # Set on failure, so that the running requests make no further attempt.
-            stop = threading.Event()
-            with self._create_executor(max_workers=self.max_workers) as executor:
-                futures: list[Future[None]] = []
-                try:
-                    with stop_retries_on(stop):
-                        for p in paths:
-                            # Tracked before it is scheduled, so that an
-                            # interrupt while it is being scheduled still lets
-                            # it be cancelled.
-                            future: Future[None] = Future()
-                            futures.append(future)
-                            executor.submit_to(
-                                future, self.chmod, p, acl, recursive=False, **kwargs
-                            )
-                    for future in as_completed(futures):
-                        future.result()
-                except BaseException:
-                    # No more ACL is applied after one failed. The updates that
-                    # have not started are cancelled before the stop, which
-                    # would let a thread that it frees take one of them.
-                    for future in futures:
-                        future.cancel()
-                    stop.set()
-                    raise
+            self._map(lambda p: self.chmod(p, acl, recursive=False, **kwargs), paths)
             if s3_path.key:
                 # A key prefix is not an object itself; only the objects
                 # below it have ACLs.
@@ -2392,30 +2402,30 @@ def clear_multipart_uploads(self, path: str) -> None:
         uploads = self.list_multipart_uploads(path)
         if not uploads:
             return
-        error: Exception | None = None
-        with self._create_executor(max_workers=self.max_workers) as executor:
-            futures = [
-                executor.submit(
-                    self.core.abort_multipart_upload,
-                    upload,
-                )
-                for upload in uploads
-            ]
-            for future in as_completed(futures):
-                try:
-                    future.result()
-                except Exception as e:
-                    cause = e.__cause__
-                    if (
-                        isinstance(e, FileNotFoundError)
-                        and isinstance(cause, botocore.exceptions.ClientError)
-                        and S3ClientError(cause).code == "NoSuchUpload"
-                    ):
-                        continue
-                    if error is None:
-                        error = e
-        if error is not None:
-            raise error
+        # Every abort runs also after another failed, and an interrupt cancels
+        # those that have not started.
+        results = self._map(
+            self.core.abort_multipart_upload,
+            uploads,
+            stop_on_failure=False,
+            counts=self._is_not_no_such_upload,
+        )
+        if results.failure is not None:
+            raise results.failure
+
+    @staticmethod
+    def _is_not_no_such_upload(error: BaseException) -> bool:
+        """Return whether an error of an abort is other than an upload that no longer exists.
+
+        Args:
+            error: The error of ``abort_multipart_upload``.
+
+        Returns:
+            False for the ``FileNotFoundError`` of a ``NoSuchUpload`` error.
+        """
+        return not (
+            isinstance(error, FileNotFoundError) and S3ClientError.code_of(error) == "NoSuchUpload"
+        )
 
     def created(self, path: str) -> datetime:
         """Return the creation time of the path.
@@ -3408,9 +3418,11 @@ def _fetch_range(self, start: int, end: int) -> bytes:
         may request a range that is empty or reaches past the end of the
         object. S3 would answer the former with the whole object and a range
         starting past the end with an ``InvalidRange`` error. A range larger
-        than the block size is read with parallel requests; when one fails,
-        those that have not started are cancelled, the running ones make no
-        further attempt, and its error is raised.
+        than the block size is read with parallel requests (see
+        :meth:`S3Executor.map() <pyathena.filesystem.s3_executor.S3Executor.map>`);
+        when one fails, those that have not started are cancelled, the running
+        ones make no further attempt, and its error is raised once they have
+        ended.
 
         Args:
             start: The offset of the first byte to read.
@@ -3428,54 +3440,17 @@ def _fetch_range(self, start: int, end: int) -> bytes:
         path = S3Path(self.bucket, self.key, self.version_id)
         request_kwargs = self._get_request_kwargs("get_object")
         if len(ranges) > 1:
-            # Set on failure, so that the running reads make no further attempt.
-            stop = threading.Event()
-            futures: list[Future[bytes]] = []
-            done: set[Future[bytes]] = set()
-            try:
-                with stop_retries_on(stop):
-                    for r in ranges:
-                        # Tracked before it is scheduled, so that an interrupt
-                        # while it is being scheduled still lets it be
-                        # cancelled; not by a variable, which would keep it in
-                        # this frame, which the traceback of an error keeps.
-                        futures.append(Future())
-                        self._executor.submit_to(
-                            futures[-1], self.fs.core.get_object, path, r, **request_kwargs
-                        )
-                done = wait(futures, return_when=FIRST_EXCEPTION).done
-                failed = next(
-                    (f for f in done if not f.cancelled() and f.exception() is not None), None
+            # Not kept by a variable of this frame, which the traceback of an
+            # interrupt keeps (F5).
+            return b"".join(
+                cast(
+                    list[bytes],
+                    self._executor.map(
+                        functools.partial(self.fs.core.get_object, path, **request_kwargs),
+                        ranges,
+                    ).outcomes,
                 )
-                if failed is not None:
-                    failed.result()
-                return b"".join(f.result() for f in futures)
-            except BaseException as e:
-                # A failed read, or an interrupt at any point, including after
-                # the wait: the reads that have not started are cancelled, and
-                # the futures of the others are dropped, so that this frame,
-                # which the traceback of the error keeps, does not keep their
-                # data, also if another interrupt stops this cleanup. The
-                # frames of the wait, which keep the futures, are cleared, and
-                # clear_frames() skips this frame, which is executing.
-                try:
-                    # Popped, not iterated, so that no variable of this frame
-                    # keeps a future.
-                    while futures:
-                        futures.pop().cancel()
-                finally:
-                    # Also if an interrupt stops the cancellation, whose error
-                    # keeps this one as its context; the frames first, as
-                    # they keep the futures of the wait.
-                    try:
-                        traceback.clear_frames(e.__traceback__)
-                    finally:
-                        done.clear()
-                        futures.clear()
-                        # After the cancellation, which a thread that it frees
-                        # could otherwise precede by taking a queued read.
-                        stop.set()
-                raise
+            )
         return self.fs.core.get_object(path, ranges[0], **request_kwargs)
 
     @staticmethod