diff --git a/benchmarks/README.md b/benchmarks/README.md index e6a323fc3..00d3546b7 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -218,9 +218,10 @@ The `[filesystem]` section of `config.toml` sets the swept values, 8 and 40 work The suite's scales are the object sizes in bytes in `[filesystem] sizes`, not the row counts in `[scales]`: 64 MiB (small), 256 MiB (medium), 1 GiB (large), and 6 GiB (xlarge) by default. The `sync` API calls the filesystem on the trial's thread, and a file's parallel requests run in an `S3ThreadPoolExecutor`. -The `aio` API runs a read or write in `asyncio.to_thread()`, as `AioS3FSCursor` reads result files, and a file's parallel requests run through `S3AioExecutor` in `asyncio.to_thread()` on the filesystem's event loop, fsspec's IO thread. -Its copy is a coroutine awaited on the trial's event loop, whose part copies run in `asyncio.to_thread()`. -Both event loops keep Python's default executor, whose `min(32, os.cpu_count() + 4)` threads bound the parallel requests of the `aio` API. +The `aio` API runs a read or write in `asyncio.to_thread()`, as `AioS3FSCursor` reads result files, and a file's parallel requests run through `S3AioExecutor`, dispatched on the filesystem's event loop, fsspec's IO thread, to a pool of `max_workers` threads of the file. +Its copy is a coroutine awaited on the trial's event loop, whose part copies run in the filesystem's pool of `max_workers` threads. +Both event loops keep Python's default executor, which runs only the trial's `asyncio.to_thread()` call. +Before #1159, the parallel requests of the `aio` API ran in the default executors, whose `min(32, os.cpu_count() + 4)` threads bound them. The source objects, of zero bytes, are uploaded with boto3 below the run's scratch prefix and recorded under `objects` in the manifest. `prepare --filesystem-scale` uploads them with the manifest, so that every fleet host reads the same objects. diff --git a/docs/aio.md b/docs/aio.md index 97599137e..9841119b7 100644 --- a/docs/aio.md +++ b/docs/aio.md @@ -249,43 +249,92 @@ See each cursor's documentation page for detailed usage examples. ## AioS3FileSystem `AioS3FileSystem` is a native asyncio filesystem interface for Amazon S3, built on -fsspec's `AsyncFileSystem`. It provides the same functionality as `S3FileSystem` but -uses `asyncio.gather` with `asyncio.to_thread` for parallel operations instead of -`ThreadPoolExecutor`. - -### Why AioS3FileSystem? - -The synchronous `S3FileSystem` uses `ThreadPoolExecutor` for parallel S3 operations -(batch deletes, multipart uploads, range reads). When used from within an asyncio -application via `AioS3FSCursor`, this creates a thread-in-thread pattern: -the cursor wraps calls in `asyncio.to_thread()`, and inside that thread -`S3FileSystem` spawns additional threads via `ThreadPoolExecutor`. - -`AioS3FileSystem` eliminates this inefficiency by dispatching all parallel -operations through the asyncio event loop. +fsspec's `AsyncFileSystem`. It provides the same functionality as `S3FileSystem`, and +its coroutines run the blocking boto3 calls in threads that PyAthena owns, sized by +`max_workers`, instead of the default executor of the event loop. | | S3FileSystem | AioS3FileSystem | |---|---|---| -| **Parallelism** | `ThreadPoolExecutor` | `asyncio.gather` + `asyncio.to_thread` | -| **File handles** | `S3File` with thread pool | `AioS3File` with `S3AioExecutor` | +| **Blocking calls** | Caller's thread | Thread pool of the filesystem | +| **Parallelism** | `ThreadPoolExecutor` | `asyncio.gather` on a thread pool | +| **File handles** | `S3File` with thread pool | `AioS3File` with `S3AioExecutor`; `AioS3StreamedFile` from `open_async()` | | **Bulk delete** | Thread pool per batch | `asyncio.gather` per batch | -| **Multipart copy** | Thread pool per part | `asyncio.gather` per part | +| **Multipart copy** | Thread pool per part | Asyncio task per part | | **Best for** | Synchronous applications | Async frameworks (FastAPI, aiohttp, etc.) | +### Threads + +The threads of `AioS3FileSystem` belong to `S3AioThreadPool` instances, not to the +default executor of the event loop, so `loop.set_default_executor()` neither limits nor +blocks them: + +- The filesystem has a pool of `max_workers` threads for the blocking calls of its + coroutines, such as `_info()`, `_cat_file()`, the requests of `_rm()`, and the part + copies of multipart copies. A multipart copy runs at most the `max_workers` of the + call of its part copies at once, and the copies together run at most the + `max_workers` of the filesystem. +- Each file opened with `open()` or `open_async()` has a pool of its own, sized by the + `max_workers` of the call. + +The pools create threads as the calls need them, and each call runs with the caller's +context variables, as with `asyncio.to_thread()`. +In a forked child process, a pool creates new threads. + ### Executor strategy `S3FileSystem` and `S3File` use a pluggable executor abstraction (`S3Executor`) for parallel operations. Two implementations are provided: - `S3ThreadPoolExecutor` — wraps `ThreadPoolExecutor` (default for sync usage) -- `S3AioExecutor` — dispatches work via `asyncio.run_coroutine_threadsafe` + `asyncio.to_thread` +- `S3AioExecutor` — dispatches work via `asyncio.run_coroutine_threadsafe` to an + `S3AioThreadPool` of its own -`AioS3FileSystem` automatically uses `S3AioExecutor` for file handles, so multipart -uploads and parallel range reads are dispatched through the event loop with -`asyncio.to_thread()` instead of a separate `ThreadPoolExecutor` per file. -At most `max_workers` of them run at once. +`AioS3FileSystem` automatically uses `S3AioExecutor` for the file handles of `open()`, +so multipart uploads and parallel range reads are dispatched through the event loop to +a pool of `max_workers` threads of each file. An instance created with `asynchronous=True` has no event loop of its own, so its file handles use `S3ThreadPoolExecutor`. +### Async file handles + +`open_async()` returns an `AioS3StreamedFile`, whose `read()`, `write()`, and `close()` +are coroutines. It supports the `rb`, `wb`, `ab`, and `xb` modes: + +```python +from pyathena.filesystem.s3_async import AioS3FileSystem + +fs = AioS3FileSystem(asynchronous=True) + +async with await fs.open_async("s3://my-bucket/data/file.bin", "wb") as f: + await f.write(data) + +async with await fs.open_async("s3://my-bucket/data/file.bin", "rb") as f: + data = await f.read() +``` + +- A read splits the requested range into ranges of the block size, and reads them in + parallel. The file has no read cache. +- A write uploads each full block as a part of a multipart upload, with at most + `max_workers` parts in flight. A write that would upload another part awaits one of + them, so the event loop keeps running. A part found failed by that wait, or a part + copied from the existing object in an append that has failed before the write uploads + another part, stops the write with its error, and the multipart upload is aborted. + A part that fails later raises its error when the file is closed. The commit cancels + the parts that have not started as soon as one fails, and waits for the running ones + before it aborts the upload and raises the error. +- Closing the file commits the written data, as `S3File` does. +- Leaving an `async with` block with an exception, or cancelling a write, closes the file + without committing it, and aborts the multipart upload after the running parts finish. + The abort completes even if the task is cancelled again. +- Once closing or the final flush has begun the commit, the commit completes even if + the task that closes or flushes the file is cancelled, and leaving an `async with` + block with an exception afterward does not abort it. +- A file opened for writing must be closed with `await f.close()` or an `async with` + block. A file that is garbage collected before it is closed or its commit has begun + is not committed: it emits a `ResourceWarning`, and its multipart upload is aborted. + A synchronous `with` block raises `TypeError`. +- `open_async()` raises `NotImplementedError` for the write modes inside a transaction, + and for `autocommit=False`, because the file commits when it is closed. + ### Usage with AioS3FSCursor `AioS3FSCursor` automatically uses `AioS3FileSystem` internally. No additional diff --git a/docs/api/filesystem.rst b/docs/api/filesystem.rst index d9d2772ee..722efcd4e 100644 --- a/docs/api/filesystem.rst +++ b/docs/api/filesystem.rst @@ -22,6 +22,9 @@ S3 FileSystem .. autoclass:: pyathena.filesystem.s3.S3File :members: +.. autoclass:: pyathena.filesystem.s3.S3FileTarget + :members: + Async S3 FileSystem ------------------- @@ -31,6 +34,9 @@ Async S3 FileSystem .. autoclass:: pyathena.filesystem.s3_async.AioS3File :members: +.. autoclass:: pyathena.filesystem.s3_async.AioS3StreamedFile + :members: + Read-ahead Stream ----------------- @@ -49,6 +55,9 @@ S3 Executor .. autoclass:: pyathena.filesystem.s3_executor.S3AioExecutor :members: +.. autoclass:: pyathena.filesystem.s3_executor.S3AioThreadPool + :members: + S3 Paths -------- diff --git a/docs/filesystem.md b/docs/filesystem.md index acc96f81f..1855bb2f9 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -625,5 +625,6 @@ for source, destination in pairing.copy_pairs(sources): ## Async filesystem `AioS3FileSystem` provides the same functionality on top of fsspec's -`AsyncFileSystem`, dispatching parallel operations through the asyncio event loop. +`AsyncFileSystem`, running its blocking calls in thread pools sized by `max_workers`. +Its `open_async()` returns a file whose reads and writes are coroutines. See {ref}`aio-s3-filesystem` for details. diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index ab3121f4c..2ad7d1b1a 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -11,6 +11,7 @@ from collections.abc import Callable, Mapping from concurrent.futures import FIRST_COMPLETED, FIRST_EXCEPTION, Future, as_completed, wait from copy import deepcopy +from dataclasses import dataclass, field from datetime import datetime from io import BytesIO from multiprocessing import cpu_count @@ -2721,71 +2722,75 @@ def _call(self, method: str | Callable[..., Any], **kwargs) -> dict[str, Any]: return self._core.call(method, **kwargs) -class S3File(AbstractBufferedFile): - """A buffered file object for reading and writing an S3 object. +@dataclass +class S3FileTarget: + """The object that a file reads or writes, resolved when it is opened. - Instances are returned by ``S3FileSystem.open()``. + In read mode, the object is looked up, and the reads are made conditional + on its ETag (``IfMatch``). In append mode, an existing object is looked + up; one smaller than ``S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE`` is also + read, so that the file rewrites it from its buffer, and a larger one is + copied by the file as the first parts of a multipart upload. In + exclusive-create mode, the object must not exist, and the upload is + committed with ``IfNoneMatch="*"``. + + Attributes: + path: The path of the file. A version that is read is carried in the + path as a ``?versionId=`` suffix, so that a reopened file reads + the same version. + bucket: The bucket of the object. + key: The key of the object. + version_id: The version that is read, if any. + request_kwargs: The parameters of the S3 requests of the file, with + the conditions above and, in append mode, the metadata of the + existing object, which the rewritten object keeps. + details: The object that is read or appended, or an empty dictionary. + size: The size of the object that is read, if known. + appended: The existing object in append mode, if any. + append_data: The data of an appended object that the file rewrites + from its buffer, if any. """ - fs: S3FileSystem - buffer: BytesIO | None + path: str + bucket: str + key: str + version_id: str | None + request_kwargs: dict[str, Any] + details: S3Object | dict[str, Any] = field(default_factory=dict) + size: int | None = None + appended: S3Object | None = None + append_data: bytes | None = None - def __init__( - self, + @classmethod + def resolve( + cls, fs: S3FileSystem, path: str, - mode: str = "rb", + mode: str, + *, + block_size: int, version_id: str | None = None, - max_workers: int = (cpu_count() or 1) * 5, - executor: S3Executor | None = None, - block_size: int = S3FileSystem.DEFAULT_BLOCK_SIZE, - cache_type: str = "bytes", - autocommit: bool = True, - cache_options: dict[Any, Any] | None = None, size: int | None = None, - s3_additional_kwargs: dict[str, Any] | None = None, - **kwargs, - ) -> None: - """Initialize the file for the path and mode. + request_kwargs: Mapping[str, Any] | None = None, + ) -> S3FileTarget: + """Validate the arguments of a file and look up its object. - In read mode, the object is looked up with ``info()`` and the reads - are made conditional on its ETag (``IfMatch``). In append mode, an - existing object smaller than ``S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE`` is - read into the write buffer; a larger one is copied with - ``UploadPartCopy`` as the first parts of a multipart upload, whatever - the block size. In exclusive-create mode, the object must not exist - when the file is opened, and the upload is committed with - ``IfNoneMatch="*"`` so that it does not replace an object created in - the meantime. + Sends the lookup requests of the mode, which block. Args: fs: The filesystem that the file belongs to. path: S3 path (s3://bucket/key) of the file. mode: The file mode: ``rb``, ``wb``, ``ab``, or ``xb``. + block_size: The block size of the file. version_id: The version ID to read. Must match the version ID in the path if both are given. A version cannot be given, in either form, for writing or appending. - max_workers: The number of parallel workers for range reads, part - copies, and part uploads. A write keeps at most - ``max_workers`` uploaded parts in flight, waiting for one of - them to finish before it submits another. - executor: The executor for parallel operations. If None, a new - ``S3ThreadPoolExecutor`` is created. - block_size: The block size for reads and writes. Must be between - ``S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE`` and - ``S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE``, inclusive, unless reading. - cache_type: The fsspec cache type for reads; see - ``S3FileSystem.__init__()``. - autocommit: Whether to commit the written data when the file is - closed. If False, :meth:`commit` must be called. - cache_options: Options for the fsspec cache. - size: The size of the object, if known. Passed to - ``fsspec.spec.AbstractBufferedFile``. - s3_additional_kwargs: Additional parameters for the S3 requests of - the file, such as ``ContentType`` or ``RequestPayer``. Each - request receives those that its operation accepts. - **kwargs: Additional parameters for the S3 requests of the file, - which take precedence over ``s3_additional_kwargs``. + size: The size of the object to read, if known. + request_kwargs: The parameters of the S3 requests of the file. A + copy is kept. + + Returns: + The resolved target. Raises: FileExistsError: If an object exists at the path in @@ -2797,36 +2802,26 @@ def __init__( between ``S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE`` and ``S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE`` for writing. """ - self.max_workers = max_workers # A new dictionary, so that the caller's is not modified. - self.s3_additional_kwargs: dict[str, Any] = {**(s3_additional_kwargs or {}), **kwargs} - - # The arguments are validated, and the objects looked up, before the - # base class initializer: a file that fails here is never opened, so - # its garbage collection does not close (flush and commit) it. + kwargs: dict[str, Any] = dict(request_kwargs or {}) s3_path = S3Path.parse(path) path_version_id = s3_path.version_id - self.bucket = s3_path.bucket if not s3_path.key: raise ValueError("The path does not contain a key.") - self.key = s3_path.key if version_id and path_version_id: if version_id != path_version_id: raise ValueError( f"The version_id: {version_id} specified in the argument and " f"the version_id: {path_version_id} specified in the path do not match." ) - self.version_id: str | None = version_id elif path_version_id: - self.version_id = path_version_id - else: - self.version_id = version_id - if self.version_id and "r" not in mode: + version_id = path_version_id + if version_id and "r" not in mode: raise ValueError("Cannot write to the file with the version specified.") - if self.version_id and not path_version_id: + if version_id and not path_version_id: # Carry the version in the path, as with the ?versionId= suffix, # so that a reopened (e.g., unpickled) file reads the same version. - path = f"{path}?versionId={self.version_id}" + path = f"{path}?versionId={version_id}" if "r" not in mode and not ( fs.core.MULTIPART_UPLOAD_MIN_PART_SIZE <= block_size @@ -2840,63 +2835,170 @@ def __init__( f"5 GiB ({fs.core.MULTIPART_UPLOAD_MAX_PART_SIZE} bytes), inclusive: {block_size}." ) - self._details: S3Object | dict[str, Any] = {} - append_info: S3Object | None = None + details: S3Object | dict[str, Any] = {} + appended: S3Object | None = None append_data: bytes | None = None # The lookups need the parameters of the file on which their # authorization depends, such as the customer-provided key. - lookup_kwargs = fs._get_lookup_kwargs(self.s3_additional_kwargs) + lookup_kwargs = fs._get_lookup_kwargs(kwargs) if "r" in mode: - # Looked up before the base class initializer, which would + # Looked up here; the base class initializer of a file would # otherwise take the size from the latest version of the object. - info = fs.info(path, version_id=self.version_id, **lookup_kwargs) + info = fs.info(path, version_id=version_id, **lookup_kwargs) if info.get("type") == S3ObjectType.S3_OBJECT_TYPE_DIRECTORY: # A prefix has no object to read. raise FileNotFoundError(path) - if fs.version_aware and not self.version_id: + if fs.version_aware and not version_id: # Pin the version observed at open time so that reads are # consistent even if the object is overwritten. info() heads # the object when the cached entry carries no version. - self.version_id = info.get("version_id") - if self.version_id: + version_id = info.get("version_id") + if version_id: # Carried in the path as an explicit version is, so that # the methods of the file use it too. - path = f"{path}?versionId={self.version_id}" + path = f"{path}?versionId={version_id}" if etag := info.get("etag"): - self.s3_additional_kwargs.update({"IfMatch": etag}) - self._details = info + kwargs.update({"IfMatch": etag}) + details = info if size is None: size = info.get("size") elif "a" in mode: # The rewritten object keeps the metadata of the existing one, # which a cached listing entry lacks, so look up the object. with contextlib.suppress(FileNotFoundError): - append_info = fs.info(path, refresh=True, **lookup_kwargs) - if ( - append_info is not None - and append_info.get("size", 0) < fs.core.MULTIPART_UPLOAD_MIN_PART_SIZE - ): - # Too small to be a part of a multipart upload: rewritten - # from the buffer. Only the lookup parameters are sent, so - # that the whole object is read. - append_data = fs.cat_file(path, **lookup_kwargs) + appended = fs.info(path, refresh=True, **lookup_kwargs) + if appended is not None: + if appended.get("size", 0) < fs.core.MULTIPART_UPLOAD_MIN_PART_SIZE: + # Too small to be a part of a multipart upload: rewritten + # from the buffer. Only the lookup parameters are sent, so + # that the whole object is read. + append_data = fs.cat_file(path, **lookup_kwargs) + kwargs.update(appended.to_api_repr()) + details = appended elif "x" in mode: # Checked up front so that no data is uploaded for an existing # object, and on commit with IfNoneMatch for one created since. if fs.exists(path, **lookup_kwargs): raise FileExistsError(path) - self.s3_additional_kwargs.update({"IfNoneMatch": "*"}) + kwargs.update({"IfNoneMatch": "*"}) + return cls( + path=path, + bucket=s3_path.bucket, + key=s3_path.key, + version_id=version_id, + request_kwargs=kwargs, + details=details, + size=size, + appended=appended, + append_data=append_data, + ) + + +class S3File(AbstractBufferedFile): + """A buffered file object for reading and writing an S3 object. + + Instances are returned by ``S3FileSystem.open()``. + """ + + fs: S3FileSystem + buffer: BytesIO | None + + def __init__( + self, + fs: S3FileSystem, + path: str, + mode: str = "rb", + version_id: str | None = None, + max_workers: int = (cpu_count() or 1) * 5, + executor: S3Executor | None = None, + block_size: int = S3FileSystem.DEFAULT_BLOCK_SIZE, + cache_type: str = "bytes", + autocommit: bool = True, + cache_options: dict[Any, Any] | None = None, + size: int | None = None, + s3_additional_kwargs: dict[str, Any] | None = None, + **kwargs, + ) -> None: + """Initialize the file for the path and mode. + + In read mode, the object is looked up with ``info()`` and the reads + are made conditional on its ETag (``IfMatch``). In append mode, an + existing object smaller than ``S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE`` is + read into the write buffer; a larger one is copied with + ``UploadPartCopy`` as the first parts of a multipart upload, whatever + the block size. In exclusive-create mode, the object must not exist + when the file is opened, and the upload is committed with + ``IfNoneMatch="*"`` so that it does not replace an object created in + the meantime. + + Args: + fs: The filesystem that the file belongs to. + path: S3 path (s3://bucket/key) of the file. + mode: The file mode: ``rb``, ``wb``, ``ab``, or ``xb``. + version_id: The version ID to read. Must match the version ID in + the path if both are given. A version cannot be given, in + either form, for writing or appending. + max_workers: The number of parallel workers for range reads, part + copies, and part uploads. A write keeps at most + ``max_workers`` uploaded parts in flight, waiting for one of + them to finish before it submits another. + executor: The executor for parallel operations. If None, a new + ``S3ThreadPoolExecutor`` is created. + block_size: The block size for reads and writes. Must be between + ``S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE`` and + ``S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE``, inclusive, unless reading. + cache_type: The fsspec cache type for reads; see + ``S3FileSystem.__init__()``. + autocommit: Whether to commit the written data when the file is + closed. If False, :meth:`commit` must be called. + cache_options: Options for the fsspec cache. + size: The size of the object, if known. Passed to + ``fsspec.spec.AbstractBufferedFile``. + s3_additional_kwargs: Additional parameters for the S3 requests of + the file, such as ``ContentType`` or ``RequestPayer``. Each + request receives those that its operation accepts. + **kwargs: Additional parameters for the S3 requests of the file, + which take precedence over ``s3_additional_kwargs``. + + Raises: + FileExistsError: If an object exists at the path in + exclusive-create mode. + FileNotFoundError: If no object exists at the path when reading, + including when the path is a prefix. + ValueError: If the path has no key, the version IDs do not match, + a version is given for writing, or the block size is not + between ``S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE`` and + ``S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE`` for writing. + """ + self.max_workers = max_workers + # The arguments are validated, and the objects looked up, before the + # base class initializer: a file that fails here is never opened, so + # its garbage collection does not close (flush and commit) it. + target = S3FileTarget.resolve( + fs, + path, + mode, + block_size=block_size, + version_id=version_id, + size=size, + request_kwargs={**(s3_additional_kwargs or {}), **kwargs}, + ) + self.bucket = target.bucket + self.key = target.key + self.version_id: str | None = target.version_id + self.s3_additional_kwargs: dict[str, Any] = target.request_kwargs + self._details: S3Object | dict[str, Any] = target.details self._executor: S3Executor = executor or S3ThreadPoolExecutor(max_workers=max_workers) super().__init__( fs=fs, - path=path, + path=target.path, mode=mode, block_size=block_size, autocommit=autocommit, cache_type=cache_type, cache_options=cache_options, - size=size, + size=target.size, ) self.append_block = False @@ -2907,15 +3009,13 @@ def __init__( self._uploads_in_flight: set[Future[S3MultipartUploadPart]] = set() # The part copies of an append that may not have finished. self._copies_in_flight: set[Future[S3MultipartUploadPart]] = set() - if append_info is not None: - if append_data is not None: - self.write(append_data) + if target.appended is not None: + if target.append_data is not None: + self.write(target.append_data) else: # Copied with UploadPartCopy as the leading part(s). self.append_block = True - self.loc = append_info.get("size", 0) - self.s3_additional_kwargs.update(append_info.to_api_repr()) - self._details = append_info + self.loc = target.appended.get("size", 0) def _get_request_kwargs(self, method: str) -> dict[str, Any]: """Select the parameters of the file that an S3 operation accepts. diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 99a86907e..1bdee1556 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -10,21 +10,30 @@ from __future__ import annotations import asyncio +import contextlib import logging import mimetypes import os +import warnings +from concurrent.futures import FIRST_EXCEPTION, Future, wait +from io import BytesIO from multiprocessing import cpu_count from typing import TYPE_CHECKING, Any, cast -from fsspec.asyn import AsyncFileSystem, sync +from fsspec.asyn import AbstractAsyncStreamedFile, AsyncFileSystem, sync from fsspec.callbacks import _DEFAULT_CALLBACK from fsspec.core import get_compression from fsspec.implementations.local import LocalFileSystem, make_path_posix from fsspec.utils import check_contained -from pyathena.filesystem.s3 import CompressedBuffer, S3File, S3FileSystem +from pyathena.filesystem.s3 import CompressedBuffer, S3File, S3FileSystem, S3FileTarget from pyathena.filesystem.s3_core import S3Core -from pyathena.filesystem.s3_executor import S3AioExecutor, S3Executor, S3ThreadPoolExecutor +from pyathena.filesystem.s3_executor import ( + S3AioExecutor, + S3AioThreadPool, + S3Executor, + S3ThreadPoolExecutor, +) from pyathena.filesystem.s3_object import ( S3CompleteMultipartUpload, S3Metadata, @@ -36,6 +45,7 @@ ) from pyathena.filesystem.s3_path import S3Path from pyathena.filesystem.s3_path_pairing import S3PathPairing +from pyathena.filesystem.s3_writer import S3MultipartWriter if TYPE_CHECKING: from datetime import datetime @@ -50,19 +60,22 @@ class AioS3FileSystem(AsyncFileSystem): """An async filesystem interface for Amazon S3 using fsspec's AsyncFileSystem. - This class wraps ``S3FileSystem`` to provide native asyncio support. Instead of - using ``ThreadPoolExecutor`` for parallel operations, it uses ``asyncio.gather`` - with ``asyncio.to_thread`` for natural integration with the asyncio event loop. + This class wraps ``S3FileSystem`` to provide native asyncio support. Its + coroutines run the blocking boto3 calls in an ``S3AioThreadPool`` with + ``max_workers`` threads that the filesystem owns, instead of the default + executor of the event loop, and run parallel requests with + ``asyncio.gather``. The implementation uses composition: an internal ``S3FileSystem`` instance handles - all boto3 calls, while this class delegates to it via ``asyncio.to_thread()``. + all boto3 calls, while this class delegates to it in the threads of its pool. This avoids diamond inheritance issues and keeps all boto3 logic in one place. File handles created by ``_open`` use ``S3AioExecutor`` so that parallel operations (range reads, multipart uploads) are dispatched through the event - loop with ``asyncio.to_thread`` instead of a ``ThreadPoolExecutor`` per file. - An instance created with ``asynchronous=True`` has no event loop of its own, - so its file handles use a ``ThreadPoolExecutor``. + loop to a pool of ``max_workers`` threads of each file. An instance created + with ``asynchronous=True`` has no event loop of its own, so its file handles + use a ``ThreadPoolExecutor``. :meth:`open_async` returns an + ``AioS3StreamedFile``, whose reads and writes are coroutines. Attributes: _sync_fs: The internal synchronous S3FileSystem instance. @@ -105,7 +118,11 @@ def __init__( connection: Passed to the internal ``S3FileSystem``. default_block_size: Passed to the internal ``S3FileSystem``. default_cache_type: Passed to the internal ``S3FileSystem``. - max_workers: Passed to the internal ``S3FileSystem``. + max_workers: The number of threads that run the blocking calls of + the filesystem, and the default number of parallel requests + of its files and copies. If None, the default of + ``ThreadPoolExecutor``: ``min(32, CPUs + 4)``. Passed to the internal + ``S3FileSystem``. s3_additional_kwargs: Passed to the internal ``S3FileSystem``. allow_bucket_creation: Passed to the internal ``S3FileSystem``. allow_bucket_deletion: Passed to the internal ``S3FileSystem``. @@ -122,6 +139,11 @@ def __init__( batch_size=batch_size, **kwargs, ) + # The blocking calls of the coroutines run in threads of this pool, + # not in the default executor of the event loop. Created first, so + # that None takes its default for the files and copies too. + self._thread_pool = S3AioThreadPool(max_workers) + max_workers = self._thread_pool.max_workers self._sync_fs = S3FileSystem( connection=connection, default_block_size=default_block_size, @@ -144,21 +166,23 @@ def core(self) -> S3Core: return self._sync_fs.core async def _info(self, path: str, **kwargs) -> S3Object: - return await asyncio.to_thread(self._sync_fs.info, path, **kwargs) + return await self._thread_pool.run(self._sync_fs.info, path, **kwargs) async def _ls(self, path: str, detail: bool = False, **kwargs) -> list[S3Object] | list[str]: - return await asyncio.to_thread(self._sync_fs.ls, path, detail=detail, **kwargs) + return await self._thread_pool.run(self._sync_fs.ls, path, detail=detail, **kwargs) async def _cat_file( self, path: str, start: int | None = None, end: int | None = None, **kwargs ) -> bytes: - return await asyncio.to_thread(self._sync_fs.cat_file, path, start=start, end=end, **kwargs) + return await self._thread_pool.run( + self._sync_fs.cat_file, path, start=start, end=end, **kwargs + ) async def _exists(self, path: str, **kwargs) -> bool: - return await asyncio.to_thread(self._sync_fs.exists, path, **kwargs) + return await self._thread_pool.run(self._sync_fs.exists, path, **kwargs) async def _rm_file(self, path: str, **kwargs) -> None: - await asyncio.to_thread(self._sync_fs.rm_file, path, **kwargs) + await self._thread_pool.run(self._sync_fs.rm_file, path, **kwargs) async def _pipe_file( self, path: str, value: bytes | bytearray | memoryview, mode: str = "overwrite", **kwargs @@ -166,9 +190,9 @@ async def _pipe_file( if self._intrans: # The transaction belongs to this filesystem, not to the internal # S3FileSystem, so write through open() to defer the commit to it. - await asyncio.to_thread(self._pipe_file_in_transaction, path, value, mode, **kwargs) + await self._thread_pool.run(self._pipe_file_in_transaction, path, value, mode, **kwargs) return - await asyncio.to_thread(self._sync_fs.pipe_file, path, value, mode=mode, **kwargs) + await self._thread_pool.run(self._sync_fs.pipe_file, path, value, mode=mode, **kwargs) def _pipe_file_in_transaction( self, path: str, value: bytes | bytearray | memoryview, mode: str, **kwargs @@ -211,11 +235,11 @@ async def _put_file( ) -> None: if self._intrans: # See _pipe_file. - await asyncio.to_thread( + await self._thread_pool.run( self._put_file_in_transaction, lpath, rpath, callback, mode, **kwargs ) return - await asyncio.to_thread( + await self._thread_pool.run( self._sync_fs.put_file, lpath, rpath, callback=callback, mode=mode, **kwargs ) @@ -274,13 +298,17 @@ def _put_file_in_transaction( self.invalidate_cache(rpath) async def _get_file(self, rpath: str, lpath: str, callback=_DEFAULT_CALLBACK, **kwargs) -> None: - await asyncio.to_thread(self._sync_fs.get_file, rpath, lpath, callback=callback, **kwargs) + await self._thread_pool.run( + self._sync_fs.get_file, rpath, lpath, callback=callback, **kwargs + ) async def _mkdir(self, path: str, create_parents: bool = True, **kwargs) -> None: - await asyncio.to_thread(self._sync_fs.mkdir, path, create_parents=create_parents, **kwargs) + await self._thread_pool.run( + self._sync_fs.mkdir, path, create_parents=create_parents, **kwargs + ) async def _makedirs(self, path: str, exist_ok: bool = False) -> None: - await asyncio.to_thread(self._sync_fs.makedirs, path, exist_ok=exist_ok) + await self._thread_pool.run(self._sync_fs.makedirs, path, exist_ok=exist_ok) async def _rm( self, @@ -293,7 +321,7 @@ async def _rm( """Delete objects with DeleteObjects requests. See :meth:`S3FileSystem.rm`. The requests run in parallel with - ``asyncio.gather`` and ``asyncio.to_thread``. + ``asyncio.gather`` in the threads of the filesystem. Args: path: S3 path (s3://bucket/key) or list of paths to delete. @@ -307,7 +335,7 @@ async def _rm( ValueError: If a path is a bucket. OSError: If S3 could not delete some of the objects. """ - paths = await asyncio.to_thread( + paths = await self._thread_pool.run( self._sync_fs._delete_paths, path, recursive=recursive, maxdepth=maxdepth ) await self._delete_objects(paths, **kwargs) @@ -328,7 +356,10 @@ async def _delete_objects(self, paths: list[str], **kwargs) -> None: """ batches, params = self._sync_fs._delete_batches(paths, **kwargs) results = await asyncio.gather( - *[asyncio.to_thread(self._sync_fs._delete_batch, batch, **params) for batch in batches], + *[ + self._thread_pool.run(self._sync_fs._delete_batch, batch, **params) + for batch in batches + ], return_exceptions=True, ) self._sync_fs._raise_delete_errors(results) @@ -356,7 +387,7 @@ async def _mv(self, path1, path2, recursive=False, maxdepth=None, **kwargs) -> N """ if path1 == path2: return - pairs = await asyncio.to_thread( + pairs = await self._thread_pool.run( self._sync_fs._move_pairs, path1, path2, recursive=recursive, maxdepth=maxdepth ) # Every copy finishes before a failure is raised, as in fsspec's @@ -418,7 +449,7 @@ async def _copy( """ sources = [path1] if isinstance(path1, (str, os.PathLike)) else path1 if isinstance(path2, str) and any(S3Path.has_version_id(p) for p in sources): - pairs = await asyncio.to_thread( + pairs = await self._thread_pool.run( self._sync_fs._copy_pairs, S3PathPairing(path1, path2, recursive=recursive, maxdepth=maxdepth), ) @@ -459,7 +490,7 @@ async def _get( sources = [rpath] if isinstance(rpath, (str, os.PathLike)) else rpath if isinstance(lpath, (str, os.PathLike)) and any(S3Path.has_version_id(p) for p in sources): root = make_path_posix(lpath) - pairs = await asyncio.to_thread( + pairs = await self._thread_pool.run( self._sync_fs._copy_pairs, S3PathPairing(rpath, root, recursive=recursive, maxdepth=maxdepth), isdir=LocalFileSystem().isdir, @@ -521,7 +552,7 @@ async def _copy_file(self, path1: str, path2: str, **kwargs) -> bool: size1 = info1.get("size", 0) try: if size1 <= self.core.MULTIPART_UPLOAD_MAX_PART_SIZE: - await asyncio.to_thread(self.core.copy_object, source, destination, **kwargs) + await self._thread_pool.run(self.core.copy_object, source, destination, **kwargs) else: await self._copy_object_with_multipart_upload( source, @@ -546,8 +577,9 @@ async def _copy_object_with_multipart_upload( """Copy an object with a multipart upload of its byte ranges. See :meth:`S3FileSystem._copy_object_with_multipart_upload`. The part - and annotation copies run in parallel as asyncio tasks with - ``asyncio.to_thread``. On a cancellation, the creation of the 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. A repeated cancellation returns without @@ -567,12 +599,14 @@ async def _copy_object_with_multipart_upload( size. """ max_workers = max_workers if max_workers else self._sync_fs.max_workers - plan = await asyncio.to_thread( + plan = await self._thread_pool.run( self.core.plan_multipart_copy, source, destination, block_size, **kwargs ) if plan.fits_single_request: # See S3FileSystem._copy_object_with_multipart_upload. - await asyncio.to_thread(self.core.copy_object, plan.source, plan.destination, **kwargs) + await self._thread_pool.run( + self.core.copy_object, plan.source, plan.destination, **kwargs + ) return multipart_upload: S3MultipartUpload @@ -586,7 +620,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 asyncio.to_thread( + return await self._thread_pool.run( self.core.upload_part_copy, upload=multipart_upload, part_number=i + 1, @@ -600,7 +634,7 @@ async def _upload_part(i: int, range_: tuple[int, int]) -> S3MultipartUploadPart raise tasks: list[asyncio.Future[S3MultipartUploadPart | None]] = [] - completion: asyncio.Task[S3CompleteMultipartUpload] | None = None + completion: asyncio.Future[S3CompleteMultipartUpload] | None = None async def _abort() -> None: await asyncio.wait([creation]) @@ -616,7 +650,7 @@ async def _abort() -> None: # The upload completed despite the cancellation, so # there is nothing to abort. return - await asyncio.to_thread( + await self._thread_pool.run( self._sync_fs._abort_multipart_upload, creation.result(), plan.abort_params, @@ -626,7 +660,7 @@ async def _abort() -> None: # after a cancellation; scheduled right before the try, so that # nothing can fail between the two. creation = asyncio.ensure_future( - asyncio.to_thread( + self._thread_pool.run( self.core.create_multipart_upload, plan.destination, **plan.create_params ) ) @@ -645,7 +679,7 @@ async def _abort() -> None: # The tasks are in part-number order. parts = [task.result() for task in tasks] completion = asyncio.ensure_future( - asyncio.to_thread( + self._thread_pool.run( self.core.complete_multipart_upload, multipart_upload, cast(list[S3MultipartUploadPart], parts), @@ -677,7 +711,7 @@ async def _copy_annotation(name: str) -> None: # S3FileSystem does. return try: - await asyncio.to_thread( + await self._thread_pool.run( self.core.copy_object_annotation, name, plan.source, @@ -705,7 +739,7 @@ async def _find( **kwargs, ) -> dict[str, S3Object] | list[str]: detail = kwargs.pop("detail", False) - files = await asyncio.to_thread( + files = await self._thread_pool.run( self._sync_fs._find, path, maxdepth=maxdepth, withdirs=withdirs, **kwargs ) if detail: @@ -807,8 +841,78 @@ def _open( **kwargs, ) + async def open_async( + self, + path: str, + mode: str = "rb", + block_size: int | None = None, + max_workers: int | None = None, + version_id: str | None = None, + s3_additional_kwargs: dict[str, Any] | None = None, + **kwargs, + ) -> AioS3StreamedFile: + """Open a file whose reads and writes are coroutines. + + The lookups of the mode, as in ``open()``, run in the threads of the + filesystem. See :class:`AioS3StreamedFile` for how the file reads, + writes, and commits. + + Args: + path: S3 path (s3://bucket/key) of the file. + mode: The file mode: ``rb``, ``wb``, ``ab``, or ``xb``. + block_size: The block size for reads and writes. Defaults to the + default block size of the filesystem. + max_workers: The number of threads of the file, which bounds its + parallel requests. Defaults to the ``max_workers`` of the + filesystem. + version_id: The version ID to read. + s3_additional_kwargs: Additional parameters for the S3 requests of + the file, which take precedence over those of the filesystem. + **kwargs: Additional parameters for the S3 requests of the file, + which take precedence over ``s3_additional_kwargs``. + + Returns: + The opened file. + + Raises: + ValueError: If the mode is not binary, compression is requested, + or the arguments are invalid, as in ``open()``. + NotImplementedError: If the mode is not supported, or the file is + opened for writing in a transaction or with + ``autocommit=False``, as the file commits when it is closed. + FileExistsError: If an object exists at the path in + exclusive-create mode. + FileNotFoundError: If no object exists at the path when reading. + """ + if "b" not in mode or kwargs.pop("compression", None): + raise ValueError("open_async() supports only binary modes without compression.") + if not kwargs.pop("autocommit", True): + raise NotImplementedError("The file of open_async() commits when it is closed.") + if mode != "rb" and self._intrans: + raise NotImplementedError("open_async() cannot write in a transaction.") + path = self._strip_protocol(path) + if block_size is None: + block_size = self._sync_fs.default_block_size + if max_workers is None: + max_workers = self._sync_fs.max_workers + s3_additional_kwargs = { + **self._sync_fs.s3_additional_kwargs, + **(s3_additional_kwargs or {}), + } + return await self._thread_pool.run( + AioS3StreamedFile, + self, + path, + mode, + version_id=version_id, + max_workers=max_workers, + block_size=block_size, + s3_additional_kwargs=s3_additional_kwargs, + **kwargs, + ) + async def _rmdir(self, path: str) -> None: - await asyncio.to_thread(self._sync_fs.rmdir, path) + await self._thread_pool.run(self._sync_fs.rmdir, path) def rmdir(self, path: str) -> None: """Remove an S3 bucket, which must be empty. @@ -1008,7 +1112,7 @@ def invalidate_cache(self, path: str | None = None) -> None: self._sync_fs.invalidate_cache(path) async def _touch(self, path: str, truncate: bool = True, **kwargs) -> dict[str, Any]: - return await asyncio.to_thread(self._sync_fs.touch, path, truncate=truncate, **kwargs) + return await self._thread_pool.run(self._sync_fs.touch, path, truncate=truncate, **kwargs) def touch(self, path: str, truncate: bool = True, **kwargs) -> dict[str, Any]: """Create an empty object with PutObject. @@ -1034,8 +1138,498 @@ class AioS3File(S3File): ``isinstance`` checks and to document the async execution model. All parallel operations (range reads, multipart uploads) are dispatched through the ``S3Executor`` interface — the ``S3AioExecutor`` - provided by ``AioS3FileSystem`` dispatches them through the event loop with - ``asyncio.to_thread`` instead of a ``ThreadPoolExecutor`` per file. + provided by ``AioS3FileSystem`` dispatches them through the event loop to + a pool of ``max_workers`` threads of the file. For an ``AioS3FileSystem`` created with ``asynchronous=True``, it is an - ``S3ThreadPoolExecutor``. + ``S3ThreadPoolExecutor``. For reads and writes that are coroutines, see + ``AioS3StreamedFile``. + """ + + +class AioS3StreamedFile(AbstractAsyncStreamedFile): + """An S3 file whose reads and writes are coroutines. + + Returned by :meth:`AioS3FileSystem.open_async`. The requests of the file + run in an ``S3AioThreadPool`` of ``max_workers`` threads of its own, and + the coroutines wait for them without blocking the event loop: + + - A read splits the requested range into ranges of the block size, read + in parallel. Nothing is cached. As with ``S3File``, the reads are + conditional on the ETag of the object, and on its version when the + filesystem is version-aware. + - A write buffers the data, and uploads each full block as a part of a + multipart upload, with at most ``max_workers`` parts in flight. A write + that would submit another part waits for one of them to finish first, + and a part found failed by that wait, or a part copy of an append + found failed before a part is submitted, stops the write with its + error. + An append copies an existing object with ``UploadPartCopy``, or + rewrites one smaller than a part from the buffer, as ``S3File`` does. + - Closing the file commits the written data, as ``S3File.close()`` does: + with PutObject if no part was uploaded, otherwise by completing the + multipart upload, which is aborted if the completion fails. + - Leaving an ``async with`` block with an exception, a failed write, and + a cancelled write close the file without committing it: the requests + that have started are waited for and the multipart upload is aborted, + so that the existing object is left unchanged. + + The commit and the abort run to their end in the threads of the file, + also if the task that waits for them is cancelled, including a commit + that still waits for a thread; once the commit has been submitted, + nothing else aborts the upload. A file opened for + writing must be closed with ``await close()`` or an ``async with`` block; + garbage collection cannot await it, so the data of a file that is + collected before it is closed or its commit has been submitted is never + committed; its multipart upload, if any, is aborted, with a + ``ResourceWarning``. """ + + fs: AioS3FileSystem + buffer: BytesIO | None + closed: bool + forced: bool + offset: int | None + + def __init__( + self, + fs: AioS3FileSystem, + path: str, + mode: str = "rb", + version_id: str | None = None, + max_workers: int = (cpu_count() or 1) * 5, + block_size: int = S3FileSystem.DEFAULT_BLOCK_SIZE, + size: int | None = None, + s3_additional_kwargs: dict[str, Any] | None = None, + **kwargs, + ) -> None: + """Initialize the file for the path and mode. + + Sends the lookup requests of the mode (see :class:`S3FileTarget`), + which block; :meth:`AioS3FileSystem.open_async` runs this in a thread. + + Args: + fs: The filesystem that the file belongs to. + path: S3 path (s3://bucket/key) of the file. + mode: The file mode: ``rb``, ``wb``, ``ab``, or ``xb``. + version_id: The version ID to read. + max_workers: The number of threads of the file, which bounds its + parallel requests. + block_size: The block size for reads and writes. Must be between + ``S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE`` and + ``S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE``, inclusive, unless + reading. + size: The size of the object to read, if known. + s3_additional_kwargs: Additional parameters for the S3 requests of + the file. Each request receives those that its operation + accepts. + **kwargs: Additional parameters for the S3 requests of the file, + which take precedence over ``s3_additional_kwargs``. + + Raises: + FileExistsError: If an object exists at the path in + exclusive-create mode. + FileNotFoundError: If no object exists at the path when reading. + NotImplementedError: If the mode is not supported. + ValueError: If ``max_workers`` is not positive, or the arguments + are invalid as for ``S3FileTarget.resolve()``. + """ + if mode not in {"rb", "wb", "ab", "xb"}: + raise NotImplementedError("File mode not supported") + # Created first, so that an invalid max_workers raises before the lookups. + self._pool = S3AioThreadPool(max_workers) + self.max_workers = self._pool.max_workers + target = S3FileTarget.resolve( + fs._sync_fs, + path, + mode, + block_size=block_size, + version_id=version_id, + size=size, + request_kwargs={**(s3_additional_kwargs or {}), **kwargs}, + ) + self.bucket = target.bucket + self.key = target.key + self.version_id = target.version_id + self.s3_additional_kwargs: dict[str, Any] = target.request_kwargs + self._details: S3Object | dict[str, Any] = target.details + super().__init__( + fs=fs, + path=target.path, + mode=mode, + block_size=block_size, + cache_type="none", + size=target.size, + ) + + self.append_block = False + # The size of the appended object, as looked up when opening. + self._append_size = 0 + self._writer: S3MultipartWriter | None = None + if mode != "rb": + self._writer = S3MultipartWriter( + fs.core, S3Path(self.bucket, self.key), block_size=block_size + ) + self._creation: Future[S3MultipartUpload] | None = None + # Set when the commit is submitted, after which it completes the + # upload, or aborts it if it fails, and nothing else aborts it. + self._committing = False + # The futures of the part uploads and copies, in part-number order. + self._parts: list[Future[S3MultipartUploadPart]] = [] + # The parts uploaded from the buffer that may not have finished. The + # part copies of an append hold no data and are not counted. + self._uploads_in_flight: set[asyncio.Future[S3MultipartUploadPart]] = set() + # The part copies of an append that may not have finished. + self._copies_in_flight: set[Future[S3MultipartUploadPart]] = set() + if target.appended is not None: + if target.append_data is not None: + cast(BytesIO, self.buffer).write(target.append_data) + else: + # Copied with UploadPartCopy as the leading part(s). + self.append_block = True + self._append_size = self.loc = target.appended.get("size", 0) + + def __del__(self) -> None: + # close() is a coroutine, which garbage collection cannot await, so + # an unclosed file whose commit was not submitted is not committed: its + # multipart upload, if any, is aborted, and its threads are released. + pool = getattr(self, "_pool", None) + if pool is None: + return + if ( + not getattr(self, "closed", True) + and self.mode != "rb" + and not getattr(self, "_committing", False) + ): + warnings.warn( + f"s3://{self.bucket}/{self.key} was not closed, and the written data " + "was not committed.", + ResourceWarning, + stacklevel=2, + ) + with contextlib.suppress(RuntimeError): + pool.submit(self._discard) + pool.shutdown(wait=False) + + def __enter__(self) -> AioS3StreamedFile: + """Reject a synchronous ``with`` block, whose exit cannot close the file. + + Raises: + TypeError: Always; use ``async with``. + """ + raise TypeError("AioS3StreamedFile is asynchronous; use 'async with'.") + + def _get_request_kwargs(self, method: str) -> dict[str, Any]: + """Select the parameters of the file that an S3 operation accepts. + + Args: + method: The name of the client method, such as ``upload_part``. + + Returns: + The parameters in ``s3_additional_kwargs`` that the operation + accepts. + """ + return self.fs.core.operation_params(method, self.s3_additional_kwargs) + + async def write(self, data: bytes) -> int: + """Write data to the buffer, uploading each full block as a part. + + Args: + data: The bytes to write. + + Returns: + The number of bytes written. + + Raises: + ValueError: If the file is not open for writing, is closed, or + has been force-flushed. + BaseException: The error of a part that the write found failed + while it waited for a part in flight, or of a part copy of an + append found failed before a part is submitted. + """ + if self.mode not in {"wb", "ab", "xb"}: + raise ValueError("File not in write mode") + if self.closed: + raise ValueError("I/O operation on closed file.") + if self.forced: + raise ValueError("This file has been force-flushed, can only close") + buffer = cast(BytesIO, self.buffer) + out = buffer.write(data) + self.loc += out + if buffer.tell() >= self.blocksize: + await self.flush() + return out + + async def flush(self, force: bool = False) -> None: + """Upload the buffered data if it fills a block, or all of it if forced. + + Args: + force: Whether this is the final flush of the file, after which + the written data is committed. + + Raises: + ValueError: If the file is closed, or force-flushed twice. + """ + if self.closed: + raise ValueError("Flush on closed file") + if force and self.forced: + raise ValueError("Force flush cannot be called more than once") + if force: + self.forced = True + if self.mode not in {"wb", "ab", "xb"}: + return + buffer = cast(BytesIO, self.buffer) + if not force and buffer.tell() < self.blocksize: + return + if self.offset is None: + self.offset = 0 + await self._initiate_upload() + if await self._upload_chunk(final=force) is not False: + self.offset += cast(BytesIO, self.buffer).seek(0, 2) + self.buffer = BytesIO() + + async def close(self) -> None: + """Close the file, committing the written data.""" + if self.closed: + return + try: + if self.mode != "rb" and not self.forced: + await self.flush(force=True) + finally: + self.closed = True + # A commit that is still running keeps its thread. + self._pool.shutdown(wait=False) + + async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: + if self.closed: + return + if exc_type is not None and self.mode != "rb": + await self._close_without_commit() + else: + await self.close() + + async def _close_without_commit(self) -> None: + """Close the file without uploading the written data. + + Drops the buffered data, and aborts the multipart upload, if any, + once its started requests have ended (see :meth:`_discard`). + """ + if self.closed: + return + self.closed = True + if self._committing: + # A cancelled commit runs on; an abort would race its completion. + self._pool.shutdown(wait=False) + return + self.buffer = None + try: + await self._pool.run_shielded(self._discard) + finally: + self._pool.shutdown(wait=False) + + async def _initiate_upload(self) -> None: + if not self.append_block and self.tell() < self.blocksize: + return + writer = cast(S3MultipartWriter, self._writer) + # Tracked before it is scheduled, so that an interrupt while it is + # being scheduled still lets the cleanup cancel it or wait for it. + self._creation = Future() + try: + self._pool.submit_to( + self._creation, + writer.initiate, + **self._get_request_kwargs("create_multipart_upload"), + ) + await asyncio.wrap_future(self._creation) + if self.append_block: + size = self._append_size + whole_source = size <= self.fs.core.MULTIPART_UPLOAD_MAX_PART_SIZE + source = S3Path(self.bucket, self.key) + copy_kwargs = self._get_request_kwargs("upload_part_copy") + for part_number, range_ in writer.iter_copy_parts(size): + # Tracked before it is scheduled; see the creation. + copy: Future[S3MultipartUploadPart] = Future() + self._parts.append(copy) + self._copies_in_flight.add(copy) + self._pool.submit_to( + copy, + writer.upload_part_copy, + part_number=part_number, + source=source, + range_=None if whole_source else range_, + **copy_kwargs, + ) + except BaseException: + await self._close_without_commit() + raise + + async def _upload_chunk(self, final: bool = False) -> bool: + """Upload the buffered data as parts of the multipart upload. + + See :meth:`S3File._upload_chunk`; a write waits for a part in flight + without blocking the event loop. + + Args: + final: Whether this is the last flush of the file, after which + the upload is committed. + + Returns: + Whether the buffer is reset afterwards. + + Raises: + BaseException: The error of an uploaded part that the write found + failed while it waited for a part in flight, or of a part copy + of an append found failed before a part is submitted. + RuntimeError: If no multipart upload is initialized. + ValueError: If the data needs more parts than a multipart upload + allows. + """ + if not self.append_block and self.tell() < self.blocksize: + # Smaller than a block: uploaded with PutObject by the commit. + if final: + await self._commit() + return not final + writer = cast(S3MultipartWriter, self._writer) + if writer.upload is None: + raise RuntimeError("Multipart upload is not initialized.") + buffer = cast(BytesIO, self.buffer) + buffer.seek(0) + upload_kwargs = self._get_request_kwargs("upload_part") + try: + for part_number, data in writer.iter_parts( + buffer, first_part_number=len(self._parts) + 1 + ): + 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() + 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() + # 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 + if final: + await self._commit() + return not final + + async def _commit(self) -> None: + """Commit the written data in a thread of the file. + + See :meth:`_complete_upload`. + """ + if self.buffer is None: + return + self._committing = True + await self._pool.run_shielded(self._complete_upload) + + def _complete_upload(self) -> None: + """Upload or complete the written data; blocks, and runs in a thread. + + Creates an empty object if nothing was written, uploads the buffered + data with PutObject if no part was uploaded, and otherwise completes + the multipart upload once its parts have ended, aborting it if a + part or the completion fails. A failed part is found as soon as it + fails, without waiting for the parts before it, and the parts that + have not started are cancelled; the running ones are waited for + before the abort. Invalidates the cache of the path afterwards. + + Raises: + FileExistsError: If an object was created at the path after the + file was opened in exclusive-create mode. + """ + buffer = cast(BytesIO, self.buffer) + writer = cast(S3MultipartWriter, self._writer) + if self.tell() == 0: + self.fs._sync_fs.touch(self.path, **self._get_request_kwargs("put_object")) + elif not self._parts: + buffer.seek(0) + self.fs.core.put_object( + S3Path(self.bucket, self.key), + buffer.read(), + **self._get_request_kwargs("put_object"), + ) + else: + try: + # The parts were submitted before this call, so they have + # 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() + parts = [part.result() for part in self._parts] + writer.complete(parts, **self._get_request_kwargs("complete_multipart_upload")) + except BaseException: + self._discard() + raise + self.fs.invalidate_cache(self.path) + + def _discard(self) -> None: + """Abort the multipart upload, if any; blocks, and runs in a thread. + + The requests that have not started are cancelled, and the running + ones are waited for first, since a part that is still uploading when + the upload is aborted may be stored after the abort. A failure is + logged instead of raised, so that it does not mask the error that the + caller is handling, also when no task waits for the abort any more. + """ + try: + if self._creation is not None and not self._creation.cancel(): + wait([self._creation]) + if self._writer is not None and self._writer.upload is not None: + wait([part for part in self._parts if not part.cancel()]) + self._writer.abort(**self._get_request_kwargs("abort_multipart_upload")) + except Exception: + _logger.exception( + f"Failed to abort the multipart upload to s3://{self.bucket}/{self.key}." + ) + + 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. + + Args: + start: The offset of the first byte to read. + end: The offset to stop reading at (exclusive). + + Returns: + The bytes read, empty if the clamped range is empty. + """ + end = min(end, self.size) + if start >= end: + return b"" + ranges = S3File._get_ranges( + start, end, max_workers=self.max_workers, worker_block_size=self.blocksize + ) + 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) diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index ae943b9a7..ce9b57112 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -10,13 +10,18 @@ from __future__ import annotations import asyncio +import contextvars +import os +import sys import threading +import weakref from abc import ABCMeta, abstractmethod from collections.abc import Callable +from concurrent import futures from concurrent.futures import Future from concurrent.futures.thread import ThreadPoolExecutor from multiprocessing import cpu_count -from typing import Any, TypeVar +from typing import Any, ClassVar, TypeVar from pyathena.util import override @@ -127,18 +132,183 @@ def shutdown(self, wait: bool = True) -> None: self._executor.shutdown(wait=wait) +class S3AioThreadPool(S3Executor): + """A thread pool that runs the blocking calls of asyncio code. + + Unlike ``asyncio.to_thread()``, which runs a call in the default executor + of the event loop, this pool runs at most ``max_workers`` calls at once in + threads of its own, so that neither the size nor the load of the default + executor limits or blocks them. The threads are created as the calls need + them, and each call runs with a copy of the caller's context variables, + as with ``asyncio.to_thread()``. In a forked child process, whose copy + has no threads, the pool creates new ones. The pool can be shared by + several event loops and threads. + + Args: + max_workers: The maximum number of calls that run at once. + + Raises: + ValueError: If ``max_workers`` is not positive. + """ + + # The pools of this process, which a forked child resets. + _instances: ClassVar[weakref.WeakSet[S3AioThreadPool]] = weakref.WeakSet() + # The executors that a forked child copied from its parent. They are + # kept, since collecting one would put to a queue whose lock a thread of + # the parent may have held. + _forked_executors: ClassVar[list[ThreadPoolExecutor]] = [] + + def __init__(self, max_workers: int | None = None) -> None: + """Initialize the pool without starting any thread. + + Args: + max_workers: The maximum number of calls that run at once. If + None, the default of ``ThreadPoolExecutor``: ``min(32, CPUs + + 4)``. + + Raises: + ValueError: If ``max_workers`` is not positive. + """ + if max_workers is None: + # As ThreadPoolExecutor, which counts the CPUs that the process + # may use since Python 3.13. + cpus = getattr(os, "process_cpu_count", os.cpu_count)() + max_workers = min(32, (cpus or 1) + 4) + if max_workers <= 0: + raise ValueError("max_workers must be greater than 0") + self.max_workers = max_workers + self._lock = threading.Lock() + self._executor: ThreadPoolExecutor | None = None + self._shutdown = False + S3AioThreadPool._instances.add(self) + + @override + def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]: + """Schedule a blocking call in a thread of the pool. + + As with ``ThreadPoolExecutor``, the returned future can be cancelled + only before the call starts. The calls start in the order they were + submitted, so that a call that waits for the calls submitted before it + does not wait for a thread that it holds. + + Args: + fn: The blocking callable. + *args: Positional arguments passed to the callable. + **kwargs: Keyword arguments passed to the callable. + + Returns: + The future of the call. + + Raises: + RuntimeError: If the pool has been shut down. + """ + context = contextvars.copy_context() + return self._get_executor().submit(context.run, fn, *args, **kwargs) + + async def run(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> T: + """Run a blocking call in a thread of the pool and wait for it. + + Cancelling the waiting task cancels a call that has not started; a + running call keeps running, as with ``asyncio.to_thread()``. + + Args: + fn: The blocking callable. + *args: Positional arguments passed to the callable. + **kwargs: Keyword arguments passed to the callable. + + Returns: + The return value of the call. + + Raises: + RuntimeError: If the pool has been shut down. + """ + return await asyncio.wrap_future(self.submit(fn, *args, **kwargs)) + + async def run_shielded(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> T: + """Run a blocking call that a cancellation of the waiting task does not stop. + + The call runs to its end, even if it has not started when the waiting + task is cancelled; its outcome is then dropped. + + Args: + fn: The blocking callable. + *args: Positional arguments passed to the callable. + **kwargs: Keyword arguments passed to the callable. + + Returns: + The return value of the call. + + Raises: + RuntimeError: If the pool has been shut down. + """ + future = asyncio.wrap_future(self.submit(fn, *args, **kwargs)) + # Retrieved, so that asyncio does not log the failure of a call whose + # waiting task was cancelled as never retrieved. + future.add_done_callback(lambda f: f.cancelled() or f.exception()) + return await asyncio.shield(future) + + @override + def shutdown(self, wait: bool = True) -> None: + """Stop accepting calls and release the threads once they are idle. + + The calls that were submitted before still run. + + Args: + wait: Whether to wait for the submitted calls to finish. + """ + with self._lock: + self._shutdown = True + executor = self._executor + if executor is not None: + executor.shutdown(wait=wait) + + @classmethod + def _reset_after_fork(cls) -> None: + """Drop the threads and locks that a forked child copied from its parent. + + The threads of the parent do not exist in the child, and a lock that + one of them held stays held, so each pool creates new ones. + """ + for pool in list(cls._instances): + if pool._executor is not None: + cls._forked_executors.append(pool._executor) + pool._lock = threading.Lock() + pool._executor = None + + def _get_executor(self) -> ThreadPoolExecutor: + """Return the thread pool of this process, creating it if needed. + + Returns: + The thread pool. + + Raises: + RuntimeError: If the pool has been shut down. + """ + if (executor := self._executor) is not None: + # Its submit() raises after the shutdown. + return executor + with self._lock: + if self._shutdown: + raise RuntimeError("cannot schedule new futures after shutdown") + if self._executor is None: + self._executor = ThreadPoolExecutor(max_workers=self.max_workers) + return self._executor + + +if hasattr(os, "register_at_fork"): + os.register_at_fork(after_in_child=S3AioThreadPool._reset_after_fork) + + class S3AioExecutor(S3Executor): """Executor that schedules work on an asyncio event loop. - Uses ``asyncio.run_coroutine_threadsafe(asyncio.to_thread(fn), loop)`` to - dispatch blocking functions onto the event loop's thread pool, returning + Uses ``asyncio.run_coroutine_threadsafe()`` to dispatch blocking functions + through the event loop to an ``S3AioThreadPool`` of its own, returning ``concurrent.futures.Future`` objects that are compatible with ``as_completed()``, ``wait()`` and ``Future.cancel()``. As with ``ThreadPoolExecutor``, a future cannot be cancelled once its function has - started. At most ``max_workers`` of the submitted functions run at once. - - This avoids thread-in-thread nesting when ``S3File`` is used from within - ``asyncio.to_thread()`` calls (the pattern used by ``AioS3FileSystem``). + started. At most ``max_workers`` of the submitted functions run at once, + whatever the size of the default executor of the event loop. Args: loop: A running asyncio event loop. @@ -164,25 +334,11 @@ def __init__( Raises: ValueError: If ``max_workers`` is not positive. """ - if max_workers <= 0: - # As ThreadPoolExecutor does; a semaphore of 0 would never run anything. - raise ValueError("max_workers must be greater than 0") self._loop = loop - self._semaphore = asyncio.Semaphore(max_workers) - - async def _run(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> T: - """Run the function in a thread once fewer than ``max_workers`` run. - - Args: - fn: The blocking function to run. - *args: Positional arguments passed to the function. - **kwargs: Keyword arguments passed to the function. - - Returns: - The return value of the function. - """ - async with self._semaphore: - return await asyncio.to_thread(fn, *args, **kwargs) + self._pool = S3AioThreadPool(max_workers) + # Orders the dispatches with the shutdown. + self._lock = threading.Lock() + self._shutdown = False @override def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]: @@ -225,7 +381,10 @@ def settle(task: Future[None]) -> None: elif future.set_running_or_notify_cancel(): future.set_exception(task.exception()) - task = asyncio.run_coroutine_threadsafe(self._run(run), self._loop) + 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 raise RuntimeError( @@ -233,7 +392,42 @@ def settle(task: Future[None]) -> None: "Use S3ThreadPoolExecutor for synchronous usage." ) + async def _shutdown_pool(self) -> None: + """Shut down the pool once the functions submitted before have reached it. + + Run as a task of the event loop, which starts its tasks in the order + they were scheduled, so the tasks that dispatch the functions + submitted before have handed them to the pool. + """ + self._pool.shutdown(wait=False) + @override def shutdown(self, wait: bool = True) -> None: - # No resources to release — work is dispatched to the event loop. - pass + loop = self._loop + with self._lock: + self._shutdown = True + closing = None + if loop is not None and loop.is_running(): + # A function reaches the pool only when the event loop runs + # its task, so the pool is shut down after the dispatches. + closing = asyncio.run_coroutine_threadsafe(self._shutdown_pool(), loop) + if loop is None or closing is None: + self._pool.shutdown(wait=wait) + return + try: + on_loop = asyncio.get_running_loop() is loop + except RuntimeError: + on_loop = False + # On the thread of the event loop, such as when a file is garbage + # collected there, waiting would keep the loop from dispatching. + if wait and not on_loop: + # The task is cancelled if the event loop shuts down before it + # runs, and never runs if the loop stops instead; the functions + # that it has not dispatched then do not reach the pool either. + while not futures.wait([closing], timeout=0.1).done: + # The thread of a loop that runs in a daemon thread, such as + # fsspec's, stops when the interpreter exits, while the loop + # still reports running. + if not loop.is_running() or sys.is_finalizing(): + break + self._pool.shutdown(wait=True) diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index 0440d35e0..80de1537d 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -8,7 +8,7 @@ import urllib.parse import urllib.request import uuid -from concurrent.futures import FIRST_COMPLETED, wait +from concurrent.futures import FIRST_COMPLETED, Future, ThreadPoolExecutor, wait from datetime import UTC, datetime from itertools import chain from pathlib import Path @@ -22,9 +22,9 @@ from fsspec import Callback from pyathena.filesystem.s3 import S3File, S3FileSystem -from pyathena.filesystem.s3_async import AioS3File, AioS3FileSystem +from pyathena.filesystem.s3_async import AioS3File, AioS3FileSystem, AioS3StreamedFile from pyathena.filesystem.s3_core import S3Core -from pyathena.filesystem.s3_executor import S3Executor +from pyathena.filesystem.s3_executor import S3AioThreadPool, S3Executor from pyathena.filesystem.s3_object import ( S3MultipartUpload, S3MultipartUploadPart, @@ -37,8 +37,10 @@ from tests import ENV from tests.pyathena.conftest import connect from tests.pyathena.util import ( + EVENT_TIMEOUT, MULTIPART_COPY_BLOCK_SIZE, MULTIPART_COPY_KWARGS, + overlapping_calls, stub_multipart_copy, ) @@ -257,6 +259,120 @@ def upload_part_copy(**kw): assert events[2:] == ["end 1", "abort"] sync_fs.core.complete_multipart_upload.assert_not_called() + @pytest.mark.asyncio + async def test_copy_object_with_multipart_upload_max_workers(self): + # GH-1159: max_workers of the copy bounds its parallel part copies, + # whatever the size of the default executor of the event loop, in + # which they used to run. + asyncio.get_running_loop().set_default_executor(ThreadPoolExecutor(max_workers=1)) + fs = AioS3FileSystem(connection=mock.MagicMock(), 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"} + ) + ) + call, state = overlapping_calls(3) + + def upload_part_copy(**kw): + call() + 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 4 parts of the default block size. + size = 4 * S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE + sync_fs._call = sync_fs._core.call = mock.MagicMock(return_value={"ContentLength": size}) + + await fs._copy_object_with_multipart_upload( + S3Path("bucket", "src"), + S3Path("bucket", "dst"), + max_workers=3, + MetadataDirective="REPLACE", + TaggingDirective="REPLACE", + AnnotationDirective="EXCLUDE", + ) + + assert sync_fs.core.upload_part_copy.call_count == 4 + assert state["peak"] == 3 + sync_fs.core.complete_multipart_upload.assert_called_once() + + @pytest.mark.asyncio + async def test_copy_object_with_multipart_upload_shares_threads(self): + # GH-1159: concurrent copies share the threads of the filesystem, so + # that a bulk copy of large objects does not start max_workers threads + # per object. + fs = AioS3FileSystem(connection=mock.MagicMock(), max_workers=2, 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"} + ) + ) + # Waits briefly for a third overlapping call, which only copies + # with threads of their own run. + call, state = overlapping_calls(3, timeout=0.5) + + def upload_part_copy(**kw): + call() + 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() + size = 2 * S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE + sync_fs._call = sync_fs._core.call = mock.MagicMock(return_value={"ContentLength": size}) + + await asyncio.gather( + *[ + fs._copy_object_with_multipart_upload( + S3Path("bucket", "src"), + S3Path("bucket", f"dst{i}"), + MetadataDirective="REPLACE", + TaggingDirective="REPLACE", + AnnotationDirective="EXCLUDE", + ) + for i in range(2) + ] + ) + + assert sync_fs.core.upload_part_copy.call_count == 4 + assert state["peak"] == 2 + + @pytest.mark.asyncio + async def test_blocking_calls_bypass_default_executor(self): + # GH-1159: the blocking calls of the coroutines run in the threads of + # the filesystem, so that a busy default executor of the event loop + # does not hold them. They used to run in it. + loop = asyncio.get_running_loop() + loop.set_default_executor(ThreadPoolExecutor(max_workers=1)) + fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + fs._sync_fs.info = mock.MagicMock(return_value={"name": "bucket/key"}) + release = threading.Event() + # Held until it is released, after the lookup has had to finish. + blocker = loop.run_in_executor(None, release.wait, 60) + try: + info = await asyncio.wait_for(fs._info("s3://bucket/key"), 5) + finally: + release.set() + await blocker + + assert info == {"name": "bucket/key"} + + @pytest.mark.asyncio + async def test_max_workers_none(self): + # None takes the default of ThreadPoolExecutor for the threads of the + # filesystem and as the default max_workers of its files and copies. + # The files and copies used to fail with it. + fs = AioS3FileSystem( + connection=mock.MagicMock(), max_workers=None, skip_instance_cache=True + ) + fs._sync_fs.info = mock.MagicMock(return_value={"name": "bucket/key"}) + + assert await fs._info("s3://bucket/key") == {"name": "bucket/key"} + # The files and copies take the same default. + cpus = getattr(os, "process_cpu_count", os.cpu_count)() + assert fs._sync_fs.max_workers == fs._thread_pool.max_workers == min(32, (cpus or 1) + 4) + @pytest.mark.parametrize("cancellations", [1, 2]) @pytest.mark.asyncio async def test_copy_object_with_multipart_upload_cancelled(self, cancellations): @@ -1525,6 +1641,43 @@ def test_append(self, fs, base, exp): assert len(actual) == len(data + extra) assert actual == data + extra + @pytest.mark.parametrize( + "size", + [ + 2**10, # < block size: PutObject + 10 * 2**20, # > block size (5 MiB): two parts of a multipart upload + ], + ) + @pytest.mark.asyncio + async def test_open_async_write_read(self, fs, size): + data = os.urandom(size) + path = ( + f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" + f"filesystem/test_async_open_async/{uuid.uuid4()}" + ) + async with await fs.open_async(path, "wb", block_size=5 * 2**20) as f: + assert isinstance(f, AioS3StreamedFile) + await f.write(data) + async with await fs.open_async(path, "rb", block_size=5 * 2**20) as f: + assert await f.read() == data + f.seek(size // 2) + assert await f.read(10) == data[size // 2 : size // 2 + 10] + + @pytest.mark.asyncio + async def test_open_async_append_exclusive(self, fs): + path = ( + f"s3://{ENV.s3_staging_bucket}/{ENV.s3_staging_key}{ENV.schema}/" + f"filesystem/test_async_open_async_append/{uuid.uuid4()}" + ) + async with await fs.open_async(path, "xb") as f: + await f.write(b"hello") + with pytest.raises(FileExistsError): + await fs.open_async(path, "xb") + async with await fs.open_async(path, "ab") as f: + await f.write(b" world") + async with await fs.open_async(path, "rb") as f: + assert await f.read() == b"hello world" + @pytest.mark.asyncio async def test_ls_buckets(self, fs): fs.invalidate_cache() @@ -2539,3 +2692,486 @@ def test_get_ranges(self, start, end, max_workers, worker_block_size, ranges): ) == ranges ) + + +class TestAioS3StreamedFile: + @pytest.fixture + def fs(self): + fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + core = fs._sync_fs.core + core.create_multipart_upload = mock.MagicMock( + return_value=S3MultipartUpload( + {"Bucket": "bucket", "Key": "key", "UploadId": "uploadid"} + ) + ) + core.upload_part = mock.MagicMock( + side_effect=lambda **kw: S3MultipartUploadPart(kw["part_number"], {"ETag": '"e"'}) + ) + core.complete_multipart_upload = mock.MagicMock() + core.abort_multipart_upload = mock.MagicMock() + core.put_object = mock.MagicMock() + # The mocked client has no operation model to select the parameters. + core.operation_params = mock.MagicMock(side_effect=lambda method, params: dict(params)) + return fs + + @pytest.mark.asyncio + async def test_write_parts_in_flight(self, fs): + # GH-1159: a write keeps at most max_workers parts in flight, and waits + # for one to finish without blocking the event loop, on which the + # parts are released here. + block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE + lock = threading.Lock() + gates: dict[int, threading.Event] = {} + active: list[int] = [] + uploads: dict[int, Future[S3MultipartUploadPart]] = {} + in_flight: list[int] = [] + submit_to = S3AioThreadPool.submit_to + + def record_upload(pool, future, fn, *args, **kwargs): + # The parts submitted and not finished, including those that + # wait for a thread with their data. + if "body" in kwargs: + uploads[kwargs["part_number"]] = future + in_flight.append(sum(not upload.done() for upload in uploads.values())) + submit_to(pool, future, fn, *args, **kwargs) + + def upload_part(**kw): + part_number = kw["part_number"] + with lock: + gate = gates.setdefault(part_number, threading.Event()) + active.append(part_number) + if not gate.wait(5): + raise TimeoutError(part_number) + with lock: + active.remove(part_number) + return S3MultipartUploadPart(part_number, {"ETag": '"e"'}) + + fs.core.upload_part = mock.MagicMock(side_effect=upload_part) + written = asyncio.Event() + + async def release(): + while True: + with lock: + if written.is_set() or len(active) >= 2: + # Only the newest part finishes until the data is + # written, so a write that waits for the oldest one + # would time out. + for part_number in active if written.is_set() else active[-1:]: + gates.setdefault(part_number, threading.Event()).set() + if written.is_set() and not active: + return + await asyncio.sleep(0.01) + + releaser = asyncio.ensure_future(release()) + f = await fs.open_async("s3://bucket/key", "wb", block_size=block_size, max_workers=2) + with mock.patch.object( + S3AioThreadPool, "submit_to", autospec=True, side_effect=record_upload + ): + await f.write(b"a" * block_size * 4) + written.set() + await f.close() + await releaser + + assert in_flight == [1, 2, 2, 2] + (_, parts), _ = fs.core.complete_multipart_upload.call_args + assert [part.part_number for part in parts] == [1, 2, 3, 4] + fs.core.abort_multipart_upload.assert_not_called() + + @pytest.mark.asyncio + async def test_write_failed_part_while_waiting(self, fs): + # A part found failed when a write waits for a part stops the write + # with its error and aborts the upload. + block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE + fs.core.upload_part = mock.MagicMock(side_effect=PermissionError("Access Denied")) + f = await fs.open_async("s3://bucket/key", "wb", block_size=block_size, max_workers=2) + + with pytest.raises(PermissionError, match="Access Denied"): + await f.write(b"a" * block_size * 3) + + assert f.closed + assert sorted(c.kwargs["part_number"] for c in fs.core.upload_part.call_args_list) == [ + 1, + 2, + ] + fs.core.abort_multipart_upload.assert_called_once() + fs.core.complete_multipart_upload.assert_not_called() + + @pytest.mark.parametrize("blocks", [0, 2]) + @pytest.mark.asyncio + async def test_exit_with_exception(self, fs, blocks): + # Leaving the block with an exception does not commit the written + # data, and aborts the multipart upload, if any. + block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE + f = await fs.open_async("s3://bucket/key", "wb", block_size=block_size) + + async def write_and_fail(): + async with f: + await f.write(b"a" * (block_size * blocks + 1)) + raise RuntimeError("failed") + + with pytest.raises(RuntimeError, match="failed"): + await write_and_fail() + + assert f.closed + fs.core.put_object.assert_not_called() + fs.core.complete_multipart_upload.assert_not_called() + assert fs.core.abort_multipart_upload.call_count == (1 if blocks else 0) + + @pytest.mark.asyncio + async def test_exit_with_exception_after_close(self, fs): + # An exception raised after the file was closed propagates unchanged. + f = await fs.open_async("s3://bucket/key", "wb") + + async def close_and_fail(): + async with f: + await f.write(b"abc") + await f.close() + raise KeyError("failed") + + with pytest.raises(KeyError, match="failed"): + await close_and_fail() + + fs.core.put_object.assert_called_once() + fs.core.abort_multipart_upload.assert_not_called() + + @pytest.mark.asyncio + async def test_sync_with(self, fs): + # A synchronous with block cannot await close(), so it is rejected + # instead of leaving the data uncommitted. + f = await fs.open_async("s3://bucket/key", "wb") + with pytest.raises(TypeError, match="async with"), f: + pass + await f.close() + + @pytest.mark.asyncio + async def test_garbage_collected_unclosed(self, fs): + # An unclosed file that is garbage collected warns, and aborts its + # multipart upload instead of leaving it incomplete. + block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE + aborted = threading.Event() + fs.core.abort_multipart_upload = mock.MagicMock(side_effect=lambda *a, **k: aborted.set()) + + async def write_unclosed(): + # Created directly: the futures of open_async() reference the + # file until the garbage collector breaks their cycle. + f = AioS3StreamedFile(fs, "bucket/key", "wb", block_size=block_size) + await f.write(b"a" * block_size * 2) + + # The file is released when the coroutine returns. + with pytest.warns(ResourceWarning, match="was not closed"): + await write_unclosed() + assert await asyncio.to_thread(aborted.wait, 5) + + fs.core.complete_multipart_upload.assert_not_called() + + @pytest.mark.asyncio + async def test_append_failed_part_copy(self, fs): + # A part copy of an append found failed before a part is submitted + # stops the write with its error, as with S3File (GH-1178). + block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE + info = S3Object( + init={"Key": "key"}, + type=S3ObjectType.S3_OBJECT_TYPE_FILE, + bucket="bucket", + key="key", + ) + info.size = block_size + fs._sync_fs.info = mock.MagicMock(return_value=info) + release = threading.Event() + + def upload_part_copy(**kwargs): + # Held until the first write returns, so that it does not see + # the failure. + release.wait(EVENT_TIMEOUT) + raise PermissionError("Access Denied") + + fs.core.upload_part_copy = mock.MagicMock(side_effect=upload_part_copy) + f = await fs.open_async("s3://bucket/key", "ab", block_size=block_size) + try: + await f.write(b"a" * block_size) + finally: + release.set() + # The copy has failed before the next write submits a part. + await asyncio.to_thread(wait, f._parts[:1], EVENT_TIMEOUT) + + with pytest.raises(PermissionError, match="Access Denied"): + await f.write(b"a" * block_size) + + assert f.closed + assert fs.core.upload_part.call_count == 1 + fs.core.abort_multipart_upload.assert_called_once() + fs.core.complete_multipart_upload.assert_not_called() + + @pytest.mark.asyncio + async def test_commit_failed_part_first(self, fs): + # The commit raises a failed part without waiting for the parts + # before it, and cancels those that have not started, as S3File does. + f = await fs.open_async("s3://bucket/key", "wb") + await f.write(b"a") + f._writer.upload = S3MultipartUpload( + {"Bucket": "bucket", "Key": "key", "UploadId": "uploadid"} + ) + failed: Future[S3MultipartUploadPart] = Future() + failed.set_exception(PermissionError("Access Denied")) + pending: list[Future[S3MultipartUploadPart]] = [Future(), Future()] + f._parts = [pending[0], failed, pending[1]] + + try: + with pytest.raises(PermissionError, match="Access Denied"): + await asyncio.wait_for(asyncio.to_thread(f._complete_upload), 5) + finally: + # Releases a commit that waits for the parts in order. + for part in pending: + if not part.done(): + part.set_result(S3MultipartUploadPart(1, {"ETag": '"e"'})) + + assert all(part.cancelled() for part in pending) + fs.core.abort_multipart_upload.assert_called_once() + fs.core.complete_multipart_upload.assert_not_called() + # Closed without committing again. + f.closed = True + f._pool.shutdown() + + @pytest.mark.parametrize("cancellations", [1, 2]) + @pytest.mark.asyncio + async def test_write_cancelled(self, fs, cancellations): + # A cancelled write waits for the running parts in the threads of the + # file and aborts the upload, also after a repeated cancellation. + block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE + events = [] + lock = threading.Lock() + started = threading.Semaphore(0) + release = threading.Event() + aborted = threading.Event() + + def upload_part(**kw): + part_number = kw["part_number"] + with lock: + events.append(f"start {part_number}") + started.release() + # The finally block of the test always releases it. + release.wait() + with lock: + events.append(f"end {part_number}") + return S3MultipartUploadPart(part_number, {"ETag": '"e"'}) + + def abort_multipart_upload(*args, **kwargs): + events.append("abort") + aborted.set() + + fs.core.upload_part = mock.MagicMock(side_effect=upload_part) + fs.core.abort_multipart_upload = mock.MagicMock(side_effect=abort_multipart_upload) + f = await fs.open_async("s3://bucket/key", "wb", block_size=block_size, max_workers=2) + task = asyncio.ensure_future(f.write(b"a" * block_size * 3)) + try: + for _ in range(2): + assert await asyncio.to_thread(started.acquire, timeout=5) + for _ in range(cancellations): + task.cancel() + # Lets the write enter its cleanup. + await asyncio.sleep(0) + finally: + release.set() + with pytest.raises(asyncio.CancelledError): + await task + assert await asyncio.to_thread(aborted.wait, 5) + + assert f.closed + assert sorted(events[:2]) == ["start 1", "start 2"] + assert sorted(events[2:4]) == ["end 1", "end 2"] + assert events[4:] == ["abort"] + fs.core.complete_multipart_upload.assert_not_called() + + @pytest.mark.asyncio + async def test_close_cancelled(self, fs): + # A commit that has started completes even if the task that closes + # the file is cancelled, and the upload is not aborted. + block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE + started = threading.Event() + release = threading.Event() + completed = threading.Event() + + def complete_multipart_upload(*args, **kwargs): + started.set() + # The finally block of the test always releases it. + release.wait() + completed.set() + + fs.core.complete_multipart_upload = mock.MagicMock(side_effect=complete_multipart_upload) + f = await fs.open_async("s3://bucket/key", "wb", block_size=block_size) + await f.write(b"a" * block_size * 2) + task = asyncio.ensure_future(f.close()) + try: + assert await asyncio.to_thread(started.wait, 5) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + finally: + release.set() + assert await asyncio.to_thread(completed.wait, 5) + + assert f.closed + fs.core.complete_multipart_upload.assert_called_once() + fs.core.abort_multipart_upload.assert_not_called() + + @pytest.mark.asyncio + async def test_flush_cancelled_then_exit_with_exception(self, fs): + # Leaving the block with an exception after a cancelled final flush + # does not abort the upload that the running commit completes. + block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE + started = threading.Event() + release = threading.Event() + completed = threading.Event() + + def complete_multipart_upload(*args, **kwargs): + started.set() + # The finally block of the test always releases it. + release.wait() + completed.set() + + fs.core.complete_multipart_upload = mock.MagicMock(side_effect=complete_multipart_upload) + f = await fs.open_async("s3://bucket/key", "wb", block_size=block_size) + + async def flush_and_exit(): + async with f: + await f.write(b"a" * block_size * 2) + flush = asyncio.ensure_future(f.flush(force=True)) + assert await asyncio.to_thread(started.wait, 5) + flush.cancel() + await flush + + try: + with pytest.raises(asyncio.CancelledError): + await flush_and_exit() + finally: + release.set() + assert await asyncio.to_thread(completed.wait, 5) + # Lets an abort that was submitted run. + await asyncio.sleep(0.1) + + assert f.closed + fs.core.complete_multipart_upload.assert_called_once() + fs.core.abort_multipart_upload.assert_not_called() + + @pytest.mark.parametrize("data", [b"", b"abc"]) + @pytest.mark.asyncio + async def test_close_small(self, fs, data): + # Data smaller than a block is uploaded with PutObject, and no data + # creates an empty object. + fs._sync_fs.touch = mock.MagicMock() + async with await fs.open_async("s3://bucket/key", "wb", ContentType="text/plain") as f: + await f.write(data) + + if data: + fs.core.put_object.assert_called_once_with( + S3Path("bucket", "key"), data, ContentType="text/plain" + ) + fs._sync_fs.touch.assert_not_called() + else: + fs._sync_fs.touch.assert_called_once_with("bucket/key", ContentType="text/plain") + fs.core.put_object.assert_not_called() + fs.core.create_multipart_upload.assert_not_called() + + @pytest.mark.asyncio + async def test_read_parallel(self, fs): + # A read runs at most max_workers range requests at once, conditional + # on the ETag of the object. + block_size = 2**10 + size = block_size * 4 + call, state = overlapping_calls(2) + + def get_object(path, range_, **kwargs): + call() + 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 = size + fs._sync_fs.info = mock.MagicMock(return_value=info) + + async with await fs.open_async( + "s3://bucket/key", "rb", block_size=block_size, max_workers=2 + ) as f: + assert await f.read() == b"a" * size + assert await f.read() == b"" + + assert fs.core.get_object.call_count == 4 + assert state["peak"] == 2 + for c in fs.core.get_object.call_args_list: + assert c.kwargs == {"IfMatch": '"etag"'} + + @pytest.mark.parametrize( + ("mode", "kwargs", "error"), + [ + ("r", {}, ValueError), + ("wb", {"compression": "gzip"}, ValueError), + ("r+b", {}, NotImplementedError), + ("wb", {"max_workers": 0}, ValueError), + ("wb", {"block_size": 2**10}, ValueError), + ("wb", {"version_id": "v"}, ValueError), + ("wb", {"autocommit": False}, NotImplementedError), + ], + ) + @pytest.mark.asyncio + async def test_open_async_invalid(self, fs, mode, kwargs, error): + with pytest.raises(error): + await fs.open_async("s3://bucket/key", mode, **kwargs) + + @pytest.mark.asyncio + async def test_open_async_in_transaction(self, fs): + # The file always commits when it is closed, so it cannot write in a + # transaction, which defers the commit. + with fs.transaction, pytest.raises(NotImplementedError, match="transaction"): + await fs.open_async("s3://bucket/key", "wb") + + @pytest.mark.parametrize("exists", [False, True]) + @pytest.mark.asyncio + async def test_exclusive_create(self, fs, exists): + fs._sync_fs.exists = mock.MagicMock(return_value=exists) + if exists: + with pytest.raises(FileExistsError): + await fs.open_async("s3://bucket/key", "xb") + return + async with await fs.open_async("s3://bucket/key", "xb") as f: + await f.write(b"abc") + + fs.core.put_object.assert_called_once_with(S3Path("bucket", "key"), b"abc", IfNoneMatch="*") + + @pytest.mark.parametrize("size", [3, S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE]) + @pytest.mark.asyncio + async def test_append(self, fs, size): + # An existing object smaller than a part is rewritten from the buffer, + # and a larger one is copied as the first part. + info = S3Object( + init={"Key": "key", "ContentType": "text/plain"}, + type=S3ObjectType.S3_OBJECT_TYPE_FILE, + bucket="bucket", + key="key", + ) + info.size = size + fs._sync_fs.info = mock.MagicMock(return_value=info) + fs._sync_fs.cat_file = mock.MagicMock(return_value=b"a" * size) + fs.core.upload_part_copy = mock.MagicMock( + side_effect=lambda **kw: S3MultipartUploadPart(kw["part_number"], {"ETag": '"c"'}) + ) + + async with await fs.open_async("s3://bucket/key", "ab") as f: + await f.write(b"b") + + if size < S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE: + fs.core.put_object.assert_called_once() + assert fs.core.put_object.call_args.args == (S3Path("bucket", "key"), b"aaab") + # The rewritten object keeps the metadata of the existing one. + assert fs.core.put_object.call_args.kwargs["ContentType"] == "text/plain" + fs.core.upload_part_copy.assert_not_called() + else: + fs.core.put_object.assert_not_called() + fs.core.upload_part_copy.assert_called_once() + assert fs.core.upload_part_copy.call_args.kwargs["range_"] is None + (_, parts), _ = fs.core.complete_multipart_upload.call_args + assert [part.part_number for part in parts] == [1, 2] diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index cb79a2a9e..b4aa849f9 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -6,13 +6,119 @@ # SPDX-License-Identifier: MIT import asyncio +import contextvars +import os import threading import time from concurrent.futures import Future, ThreadPoolExecutor, wait +from unittest import mock import pytest -from pyathena.filesystem.s3_executor import S3AioExecutor, S3ThreadPoolExecutor +from pyathena.filesystem.s3_executor import S3AioExecutor, S3AioThreadPool, S3ThreadPoolExecutor +from tests.pyathena.util import EVENT_TIMEOUT, overlapping_calls + + +class TestS3AioThreadPool: + def test_init(self): + # None takes the default of ThreadPoolExecutor. + cpus = getattr(os, "process_cpu_count", os.cpu_count)() + assert S3AioThreadPool(None).max_workers == min(32, (cpus or 1) + 4) + with pytest.raises(ValueError, match="max_workers must be greater than 0"): + S3AioThreadPool(0) + + def test_run_beyond_default_executor(self): + # GH-1159: the calls run in threads of the pool, so that max_workers, + # not the default executor of the event loop, bounds them. + call, state = overlapping_calls(4) + pool = S3AioThreadPool(4) + + async def main(): + asyncio.get_running_loop().set_default_executor(ThreadPoolExecutor(max_workers=1)) + await asyncio.gather(*[pool.run(call) for _ in range(6)]) + + asyncio.run(main()) + pool.shutdown() + + assert state["peak"] == 4 + + def test_run_context(self): + # The call sees the context variables of the caller, as with + # asyncio.to_thread(). + variable: contextvars.ContextVar[str] = contextvars.ContextVar("variable") + pool = S3AioThreadPool(1) + + async def main(): + variable.set("caller") + return await pool.run(variable.get) + + assert asyncio.run(main()) == "caller" + pool.shutdown() + + @pytest.mark.parametrize("shielded", [False, True]) + def test_run_cancelled_before_start(self, shielded): + # Cancelling the waiting task drops a call that has not started, + # unless it runs shielded. + events = [] + started = threading.Event() + release = threading.Event() + pool = S3AioThreadPool(1) + + def work(): + started.set() + release.wait(EVENT_TIMEOUT) + + async def main(): + running = pool.submit(work) + run = pool.run_shielded if shielded else pool.run + pending = asyncio.ensure_future(run(events.append, "pending finished")) + await asyncio.to_thread(started.wait, EVENT_TIMEOUT) + # Lets the task submit the call, which waits for the thread. + while not pending.done() and not pool._get_executor()._work_queue.qsize(): + await asyncio.sleep(0.01) + pending.cancel() + with pytest.raises(asyncio.CancelledError): + await pending + release.set() + await asyncio.wrap_future(running) + + asyncio.run(main()) + pool.shutdown() + + assert events == (["pending finished"] if shielded else []) + + def test_shutdown(self): + # The calls submitted before the shutdown still run. + events = [] + release = threading.Event() + pool = S3AioThreadPool(1) + pool.submit(release.wait, EVENT_TIMEOUT) + pending = pool.submit(events.append, "pending finished") + pool.shutdown(wait=False) + with pytest.raises(RuntimeError, match="after shutdown"): + pool.submit(sum, [1, 2]) + release.set() + pending.result(timeout=EVENT_TIMEOUT) + + assert events == ["pending finished"] + + def test_fork(self): + # A forked child, whose copy of the pool has no threads, creates new + # ones, also if a thread of the parent held the lock of the pool. + pool = S3AioThreadPool(1) + parent = pool._get_executor() + pool._lock.acquire() + # As os.register_at_fork() runs it in a forked child. + S3AioThreadPool._reset_after_fork() + # The copied executor is kept, so that its collection does not touch + # the locks that the parent's threads may have held. + assert parent in S3AioThreadPool._forked_executors + child = pool._get_executor() + assert pool.submit(sum, [1, 2]).result(timeout=EVENT_TIMEOUT) == 3 + pool.shutdown() + assert child is not parent + S3AioThreadPool._forked_executors.remove(parent) + parent.shutdown() class TestS3AioExecutor: @@ -36,6 +142,154 @@ async def main(): with pytest.raises(ValueError, match="invalid literal"): failed.result() + def test_submit_beyond_default_executor(self): + # GH-1159: max_workers functions run at once, whatever the size of the + # default executor of the event loop. They used to run in it. + call, state = overlapping_calls(3) + + async def main(): + loop = asyncio.get_running_loop() + loop.set_default_executor(ThreadPoolExecutor(max_workers=1)) + executor = S3AioExecutor(loop=loop, max_workers=3) + futures = [executor.submit(call) for _ in range(3)] + await asyncio.gather(*[asyncio.wrap_future(f) for f in futures]) + executor.shutdown() + + asyncio.run(main()) + + assert state["peak"] == 3 + + def test_submit_from_default_executor(self): + # GH-1159: a function submitted from a thread of the default executor + # of the event loop does not wait for that thread. With a default + # executor of one thread, this used to deadlock a multipart write + # through the loop's asyncio.to_thread(). + async def main(): + loop = asyncio.get_running_loop() + loop.set_default_executor(ThreadPoolExecutor(max_workers=1)) + executor = S3AioExecutor(loop=loop) + # Waits in the only thread of the default executor, as a + # multipart write waits for the creation of its upload. + result = await asyncio.to_thread( + lambda: executor.submit(sum, [1, 2]).result(EVENT_TIMEOUT) + ) + executor.shutdown() + return result + + assert asyncio.run(main()) == 3 + + def test_shutdown_before_dispatch(self): + # GH-1159: a function submitted before the shutdown runs, although it + # reaches the pool of the executor only when the event loop runs its + # task. An AioS3File in a transaction shuts its executor down when it + # is closed, before its parts may have reached the pool. + async def main(): + executor = S3AioExecutor(loop=asyncio.get_running_loop()) + future = executor.submit(sum, [1, 2]) + executor.shutdown(wait=False) + with pytest.raises(RuntimeError, match="after shutdown"): + executor.submit(sum, [3, 4]) + return await asyncio.wrap_future(future) + + assert asyncio.run(main()) == 3 + + def test_shutdown_on_loop_thread(self): + # Shutting down with wait=True on the thread of the event loop, such + # as when a file is garbage collected there, does not wait for the + # functions that only the loop dispatches; they still run. + async def main(): + executor = S3AioExecutor(loop=asyncio.get_running_loop()) + future = executor.submit(sum, [1, 2]) + executor.shutdown(wait=True) + assert not future.done() + return await asyncio.wrap_future(future) + + assert asyncio.run(main()) == 3 + + def test_shutdown_loop_stopped(self): + # Shutting down with wait=True returns when the event loop stops + # before it runs the task that shuts the pool down. + loop = asyncio.new_event_loop() + busy = threading.Event() + release = threading.Event() + + def block(): + busy.set() + release.wait(EVENT_TIMEOUT) + # The loop exits after this callback, without running the task + # that the shutdown schedules meanwhile. + loop.stop() + + scheduled = threading.Event() + run_coroutine_threadsafe = asyncio.run_coroutine_threadsafe + + def schedule(coro, loop): + future = run_coroutine_threadsafe(coro, loop) + scheduled.set() + return future + + thread = threading.Thread(target=loop.run_forever) + thread.start() + try: + loop.call_soon_threadsafe(block) + assert busy.wait(EVENT_TIMEOUT) + executor = S3AioExecutor(loop=loop) + shut_down = threading.Event() + waiter = threading.Thread( + target=lambda: (executor.shutdown(wait=True), shut_down.set()) + ) + with mock.patch("asyncio.run_coroutine_threadsafe", side_effect=schedule): + waiter.start() + # The loop stops only after the task has been scheduled, so + # that the shutdown waits for a task that never runs. + assert scheduled.wait(EVENT_TIMEOUT) + assert not shut_down.is_set() + release.set() + assert shut_down.wait(EVENT_TIMEOUT) + waiter.join() + finally: + release.set() + thread.join() + # Runs the pending task, which would otherwise be dropped + # unawaited when the loop is closed. + loop.run_until_complete(asyncio.sleep(0)) + loop.close() + + def test_shutdown_interpreter_finalizing(self): + # Shutting down with wait=True while the interpreter exits returns, + # although the loop, whose daemon thread no longer runs, still + # reports running. + loop = asyncio.new_event_loop() + busy = threading.Event() + release = threading.Event() + + def block(): + busy.set() + # Held until the shutdown has returned, so that the loop cannot + # run the task that the shutdown waits for. + release.wait(6 * EVENT_TIMEOUT) + + thread = threading.Thread(target=loop.run_forever) + thread.start() + try: + loop.call_soon_threadsafe(block) + assert busy.wait(EVENT_TIMEOUT) + executor = S3AioExecutor(loop=loop) + shut_down = threading.Event() + waiter = threading.Thread( + target=lambda: (executor.shutdown(wait=True), shut_down.set()) + ) + with mock.patch("pyathena.filesystem.s3_executor.sys.is_finalizing", return_value=True): + waiter.start() + assert shut_down.wait(EVENT_TIMEOUT) + waiter.join() + finally: + release.set() + loop.call_soon_threadsafe(loop.stop) + thread.join() + loop.run_until_complete(asyncio.sleep(0)) + loop.close() + def test_submit_without_running_loop(self): with pytest.raises(RuntimeError, match="requires a running event loop"): S3AioExecutor().submit(sum, [1, 2]) @@ -66,10 +320,9 @@ def cancel(executor: S3AioExecutor) -> tuple[Future[None], Future[None]]: return running, pending async def main(): - loop = asyncio.get_running_loop() - loop.set_default_executor(ThreadPoolExecutor(max_workers=2)) - executor = S3AioExecutor(loop=loop) - # The function runs in one of the two threads, and cancel() in the other. + # The function runs in the only thread of the executor, so that + # the pending one waits for it, and cancel() in another thread. + executor = S3AioExecutor(loop=asyncio.get_running_loop(), max_workers=1) return await asyncio.to_thread(cancel, executor) running, pending = asyncio.run(main()) @@ -79,14 +332,7 @@ async def main(): assert not running.cancelled() assert pending.cancelled() - @pytest.mark.parametrize( - ("threads", "max_workers"), - [ - (1, 5), # The pending function waits for a thread. - (2, 1), # The pending function waits for a permit. - ], - ) - def test_loop_shutdown(self, threads, max_workers): + def test_loop_shutdown(self): # A function that has not started when the event loop shuts down is # never run, and its future is cancelled and settled, so that wait() # returns, instead of left pending. @@ -102,9 +348,8 @@ def work(): events.append("finished") async def main(): - loop = asyncio.get_running_loop() - loop.set_default_executor(ThreadPoolExecutor(max_workers=threads)) - executor = S3AioExecutor(loop=loop, max_workers=max_workers) + # The pending function waits for the only thread of the executor. + executor = S3AioExecutor(loop=asyncio.get_running_loop(), max_workers=1) running = executor.submit(work) pending = executor.submit(events.append, "pending finished") pending.add_done_callback(lambda _: settled.set()) @@ -123,11 +368,10 @@ async def main(): def test_executor_shut_down(self): # An error raised before the function starts is set on the future. async def main(): - loop = asyncio.get_running_loop() - default_executor = ThreadPoolExecutor(max_workers=1) - loop.set_default_executor(default_executor) - default_executor.shutdown() - future = S3AioExecutor(loop=loop).submit(sum, [1, 2]) + executor = S3AioExecutor(loop=asyncio.get_running_loop()) + # Fails the dispatch of the function to the pool. + executor._pool.shutdown() + future = executor.submit(sum, [1, 2]) with pytest.raises(RuntimeError, match="after shutdown"): await asyncio.wrap_future(future) return future @@ -188,9 +432,7 @@ def work(): pending.add_done_callback(lambda _: settled.set()) async def main(): - loop = asyncio.get_running_loop() - loop.set_default_executor(ThreadPoolExecutor(max_workers=1)) - executor = S3AioExecutor(loop=loop) + executor = S3AioExecutor(loop=asyncio.get_running_loop(), max_workers=1) running = executor.submit(work) executor.submit_to(pending, events.append, "pending finished") while not started.is_set(): @@ -210,11 +452,10 @@ def test_submit_to_executor_shut_down(self): future: Future[int] = Future() async def main(): - loop = asyncio.get_running_loop() - default_executor = ThreadPoolExecutor(max_workers=1) - loop.set_default_executor(default_executor) - default_executor.shutdown() - S3AioExecutor(loop=loop).submit_to(future, sum, [1, 2]) + executor = S3AioExecutor(loop=asyncio.get_running_loop()) + # Fails the dispatch of the callable to the pool. + executor._pool.shutdown() + executor.submit_to(future, sum, [1, 2]) with pytest.raises(RuntimeError, match="after shutdown"): await asyncio.wrap_future(future) @@ -229,10 +470,10 @@ def test_submit_to_cancelled_executor_shut_down(self): async def main(): loop = asyncio.get_running_loop() - default_executor = ThreadPoolExecutor(max_workers=1) - loop.set_default_executor(default_executor) - default_executor.shutdown() - S3AioExecutor(loop=loop).submit_to(future, sum, [1, 2]) + executor = S3AioExecutor(loop=loop) + # Fails the dispatch of the callable to the pool. + executor._pool.shutdown() + executor.submit_to(future, sum, [1, 2]) assert future.cancel() with ThreadPoolExecutor(max_workers=1) as waiter: return await loop.run_in_executor(waiter, lambda: wait([future], timeout=5)) diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index 50f6712e5..2b2784f6d 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -6,6 +6,7 @@ # SPDX-License-Identifier: MIT import io +import threading import time from concurrent.futures import Future, wait from datetime import UTC, datetime, timedelta, timezone @@ -30,6 +31,34 @@ ) +def overlapping_calls(calls, timeout=5): + """Return a callable that records the peak number of its overlapping calls. + + Each call waits until ``calls`` calls overlap, or until the timeout. + + Args: + calls: The number of calls to wait for. + timeout: The longest wait of a call, in seconds. + + Returns: + The callable, which accepts and ignores any arguments, and a + dictionary whose ``peak`` item is the peak number. + """ + condition = threading.Condition() + state = {"active": 0, "peak": 0} + + def call(*args, **kwargs): + with condition: + state["active"] += 1 + state["peak"] = max(state["peak"], state["active"]) + condition.notify_all() + condition.wait_for(lambda: state["active"] >= calls, timeout=timeout) + with condition: + state["active"] -= 1 + + return call, state + + def read_query(name, **kwargs): template = _queries.get_template(name) return [q.strip() for q in template.render(**kwargs).split(";") if q and q.strip()]