diff --git a/benchmarks/README.md b/benchmarks/README.md index d0c67b7c9..4edc602eb 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -238,8 +238,8 @@ 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. A `get` also writes the object to the trial's temporary directory, which the parent process measures as the peak temp dir and removes; on the Amazon Linux 2023 hosts, that directory is on the tmpfs `/tmp`, so the object stays in memory outside the process's RSS. A `put` reads a sparse file, which no disk bounds. -The `sync` API calls the filesystem on the trial's thread, and a file's parallel requests run in an `S3ThreadPool`. -The `aio` API runs a read, get, write, or put in `asyncio.to_thread()`, as `AioS3FSCursor` reads result files, and a file's parallel requests run in a pool of `max_workers` threads of the file, as in the `sync` API; its `get_file()` and `put_file()` run the synchronous filesystem's in the filesystem's thread pool. +The `sync` API calls the filesystem on the trial's thread, and a file's parallel requests run in the request pool of the filesystem, which its files and operations share. +The `aio` API runs a read, get, write, or put in `asyncio.to_thread()`, as `AioS3FSCursor` reads result files, and a file's parallel requests run in the request pool of the filesystem, as in the `sync` API; its `get_file()` and `put_file()` run the synchronous filesystem's in the filesystem's thread pool. 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. diff --git a/benchmarks/pyathena_bench/filesystem.py b/benchmarks/pyathena_bench/filesystem.py index d10d7a1bc..a88203af7 100644 --- a/benchmarks/pyathena_bench/filesystem.py +++ b/benchmarks/pyathena_bench/filesystem.py @@ -136,7 +136,7 @@ async def transfer( for each operation, a coroutine on fsspec's IO loop. As ``AioS3FSCursor`` reads result files, the asyncio filesystem's file and local-file transfers run in ``asyncio.to_thread()``: its files run their range requests and - part uploads in a pool of ``max_workers`` threads of the file, and its + part uploads in the request pool of the filesystem, and its ``get_file()`` and ``put_file()`` run the synchronous filesystem's in the filesystem's thread pool. Its copy is a coroutine awaited on the trial's event loop. diff --git a/docs/aio.md b/docs/aio.md index 6bbab094b..5bf72b50b 100644 --- a/docs/aio.md +++ b/docs/aio.md @@ -294,23 +294,29 @@ default executor of the event loop, so `loop.set_default_executor()` neither lim 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. + coroutines, such as `_info()`, `_cat_file()`, `_put_file()`, and the completion and + cleanup of a multipart copy, and for the commit and abort of the files of + `open_async()`. +- Its internal `S3FileSystem` has a request pool of `max_workers` threads for the + requests that these calls and the files run in parallel: the part copies of multipart + copies and their annotation copies, the DeleteObjects batches of `_rm()`, and the + range reads and parts of the files of `open()` and `open_async()`. A file or a copy + runs at most its own `max_workers` requests at once, and all of them together run at + most the `max_workers` of the filesystem. + +`max_workers` defaults to the `max_pool_connections` of the S3 client, 10 unless +configured, as for `S3FileSystem`. 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. -`close()` releases the threads of the filesystem once no call runs or waits in them, and +`close()` releases the threads of both pools once no call runs or waits in them, and so does the end of a `with` or `async with` block. Until then, calls, including those made after `close()`, run in those threads within `max_workers`, and the calls in progress, including the cleanup of a multipart copy, run to their end. The filesystem stays usable, also when other callers share it through fsspec's instance cache: a call after the release -creates new threads. Files opened before keep their own threads. Without `close()`, the +creates new threads. The requests of open files run to their end. Without `close()`, the threads are released when the filesystem is garbage collected, which a filesystem in fsspec's instance cache is not until `clear_instance_cache()`. @@ -339,7 +345,8 @@ creates new threads. `S3ThreadPoolExecutor` and `S3AioThreadPool` are aliases of `loop` argument, with a `DeprecationWarning` if one is given. The file handles of `AioS3FileSystem.open()` run their multipart uploads and parallel -range reads in a pool of `max_workers` threads of each file, as `S3File` does. +range reads in the request pool of the filesystem, at most the `max_workers` of the +file at once, as `S3File` does. In a forked child process, create a new filesystem. A filesystem created in the parent keeps the parent's boto3 client, which is not meant to be shared across processes, and @@ -386,7 +393,10 @@ async with await fs.open_async("s3://my-bucket/data/file.bin", "rb") as f: upload; leaving the block with an exception after a cancelled final flush, while that abort runs, waits for it. An interrupt while the commit is being queued, such as a repeated Ctrl+C under `asyncio.run()`, closes the file without committing it and - aborts its multipart upload, if any. + aborts its multipart upload, if any. A further interrupt during that cleanup can leave + the multipart upload incomplete; the close still does not wait forever. +- Closing a file, or leaving its `async with` block, waits for its range reads, commit, + and abort, also if the task is cancelled, and raises the cancellation after they end. - 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. diff --git a/docs/filesystem.md b/docs/filesystem.md index 307f4852b..15f58f551 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -100,6 +100,33 @@ transaction commits, and the multipart upload is aborted. `max_workers` is an argument of `S3FileSystem`, `AioS3FileSystem`, `open`, `put`, and `pipe`. +The requests that the files and operations of a filesystem run in parallel (range +reads, parts, part copies, the ACLs of a recursive `chmod`, the DeleteObjects batches of +`rm`, and the aborts of `clear_multipart_uploads`) share one request pool of the +filesystem, with `max_workers` threads. Each file or operation runs at most its own +`max_workers` requests at once in that pool. The `max_workers` of the filesystem +defaults to the `max_pool_connections` of its S3 client: 10 unless configured, or that +of the client of the connection given with `connection`. A client that the filesystem +builds takes its `max_pool_connections` from `max_workers` unless `config_kwargs` sets +it. For more parallel requests, raise `max_workers`, and with `connection`, also the +`max_pool_connections` of its S3 client, such as with its `s3_config`. `close()`, or the end of a `with` block, releases the threads of the pool once they +are idle, and the filesystem stays usable. + +In a forked child process, create a new filesystem. A file that the parent opened raises +`RuntimeError` there from its reads, writes, and flushes, `commit()`, and `discard()`, +and closing it only marks it closed, without uploading or aborting anything. + +An interrupt, such as Ctrl+C, while a file reads ranges in parallel cancels the range +reads that have not started. Those that are running finish, and their data is released +once they end. Closing the file waits for its running requests; an interrupt of that +wait stops it, and the requests still finish. An interrupt while a +request is being submitted can let one more request than `max_workers` run at once +until the request pool is idle. These guarantees hold for one interrupt at a time. If another interrupt arrives while the cleanup of an interrupt or +an error runs, the call still does not wait forever, but the cleanup may stop early: +the data of finished requests can stay referenced until the garbage collector frees it, +and a multipart upload can be left incomplete. `clear_multipart_uploads` removes such +uploads. + A file returned by `open` whose write or commit has failed, or which has been discarded, cannot be committed: its `commit()` raises `S3MultipartStateError`. Inside a transaction, a write error that the caller catches therefore makes the transaction fail @@ -642,6 +669,7 @@ for source, destination in pairing.copy_pairs(sources): ## Async filesystem `AioS3FileSystem` provides the same functionality on top of fsspec's -`AsyncFileSystem`, running its blocking calls in thread pools sized by `max_workers`. +`AsyncFileSystem`, running its blocking calls in a thread pool sized by `max_workers`, +and its parallel requests in the request pool of its internal `S3FileSystem`. Its `open_async()` returns a file whose reads and writes are coroutines. See {ref}`aio-s3-filesystem` for details. diff --git a/docs/usage.md b/docs/usage.md index 5f072117b..2d79926fd 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -59,11 +59,10 @@ The reader rejects an unknown name when it reads the query results, which PyAthe ## Worker settings in PyAthena 4.0 -The pandas and Polars cursors take the number of threads that read each result file from S3 as `s3_max_workers`. +The pandas and Polars cursors take, as `s3_max_workers`, the size of the request pool of the S3 filesystem that reads the result files; a file reads at most that many ranges at once. It can be set in the cursor constructor and overridden for one query in `execute()`. -The default is `(cpu_count() or 1) * 5`. -Polars uses it only for CSV results read without `chunksize`. -Chunked CSV reads and UNLOAD results use Polars' native readers, which do not use this setting. +The default is None, which takes the `max_pool_connections` of the connection's S3 client: 10, unless set with `s3_config`. +Polars uses it for CSV results read without `chunksize`, and for UNLOAD results read in chunks through PyArrow without read options; its other readers do not use this setting. The thread-pool cursors, such as `AsyncCursor`, `AsyncPandasCursor`, and `AsyncPolarsCursor`, keep `max_workers` for the size of the thread pool that waits for queries and collects their results. diff --git a/pyathena/aio/pandas/cursor.py b/pyathena/aio/pandas/cursor.py index e95ee2599..8396a86aa 100644 --- a/pyathena/aio/pandas/cursor.py +++ b/pyathena/aio/pandas/cursor.py @@ -6,7 +6,6 @@ import logging import operator from collections.abc import AsyncGenerator, Callable, Iterable -from multiprocessing import cpu_count from typing import ( TYPE_CHECKING, Any, @@ -65,7 +64,7 @@ def __init__( chunksize: int | None = None, block_size: int | None = None, cache_type: str | None = None, - s3_max_workers: int = (cpu_count() or 1) * 5, + s3_max_workers: int | None = None, result_reuse_enable: bool = False, result_reuse_minutes: int = CursorIterator.DEFAULT_RESULT_REUSE_MINUTES, auto_optimize_chunksize: bool = False, @@ -95,7 +94,10 @@ def __init__( the next block while pandas parses the current one. Without ``cache_type``, PyAthena's filesystem reads a Parquet file of a chunked UNLOAD result that is larger than 32 MiB without a cache. - s3_max_workers: Maximum number of workers of the S3 filesystem. + s3_max_workers: The size of the request pool of the S3 filesystem + that reads the result files. If None, the + ``max_pool_connections`` of the connection's S3 client (10 + unless set with ``s3_config``). result_reuse_enable: Whether to enable Athena query result reuse. result_reuse_minutes: Maximum age of a reused query result in minutes. auto_optimize_chunksize: Whether to choose a chunk size from the size of the diff --git a/pyathena/aio/polars/cursor.py b/pyathena/aio/polars/cursor.py index 5ff234a6e..5c2d7bffd 100644 --- a/pyathena/aio/polars/cursor.py +++ b/pyathena/aio/polars/cursor.py @@ -6,7 +6,6 @@ import logging import operator from collections.abc import AsyncGenerator, Callable -from multiprocessing import cpu_count from typing import TYPE_CHECKING, Any from pyathena.aio.common import WithAsyncFetch @@ -62,7 +61,7 @@ def __init__( result_reuse_minutes: int = CursorIterator.DEFAULT_RESULT_REUSE_MINUTES, block_size: int | None = None, cache_type: str | None = None, - s3_max_workers: int = (cpu_count() or 1) * 5, + s3_max_workers: int | None = None, chunksize: int | None = None, **kwargs, ) -> None: @@ -87,7 +86,10 @@ def __init__( ``"background"``; see ``S3FileSystem``. Without ``cache_type``, PyAthena's filesystem reads a Parquet file of a chunked UNLOAD result that is larger than 32 MiB without a cache. - s3_max_workers: Maximum number of workers of the S3 filesystem. + s3_max_workers: The size of the request pool of the S3 filesystem + that reads the result files. If None, the + ``max_pool_connections`` of the connection's S3 client (10 + unless set with ``s3_config``). chunksize: Number of rows per chunk. If set, result files in S3 are read lazily in chunks of this size. **kwargs: Other cursor arguments, such as ``connection`` and ``arraysize``, diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index e79b2a663..480f377b9 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -16,7 +16,6 @@ from dataclasses import dataclass, field from datetime import datetime from io import BytesIO -from multiprocessing import cpu_count from typing import Any, BinaryIO, cast from urllib.parse import unquote_plus @@ -36,7 +35,7 @@ from pyathena.connection import Connection from pyathena.filesystem.s3_core import S3Core, S3DeleteBatch, S3DeleteResult from pyathena.filesystem.s3_errors import S3ClientError -from pyathena.filesystem.s3_executor import S3Executor, S3ThreadPool +from pyathena.filesystem.s3_executor import S3Executor, S3ThreadPool, _S3PoolScope from pyathena.filesystem.s3_object import ( S3Metadata, S3MultipartUpload, @@ -170,7 +169,7 @@ def __init__( connection: Connection[Any] | None = None, default_block_size: int | None = None, default_cache_type: str | None = None, - max_workers: int = (cpu_count() or 1) * 5, + max_workers: int | None = None, s3_additional_kwargs=None, allow_bucket_creation: bool = False, allow_bucket_deletion: bool = False, @@ -197,7 +196,13 @@ def __init__( read; ``"blockcache"``, which requests whole blocks and keeps the most recently used ones; and ``"background"``, which is like ``"blockcache"`` and also reads the next block in a thread. - max_workers: The number of threads for parallel transfers. + max_workers: The number of threads of the request pool of the + filesystem, which its files and parallel operations share, + and the default number of parallel requests of a file. If + None, the ``max_pool_connections`` of the S3 client (10 + unless configured). A client that the filesystem builds gets + ``max_pool_connections`` from ``max_workers`` unless + ``config_kwargs`` sets it. s3_additional_kwargs: Extra arguments for the object requests of ``open()`` and ``pipe_file()``; listings and other requests do not use them. Each request receives those that its operation @@ -218,13 +223,18 @@ def __init__( client = connection.s3_client retry_config = connection.retry_config else: - client = self._get_client_compatible_with_s3fs(**kwargs) + client = self._get_client_compatible_with_s3fs(max_workers=max_workers, **kwargs) retry_config = RetryConfig() self.default_block_size = ( default_block_size if default_block_size else self.DEFAULT_BLOCK_SIZE ) self.default_cache_type = default_cache_type if default_cache_type else "bytes" - self.max_workers = max_workers + # The requests of the files and parallel operations of the + # filesystem run in this pool, at most max_workers at once. + self._request_pool = S3ThreadPool( + max_workers if max_workers is not None else client.meta.config.max_pool_connections + ) + self.max_workers: int = self._request_pool.max_workers self.s3_additional_kwargs = s3_additional_kwargs if s3_additional_kwargs else {} self.allow_bucket_creation = allow_bucket_creation self.allow_bucket_deletion = allow_bucket_deletion @@ -261,7 +271,26 @@ def request_kwargs(self) -> dict[str, Any]: """ return self._core.request_kwargs - def _get_client_compatible_with_s3fs(self, **kwargs) -> BaseClient: + def close(self) -> None: + """Release the threads of the request pool, which stays usable, once they are idle. + + See :meth:`S3ThreadPool.release() + `. The calls in + progress, including the requests of open files, run to their end, and + a call after the release creates new threads. + """ + self._request_pool.release() + + def __enter__(self) -> S3FileSystem: + """Return the filesystem, whose threads the end of the block releases.""" + return self + + def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: + self.close() + + def _get_client_compatible_with_s3fs( + self, max_workers: int | None = None, **kwargs + ) -> BaseClient: """Build a boto3 S3 client from s3fs-compatible constructor arguments. Accepts the constructor arguments that s3fs users pass through fsspec @@ -273,6 +302,9 @@ def _get_client_compatible_with_s3fs(self, **kwargs) -> BaseClient: is used as ``profile_name`` when ``profile_name`` is not given. Args: + max_workers: The size of the request pool of the filesystem, which + becomes the ``max_pool_connections`` of the client unless + ``config_kwargs`` sets it; None keeps botocore's default. **kwargs: The filesystem constructor arguments. Returns: @@ -280,6 +312,8 @@ def _get_client_compatible_with_s3fs(self, **kwargs) -> BaseClient: """ config_kwargs = deepcopy(kwargs.pop("config_kwargs", {})) client_kwargs = deepcopy(kwargs.pop("client_kwargs", {})) + if max_workers is not None: + config_kwargs.setdefault("max_pool_connections", max_workers) user_agent_extra = config_kwargs.pop("user_agent_extra", None) if user_agent_extra and pyathena.user_agent_extra not in user_agent_extra: @@ -1098,20 +1132,27 @@ def _delete_paths( unversioned_paths, recursive=recursive, maxdepth=maxdepth ) - def _create_executor(self, max_workers: int) -> S3Executor: + def _create_executor(self, max_workers: int) -> _S3PoolScope: """Create the executor of the parallel requests of a file or an operation. The files of ``S3FileSystem`` and ``AioS3FileSystem`` and the - parallel operations of ``S3FileSystem`` take their executors from - this method. + parallel operations of both take their executors from this method. + The requests run in the request pool of the filesystem, which they + share, at most ``max_workers`` of those of the executor at once; + shutting the executor down waits for its own requests only, and + leaves the pool running. Args: - max_workers: Maximum number of parallel workers. + max_workers: The maximum number of requests of the executor in + the request pool at once. Returns: - A new ``S3ThreadPool``. + An executor on the request pool of the filesystem. + + Raises: + ValueError: If ``max_workers`` is not positive. """ - return S3ThreadPool(max_workers=max_workers) + return _S3PoolScope(self._request_pool, max_workers) def _delete_objects(self, paths: list[str], max_workers: int | None = None, **kwargs) -> None: """Delete objects with DeleteObjects requests grouped by bucket. @@ -2848,7 +2889,7 @@ def __init__( path: str, mode: str = "rb", version_id: str | None = None, - max_workers: int = (cpu_count() or 1) * 5, + max_workers: int | None = None, executor: S3Executor | None = None, block_size: int = S3FileSystem.DEFAULT_BLOCK_SIZE, cache_type: str = "bytes", @@ -2877,13 +2918,13 @@ def __init__( 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: The number of parallel requests of the file: range + reads, part copies, and part uploads. A write keeps at most ``max_workers`` parts, including the part copies of an append, in flight, waiting for one of them to finish before it submits - another. - executor: The executor for parallel operations. If None, a new - ``S3ThreadPool`` is created. + another. If None, the ``max_workers`` of the filesystem. + executor: The executor for parallel operations. If None, one on + the request pool of the filesystem (``fs._create_executor()``). 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. @@ -2908,8 +2949,15 @@ def __init__( 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. + ``S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE`` for writing, or + ``max_workers`` is not positive. """ + if max_workers is None: + max_workers = fs.max_workers + if max_workers <= 0: + raise ValueError("max_workers must be greater than 0") + # The process that opened the file; a forked child rejects it. + self._pid = os.getpid() 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 @@ -2929,7 +2977,7 @@ def __init__( self.s3_additional_kwargs: dict[str, Any] = target.request_kwargs self._details: S3Object | dict[str, Any] = target.details - self._executor: S3Executor = executor or S3ThreadPool(max_workers=max_workers) + self._executor: S3Executor = executor or fs._create_executor(max_workers=max_workers) # The requests and the state of the write; None when reading. self._session: _S3WriteSession | None = None if mode != "rb": @@ -2976,13 +3024,78 @@ def _get_request_kwargs(self, method: str) -> dict[str, Any]: return self.fs.core.operation_params(method, self.s3_additional_kwargs) def close(self) -> None: - """Close the file, flushing any written data, and shut down its executor.""" + """Close the file, flushing any written data, and shut down its executor. + + In a forked child process, a file that the parent opened is only + marked closed: its data is neither flushed nor committed, and its + upload is not aborted, since they belong to the parent. + """ + if self._inherited(): + self.closed = True + return try: super().close() finally: # The executor is shut down even if the final flush fails. self._executor.shutdown() + def read(self, length: int = -1) -> bytes: + """Read bytes from the file. + + Args: + length: The number of bytes to read; -1 reads to the end. + + Returns: + The bytes read. + + Raises: + RuntimeError: In a forked child process of the process that + opened the file. + ValueError: If the file is closed or not open for reading. + """ + self._check_process() + data: bytes = super().read(length) + return data + + def flush(self, force: bool = False) -> None: + """Upload the buffered data as a part once a block is full, or all of it with ``force``. + + Args: + force: Whether to upload the remaining data and commit the file, + after which it cannot be written. + + Raises: + RuntimeError: In a forked child process of the process that + opened the file. + ValueError: If ``force`` is given again after a forced flush. + """ + self._check_process() + super().flush(force=force) + + def _inherited(self) -> bool: + """Whether this process is a forked child of the process that opened the file. + + Returns: + True in a forked child process of the process that opened the + file. + """ + # Compared with the process ID, not a count that an at-fork hook + # updates, since hooks registered before it may close the file. + return getattr(self, "_pid", os.getpid()) != os.getpid() + + def _check_process(self) -> None: + """Reject the use of a file that the parent of a forked process opened. + + Raises: + RuntimeError: In a forked child process of the process that + opened the file. + """ + if self._inherited(): + raise RuntimeError( + f"s3://{self.bucket}/{self.key} was opened by the parent of this forked " + "process and cannot be used here; open it again in this process." + ) + def write(self, data: bytes | bytearray | memoryview) -> int: """Write data to the buffer, uploading each full block as a part. @@ -2996,9 +3109,12 @@ def write(self, data: bytes | bytearray | memoryview) -> int: BaseException: The error of a request of the write that failed since the last write, after which the file is closed without committing it. + RuntimeError: In a forked child process of the process that + opened the file. S3MultipartStateError: If the write was discarded. ValueError: If the file is not open for writing. """ + self._check_process() if self._session is not None and not self.closed: try: self._session.check() @@ -3186,9 +3302,12 @@ def commit(self) -> None: Raises: FileExistsError: If an object was created at the path after the file was opened in exclusive-create mode. + RuntimeError: In a forked child process of the process that + opened the file. S3MultipartStateError: If the write was discarded, or the write or an earlier commit failed. """ + self._check_process() session = cast(_S3WriteSession, self._session) if not session.multipart: body = None @@ -3218,7 +3337,12 @@ def discard(self) -> None: abort. Afterwards, :meth:`commit` raises :class:`~pyathena.filesystem.s3_writer.S3MultipartStateError`. Does nothing for a committed file or a file opened for reading. + + Raises: + RuntimeError: In a forked child process of the process that + opened the file. """ + self._check_process() if self._session is not None: self._session.abort() @@ -3307,6 +3431,7 @@ def _fetch_range(self, start: int, end: int) -> bytes: # Set on failure, so that the running reads make no further attempt. stop = threading.Event() futures: list[Future[bytes]] = [] + done: set[Future[bytes]] = set() try: with stop_retries_on(stop): for r in ranges: @@ -3319,30 +3444,38 @@ def _fetch_range(self, start: int, end: int) -> bytes: futures[-1], self.fs.core.get_object, path, r, **request_kwargs ) done = wait(futures, return_when=FIRST_EXCEPTION).done + failed = next( + (f for f in done if not f.cancelled() and f.exception() is not None), None + ) + if failed is not None: + failed.result() + return b"".join(f.result() for f in futures) except BaseException as e: - # The reads are cancelled with the read of the range, as below, - # and the frames of the wait, which keep their futures, are - # cleared; clear_frames() skips this frame, which is executing. - traceback.clear_frames(e.__traceback__) - while futures: - futures.pop().cancel() - # After the cancellation, which a thread that it frees could - # otherwise precede by taking a queued read. - stop.set() + # A failed read, or an interrupt at any point, including after + # the wait: the reads that have not started are cancelled, and + # the futures of the others are dropped, so that this frame, + # which the traceback of the error keeps, does not keep their + # data, also if another interrupt stops this cleanup. The + # frames of the wait, which keep the futures, are cleared, and + # clear_frames() skips this frame, which is executing. + try: + # Popped, not iterated, so that no variable of this frame + # keeps a future. + while futures: + futures.pop().cancel() + finally: + # Also if an interrupt stops the cancellation, whose error + # keeps this one as its context; the frames first, as + # they keep the futures of the wait. + try: + traceback.clear_frames(e.__traceback__) + finally: + done.clear() + futures.clear() + # After the cancellation, which a thread that it frees + # could otherwise precede by taking a queued read. + stop.set() raise - failed = next( - (f for f in done if not f.cancelled() and f.exception() is not None), None - ) - if failed is not None: - # The reads that have not started are cancelled, and the - # futures of the others are dropped, so that this frame, which - # the traceback of the error keeps, does not keep their data. - while futures: - futures.pop().cancel() - stop.set() - done.clear() - failed.result() - return b"".join(f.result() for f in futures) return self.fs.core.get_object(path, ranges[0], **request_kwargs) @staticmethod diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 029054cab..8a20ce546 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -20,7 +20,6 @@ import warnings from concurrent.futures import Future, wait from io import BytesIO -from multiprocessing import cpu_count from typing import TYPE_CHECKING, Any, cast from fsspec.asyn import AbstractAsyncStreamedFile, AsyncFileSystem, sync @@ -61,18 +60,21 @@ class AioS3FileSystem(AsyncFileSystem): coroutines run the blocking boto3 calls in an ``S3ThreadPool`` with ``max_workers`` threads that the filesystem owns, instead of the default executor of the event loop, and run parallel requests with - ``asyncio.gather``. + ``asyncio.gather``. The single requests that these calls run in + parallel, such as the parts of a file or a copy and the DeleteObjects + batches of ``_rm()``, run in the request pool of the wrapped + ``S3FileSystem``, which its files share. The implementation uses composition: an internal ``S3FileSystem`` instance handles 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`` run their parallel operations (range - reads, multipart uploads) in an ``S3ThreadPool`` of ``max_workers`` - threads of each file, as ``S3File`` does. :meth:`open_async` returns an - ``AioS3StreamedFile``, whose reads and writes are coroutines. - :meth:`close`, or the end of a ``with`` or ``async with`` block, releases - the threads of the filesystem, which stays usable. + reads, multipart uploads) in that request pool, as ``S3File`` does. + :meth:`open_async` returns an ``AioS3StreamedFile``, whose reads and + writes are coroutines. :meth:`close`, or the end of a ``with`` or + ``async with`` block, releases the threads of the filesystem, which + stays usable. Attributes: _sync_fs: The internal synchronous S3FileSystem instance. @@ -99,7 +101,7 @@ def __init__( connection: Connection[Any] | None = None, default_block_size: int | None = None, default_cache_type: str | None = None, - max_workers: int = (cpu_count() or 1) * 5, + max_workers: int | None = None, s3_additional_kwargs: dict[str, Any] | None = None, allow_bucket_creation: bool = False, allow_bucket_deletion: bool = False, @@ -116,10 +118,11 @@ def __init__( default_block_size: Passed to the internal ``S3FileSystem``. default_cache_type: 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``. + the filesystem, and, passed to the internal ``S3FileSystem``, + the size of its request pool and the default number of + parallel requests of its files and copies. If None, the + ``max_pool_connections`` of the S3 client (10 unless + configured), as for ``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``. @@ -136,11 +139,6 @@ 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 = S3ThreadPool(max_workers) - max_workers = self._thread_pool.max_workers self._sync_fs = S3FileSystem( connection=connection, default_block_size=default_block_size, @@ -154,6 +152,11 @@ def __init__( skip_instance_cache=True, **kwargs, ) + # The blocking calls of the coroutines run in threads of this pool, + # not in the default executor of the event loop, and wait for the + # requests that they run in parallel in the request pool of the + # internal S3FileSystem, which runs no call that waits. + self._thread_pool = S3ThreadPool(self._sync_fs.max_workers) # Share dircache for cache coherence between async and sync instances self.dircache = self._sync_fs.dircache @@ -165,11 +168,15 @@ def core(self) -> S3Core: def close(self) -> None: """Release the threads of the filesystem, which stays usable, once they are idle. - See :meth:`S3ThreadPool.release() - `. Files opened - before have threads of their own and are not affected. + Releases the threads of the blocking calls and those of the request + pool of the internal ``S3FileSystem`` (see + :meth:`S3FileSystem.close() ` + and :meth:`S3ThreadPool.release() + `). The calls in + progress, including the requests of open files, run to their end. """ self._thread_pool.release() + self._sync_fs.close() def __enter__(self) -> AioS3FileSystem: """Return the filesystem, whose threads the end of the block releases.""" @@ -342,8 +349,9 @@ async def _rm( ) -> None: """Delete objects with DeleteObjects requests. - See :meth:`S3FileSystem.rm`. The requests run in parallel with - ``asyncio.gather`` in the threads of the filesystem. + See :meth:`S3FileSystem.rm`. The DeleteObjects requests run in + parallel with ``asyncio.gather``, at most ``max_workers`` at once, in + the request pool of the internal ``S3FileSystem``. Args: path: S3 path (s3://bucket/key) or list of paths to delete. @@ -377,11 +385,12 @@ async def _delete_objects(self, paths: list[str], **kwargs) -> None: OSError: If S3 could not delete some of the objects. """ batches, params = self._sync_fs._delete_batches(paths, **kwargs) + # At most max_workers requests at once, as S3FileSystem.rm() sends, + # so that a large deletion does not queue ahead of the requests of + # the files that share the request pool. + requests = self._sync_fs._create_executor(max_workers=self._sync_fs.max_workers) results = await asyncio.gather( - *[ - self._thread_pool.run(self._sync_fs._delete_batch, batch, **params) - for batch in batches - ], + *[requests.run(self._sync_fs._delete_batch, batch, **params) for batch in batches], return_exceptions=True, ) self._sync_fs._raise_delete_errors(results) @@ -599,19 +608,20 @@ 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 - copies run in the threads of the filesystem, which bound the requests - of all its copies together, with at most ``max_workers`` of them in - flight; each one is sent once an earlier one has finished, and the - upload is completed once all of them have. The annotation copies then - run in parallel as asyncio tasks, at most ``max_workers`` at once, in - the same threads. On a cancellation, the requests that have not - started, including the completion, are cancelled, the running ones - are waited for in a thread of the filesystem, the upload is aborted - if it was created and has not completed, and the cancellation is - re-raised. On a failure or cancellation, the running part copies make - no further attempt. This also holds when the shutdown of the event loop, such as - at the end of ``asyncio.run()``, cancels the copy, also after its - cleanup has started. A repeated cancellation returns without stopping + copies run in the request pool of the internal ``S3FileSystem``, which + bounds the requests of all its copies and files together, with at most + ``max_workers`` of them in flight; each one is sent once an earlier + one has finished, and the upload is completed once all of them have. + The annotation copies then run in parallel as asyncio tasks, at most + ``max_workers`` at once, in the same pool. On a cancellation, the + requests that have not started, including the completion, are + cancelled, the running ones are waited for in a thread of the + filesystem, the upload is aborted if it was created and has not + completed, and the cancellation is re-raised. On a failure or + cancellation, the running part copies make no further attempt. This + also holds when the shutdown of the event loop, such as at the end of + ``asyncio.run()``, cancels the copy, also after its cleanup has + started. A repeated cancellation returns without stopping this cleanup. A failed abort is logged with the upload ID. Args: @@ -638,13 +648,14 @@ async def _copy_object_with_multipart_upload( ) return - # The requests run in the threads of the filesystem, not as tasks, - # which the shutdown of the event loop would cancel while the - # requests keep running. + # The requests run in the request pool, not as tasks, which the + # shutdown of the event loop would cancel while the requests keep + # running. The completion and the abort, which wait for them, run in + # the threads of the filesystem. session = _S3WriteSession( self.core, plan.destination, - self._thread_pool, + self._sync_fs._create_executor(max_workers=max_workers), block_size=block_size if block_size else self.core.MULTIPART_UPLOAD_MAX_PART_SIZE, max_workers=max_workers, prepared_params={ @@ -685,9 +696,8 @@ def _abort() -> None: # A cancellation of this task cancels only its wait. await asyncio.wrap_future(session.ready()) session.submit_copy(part_number, plan.source, range_) - # Here, not in the thread of the completion, which may wait for - # a thread behind the part copies, so that a failed part stops - # the others at once. + # Here, not in the thread of the completion, so that a failed + # part stops the others at once. await asyncio.wrap_future(session.settled()) # Reserved before it is queued, so that a cancellation of this # task can still cancel a completion that has not started. @@ -710,10 +720,10 @@ def _abort() -> None: await self._thread_pool.run_shielded(_abort) raise - semaphore = asyncio.Semaphore(max_workers) - failed = False + # At most max_workers copies at once in the request pool. + requests = self._sync_fs._create_executor(max_workers=max_workers) # Set when an annotation copy fails, so that the running ones make no - # further attempt. + # further attempt, and the waiting ones do not start. stop = threading.Event() # The index of the copy that failed first, whose error is raised # instead of those of the copies that its failure stopped. Recorded @@ -730,7 +740,7 @@ def _copy_annotation_in_thread(i: int, name: str) -> None: """ nonlocal first_failure if stop.is_set(): - # Waited for a thread of the filesystem until a copy failed. + # Waited for a thread of the request pool until a copy failed. return try: with stop_retries_on(stop): @@ -751,19 +761,13 @@ def _copy_annotation_in_thread(i: int, name: str) -> None: raise async def _copy_annotation(i: int, name: str) -> None: - nonlocal failed - async with semaphore: - if failed: - # Do not start more copies after one failed, as - # S3FileSystem does. - return - try: - await self._thread_pool.run(_copy_annotation_in_thread, i, name) - except Exception: - # The running copies make no further attempt. - failed = True - stop.set() - raise + try: + await requests.run(_copy_annotation_in_thread, i, name) + except Exception: + # The running copies make no further attempt, and the waiting + # ones do not start, as in S3FileSystem. + stop.set() + raise results = await asyncio.gather( *[_copy_annotation(i, name) for i, name in enumerate(plan.annotations)], @@ -891,9 +895,9 @@ async def open_async( 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. + max_workers: The number of parallel requests of the file, in the + request pool of the filesystem. 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. @@ -1168,8 +1172,9 @@ class AioS3File(S3File): """The file handle that ``AioS3FileSystem.open()`` returns. Functionally identical to ``S3File``, whose parallel operations (range - reads, multipart uploads) it runs in a pool of ``max_workers`` threads of - the file; it exists as a distinct type for ``isinstance`` checks. For + reads, multipart uploads) it runs in the request pool of the filesystem, + at most ``max_workers`` of them at once; it exists as a distinct type for + ``isinstance`` checks. For reads and writes that are coroutines, see ``AioS3StreamedFile``. """ @@ -1178,8 +1183,9 @@ 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 ``S3ThreadPool`` of ``max_workers`` threads of its own, and - the coroutines wait for them without blocking the event loop: + run in the request pool of the filesystem, at most ``max_workers`` of + them at once, 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 @@ -1201,10 +1207,13 @@ class AioS3StreamedFile(AbstractAsyncStreamedFile): 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 + The commit and the abort, which wait for the requests, run to their end + in the threads of the filesystem, 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. In a forked + child process, a file that the parent opened raises ``RuntimeError`` + from its reads, writes, and flushes, and closing it only marks it closed. + 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 @@ -1224,7 +1233,7 @@ def __init__( path: str, mode: str = "rb", version_id: str | None = None, - max_workers: int = (cpu_count() or 1) * 5, + max_workers: int | None = None, block_size: int = S3FileSystem.DEFAULT_BLOCK_SIZE, size: int | None = None, s3_additional_kwargs: dict[str, Any] | None = None, @@ -1240,8 +1249,9 @@ def __init__( 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. + max_workers: The number of parallel requests of the file, in the + request pool of the filesystem. If None, the ``max_workers`` + of the filesystem. 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 @@ -1263,9 +1273,20 @@ def __init__( """ 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 = S3ThreadPool(max_workers) - self.max_workers = self._pool.max_workers + # The requests of the file, at most max_workers at once in the + # request pool of the filesystem. Created first, so that an invalid + # max_workers raises before the lookups. + if max_workers is None: + max_workers = fs._sync_fs.max_workers + self._requests = fs._sync_fs._create_executor(max_workers=max_workers) + # The scheduled commit and abort, which close() waits for, and those of + # them whose queuing has returned. + self._commit_future: Future[None] | None = None + self._discard_future: Future[None] | None = None + self._queued: set[Future[None]] = set() + self.max_workers = max_workers + # The process that opened the file; a forked child rejects it. + self._pid = os.getpid() target = S3FileTarget.resolve( fs._sync_fs, path, @@ -1298,7 +1319,7 @@ def __init__( self._session = _S3WriteSession( fs.core, S3Path(self.bucket, self.key), - self._pool, + self._requests, block_size=block_size, max_workers=self.max_workers, request_kwargs=self.s3_additional_kwargs, @@ -1314,9 +1335,10 @@ def __init__( 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: + # multipart upload, if any, is aborted. + requests = getattr(self, "_requests", None) + if requests is None or self._inherited(): + # The upload of a file inherited across a fork is the parent's. return session = getattr(self, "_session", None) if ( @@ -1331,8 +1353,30 @@ def __del__(self) -> None: stacklevel=2, ) with contextlib.suppress(RuntimeError): - pool.submit(self._discard) - pool.shutdown(wait=False) + self.fs._thread_pool.submit(self._discard) + requests.shutdown(wait=False) + + def _inherited(self) -> bool: + """Whether this process is a forked child of the process that opened the file. + + Returns: + True in a forked child process of the process that opened the + file. + """ + return getattr(self, "_pid", os.getpid()) != os.getpid() + + def _check_process(self) -> None: + """Reject the use of a file that the parent of a forked process opened. + + Raises: + RuntimeError: In a forked child process of the process that + opened the file. + """ + if self._inherited(): + raise RuntimeError( + f"s3://{self.bucket}/{self.key} was opened by the parent of this forked " + "process and cannot be used here; open it again in this process." + ) def __enter__(self) -> AioS3StreamedFile: """Reject a synchronous ``with`` block, whose exit cannot close the file. @@ -1354,6 +1398,24 @@ def _get_request_kwargs(self, method: str) -> dict[str, Any]: """ return self.fs.core.operation_params(method, self.s3_additional_kwargs) + async def read(self, length: int = -1) -> bytes: + """Read bytes from the file. + + Args: + length: The number of bytes to read; -1 reads to the end. + + Returns: + The bytes read. + + Raises: + RuntimeError: In a forked child process of the process that + opened the file. + ValueError: If the file is closed or not open for reading. + """ + self._check_process() + data: bytes = await super().read(length) + return data + async def write(self, data: bytes) -> int: """Write data to the buffer, uploading each full block as a part. @@ -1364,12 +1426,15 @@ async def write(self, data: bytes) -> int: The number of bytes written. Raises: + RuntimeError: In a forked child process of the process that + opened the file. ValueError: If the file is not open for writing, is closed, or has been force-flushed. BaseException: The error of a request of the write that failed since the last write, or before a part is submitted, after which the file is closed without committing it. """ + self._check_process() if self.mode not in {"wb", "ab", "xb"}: raise ValueError("File not in write mode") if self.closed: @@ -1396,8 +1461,11 @@ async def flush(self, force: bool = False) -> None: the written data is committed. Raises: + RuntimeError: In a forked child process of the process that + opened the file. ValueError: If the file is closed, or force-flushed twice. """ + self._check_process() if self.closed: raise ValueError("Flush on closed file") if force and self.forced: @@ -1417,16 +1485,60 @@ async def flush(self, force: bool = False) -> None: self.buffer = BytesIO() async def close(self) -> None: - """Close the file, committing the written data.""" + """Close the file, committing the written data. + + Returns once the requests of the file and its scheduled commit have + ended, also when the close or an earlier read or flush was + cancelled; the event loop runs meanwhile. In a forked child process, + a file that the parent opened is only marked closed, without + committing or aborting anything. + """ if self.closed: return + if self._inherited(): + self.closed = True + 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) + await self._drain() + + async def _drain(self) -> None: + """Wait, without blocking the event loop, for the requests, commit, and abort of the file. + + A cancellation of the waiting task, also a repeated one, does not stop + the wait; it is re-raised once they have ended, so that nothing of + the file runs after its close returns (I1). A commit or abort whose + queuing did not return, as when interrupts stopped it and its + cleanup, is cancelled unless it has started, so that the wait does + not last forever for one that was never queued. + """ + self._requests.shutdown(wait=False) + pending = self._requests.pending() + for future in (self._commit_future, self._discard_future): + if future is not None and (future in self._queued or not future.cancel()): + pending.append(future) + if not pending: + return + waiters = [asyncio.wrap_future(future) for future in pending] + for waiter in waiters: + # The outcomes belong to the reads, the commit, and the abort that + # awaited them; retrieved here, so that asyncio does not log them. + waiter.add_done_callback(lambda f: f.cancelled() or f.exception()) + cancelled = False + # The states are also checked every second, since an interrupt of a + # cancellation can skip the callbacks that the waiters rely on. + while not all(future.done() for future in pending): + try: + # Unlike gather(), wait() neither raises their outcomes nor + # cancels them when this task is cancelled. + await asyncio.wait(waiters, timeout=1.0) + except asyncio.CancelledError: + cancelled = True + if cancelled: + raise asyncio.CancelledError async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: if self.closed: @@ -1445,11 +1557,21 @@ async def _close_without_commit(self) -> None: """ if self.closed: return + if self._inherited(): + self.closed = True + return try: if self._begin_close_without_commit(): - await self._pool.run_shielded(self._discard) + # Tracked before it is queued, so that the drain below waits + # for it also after a cancellation of this task. + discard: Future[None] = Future() + self._discard_future = discard + self.fs._thread_pool.submit_to(discard, self._discard) + self._queued.add(discard) + await self.fs._thread_pool.wait_shielded(discard) finally: - self._pool.shutdown(wait=False) + # A commit that was scheduled before runs on, and is waited for. + await self._drain() def _close_without_commit_blocking(self) -> None: """Close the file without committing it; blocks. @@ -1464,7 +1586,7 @@ def _close_without_commit_blocking(self) -> None: if self._begin_close_without_commit(): self._discard() finally: - self._pool.shutdown(wait=False) + self._requests.shutdown(wait=False) def _begin_close_without_commit(self) -> bool: """Mark the file closed, drop the buffered data, and begin the abort. @@ -1472,7 +1594,7 @@ def _begin_close_without_commit(self) -> bool: The part of :meth:`_close_without_commit` and :meth:`_close_without_commit_blocking` that runs before the discard. The abort begins here, so that the creation and the running parts stop - even while they hold every thread of the file, ahead of + even while they hold every thread of the request pool, ahead of :meth:`_discard`. Returns: @@ -1573,20 +1695,25 @@ async def _commit(self) -> None: committing it. """ session = cast(_S3WriteSession, self._session) - commit: Future[None] | None = None + # Tracked before it is queued, so that close() also waits for a + # commit that an interrupt of the queuing has let start. + commit: Future[None] = Future() + self._commit_future = commit try: session.begin_complete() # Passed, as closing the file without committing it drops the buffer. - commit = self._pool.submit(self._complete_upload, cast(BytesIO, self.buffer)) + self.fs._thread_pool.submit_to( + commit, self._complete_upload, cast(BytesIO, self.buffer) + ) + self._queued.add(commit) except BaseException: - if commit is None: - # Not queued, or queued by a submit() that an interrupt - # stopped: such a commit either has taken the reservation - # and runs on, or finds the write failed. + if commit.cancel(): + # Not started, and now never: a commit that was queued by a + # submission that an interrupt stopped finds itself cancelled. session.withdraw_complete() await self._close_without_commit() raise - await self._pool.wait_shielded(commit) + await self.fs._thread_pool.wait_shielded(commit) def _complete_upload(self, buffer: BytesIO) -> None: """Upload or complete the written data; blocks, and runs in a thread. @@ -1641,10 +1768,10 @@ 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. When a read fails, - the reads that have not finished are cancelled, which stops the retries - of the running ones, and its error is raised, as - :meth:`S3File._fetch_range` does. + of the block size, read at most ``max_workers`` at once in the + request pool of the filesystem. When a read fails, the reads that have + not finished are cancelled, which stops the retries of the running + ones, and its error is raised, as :meth:`S3File._fetch_range` does. Args: start: The offset of the first byte to read. @@ -1652,7 +1779,12 @@ async def _fetch_range(self, start: int, end: int) -> bytes: Returns: The bytes read, empty if the clamped range is empty. + + Raises: + RuntimeError: In a forked child process of the process that + opened the file. """ + self._check_process() end = min(end, self.size) if start >= end: return b"" @@ -1663,7 +1795,7 @@ async def _fetch_range(self, start: int, end: int) -> bytes: request_kwargs = self._get_request_kwargs("get_object") reads = [ asyncio.ensure_future( - self._pool.run(self.fs.core.get_object, path, r, **request_kwargs) + self._requests.run(self.fs.core.get_object, path, r, **request_kwargs) ) for r in ranges ] diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index 2da2b57e7..a5f43b9ef 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -17,6 +17,7 @@ import warnings import weakref from abc import ABCMeta, abstractmethod +from collections import deque from collections.abc import Callable from concurrent.futures import Future from concurrent.futures.thread import ThreadPoolExecutor @@ -141,6 +142,93 @@ def shutdown(self, wait: bool = True) -> None: """Shut down the executor, freeing any resources.""" ... + async def run(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> T: + """Run a blocking call in the executor and wait for it. + + Cancelling the waiting task cancels a call that has not started, and + stops the retries of the AWS API calls that a running call makes, as + :func:`~pyathena.aio.util.run_in_thread` does; the call itself keeps + running until it returns or raises. + + 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 executor has been shut down. + """ + stop = threading.Event() + # Orders the end of the call with a cancellation, after which work + # that the call left running, such as the parts of a write in a + # transaction, is not stopped. + lock = threading.Lock() + returned = False + + def call() -> T: + """Run the call in a scope that the cancellation stops until it returns.""" + nonlocal returned + try: + with stop_retries_on(stop): + return fn(*args, **kwargs) + finally: + with lock: + returned = True + + try: + return await asyncio.wrap_future(self.submit(call)) + except asyncio.CancelledError: + with lock: + if not returned: + stop.set() + raise + + 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, and its retries are not stopped; its outcome is then + dropped. See :meth:`wait_shielded`. + + 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 executor has been shut down. + """ + return await self.wait_shielded(self.submit(fn, *args, **kwargs)) + + async def wait_shielded(self, future: Future[T]) -> T: + """Wait for a submitted call that a cancellation of the waiting task does not cancel. + + Unlike awaiting ``asyncio.wrap_future(future)``, a cancellation of + the waiting task leaves the future to the caller, which can still + wait for it or cancel it, and its outcome is retrieved, so that + asyncio does not log a failure that ends after the cancellation as + never retrieved. + + Args: + future: The future of a call submitted to this executor. + + Returns: + The return value of the call. + """ + waiter = asyncio.wrap_future(future) + waiter.add_done_callback(lambda f: f.cancelled() or f.exception()) + # Unlike shield(), which on Python 3.14 and later also logs the failure + # of a future that ends after a cancellation, wait() leaves the future + # alone when the waiting task is cancelled. + await asyncio.wait([waiter]) + return waiter.result() + def submit_to(self, future: Future[T], fn: Callable[..., T], *args: Any, **kwargs: Any) -> None: """Submit a callable whose outcome resolves a future of the caller. @@ -176,6 +264,10 @@ def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: self.shutdown(wait=True) +# The pool, weakly referenced, whose thread this is. +_pool_thread = threading.local() + + class S3ThreadPool(S3Executor): """A thread pool for the requests of the S3 filesystems. @@ -278,93 +370,6 @@ def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]: future.add_done_callback(self._end_call_callback) return future - 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, and - stops the retries of the AWS API calls that a running call makes, as - :func:`~pyathena.aio.util.run_in_thread` does; the call itself keeps - running until it returns or raises. - - 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. - """ - stop = threading.Event() - # Orders the end of the call with a cancellation, after which work - # that the call left running, such as the parts of a write in a - # transaction, is not stopped. - lock = threading.Lock() - returned = False - - def call() -> T: - """Run the call in a scope that the cancellation stops until it returns.""" - nonlocal returned - try: - with stop_retries_on(stop): - return fn(*args, **kwargs) - finally: - with lock: - returned = True - - try: - return await asyncio.wrap_future(self.submit(call)) - except asyncio.CancelledError: - with lock: - if not returned: - stop.set() - raise - - 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, and its retries are not stopped; its outcome is then - dropped. See :meth:`wait_shielded`. - - 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 self.wait_shielded(self.submit(fn, *args, **kwargs)) - - async def wait_shielded(self, future: Future[T]) -> T: - """Wait for a submitted call that a cancellation of the waiting task does not cancel. - - Unlike awaiting ``asyncio.wrap_future(future)``, a cancellation of - the waiting task leaves the future to the caller, which can still - wait for it or cancel it, and its outcome is retrieved, so that - asyncio does not log a failure that ends after the cancellation as - never retrieved. - - Args: - future: The future of a call submitted to this pool. - - Returns: - The return value of the call. - """ - waiter = asyncio.wrap_future(future) - waiter.add_done_callback(lambda f: f.cancelled() or f.exception()) - # Unlike shield(), which on Python 3.14 and later also logs the failure - # of a future that ends after a cancellation, wait() leaves the future - # alone when the waiting task is cancelled. - await asyncio.wait([waiter]) - return waiter.result() - @override def shutdown(self, wait: bool = True) -> None: """Stop accepting calls and release the threads once they are idle. @@ -424,6 +429,15 @@ def _take_idle_executor(self) -> ThreadPoolExecutor | None: executor, self._executor = self._executor, None return executor + def owns_current_thread(self) -> bool: + """Whether the calling thread is a thread of this pool. + + Returns: + True in a thread that this pool started. + """ + pool = getattr(_pool_thread, "pool", None) + return pool is not None and pool() is self + @classmethod def _reset_after_fork(cls) -> None: """Drop the threads and locks that a forked child copied from its parent. @@ -453,7 +467,11 @@ def _get_executor(self) -> ThreadPoolExecutor: if self._shutdown: raise RuntimeError("cannot schedule new futures after shutdown") if self._executor is None: - self._executor = ThreadPoolExecutor(max_workers=self.max_workers) + pool = weakref.ref(self) + self._executor = ThreadPoolExecutor( + max_workers=self.max_workers, + initializer=lambda: setattr(_pool_thread, "pool", pool), + ) return self._executor @@ -461,6 +479,291 @@ def _get_executor(self) -> ThreadPoolExecutor: os.register_at_fork(after_in_child=S3ThreadPool._reset_after_fork) +class _S3PoolScope(S3Executor): + """The calls of one file or operation in a pool that others share. + + At most ``max_workers`` calls of the scope are in the pool at once; the + others wait in the scope, in the order they were submitted, and enter the + pool as earlier ones end, from the thread of the call that ends. Each call + runs with a copy of the context of its submitter. :meth:`shutdown` stops + accepting calls and, with ``wait``, waits for those submitted, without + shutting the pool down, so that a file or an operation drains its own + requests while others keep using the pool. + + A call counts toward ``max_workers`` once the pool has accepted it, so + that an interrupt of the submitter at any point leaves no place taken. + A call whose submission an interrupt stopped after the pool accepted it + may then run without counting, so that each such interrupt can let one + more call run at once. + + Args: + pool: The shared pool. + max_workers: The maximum number of calls of the scope in the pool. + + Raises: + ValueError: If ``max_workers`` is not positive. + """ + + def __init__(self, pool: S3ThreadPool, max_workers: int) -> None: + """Initialize the scope on a shared pool. + + Args: + pool: The shared pool. + max_workers: The maximum number of calls of the scope in the pool. + + Raises: + ValueError: If ``max_workers`` is not positive. + """ + if max_workers <= 0: + raise ValueError("max_workers must be greater than 0") + self.max_workers = max_workers + self._pool = pool + self._lock = threading.Lock() + # The calls that were submitted and have not ended, by future. + self._futures: dict[Future[Any], _S3ScopeCall] = {} + # The futures of the calls that threads of the pool run themselves. + self._inline: set[Future[Any]] = set() + # The calls that wait for a place in the pool, and those that the + # pool has accepted and that have not ended. + self._waiting: deque[_S3ScopeCall] = deque() + # The waiting calls that have not been cancelled, by future. + self._waiting_calls: dict[Future[Any], _S3ScopeCall] = {} + self._entered: set[_S3ScopeCall] = set() + self._shutdown = False + + @override + def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]: + """Submit a call, which enters the pool once fewer than ``max_workers`` calls are there. + + As with ``ThreadPoolExecutor``, the returned future can be cancelled + only before the call starts. An interrupt of the submission cancels + the call unless it has started. + + Args: + fn: The callable to run. + *args: Positional arguments passed to the callable. + **kwargs: Keyword arguments passed to the callable. + + Returns: + The future of the call. + + Raises: + RuntimeError: If the scope or the pool has been shut down. + """ + future: Future[T] = Future() + # The context is captured here, since a waiting call enters the pool + # from the thread of the call that ends before it. + call = _S3ScopeCall( + contextvars.copy_context(), _S3WorkItem(future, fn, args, kwargs), future + ) + if self._pool.owns_current_thread(): + # A thread of the pool, such as one in which garbage collection + # closes a file, runs the call itself instead of waiting for a + # place in the pool that it holds (T1). It does not count toward + # max_workers, and a shutdown waits for it. + with self._lock: + if self._shutdown: + raise RuntimeError("cannot schedule new futures after shutdown") + self._futures[future] = call + self._inline.add(future) + future.add_done_callback(self._untrack) + call.context.run(call.item.run) + return future + try: + with self._lock: + if self._shutdown: + raise RuntimeError("cannot schedule new futures after shutdown") + self._futures[future] = call + # Behind any waiting call, so that the calls enter in order, + # also when a place frees before a waiting call takes it. + enter = not self._waiting_calls and len(self._entered) < self.max_workers + if enter: + # Counted only once the pool has accepted it. + self._pool.submit(self._run, call) + self._entered.add(call) + else: + self._waiting.append(call) + self._waiting_calls[future] = call + future.add_done_callback(self._untrack) + except BaseException: + # A call that the pool has not taken, or that has not started, + # never runs, and nothing waits for it; one that has started + # runs to its end. The work item marks the cancelled future as + # notified once, also when the pool or the scope runs it later. + if future.cancel(): + call.item.run() + with self._lock: + self._futures.pop(future, None) + raise + return future + + @override + def shutdown(self, wait: bool = True) -> None: + """Stop accepting calls and, with ``wait``, wait for those submitted. + + The pool keeps running, also for the calls of this scope that have + not entered it yet. + + Args: + wait: Whether to wait for the calls submitted through this scope. + """ + with self._lock: + self._shutdown = True + pending = self.pending() + if wait and pending: + # Waits through done callbacks, which a cancellation runs at once, + # not through futures.wait(), which waits for a cancelled future + # until the pool reaches its call and marks it as notified. The + # states are also checked every second, since an interrupt of a + # cancellation can skip its callbacks. + ended = threading.Event() + for future in pending: + future.add_done_callback(lambda _: ended.set()) + while not all(future.done() for future in pending): + ended.wait(1.0) + ended.clear() + + def pending(self) -> list[Future[Any]]: + """Return the futures of the calls that were submitted and have not ended. + + A call that has neither entered the pool, nor waits for a place, nor + started, is one whose submission interrupts stopped before its + cleanup ended. It is cancelled instead, and a cancelled call is + left out, so that nothing waits forever for a call that never runs. + + Returns: + The futures, which an asyncio caller can wait for without + blocking the event loop, such as with :meth:`wait_shielded`. + """ + with self._lock: + tracked = {call.future for call in self._entered} + tracked.update(self._waiting_calls, self._inline) + calls = list(self._futures.items()) + pending = [] + for future, call in calls: + if future.cancelled(): + # Never runs, also if an interrupt stopped its notification. + continue + if future in tracked or not future.cancel(): + pending.append(future) + else: + # Marks the cancelled future as notified, so that wait() + # counts it as done; the call does not run. + call.item.run() + return pending + + def _run(self, call: _S3ScopeCall) -> None: + """Run a call in the pool, then let the next waiting calls enter it. + + The next calls enter the pool from this thread, before this call ends + there, so that a release of the pool waits for the calls that wait + here. If the pool takes no more calls, as when the interpreter exits, + this thread runs the waiting calls itself. + + Args: + call: The call. + """ + try: + call.context.run(call.item.run) + finally: + self._next(call) + + def _next(self, ended: _S3ScopeCall) -> None: + """Free the place of a call that has ended, and fill the free places. + + Calls that were cancelled while they waited are settled without + entering the pool. + + Args: + ended: The call that has ended. + """ + freed = False + while True: + with self._lock: + if not freed: + # Freed under the lock that hands the place to the oldest + # waiting call, so that a new submitter cannot take it. + self._entered.discard(ended) + freed = True + if not self._waiting or len(self._entered) >= self.max_workers: + return + call = self._waiting.popleft() + if self._waiting_calls.pop(call.future, None) is None: + # Cancelled while it waited, and already settled. + continue + error: Exception | None = None + try: + self._pool.submit(self._run, call) + except RuntimeError: + # The pool takes no more calls, as when the interpreter + # exits: this thread runs the call, which a shutdown + # waits for. + run_inline = True + self._inline.add(call.future) + except Exception as e: + # Such as a failure to start a thread: the call fails. + run_inline, error = False, e + else: + run_inline = False + self._entered.add(call) + if run_inline: + call.context.run(call.item.run) + elif error is not None: + rejected: Future[None] = Future() + rejected.set_running_or_notify_cancel() + rejected.set_exception(error) + call.item.settle(rejected) + + def _untrack(self, future: Future[Any]) -> None: + """Stop tracking a call that has ended or was cancelled. + + A call cancelled while it waits leaves the queue at once, so that its + arguments and context are not kept until a place frees. + + Args: + future: The future of the call. + """ + with self._lock: + self._futures.pop(future, None) + self._inline.discard(future) + dropped = self._waiting_calls.pop(future, None) if future.cancelled() else None + if dropped is not None and len(self._waiting) > 2 * len(self._waiting_calls) + 16: + # Drops the entries of cancelled calls, so that calls submitted + # and cancelled behind a long call do not accumulate. + self._waiting = deque(c for c in self._waiting if c.future in self._waiting_calls) + if dropped is not None: + # Marks the future as notified and drops the arguments, without + # running the call; the entry left in the queue is skipped. + dropped.item.run() + dropped.context = contextvars.Context() + + +class _S3ScopeCall: + """A call of an :class:`_S3PoolScope`, with the context of its submitter. + + Args: + context: The context of the submitter. + item: The work item of the call. + future: The future of the call. + """ + + __slots__ = ("context", "future", "item") + + def __init__( + self, context: contextvars.Context, item: _S3WorkItem[Any], future: Future[Any] + ) -> None: + """Initialize the call. + + Args: + context: The context of the submitter. + item: The work item of the call. + future: The future of the call. + """ + self.context = context + self.item = item + self.future = future + + class S3AioExecutor(S3ThreadPool): """A thread pool kept for compatibility, an :class:`S3ThreadPool`. diff --git a/pyathena/pandas/async_cursor.py b/pyathena/pandas/async_cursor.py index 777a98db3..00c73574c 100644 --- a/pyathena/pandas/async_cursor.py +++ b/pyathena/pandas/async_cursor.py @@ -85,7 +85,7 @@ def __init__( block_size: int | None = None, cache_type: str | None = None, auto_optimize_chunksize: bool = False, - s3_max_workers: int = (cpu_count() or 1) * 5, + s3_max_workers: int | None = None, **kwargs, ) -> None: """Initialize an AsyncPandasCursor. @@ -119,7 +119,10 @@ def __init__( UNLOAD result that is larger than 32 MiB without a cache. auto_optimize_chunksize: Whether to choose a chunk size from the size of the CSV result file when ``chunksize`` is None. - s3_max_workers: Maximum worker threads per S3 file reader. + s3_max_workers: The size of the request pool of the S3 filesystem + that reads the result files. If None, the + ``max_pool_connections`` of the connection's S3 client (10 + unless set with ``s3_config``). **kwargs: Other cursor arguments, such as ``connection`` and ``converter``, passed to ``AsyncCursor.__init__``. """ diff --git a/pyathena/pandas/cursor.py b/pyathena/pandas/cursor.py index 3e428d00f..9289c218f 100644 --- a/pyathena/pandas/cursor.py +++ b/pyathena/pandas/cursor.py @@ -4,7 +4,6 @@ import logging from collections.abc import Callable, Generator, Iterable -from multiprocessing import cpu_count from typing import ( TYPE_CHECKING, Any, @@ -81,7 +80,7 @@ def __init__( chunksize: int | None = None, block_size: int | None = None, cache_type: str | None = None, - s3_max_workers: int = (cpu_count() or 1) * 5, + s3_max_workers: int | None = None, result_reuse_enable: bool = False, result_reuse_minutes: int = CursorIterator.DEFAULT_RESULT_REUSE_MINUTES, auto_optimize_chunksize: bool = False, @@ -110,7 +109,10 @@ def __init__( the next block while pandas parses the current one. Without ``cache_type``, PyAthena's filesystem reads a Parquet file of a chunked UNLOAD result that is larger than 32 MiB without a cache. - s3_max_workers: Maximum worker threads per S3 file reader. + s3_max_workers: The size of the request pool of the S3 filesystem + that reads the result files. If None, the + ``max_pool_connections`` of the connection's S3 client (10 + unless set with ``s3_config``). result_reuse_enable: Enable query result reuse. result_reuse_minutes: Result reuse duration in minutes. auto_optimize_chunksize: Enable automatic chunksize determination for diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index 34c183acb..548cb7604 100644 --- a/pyathena/pandas/result_set.py +++ b/pyathena/pandas/result_set.py @@ -16,7 +16,6 @@ from functools import partial from io import BufferedReader, BytesIO, IOBase, StringIO, TextIOWrapper from itertools import chain, islice -from multiprocessing import cpu_count from types import GeneratorType from typing import ( TYPE_CHECKING, @@ -868,7 +867,7 @@ def __init__( chunksize: int | None = None, block_size: int | None = None, cache_type: str | None = None, - max_workers: int = (cpu_count() or 1) * 5, + max_workers: int | None = None, **kwargs: Any, ) -> None: """Initialize the reader and the filesystem of the result files; it sends no request. @@ -884,7 +883,9 @@ def __init__( chunksize: The number of rows per chunk, if the result is read in chunks. block_size: The S3 read block size of the filesystem. cache_type: The default fsspec cache type of the filesystem. - max_workers: The maximum number of threads of the filesystem. + max_workers: The size of the request pool of the S3 filesystem + that reads the result files. If None, the + ``max_pool_connections`` of the connection's S3 client. **kwargs: The options of ``pandas.read_csv()`` or ``pandas.read_parquet()``. ``coerce_int96_timestamp_unit=None`` is dropped for an UNLOAD result. """ @@ -1216,7 +1217,7 @@ def __init__( chunksize: int | None = None, block_size: int | None = None, cache_type: str | None = None, - max_workers: int = (cpu_count() or 1) * 5, + max_workers: int | None = None, auto_optimize_chunksize: bool = False, result_set_type_hints: dict[str | int, str] | None = None, **kwargs, @@ -1248,7 +1249,9 @@ def __init__( the next block while pandas parses the current one. Without ``cache_type``, PyAthena's filesystem reads a Parquet file of a chunked UNLOAD result that is larger than 32 MiB without a cache. - max_workers: Maximum worker threads for parallel operations. + max_workers: The size of the request pool of the S3 filesystem + that reads the result files. If None, the + ``max_pool_connections`` of the connection's S3 client. auto_optimize_chunksize: Enable automatic chunksize determination for large files when chunksize is None. result_set_type_hints: Athena type signatures for complex-type columns, @@ -2426,6 +2429,10 @@ def close(self) -> None: super().close() self._closed_event.set() self._df_iter.close() + # Ends the idle threads of the request pool of the filesystem, which + # a read that is still running keeps until it ends. + if (fs := getattr(self, "_fs", None)) is not None: + fs.close() self._df = pd.DataFrame() self._df_iter = PandasDataFrameIterator(self._df, _no_trunc_date) self._row_iter = iter(()) diff --git a/pyathena/polars/async_cursor.py b/pyathena/polars/async_cursor.py index 7f893d730..55e3230f8 100644 --- a/pyathena/polars/async_cursor.py +++ b/pyathena/polars/async_cursor.py @@ -80,7 +80,7 @@ def __init__( block_size: int | None = None, cache_type: str | None = None, chunksize: int | None = None, - s3_max_workers: int = (cpu_count() or 1) * 5, + s3_max_workers: int | None = None, **kwargs, ) -> None: """Initialize an AsyncPolarsCursor. @@ -108,7 +108,10 @@ def __init__( chunksize: Number of rows per chunk for memory-efficient processing. If specified, data is loaded lazily in chunks for all data access methods including fetchone(), fetchmany(), and iter_chunks(). - s3_max_workers: Maximum worker threads per S3 file reader. + s3_max_workers: The size of the request pool of the S3 filesystem + that reads the result files. If None, the + ``max_pool_connections`` of the connection's S3 client (10 + unless set with ``s3_config``). **kwargs: Additional connection parameters. Example: diff --git a/pyathena/polars/cursor.py b/pyathena/polars/cursor.py index 6ef1ae4c9..cd4c1c8fa 100644 --- a/pyathena/polars/cursor.py +++ b/pyathena/polars/cursor.py @@ -4,7 +4,6 @@ import logging from collections.abc import Callable, Iterator -from multiprocessing import cpu_count from typing import ( TYPE_CHECKING, Any, @@ -81,7 +80,7 @@ def __init__( result_reuse_minutes: int = CursorIterator.DEFAULT_RESULT_REUSE_MINUTES, block_size: int | None = None, cache_type: str | None = None, - s3_max_workers: int = (cpu_count() or 1) * 5, + s3_max_workers: int | None = None, chunksize: int | None = None, **kwargs, ) -> None: @@ -105,7 +104,10 @@ def __init__( ``"background"``; see ``S3FileSystem``. Without ``cache_type``, PyAthena's filesystem reads a Parquet file of a chunked UNLOAD result that is larger than 32 MiB without a cache. - s3_max_workers: Maximum worker threads for parallel S3 operations. + s3_max_workers: The size of the request pool of the S3 filesystem + that reads the result files. If None, the + ``max_pool_connections`` of the connection's S3 client (10 + unless set with ``s3_config``). chunksize: Number of rows per chunk for memory-efficient processing. If specified, data is loaded lazily in chunks for all data access methods including fetchone(), fetchmany(), and iter_chunks(). diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index fdd672355..4412f9b54 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -20,7 +20,6 @@ from functools import partial from io import BytesIO, StringIO from itertools import islice, repeat -from multiprocessing import cpu_count from types import GeneratorType from typing import ( TYPE_CHECKING, @@ -474,7 +473,7 @@ def __init__( unload_location: str | None = None, block_size: int | None = None, cache_type: str | None = None, - max_workers: int = (cpu_count() or 1) * 5, + max_workers: int | None = None, chunksize: int | None = None, result_set_type_hints: dict[str | int, str] | None = None, **kwargs, @@ -495,7 +494,9 @@ def __init__( ``"background"``; see ``S3FileSystem``. Without ``cache_type``, PyAthena's filesystem reads a Parquet file of a chunked UNLOAD result that is larger than 32 MiB without a cache. - max_workers: Maximum number of worker threads. + max_workers: The size of the request pool of the S3 filesystem + that reads the result files. If None, the + ``max_pool_connections`` of the connection's S3 client. chunksize: Number of rows per chunk for memory-efficient processing. If specified, data is loaded lazily in chunks for all data access methods including fetchone(), fetchmany(), and iter_chunks(). diff --git a/pyathena/s3fs/result_set.py b/pyathena/s3fs/result_set.py index bed142ea5..b6518e268 100644 --- a/pyathena/s3fs/result_set.py +++ b/pyathena/s3fs/result_set.py @@ -255,3 +255,6 @@ def close(self) -> None: if self._csv_reader: self._csv_reader.close() self._csv_reader = None + if isinstance(fs := getattr(self, "_fs", None), (S3FileSystem, AioS3FileSystem)): + # Ends the idle threads of the request pool of the filesystem. + fs.close() diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 50038d249..3808fcc11 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -9,7 +9,6 @@ import lzma import os import re -import signal import sys import tempfile import threading @@ -42,7 +41,7 @@ from pyathena.filesystem.s3 import CompressedBuffer, S3File, S3FileSystem from pyathena.filesystem.s3_core import S3Core, S3DeleteBatch, S3MultipartCopyPlan from pyathena.filesystem.s3_errors import S3ClientError -from pyathena.filesystem.s3_executor import S3ThreadPool +from pyathena.filesystem.s3_executor import S3ThreadPool, _S3PoolScope from pyathena.filesystem.s3_object import S3MultipartUpload, S3Object, S3ObjectType, S3StorageClass from pyathena.filesystem.s3_path import S3Path from pyathena.filesystem.s3_writer import S3MultipartStateError, _S3WriteSession @@ -55,7 +54,14 @@ MULTIPART_COPY_KWARGS, MULTIPART_COPY_SIZE, ThrottledOperation, + forking, + interrupt_each_opcode, + mock_connection, + overlapping_calls, + poll_until, retried, + run_in_forked_child, + sigint_when, stub_multipart_copy, ) @@ -86,6 +92,7 @@ def _make_fs(): # The requests of the core and of the filesystem go to one mock. fs._call = fs._core.call = mock.MagicMock() fs.max_workers = 4 + fs._request_pool = S3ThreadPool(fs.max_workers) fs.default_block_size = S3FileSystem.DEFAULT_BLOCK_SIZE fs.allow_bucket_creation = False fs.allow_bucket_deletion = False @@ -2680,6 +2687,137 @@ def test_open_max_workers(self): with fs.open("s3://bucket/key", "wb") as f: assert f.max_workers == fs.max_workers + @pytest.mark.parametrize( + ("max_workers", "config_kwargs", "pool_size", "pool_connections"), + [ + (None, {}, 10, 10), + (None, {"max_pool_connections": 25}, 25, 25), + (4, {}, 4, 4), + (4, {"max_pool_connections": 25}, 4, 25), + ], + ) + def test_request_pool_size(self, max_workers, config_kwargs, pool_size, pool_connections): + # GH-1332: the request pool defaults to the connection pool of the S3 + # client, and a client that the filesystem builds takes its + # connection pool from max_workers unless config_kwargs sets it. + # max_workers used to default to cpu_count() * 5. + fs = S3FileSystem( + max_workers=max_workers, + config_kwargs=config_kwargs, + key="dummy", + secret="dummy", + region_name="us-east-1", + skip_instance_cache=True, + ) + assert fs.max_workers == pool_size + assert fs._request_pool.max_workers == pool_size + assert fs._client.meta.config.max_pool_connections == pool_connections + + def test_request_pool_size_of_connection(self): + # GH-1332: with a connection, the request pool takes the connection + # pool of its S3 client. + fs = S3FileSystem(connection=mock_connection(7), skip_instance_cache=True) + assert fs.max_workers == 7 + + def test_files_share_request_pool(self): + # GH-1332: the reads of the files of a filesystem run in its request + # pool, which bounds them together. Each file used to have a pool of + # max_workers threads of its own. + fs = self._make_fs() + fs._request_pool = S3ThreadPool(2) + fs.info = mock.MagicMock(return_value=self._file_object("key")) + fs.info.return_value.size = 64 + call, state = overlapping_calls(2) + + def get_object(path, range_, **kwargs): + call() + return bytes(range_[1] - range_[0]) + + fs.core.get_object = get_object + + def read(_): + with fs.open( + "s3://bucket/key", "rb", cache_type="none", block_size=16, max_workers=2 + ) as f: + return f.read() + + with ThreadPoolExecutor(max_workers=2) as readers: + assert list(readers.map(read, range(2))) == [bytes(64)] * 2 + assert state["peak"] == 2 + # Closing the files leaves the pool running. + assert fs._request_pool.submit(sum, [1, 2]).result(EVENT_TIMEOUT) == 3 + fs._request_pool.shutdown() + + def test_file_without_executor_shares_request_pool(self): + # GH-1332 (F1): a file constructed without an executor or max_workers + # takes the max_workers of the filesystem and runs its reads in the + # request pool of the filesystem; it used to start a pool of + # cpu_count() * 5 threads of its own. + fs = self._make_fs() + fs._request_pool = S3ThreadPool(2) + fs.info = mock.MagicMock(return_value=self._file_object("key")) + fs.info.return_value.size = 64 + call, state = overlapping_calls(2) + + def get_object(path, range_, **kwargs): + call() + return bytes(range_[1] - range_[0]) + + fs.core.get_object = get_object + with S3File(fs, "s3://bucket/key", "rb", cache_type="none", block_size=16) as f: + assert f.max_workers == fs.max_workers + assert f.read() == bytes(64) + # The file's max_workers is 4, but the pool of the filesystem has 2. + assert state["peak"] == 2 + fs._request_pool.shutdown() + + @pytest.mark.parametrize("exit_with", ["close", "with"]) + def test_close_releases_threads(self, exit_with): + # GH-1332: close(), also at the end of a with block, ends the threads + # of the request pool, and the filesystem stays usable with new ones. + fs = self._make_fs() + threads = [] + + def record(): + threads.append(threading.current_thread()) + + if exit_with == "close": + fs._request_pool.submit(record).result(EVENT_TIMEOUT) + fs.close() + else: + with fs: + fs._request_pool.submit(record).result(EVENT_TIMEOUT) + threads[0].join(EVENT_TIMEOUT) + assert not threads[0].is_alive() + + fs._request_pool.submit(record).result(EVENT_TIMEOUT) + assert threads[1] is not threads[0] + fs.close() + + @forking + def test_file_inherited_across_fork(self): + # GH-1332: in a forked child, a file that the parent opened raises on + # use, and closing it, also at garbage collection, only marks it + # closed, without uploading or aborting anything: the upload belongs + # to the parent. + fs = self._make_fs() + fs.default_cache_type = "bytes" + f = fs.open("s3://bucket/key", "wb") + f.write(b"x") + + def child(): + for use in (lambda: f.write(b"y"), f.flush, f.commit, f.discard, f.read): + with pytest.raises(RuntimeError, match="parent of this forked process"): + use() + requests = fs._call.call_count + len(fs._core.client.method_calls) + f.close() + return f.closed and fs._call.call_count + len(fs._core.client.method_calls) == requests + + try: + assert run_in_forked_child(child) == 0 + finally: + f._close_without_commit() + def test_open_version_id(self): fs = self._make_fs() fs.default_cache_type = "bytes" @@ -3495,6 +3633,138 @@ def interrupted_submit(fn, *args, **kwargs): assert sorted(read) == [0, 16] assert operation.calls == 2 + def test_read_interrupted_at_each_opcode_frees_ranges(self): + # GH-1332 (F5): an interrupt at any point of a parallel read, including + # after its wait returns, cancels the reads that have not started and + # frees the data of the others while the caller holds the interrupt. + # The data used to stay in the frame that the traceback keeps after + # the wait. + def case(interrupt): + fs = self._make_fs() + fs.info = mock.MagicMock(return_value=self._file_object("key")) + fs.info.return_value.size = 64 + read = [] + + def get_object(path, range_, **kwargs): + # Unlike bytes, an array can be weakly referenced. + data = array.array("B", bytes(range_[1] - range_[0])) + read.append(weakref.ref(data)) + return data + + fs.core.get_object = get_object + f = fs.open("s3://bucket/key", "rb", cache_type="none", block_size=16, max_workers=4) + try: + result = interrupt(f._fetch_range, 0, 64) + f.close() + if not interrupt.fired: + assert result == bytes(64) + return + assert interrupt.error.__traceback__ is not None + # The threads leave the frames of the reads right after the + # close has waited for them. + poll_until( + lambda: all(ref() is None for ref in read), "the data of a read was kept" + ) + finally: + f.close() + fs._request_pool.shutdown() + + assert interrupt_each_opcode(S3File._fetch_range, case) > 0 + + def test_read_interrupted_twice_frees_ranges(self): + # GH-1332 (F5): a second interrupt at any point of the cleanup after + # an interrupt of the wait, whose error then keeps the first as its + # context, also frees the data of the ranges that were read. The + # frame of the wait keeps the futures, as concurrent.futures.wait() + # does. + def interrupted_wait(futures, **kwargs): + # A set of its own, as concurrent.futures.wait() makes. + held = set(futures) + wait(held) + raise KeyboardInterrupt + + def case(interrupt): + fs = self._make_fs() + fs.info = mock.MagicMock(return_value=self._file_object("key")) + fs.info.return_value.size = 64 + read = [] + + def get_object(path, range_, **kwargs): + # Unlike bytes, an array can be weakly referenced. + data = array.array("B", bytes(range_[1] - range_[0])) + read.append(weakref.ref(data)) + return data + + fs.core.get_object = get_object + f = fs.open("s3://bucket/key", "rb", cache_type="none", block_size=16, max_workers=4) + try: + with mock.patch("pyathena.filesystem.s3.wait", interrupted_wait): + try: + interrupt(f._fetch_range, 0, 64) + error = interrupt.error + except KeyboardInterrupt as e: + error = e + f.close() + # Held, as a caller that handles the interrupt does. + assert error.__traceback__ is not None + poll_until( + lambda: all(ref() is None for ref in read), "the data of a read was kept" + ) + finally: + f.close() + fs._request_pool.shutdown() + + assert interrupt_each_opcode(S3File._fetch_range, case) > 0 + + @pytest.mark.skipif( + threading.current_thread() is not threading.main_thread(), + reason="SIGINT interrupts the main thread.", + ) + def test_read_interrupted_by_signal_frees_ranges(self): + # GH-1332 (F5): a real interrupt of a parallel read while it waits for + # its ranges cancels those that have not started, and frees the data + # of the others once they end, while the caller holds the interrupt. + # concurrent.futures.wait() leaves its waiter on the futures when it + # is interrupted, which forms a reference cycle with them, so the + # data is freed by the cyclic garbage collector, not when they end. + fs = self._make_fs() + fs.info = mock.MagicMock(return_value=self._file_object("key")) + fs.info.return_value.size = 64 + read = [] + running = threading.Semaphore(0) + release = threading.Event() + + def get_object(path, range_, **kwargs): + running.release() + release.wait(EVENT_TIMEOUT) + # Unlike bytes, an array can be weakly referenced. + data = array.array("B", bytes(range_[1] - range_[0])) + read.append(weakref.ref(data)) + return data + + fs.core.get_object = get_object + f = fs.open("s3://bucket/key", "rb", cache_type="none", block_size=16, max_workers=2) + try: + with ( + sigint_when(lambda: running._value >= 2), + pytest.raises(KeyboardInterrupt) as info, + ): + f.read() + release.set() + f.close() + assert info.value.__traceback__ is not None + poll_until( + lambda: all(ref() is None for ref in read), + "the data of a read was kept", + collect=True, + ) + # The two ranges that had not started were cancelled. + assert len(read) == 2 + finally: + release.set() + f.close() + fs._request_pool.shutdown() + def test_interrupted_read_frees_ranges(self): # GH-1186: an interrupted parallel read cancels the reads that have # not started and frees the data of the others, also while the @@ -3589,12 +3859,11 @@ def test_copy_object_with_multipart_upload_interrupted_creation(self): fs = self._make_fs() started = threading.Event() waiting = threading.Event() - interrupted = threading.Event() def create_multipart_upload(*args, **kw): started.set() # Still running when the interrupt arrives, which releases it. - interrupted.wait(30) + poll_until(lambda: signals, "the interrupt was not raised", timeout=30) return S3MultipartUpload({"Bucket": "bucket", "Key": "dst", "UploadId": "uploadid"}) fs.core.create_multipart_upload = mock.MagicMock(side_effect=create_multipart_upload) @@ -3610,52 +3879,20 @@ def wait_for_ready(session): waiting.set() return waiter - handled = [] - - def handle_interrupt(signum, frame): - # Only the first interrupt is raised; repeated signals are ignored. - # The list is checked instead of the event, whose set() takes a - # lock that a nested call of this handler would wait for forever. - if not handled: - handled.append(signum) - interrupted.set() - raise KeyboardInterrupt - - def interrupt(): + with ( + mock.patch.object(_S3WriteSession, "ready", wait_for_ready), # Sent only while the creation is running and the copy waits for # it, which the creation cannot stop doing before the interrupt. - if started.wait(5) and waiting.wait(5): - # A signal that arrives just before the main thread blocks on - # the creation can miss waking it, so it is repeated until handled. - while True: - signal.pthread_kill(threading.main_thread().ident, signal.SIGINT) - if interrupted.wait(0.1): - break - - thread = threading.Thread(target=interrupt, daemon=True) - previous_handler = signal.signal(signal.SIGINT, handle_interrupt) - try: - thread.start() - with ( - mock.patch.object(_S3WriteSession, "ready", wait_for_ready), - pytest.raises(KeyboardInterrupt), - ): - fs._copy_object_with_multipart_upload( - S3Path("bucket", "src"), - S3Path("bucket", "dst"), - MetadataDirective="REPLACE", - TaggingDirective="REPLACE", - AnnotationDirective="EXCLUDE", - ) - finally: - # A late interrupt is ignored instead of reaching a later test. - signal.signal(signal.SIGINT, lambda signum, frame: None) - started.set() - waiting.set() - interrupted.set() - if thread.ident is not None: - thread.join() - signal.signal(signal.SIGINT, previous_handler) + sigint_when(lambda: started.is_set() and waiting.is_set()) as signals, + pytest.raises(KeyboardInterrupt), + ): + fs._copy_object_with_multipart_upload( + S3Path("bucket", "src"), + S3Path("bucket", "dst"), + MetadataDirective="REPLACE", + TaggingDirective="REPLACE", + AnnotationDirective="EXCLUDE", + ) fs.core.abort_multipart_upload.assert_called_once() (upload,) = fs.core.abort_multipart_upload.call_args.args @@ -6237,6 +6474,12 @@ def _make_mock_fs(): ): setattr(fs.core, name, mock.MagicMock()) fs._get_lookup_kwargs.side_effect = S3FileSystem._get_lookup_kwargs + # The request pool, whose scopes the files take by default. + fs.max_workers = 4 + fs._request_pool = S3ThreadPool(fs.max_workers) + fs._create_executor.side_effect = lambda max_workers: _S3PoolScope( + fs._request_pool, max_workers + ) return fs @staticmethod diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index f1bf51820..f6f5eaee8 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -23,7 +23,7 @@ from pyathena.filesystem.s3 import S3FileSystem from pyathena.filesystem.s3_async import AioS3File, AioS3FileSystem, AioS3StreamedFile from pyathena.filesystem.s3_core import S3Core, S3MultipartCopyPlan -from pyathena.filesystem.s3_executor import S3ThreadPool +from pyathena.filesystem.s3_executor import S3ThreadPool, _S3PoolScope from pyathena.filesystem.s3_object import ( S3Metadata, S3MultipartUpload, @@ -45,6 +45,7 @@ await_bounded, cancel_when_started, forking, + mock_connection, overlapping_calls, retried, run_bounded, @@ -135,7 +136,7 @@ def test_read_in_forked_child(self): # parent opens reads its ranges in parallel. Its open() raised there # since GH-1336, and before that, its reads waited forever for the # event loop of the parent. - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) block_size = 2**10 size = block_size * 4 fs.core.get_object = mock.MagicMock( @@ -229,7 +230,7 @@ async def test_copy_object_with_multipart_upload_failed_annotation(self, release @pytest.mark.asyncio async def test_cp_file_failed_multipart_copy_invalidates_cache(self): # GH-973: see TestS3FileSystem.test_cp_file_failed_multipart_copy_invalidates_cache. - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) fs._info = mock.AsyncMock( return_value=S3Object( init={"ContentLength": S3Core.MULTIPART_UPLOAD_MAX_PART_SIZE + 1}, @@ -251,7 +252,7 @@ async def test_copy_object_with_multipart_upload_max_workers(self): # 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) + fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) sync_fs = fs._sync_fs sync_fs.core.create_multipart_upload = mock.MagicMock( return_value=S3MultipartUpload( @@ -306,7 +307,7 @@ 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) + fs = AioS3FileSystem(connection=mock_connection(), max_workers=2, skip_instance_cache=True) sync_fs = fs._sync_fs sync_fs.core.create_multipart_upload = mock.MagicMock( return_value=S3MultipartUpload( @@ -349,7 +350,7 @@ async def test_blocking_calls_bypass_default_executor(self): # 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 = AioS3FileSystem(connection=mock_connection(), 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. @@ -628,13 +629,15 @@ def test_invalidate_cache(self, path): @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. + # None takes the connection pool of the S3 client for the threads of + # the filesystem, its request pool, and as the default max_workers of + # its files and copies (GH-1332; it used to take the default of + # ThreadPoolExecutor). The files and copies used to fail with it. fs = AioS3FileSystem( - connection=mock.MagicMock(), max_workers=None, skip_instance_cache=True + connection=mock_connection(7), max_workers=None, skip_instance_cache=True ) sync_fs = fs._sync_fs + assert fs._thread_pool.max_workers == sync_fs._request_pool.max_workers == 7 sync_fs.info = mock.MagicMock(return_value={"name": "bucket/key"}) assert await fs._info("s3://bucket/key") == {"name": "bucket/key"} @@ -677,7 +680,7 @@ async def test_copy_object_with_multipart_upload_cancelled(self, cancellations): # aborts the upload, and is re-raised, as S3FileSystem does on an # interrupt. A repeated cancellation returns without stopping the # cleanup. - fs = AioS3FileSystem(connection=mock.MagicMock(), max_workers=2, skip_instance_cache=True) + fs = AioS3FileSystem(connection=mock_connection(), max_workers=2, skip_instance_cache=True) sync_fs = fs._sync_fs sync_fs.core.create_multipart_upload = mock.MagicMock( return_value=S3MultipartUpload( @@ -755,7 +758,7 @@ def abort_multipart_upload(*args, **kwargs): async def test_close_releases_threads(self, exit_with): # GH-1209: close(), also at the end of a with block, ends the threads # of the filesystem, which stays usable with new ones. - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) threads = [] def exists(path, **kwargs): @@ -779,12 +782,24 @@ def exists(path, **kwargs): assert threads[1] is not threads[0] fs.close() + def test_close_releases_request_threads(self): + # GH-1332: close() also ends the threads of the request pool of the + # internal S3FileSystem, which the files and parallel requests share. + fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) + threads = [] + fs._sync_fs._request_pool.submit(lambda: threads.append(threading.current_thread())).result( + EVENT_TIMEOUT + ) + fs.close() + threads[0].join(EVENT_TIMEOUT) + assert not threads[0].is_alive() + @pytest.mark.parametrize("cancel", [False, True]) @pytest.mark.asyncio async def test_close_during_multipart_copy(self, cancel): # GH-1209: a copy that runs at close() still schedules its remaining # requests and its completion or abort. - fs = AioS3FileSystem(connection=mock.MagicMock(), max_workers=1, skip_instance_cache=True) + fs = AioS3FileSystem(connection=mock_connection(), max_workers=1, skip_instance_cache=True) sync_fs = fs._sync_fs sync_fs.core.create_multipart_upload = mock.MagicMock( return_value=S3MultipartUpload( @@ -847,7 +862,7 @@ def upload_part_copy(**kw): async def test_copy_object_with_multipart_upload_cancelled_completion(self, completion_fails): # GH-1046: a cancellation during CompleteMultipartUpload waits for it, # aborts the upload only if it failed, and is re-raised. - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) sync_fs = fs._sync_fs sync_fs.core.create_multipart_upload = mock.MagicMock( return_value=S3MultipartUpload( @@ -926,7 +941,7 @@ async def test_copy_object_with_multipart_upload_cancelled_creation(self, creati # A cancellation during CreateMultipartUpload waits for it, aborts # the upload that it created before any part is copied, and is # re-raised. The created upload used to be left incomplete. - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) sync_fs = fs._sync_fs events = [] started = threading.Event() @@ -996,7 +1011,7 @@ async def test_copy_object_with_multipart_upload_cancelled_queued_completion(sel # GH-1186: a cancellation cancels a CompleteMultipartUpload request # that waits for a thread of the filesystem, and aborts the upload. The # request used to run after the cancellation and complete the copy. - fs = AioS3FileSystem(connection=mock.MagicMock(), max_workers=1, skip_instance_cache=True) + fs = AioS3FileSystem(connection=mock_connection(), max_workers=1, skip_instance_cache=True) sync_fs = fs._sync_fs sync_fs.core.create_multipart_upload = mock.MagicMock( return_value=S3MultipartUpload( @@ -1080,7 +1095,7 @@ def test_copy_object_with_multipart_upload_interrupted_while_scheduling_completi # completion is being scheduled, before or after it is queued, keeps it # from running and aborts the upload, instead of leaving the upload # incomplete with a completion that nothing sends. - fs = AioS3FileSystem(connection=mock.MagicMock(), max_workers=1, skip_instance_cache=True) + fs = AioS3FileSystem(connection=mock_connection(), max_workers=1, skip_instance_cache=True) sync_fs = fs._sync_fs sync_fs.core.create_multipart_upload = mock.MagicMock( return_value=S3MultipartUpload( @@ -1137,10 +1152,10 @@ def interrupt_completion(pool, future, fn, *args, **kwargs): async def test_copy_object_with_multipart_upload_failed_part_with_threads_taken(self): # GH-1230: a part copy that fails after the last one was submitted # stops the copy at once, also when a part copy that it frees the only - # thread of the filesystem for retries; a completion queued behind it + # thread of the request pool for retries; a completion queued behind it # would notice the failure only after its backoff wait of at least # EVENT_TIMEOUT seconds. - fs = AioS3FileSystem(connection=mock.MagicMock(), max_workers=1, skip_instance_cache=True) + fs = AioS3FileSystem(connection=mock_connection(), max_workers=1, skip_instance_cache=True) sync_fs = fs._sync_fs plan = S3MultipartCopyPlan( source=S3Path("bucket", "src"), @@ -1157,10 +1172,10 @@ async def test_copy_object_with_multipart_upload_failed_part_with_threads_taken( operation = ThrottledOperation(operation_name="UploadPartCopy") throttled = retried(operation, RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT)) submitted = threading.Event() - submit_to = S3ThreadPool.submit_to + submit_to = _S3PoolScope.submit_to - def record_part(pool, future, fn, *args, **kwargs): - submit_to(pool, future, fn, *args, **kwargs) + def record_part(scope, future, fn, *args, **kwargs): + submit_to(scope, future, fn, *args, **kwargs) if kwargs.get("part_number") == 2: submitted.set() @@ -1177,7 +1192,7 @@ def upload_part_copy(part_number, **kwargs): sync_fs.core.upload_part_copy = mock.MagicMock(side_effect=upload_part_copy) with ( - mock.patch.object(S3ThreadPool, "submit_to", autospec=True, side_effect=record_part), + mock.patch.object(_S3PoolScope, "submit_to", autospec=True, side_effect=record_part), pytest.raises(PermissionError, match="Access Denied"), ): # The retries stop, or part 2 is cancelled, before the bound. @@ -1196,7 +1211,7 @@ async def test_copy_object_with_multipart_upload_failed_annotation_stops_running # GH-1203: a failed annotation copy stops the retries of the running # ones, whose backoff wait would take at least EVENT_TIMEOUT seconds, # and its error is raised, whichever annotation comes first. - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) sync_fs = fs._sync_fs plan = S3MultipartCopyPlan( source=S3Path("bucket", "src"), @@ -1236,7 +1251,7 @@ def copy_object_annotation(name, *args, **kwargs): async def test_copy_object_with_multipart_upload_failed_annotation_skips_queued_copies(self): # GH-1203: the annotation copies that wait for a thread of the filesystem # when one fails are not sent. - fs = AioS3FileSystem(connection=mock.MagicMock(), max_workers=1, skip_instance_cache=True) + fs = AioS3FileSystem(connection=mock_connection(), max_workers=1, skip_instance_cache=True) sync_fs = fs._sync_fs plan = S3MultipartCopyPlan( source=S3Path("bucket", "src"), @@ -1276,10 +1291,11 @@ def copy_object_annotation(name, *args, **kwargs): @pytest.mark.asyncio async def test_copy_object_with_multipart_upload_cancelled_queued_parts(self): # GH-1186: a cancellation cancels the part copies that wait for a - # thread of the filesystem, whose threads bound all its copies, before - # it waits for the running ones and aborts the upload. The queued - # part copies used to run before the abort. - fs = AioS3FileSystem(connection=mock.MagicMock(), max_workers=1, skip_instance_cache=True) + # thread of the request pool, whose threads bound all the copies and + # files of the filesystem, before it waits for the running ones and + # aborts the upload. The queued part copies used to run before the + # abort. + fs = AioS3FileSystem(connection=mock_connection(), max_workers=1, skip_instance_cache=True) sync_fs = fs._sync_fs sync_fs.core.create_multipart_upload = mock.MagicMock( return_value=S3MultipartUpload( @@ -1306,16 +1322,16 @@ def upload_part_copy(**kw): sync_fs.core.abort_multipart_upload = mock.MagicMock( side_effect=lambda upload, **params: events.append("abort") ) - submit_to = fs._thread_pool.submit_to + submit_to = _S3PoolScope.submit_to queued = threading.Event() - def record_part(future, fn, *args, **kwargs): - submit_to(future, fn, *args, **kwargs) + def record_part(scope, future, fn, *args, **kwargs): + submit_to(scope, future, fn, *args, **kwargs) if kwargs.get("part_number") == 2: queued.set() - fs._thread_pool.submit_to = record_part - + patch = mock.patch.object(_S3PoolScope, "submit_to", autospec=True, side_effect=record_part) + patch.start() task = asyncio.ensure_future( fs._copy_object_with_multipart_upload( S3Path("bucket", "src"), @@ -1338,6 +1354,7 @@ def record_part(future, fn, *args, **kwargs): raise finally: release.set() + patch.stop() with pytest.raises(asyncio.CancelledError): await await_bounded(task) @@ -1361,7 +1378,7 @@ def test_copy_object_with_multipart_upload_loop_shutdown(self, stage, expected, # also cancelled the tasks of the requests and of the cleanup, so the # cleanup left an upload created during the cancellation, aborted # before the running parts ended, or aborted during the completion. - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) sync_fs = fs._sync_fs events = [] lock = threading.Lock() @@ -1470,7 +1487,7 @@ def write(): def test_transaction_pipe_put_file_create_existing(self, tmp_path): # GH-972: in a transaction, put_file(mode="create") also raises, when # the file is opened, for an existing object. - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) fs._sync_fs.exists = mock.MagicMock(return_value=True) fs._sync_fs._call = fs._sync_fs._core.call = mock.MagicMock() local = tmp_path / "local" @@ -1495,7 +1512,7 @@ def test_transaction_pipe_file_compression(self, path, compression): # GH-1037: the value is compressed before it is uploaded, also in a # transaction of this filesystem, whose pipe_file() is its own; outside # one, it is the synchronous filesystem's. - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) put_object = fs._sync_fs.core.put_object = mock.MagicMock() with fs.transaction: @@ -1510,7 +1527,7 @@ def test_transaction_pipe_file_write(self): # and a failed write does not replace the object with an empty one # when the transaction commits. GH-1211: the commit of the failed # write raises, and fsspec discards the rest of the transaction. - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) put_object = fs._sync_fs.core.put_object = mock.MagicMock() def write_in_transaction(): @@ -1539,7 +1556,7 @@ def test_transaction_put_file_failed_write(self, tmp_path, error): # GH-1211: the commit of the failed write raises, and fsspec discards # the rest of the transaction. A local file that cannot be read opens # no remote file. - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) put_object = fs._sync_fs.core.put_object = mock.MagicMock() local = tmp_path / "local" local.write_bytes(b"a") @@ -1574,7 +1591,7 @@ def test_transaction_pipe_put_file_exceeding_max_parts(self, tmp_path, kwargs): # GH-953: in a transaction, as outside one, pipe_file() and put_file() # reject data that does not fit in the maximum number of parts before # opening the file. - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) fs._sync_fs.core.MULTIPART_UPLOAD_MAX_PARTS = 3 fs._sync_fs.default_block_size = 4 fs._sync_fs._call = fs._sync_fs._core.call = mock.MagicMock() @@ -1595,7 +1612,7 @@ async def test_rm_requests(self): # GH-971: _rm() sent the keys of every bucket to the bucket of the # first path and dropped the DeleteObjects parameters. The async # filesystem sends the request of each bucket in parallel. - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) sync_fs = fs._sync_fs sync_fs._call = sync_fs._core.call = mock.MagicMock(return_value={}) @@ -1610,7 +1627,7 @@ async def test_rm_requests(self): @pytest.mark.asyncio async def test_rm_request_error_keeps_errors(self): - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) def call(method, **request): if request["Bucket"] == "b1": @@ -1628,7 +1645,7 @@ def call(method, **request): async def test_rm_cancel_invalidates_cache_after_each_request(self): # The request threads keep running after _rm() is cancelled, so each # request invalidates the cache of its objects when it finishes. - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) started = {"b1": threading.Event(), "b2": threading.Event()} release = {"b1": threading.Event(), "b2": threading.Event()} @@ -1705,7 +1722,7 @@ def test_put_file_in_transaction_open_parameters(self, tmp_path, mode, open_mode # GH-969: in a transaction, the open() parameters of put_file() go to # open(), and the other parameters, also in s3_additional_kwargs, to # S3. GH-972: fsspec's mode argument selects the mode of the file. - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) fs.open = mock.MagicMock() lpath = tmp_path / "data.csv" lpath.write_bytes(b"a") @@ -1743,7 +1760,7 @@ def test_put_file_in_transaction_open_parameters(self, tmp_path, mode, open_mode async def test_cp_file_directory(self): # GH-1008: recursive copy() passes the directories, which used to be # sent to CopyObject and fail with NoSuchKey. - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) fs._info = mock.AsyncMock(return_value=S3FileSystem._directory_object("bucket", "src")) fs._sync_fs._call = fs._sync_fs._core.call = mock.MagicMock() @@ -1774,7 +1791,7 @@ async def test_copy_version(self): async def test_expand_path_glob_lists_stem_prefix(self): # Unversioned globs keep fsspec's async expansion, which lists only # the keys that start with the stem before the first wildcard. - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) sync_fs = fs._sync_fs sync_fs._call = sync_fs._core.call = mock.MagicMock( side_effect=lambda method, **request: {"Contents": [], "KeyCount": 0} @@ -1793,7 +1810,7 @@ async def test_expand_path_glob_lists_stem_prefix(self): async def test_expand_path_version_maxdepth(self): # A maxdepth below 1 is rejected, also for a version, which fsspec's # expansion does not see. - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) with pytest.raises(ValueError): # noqa: PT011 await fs._expand_path("s3://bucket/b?versionId=v1", maxdepth=0) @@ -1812,7 +1829,7 @@ async def test_expand_path_version_maxdepth(self): ) @pytest.mark.asyncio async def test_get_version(self, tmp_path, destination, expected): - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) fs._sync_fs.isdir = mock.MagicMock(return_value=False) fs._get_file = mock.AsyncMock() @@ -1827,7 +1844,7 @@ async def test_get_version(self, tmp_path, destination, expected): async def test_mv(self): # GH-1008: the files are copied in parallel, the directories are # skipped, and only the copied sources are deleted. - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) started = asyncio.Event() copies = [] @@ -1863,7 +1880,7 @@ async def copy_file(path1, path2, **kwargs): async def test_mv_copy_failure(self): # A failed copy is raised after the other copies have finished, and # nothing is deleted. - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) finished = [] async def copy_file(path1, path2, **kwargs): @@ -1963,7 +1980,7 @@ def upload_part_copy(**kw): def test_internal_file_system_not_cached(self): # GH-978: the internal S3FileSystem was kept in the fsspec instance # cache, so skip_instance_cache=True instances shared it. - connection = mock.MagicMock() + connection = mock_connection() fs1 = AioS3FileSystem(connection=connection, skip_instance_cache=True) fs2 = AioS3FileSystem(connection=connection, skip_instance_cache=True) assert fs1._sync_fs is not fs2._sync_fs @@ -2105,7 +2122,7 @@ class TestAioS3File: def test_open_max_workers(self, max_workers): # GH-1158: None takes the default of the filesystem. It used to raise # TypeError when the executor of the file was created. - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) with fs.open("s3://bucket/key", "wb", max_workers=max_workers) as f: assert isinstance(f, AioS3File) assert f.max_workers == (max_workers or fs._sync_fs.max_workers) @@ -2117,7 +2134,7 @@ async def test_open_parallel_requests(self, asynchronous): # reads. GH-977: a filesystem created with asynchronous=True has no # event loop of its own and used to fail to run them. fs = AioS3FileSystem( - connection=mock.MagicMock(), asynchronous=asynchronous, skip_instance_cache=True + connection=mock_connection(), asynchronous=asynchronous, skip_instance_cache=True ) block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE size = block_size * 4 @@ -2182,14 +2199,14 @@ def read(): assert state["peak"] == 2 def test_open_invalid_max_workers(self): - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) with pytest.raises(ValueError, match="max_workers must be greater than 0"): fs.open("s3://bucket/key", "wb", max_workers=0) def test_open_lookup_parameters(self): # GH-1004: the lookup of the file sends its lookup parameters, and # looks up the version that is opened. - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) fs._sync_fs.info = mock.MagicMock( return_value=S3Object( init={"ContentLength": 4}, @@ -2214,7 +2231,7 @@ def test_open_lookup_parameters(self): class TestAioS3StreamedFile: @pytest.fixture def fs(self): - fs = AioS3FileSystem(connection=mock.MagicMock(), skip_instance_cache=True) + fs = AioS3FileSystem(connection=mock_connection(), skip_instance_cache=True) core = fs._sync_fs.core core.create_multipart_upload = mock.MagicMock( return_value=S3MultipartUpload( @@ -2231,6 +2248,46 @@ def fs(self): core.operation_params = mock.MagicMock(side_effect=lambda method, params: dict(params)) return fs + @forking + def test_file_inherited_across_fork(self, fs): + # GH-1332: in a forked child, a file that the parent opened raises on + # use, and closing it only marks it closed, without uploading or + # aborting anything: the upload belongs to the parent. + core = fs._sync_fs.core + f = asyncio.run(fs.open_async("s3://bucket/key", "wb")) + asyncio.run(f.write(b"x")) + + async def use(): + for call in (lambda: f.write(b"y"), f.flush, lambda: f.read(0)): + with pytest.raises(RuntimeError, match="parent of this forked process"): + await call() + await f.close() + + def child(): + asyncio.run(use()) + return ( + f.closed + and not core.put_object.called + and not core.complete_multipart_upload.called + and not core.abort_multipart_upload.called + ) + + try: + assert run_in_forked_child(child) == 0 + finally: + asyncio.run(f._close_without_commit()) + core.put_object.assert_not_called() + + @pytest.mark.asyncio + async def test_default_max_workers(self, fs): + # GH-1332: a streamed file constructed without max_workers takes that + # of the filesystem, as open_async() passes it; None used to raise. + f = await asyncio.to_thread(AioS3StreamedFile, fs, "s3://bucket/key", "wb") + try: + assert f.max_workers == fs._sync_fs.max_workers + finally: + await f._close_without_commit() + @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 @@ -2242,15 +2299,15 @@ async def test_write_parts_in_flight(self, fs): active: list[int] = [] uploads: dict[int, Future[S3MultipartUploadPart]] = {} in_flight: list[int] = [] - submit_to = S3ThreadPool.submit_to + submit_to = _S3PoolScope.submit_to - def record_upload(pool, future, fn, *args, **kwargs): + def record_upload(scope, 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) + submit_to(scope, future, fn, *args, **kwargs) def upload_part(**kw): part_number = kw["part_number"] @@ -2283,7 +2340,7 @@ async def release(): 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(S3ThreadPool, "submit_to", autospec=True, side_effect=record_upload): + with mock.patch.object(_S3PoolScope, "submit_to", autospec=True, side_effect=record_upload): await f.write(b"a" * block_size * 4) written.set() await f.close() @@ -2396,8 +2453,8 @@ def blocking_init(self, *args, **kwargs): with mock.patch.object(AioS3StreamedFile, "__init__", blocking_init): task = asyncio.create_task(fs.open_async("s3://bucket/key", mode)) await cancel_when_started(task, started, release) - # The pool is shut down last. - await wait_until(lambda: files[0]._pool._shutdown) + # The requests of the file are shut down last. + await wait_until(lambda: files[0]._requests._shutdown) assert files[0].closed with warnings.catch_warnings(record=True) as caught: @@ -2503,12 +2560,12 @@ async def test_commit_failed_part_first(self, fs): } queued = threading.Event() - submit = S3ThreadPool.submit + submit_to = S3ThreadPool.submit_to - def record_commit(pool, fn, *args, **kwargs): + def record_commit(pool, future, fn, *args, **kwargs): if getattr(fn, "__name__", None) == "_complete_upload": queued.set() - return submit(pool, fn, *args, **kwargs) + return submit_to(pool, future, fn, *args, **kwargs) def upload_part(part_number, **kwargs): try: @@ -2531,7 +2588,7 @@ def upload_part(part_number, **kwargs): await f.write(b"a" * block_size * 3) started = time.monotonic() with ( - mock.patch.object(S3ThreadPool, "submit", autospec=True, side_effect=record_commit), + mock.patch.object(S3ThreadPool, "submit_to", autospec=True, side_effect=record_commit), pytest.raises(PermissionError, match="Access Denied"), ): await f.close() @@ -2629,7 +2686,9 @@ def abort_multipart_upload(*args, **kwargs): async def test_commit_cancelled(self, fs, cancelled): # A commit that has started completes, and the upload is not aborted, # when the task that closes the file is cancelled, or when the block - # is left with the exception of a cancelled final flush. + # is left with the exception of a cancelled final flush. GH-1332 (F2): + # the close waits for the commit before it raises the cancellation; + # it used to return while the commit ran. block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE started = threading.Event() release = threading.Event() @@ -2660,71 +2719,123 @@ async def flush_and_exit(): flush.cancel() await flush + scenario = asyncio.ensure_future(close() if cancelled == "close" else flush_and_exit()) try: + assert await asyncio.to_thread(started.wait, EVENT_TIMEOUT) + # The cancelled close waits for the running commit. + await asyncio.sleep(0.1) + assert not scenario.done() + release.set() with pytest.raises(asyncio.CancelledError): - await await_bounded(close() if cancelled == "close" else flush_and_exit()) + await await_bounded(scenario) finally: release.set() - # Waits for the commit, and for an abort queued after it. - await asyncio.to_thread(shutdown_bounded, f._pool) + scenario.cancel() assert f.closed assert released == [True] fs.core.complete_multipart_upload.assert_called_once() fs.core.abort_multipart_upload.assert_not_called() + @pytest.mark.asyncio + async def test_close_waits_for_cancelled_read(self, fs): + # GH-1332 (F2): closing a file whose read was cancelled while its + # request ran waits for the request, which used to run on after the + # close returned. + block_size = 2**10 + started = threading.Event() + release = threading.Event() + finished = [] + + def get_object(path, range_, **kwargs): + started.set() + release.wait(EVENT_TIMEOUT) + finished.append(range_) + return bytes(range_[1] - range_[0]) + + fs.core.get_object = mock.MagicMock(side_effect=get_object) + info = S3Object( + init={"Key": "key", "ETag": '"etag"'}, + type=S3ObjectType.S3_OBJECT_TYPE_FILE, + bucket="bucket", + key="key", + ) + info.size = block_size + fs._sync_fs.info = mock.MagicMock(return_value=info) + f = await fs.open_async("s3://bucket/key", "rb", block_size=block_size) + read = asyncio.ensure_future(f.read()) + close = None + try: + assert await asyncio.to_thread(started.wait, EVENT_TIMEOUT) + read.cancel() + with pytest.raises(asyncio.CancelledError): + await read + close = asyncio.ensure_future(f.close()) + await asyncio.sleep(0.1) + assert not close.done() + # A cancellation of the close does not stop it waiting either. + close.cancel() + await asyncio.sleep(0.1) + assert not close.done() + release.set() + with pytest.raises(asyncio.CancelledError): + await await_bounded(close) + assert finished == [(0, block_size)] + finally: + release.set() + if close is not None: + close.cancel() + @pytest.mark.asyncio async def test_flush_cancelled_while_queued_then_exit_with_exception(self, fs): # GH-1211: the final flush reserves the commit before it is queued, so # leaving the block with an exception after the flush is cancelled # while its commit waits for a thread does not abort the upload. block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE - started = threading.Event() release = threading.Event() released = [] - respond = fs.core.upload_part.side_effect - def upload_part(**kwargs): - started.set() + def hold(): # The finally block of the test releases it. The wait is bounded, # and its timeout recorded, so that a regression that blocks the # event loop fails instead of hanging or passing. released.append(release.wait(EVENT_TIMEOUT)) - return respond(**kwargs) - fs.core.upload_part = mock.MagicMock(side_effect=upload_part) queued = threading.Event() - submit = S3ThreadPool.submit + submit_to = S3ThreadPool.submit_to - def record_commit(pool, fn, *args, **kwargs): + def record_commit(pool, future, fn, *args, **kwargs): if getattr(fn, "__name__", None) == "_complete_upload": queued.set() - return submit(pool, fn, *args, **kwargs) + return submit_to(pool, future, fn, *args, **kwargs) - # The part holds the only thread, for which the commit waits. - f = await fs.open_async("s3://bucket/key", "wb", block_size=block_size, max_workers=1) + fs._thread_pool = S3ThreadPool(1) + 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) - assert await asyncio.to_thread(started.wait, EVENT_TIMEOUT) + # The commit waits for the only thread of the filesystem, + # which this call holds. + fs._thread_pool.submit(hold) flush = asyncio.ensure_future(f.flush(force=True)) await wait_until(queued.is_set) flush.cancel() await flush - try: - with ( - mock.patch.object(S3ThreadPool, "submit", autospec=True, side_effect=record_commit), - pytest.raises(asyncio.CancelledError), - ): - # Bounded, so that a regression whose exit waits for the - # held part fails here and releases it. - await await_bounded(flush_and_exit()) - finally: - release.set() - # Waits for the commit, and for an abort queued after it. - await asyncio.to_thread(shutdown_bounded, f._pool) + with mock.patch.object(S3ThreadPool, "submit_to", autospec=True, side_effect=record_commit): + scenario = asyncio.ensure_future(flush_and_exit()) + try: + await await_bounded(asyncio.to_thread(queued.wait, EVENT_TIMEOUT)) + # GH-1332 (F2): leaving the block waits for the queued commit. + await asyncio.sleep(0.1) + assert not scenario.done() + release.set() + with pytest.raises(asyncio.CancelledError): + await await_bounded(scenario) + finally: + release.set() + scenario.cancel() assert f.closed assert released == [True] @@ -2737,12 +2848,12 @@ def test_close_interrupted_while_queuing_commit(self, fs): # committing it and aborts the upload. The upload used to be left # incomplete with a commit that nothing sends. block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE - submit = S3ThreadPool.submit + submit_to = S3ThreadPool.submit_to - def interrupt_commit(pool, fn, *args, **kwargs): + def interrupt_commit(pool, future, fn, *args, **kwargs): if getattr(fn, "__name__", None) == "_complete_upload": raise KeyboardInterrupt - return submit(pool, fn, *args, **kwargs) + return submit_to(pool, future, fn, *args, **kwargs) async def write_and_close(): async with asyncio.timeout(EVENT_TIMEOUT): @@ -2753,7 +2864,7 @@ async def write_and_close_file(): await f.write(b"a" * block_size) try: with mock.patch.object( - S3ThreadPool, "submit", autospec=True, side_effect=interrupt_commit + S3ThreadPool, "submit_to", autospec=True, side_effect=interrupt_commit ): await f.close() finally: @@ -2769,6 +2880,88 @@ async def write_and_close_file(): fs.core.complete_multipart_upload.assert_not_called() fs.core.abort_multipart_upload.assert_called_once() + def test_close_interrupted_after_queuing_commit(self, fs): + # GH-1332 (F2, I3): an interrupt of the queuing of the commit after the + # commit has started lets the close return only after the commit has + # ended. The commit used to be tracked only once its queuing returned. + block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE + submit_to = S3ThreadPool.submit_to + started, interrupted = threading.Event(), threading.Event() + events = [] + + def complete_multipart_upload(*args, **kwargs): + started.set() + # Until the interrupt has been raised and the close waits. + assert interrupted.wait(EVENT_TIMEOUT) + time.sleep(0.1) + events.append("committed") + + fs.core.complete_multipart_upload = mock.MagicMock(side_effect=complete_multipart_upload) + + def interrupt_commit(pool, future, fn, *args, **kwargs): + submit_to(pool, future, fn, *args, **kwargs) + if getattr(fn, "__name__", None) == "_complete_upload": + assert started.wait(EVENT_TIMEOUT) + try: + raise KeyboardInterrupt + finally: + interrupted.set() + + async def write_and_close(): + async with asyncio.timeout(EVENT_TIMEOUT): + f = await fs.open_async("s3://bucket/key", "wb", block_size=block_size) + await f.write(b"a" * block_size) + try: + with mock.patch.object( + S3ThreadPool, "submit_to", autospec=True, side_effect=interrupt_commit + ): + await f.close() + finally: + events.append("closed") + + with pytest.raises(KeyboardInterrupt): + asyncio.run(write_and_close()) + + assert events == ["committed", "closed"] + fs.core.abort_multipart_upload.assert_not_called() + + def test_exit_interrupted_while_queuing_abort(self, fs): + # GH-1332 (F2, I1): an interrupt that stops the queuing of the abort + # of an exit with an exception is raised; the close used to wait + # forever for the abort, which nothing runs. + block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE + submit_to = S3ThreadPool.submit_to + + def interrupt_abort(pool, future, fn, *args, **kwargs): + if getattr(fn, "__name__", None) == "_discard": + raise KeyboardInterrupt + return submit_to(pool, future, fn, *args, **kwargs) + + async def write_and_fail(): + f = await fs.open_async("s3://bucket/key", "wb", block_size=block_size) + await f.write(b"a" * block_size * 2) + with mock.patch.object( + S3ThreadPool, "submit_to", autospec=True, side_effect=interrupt_abort + ): + async with f: + raise ValueError("failed") + + errors = [] + + def run(): + try: + asyncio.run(write_and_fail()) + except BaseException as e: + errors.append(e) + + # In a thread, which a regression leaves waiting instead of hanging + # the test. + thread = threading.Thread(target=run, daemon=True) + thread.start() + thread.join(EVENT_TIMEOUT) + assert not thread.is_alive() + assert [type(e) for e in errors] == [KeyboardInterrupt] + @pytest.mark.asyncio async def test_flush_cancelled_then_exit_joins_abort(self, fs): # GH-1211: leaving the block with an exception after a cancelled final @@ -2811,6 +3004,41 @@ async def flush_and_exit(): fs.core.complete_multipart_upload.assert_called_once() fs.core.abort_multipart_upload.assert_called_once() + @pytest.mark.asyncio + async def test_exit_cancelled_while_aborting(self, fs): + # GH-1332 (F2): a cancellation of the exit that aborts the upload is + # raised only after the abort has ended; the abort used to run on + # after the cancelled exit returned. + block_size = S3Core.MULTIPART_UPLOAD_MIN_PART_SIZE + aborting, release = threading.Event(), threading.Event() + aborted = [] + + def abort_multipart_upload(*args, **kwargs): + aborting.set() + release.wait(EVENT_TIMEOUT) + aborted.append(True) + + 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) + await f.write(b"a" * block_size * 2) + + async def write_and_fail(): + async with f: + raise ValueError("failed") + + exit_ = asyncio.ensure_future(write_and_fail()) + try: + assert await asyncio.to_thread(aborting.wait, EVENT_TIMEOUT) + exit_.cancel() + await asyncio.sleep(0.1) + assert not exit_.done() + release.set() + with pytest.raises(asyncio.CancelledError): + await await_bounded(exit_) + assert aborted == [True] + finally: + release.set() + @pytest.mark.parametrize("data", [b"", b"abc"]) @pytest.mark.asyncio async def test_close_small(self, fs, data): @@ -2857,13 +3085,12 @@ def get_object(path, range_, **kwargs): with pytest.raises(PermissionError, match="Access Denied") as info_: await f.read() await f.close() - # The threads have set the results once the pool of the file is shut - # down, and the event loop runs their callbacks on the next iterations. - await asyncio.to_thread(f._pool.shutdown, True) - for _ in range(3): - await asyncio.sleep(0) + # The threads have set the results once the requests of the file are + # shut down, and leave the frames of the calls, and the event loop + # runs their callbacks, right after. + await asyncio.to_thread(f._requests.shutdown, True) assert len(read) == 3 - assert [ref() for ref in read] == [None] * 3 + await wait_until(lambda: [ref() for ref in read] == [None] * 3) assert info_.value.__traceback__ is not None @pytest.mark.asyncio @@ -2895,7 +3122,7 @@ def get_object(path, range_, **kwargs): await f.read() await f.close() # Waits for the thread of the running read. - await asyncio.to_thread(f._pool.shutdown, True) + await asyncio.to_thread(f._requests.shutdown, True) assert time.monotonic() - started < EVENT_TIMEOUT assert operation.calls == 1 @@ -2968,7 +3195,7 @@ def running_reads(): return [ t for t in asyncio.all_tasks() - if t.get_coro().__qualname__ == "S3ThreadPool.run" and not t.done() + if t.get_coro().__qualname__ == "S3Executor.run" and not t.done() ] # At least the held read is found. @@ -3014,14 +3241,14 @@ def get_object(path, range_, **kwargs): with pytest.raises(asyncio.CancelledError) as info_: await asyncio.wait_for(task, EVENT_TIMEOUT) await f.close() - await asyncio.wait_for(asyncio.to_thread(f._pool.shutdown, True), EVENT_TIMEOUT) + await asyncio.wait_for(asyncio.to_thread(f._requests.shutdown, True), EVENT_TIMEOUT) # The retries of the held range stopped. assert time.monotonic() - started < EVENT_TIMEOUT assert operation.calls == 1 - for _ in range(3): - await asyncio.sleep(0) assert len(read) == 4 - assert [ref() for ref in read] == [None] * 4 + # The threads leave the frames of the calls, and the event loop runs + # their callbacks, right after the results are set. + await wait_until(lambda: [ref() for ref in read] == [None] * 4) assert info_.value.__traceback__ is not None @pytest.mark.asyncio diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index 2402d6a4e..e774599c9 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -6,8 +6,10 @@ # SPDX-License-Identifier: MIT import asyncio +import contextlib import contextvars import gc +import sys import threading import time import traceback @@ -19,18 +21,76 @@ import pytest from botocore.exceptions import ClientError -from pyathena.filesystem.s3_executor import S3AioExecutor, S3Executor, S3ThreadPool +from pyathena.filesystem.s3_executor import S3AioExecutor, S3Executor, S3ThreadPool, _S3PoolScope from pyathena.util import RetryConfig, override, retry_api_call, stop_retries_on from tests.pyathena.util import ( EVENT_TIMEOUT, + OpcodeInterrupter, ThrottledOperation, forking, + interrupt_each_opcode, overlapping_calls, run_in_forked_child, + shutdown_bounded, + sigint_when, ) class TestS3ThreadPool: + def test_submit_interrupted_at_each_opcode(self): + # GH-1332 (P6, I2): an interrupt at any opcode of submit() leaves + # nothing to wait for forever: the next call runs, and the shutdown + # returns. + def case(interrupt): + pool = S3ThreadPool(1) + hold = threading.Event() + lock = threading.Lock() + state = {"active": 0, "peak": 0} + events = [] + + overlap = threading.Condition(lock) + + def call(name, timeout): + with overlap: + state["active"] += 1 + state["peak"] = max(state["peak"], state["active"]) + overlap.notify_all() + if name == "other": + # Waits for a third overlapping call, which only a + # pool that exceeds the bound starts. + overlap.wait_for(lambda: state["peak"] > 2, timeout=timeout) + try: + if name == "call": + hold.wait(timeout) + finally: + with overlap: + state["active"] -= 1 + events.append(name) + + try: + # Running until the others have had the time to overlap it. + future = interrupt(pool.submit, call, "call", EVENT_TIMEOUT) + pool.release() + others = [pool.submit(call, "other", 0.1) for _ in range(3)] + # The others run beside the interrupted call only if the + # release has left it untracked; otherwise they wait for it. + wait(others, timeout=0.5) + hold.set() + for other in others: + other.result(EVENT_TIMEOUT) + # At most one call more than max_workers runs at once (I2). + assert state["peak"] <= 2 + shutdown_bounded(pool) + if future is not None: + assert future.done() + assert events.count("other") == 3 + assert events.count("call") <= 1 + finally: + hold.set() + pool.shutdown() + + assert interrupt_each_opcode(S3ThreadPool.submit, case) > 0 + def test_default_max_workers(self): # None takes the default of ThreadPoolExecutor, as documented. with ThreadPoolExecutor() as executor: @@ -310,6 +370,558 @@ def hold_lock(): assert code == 0 +class TestS3PoolScope: + @pytest.mark.parametrize("busy", [False, True], ids=["entering", "waiting"]) + def test_submit_interrupted_at_each_opcode(self, busy): + # GH-1332 (S1, S2, S5, S8, S9): an interrupt at any opcode of submit(), + # for a call that enters the pool at once or waits for a place, leaves + # no place taken and nothing to wait for forever: the next call runs, + # the shutdown returns, and the interrupted call runs at most once. + def case(interrupt): + pool = S3ThreadPool(2) + scope = _S3PoolScope(pool, 1) + release = threading.Event() + events = [] + try: + if busy: + scope.submit(release.wait, EVENT_TIMEOUT) + future = interrupt(scope.submit, events.append, "call") + release.set() + # A place left taken would keep the next call waiting forever. + assert scope.submit(sum, [1, 2]).result(EVENT_TIMEOUT) == 3 + shutdown_bounded(scope) + if future is not None: + assert future.done() + if interrupt.fired and busy: + # A waiting call cannot have started, so it was cancelled. + assert events == [] + else: + assert events in ([], ["call"]) + finally: + release.set() + pool.shutdown() + + assert interrupt_each_opcode(_S3PoolScope.submit, case) > 0 + + def test_submit_to_interrupted_at_each_opcode(self): + # GH-1332 (I3): an interrupt at any opcode of submit_to() through a + # scope leaves the caller's future cancellable by the caller, as the + # write session does, after which nothing waits forever and the call + # runs at most once. + def case(interrupt): + pool = S3ThreadPool(2) + scope = _S3PoolScope(pool, 1) + release = threading.Event() + events = [] + future: Future[None] = Future() + try: + scope.submit(release.wait, EVENT_TIMEOUT) + interrupt(scope.submit_to, future, events.append, "call") + if interrupt.fired: + # As the write session does with its tracked request. + future.cancel() + release.set() + assert scope.submit(sum, [1, 2]).result(EVENT_TIMEOUT) == 3 + shutdown_bounded(scope) + assert events == ([] if future.cancelled() else ["call"]) + finally: + release.set() + pool.shutdown() + + assert interrupt_each_opcode(S3Executor.submit_to, case) > 0 + + def test_rejected_submit_interrupted_at_each_opcode(self): + # GH-1332 (I1, S5): an interrupt at any point of a submission that the + # pool rejects, including its cleanup, leaves nothing for the + # shutdown of the scope to wait for forever. + def case(interrupt): + pool = S3ThreadPool(1) + pool.shutdown() + scope = _S3PoolScope(pool, 1) + with contextlib.suppress(RuntimeError): + interrupt(scope.submit, sum, [1, 2]) + shutdown_bounded(scope) + + assert interrupt_each_opcode(_S3PoolScope.submit, case) > 0 + + @pytest.mark.skipif(sys.version_info < (3, 12), reason="Interleaving needs sys.monitoring.") + def test_shutdown_during_hand_on_keeps_waiting_call(self): + # GH-1332 (S7): a shutdown of the scope at any point of the hand-on + # of a place keeps the waiting call, which then runs, also when the + # pool takes no more calls and the ending call's thread runs it. + def case(n): + pool = S3ThreadPool(1) + scope = _S3PoolScope(pool, 1) + release = threading.Event() + shutdowns = [] + errors = [] + + def shut_down(): + # From another thread, as a close would; this thread may hold + # the lock of the scope, so it waits only briefly. + thread = threading.Thread(target=scope.shutdown, args=(False,), daemon=True) + try: + thread.start() + except BaseException as e: + errors.append(e) + raise + shutdowns.append(thread) + thread.join(0.05) + + interleave = OpcodeInterrupter(_S3PoolScope._next.__code__, n, action=shut_down) + try: + scope.submit(release.wait, EVENT_TIMEOUT) + waiting = scope.submit(sum, [1, 2]) + pool.shutdown(wait=False) + + def end_first_call(): + release.set() + return waiting.result(EVENT_TIMEOUT) + + assert interleave(end_first_call) == 3 + for thread in shutdowns: + thread.join(EVENT_TIMEOUT) + assert errors == [] + finally: + release.set() + pool.shutdown() + return interleave.fired + + n = 1 + while case(n): + n += 1 + assert n <= 2000 + assert n > 1 + + def test_shutdown_returns_once_calls_are_cancelled(self): + # GH-1332 (S5, I1): a shutdown that waits for a call returns once the + # call is cancelled, without waiting for the pool to reach it. A + # wait for its notification by the pool, which an interrupt can stop, + # used to keep the shutdown waiting. + pool = S3ThreadPool(1) + other = _S3PoolScope(pool, 1) + scope = _S3PoolScope(pool, 1) + release = threading.Event() + try: + # Longer than the waits below, so that the pool cannot reach the + # cancelled call before the test ends. + other.submit(release.wait, 3 * EVENT_TIMEOUT) + # Accepted by the pool, behind the call of the other scope. + future = scope.submit(sum, [1, 2]) + shut_down = threading.Event() + thread = threading.Thread( + target=lambda: (scope.shutdown(wait=True), shut_down.set()), daemon=True + ) + thread.start() + assert not shut_down.wait(0.1) + assert future.cancel() + assert shut_down.wait(EVENT_TIMEOUT) + assert not release.is_set() + finally: + release.set() + pool.shutdown() + + def test_shutdown_after_interrupted_cancel(self): + # GH-1332 (S5, I1): an interrupt at any point of the cancellation of a + # call that a shutdown waits for, including one that skips the + # callbacks of the future, does not keep the shutdown waiting. + def case(interrupt): + pool = S3ThreadPool(1) + other = _S3PoolScope(pool, 1) + scope = _S3PoolScope(pool, 1) + release = threading.Event() + try: + # Longer than the waits below, so that the pool cannot reach + # the cancelled call before the test ends. + other.submit(release.wait, 3 * EVENT_TIMEOUT) + future = scope.submit(sum, [1, 2]) + waiting = threading.Event() + add_done_callback = future.add_done_callback + + def watch(fn): + add_done_callback(fn) + waiting.set() + + future.add_done_callback = watch + shut_down = threading.Event() + thread = threading.Thread( + target=lambda: (scope.shutdown(wait=True), shut_down.set()), daemon=True + ) + thread.start() + assert waiting.wait(EVENT_TIMEOUT) + interrupt(future.cancel) + if not future.cancelled(): + # Interrupted before it changed the state. + assert future.cancel() + assert shut_down.wait(EVENT_TIMEOUT) + finally: + release.set() + pool.shutdown() + + assert interrupt_each_opcode(Future.cancel, case) > 0 + + def test_waiting_calls_enter_in_order(self): + # GH-1332 (S3): the calls that wait for a place enter the pool in the + # order they were submitted. + pool = S3ThreadPool(2) + scope = _S3PoolScope(pool, 1) + release = threading.Event() + order = [] + try: + scope.submit(release.wait, EVENT_TIMEOUT) + waiting = [scope.submit(order.append, name) for name in ("first", "second")] + release.set() + for future in waiting: + future.result(EVENT_TIMEOUT) + shutdown_bounded(scope) + finally: + release.set() + pool.shutdown() + assert order == ["first", "second"] + + def test_shutdown_waits_for_call_from_pool_thread(self): + # GH-1332 (S5): the shutdown of a scope also waits for a call that a + # thread of its pool runs itself (T1). + pool = S3ThreadPool(2) + scope = _S3PoolScope(pool, 1) + started, release, ended = threading.Event(), threading.Event(), threading.Event() + + def call(): + started.set() + release.wait(EVENT_TIMEOUT) + ended.set() + + try: + nested = pool.submit(lambda: scope.submit(call).result(EVENT_TIMEOUT)) + assert started.wait(EVENT_TIMEOUT) + # Released from another thread while this one shuts the scope down. + threading.Timer(0.1, release.set).start() + scope.shutdown(wait=True) + assert ended.is_set() + nested.result(EVENT_TIMEOUT) + finally: + release.set() + pool.shutdown() + + def test_cancelled_waiting_calls_do_not_accumulate(self): + # GH-1332 (S8): calls submitted and cancelled while one long call + # holds the only place are freed before that call ends. + pool = S3ThreadPool(1) + scope = _S3PoolScope(pool, 1) + release = threading.Event() + refs = [] + try: + scope.submit(release.wait, EVENT_TIMEOUT) + for _ in range(100): + future = scope.submit(sum, [1, 2]) + assert future.cancel() + refs.append(weakref.ref(future)) + del future + gc.collect() + assert sum(ref() is not None for ref in refs) < 50 + finally: + release.set() + pool.shutdown() + + def test_cancelled_waiting_call_frees_arguments(self): + # GH-1332 (S8): a call cancelled while it waits leaves the queue at + # once, so that its arguments are not kept until a place frees. + class Body: + pass + + pool = S3ThreadPool(1) + scope = _S3PoolScope(pool, 1) + release = threading.Event() + body = Body() + body_ref = weakref.ref(body) + try: + scope.submit(release.wait, EVENT_TIMEOUT) + waiting = scope.submit(lambda b: None, body) + del body + assert waiting.cancel() + gc.collect() + assert body_ref() is None + finally: + release.set() + shutdown_bounded(scope) + pool.shutdown() + + @pytest.mark.skipif( + threading.current_thread() is not threading.main_thread(), + reason="SIGINT interrupts the main thread.", + ) + def test_shutdown_interrupted_while_waiting(self): + # GH-1332 (I1, S5): a real interrupt of a shutdown that waits for a + # running call leaves the scope consistent: the call runs to its end, + # and a later shutdown waits for it and returns. + pool = S3ThreadPool(2) + scope = _S3PoolScope(pool, 1) + started, release = threading.Event(), threading.Event() + + def call(): + started.set() + return release.wait(EVENT_TIMEOUT) + + try: + running = scope.submit(call) + assert started.wait(EVENT_TIMEOUT) + (pending,) = scope.pending() + # Sent once the shutdown waits for the call, which it does through + # a done callback of the future. + waiting = [] + add_done_callback = pending.add_done_callback + + def watch(fn): + add_done_callback(fn) + waiting.append(fn) + + pending.add_done_callback = watch + with sigint_when(lambda: waiting), pytest.raises(KeyboardInterrupt): + scope.shutdown(wait=True) + assert not running.done() + release.set() + shutdown_bounded(scope) + assert running.result(EVENT_TIMEOUT) + finally: + release.set() + pool.shutdown() + + def test_submit_from_pool_thread(self): + # GH-1332 (T1): a call submitted to a scope from a thread of its own + # pool, as when garbage collection there closes a file, runs in that + # thread instead of waiting for a place in the pool that it holds. + pool = S3ThreadPool(1) + scope = _S3PoolScope(pool, 1) + try: + nested = pool.submit(lambda: scope.submit(sum, [1, 2]).result(EVENT_TIMEOUT)) + assert nested.result(2 * EVENT_TIMEOUT) == 3 + shutdown_bounded(scope) + finally: + pool.shutdown() + + def test_waiting_call_rejected_by_pool(self): + # GH-1332 (S5): a waiting call that the pool fails to take, such as + # when a thread cannot start, fails instead of being lost, so that + # the shutdown of the scope returns. + pool = S3ThreadPool(2) + scope = _S3PoolScope(pool, 1) + release = threading.Event() + submit = pool.submit + + def failing_submit(fn, *args, **kwargs): + if threading.current_thread() is not threading.main_thread(): + raise OSError("can't start new thread") + return submit(fn, *args, **kwargs) + + try: + running = scope.submit(release.wait, EVENT_TIMEOUT) + with mock.patch.object(pool, "submit", side_effect=failing_submit): + waiting = scope.submit(sum, [1, 2]) + release.set() + with pytest.raises(OSError, match="can't start new thread"): + waiting.result(EVENT_TIMEOUT) + assert running.result(EVENT_TIMEOUT) + shutdown_bounded(scope) + finally: + release.set() + pool.shutdown() + + def test_window(self): + # GH-1332: at most max_workers calls of a scope are in the shared + # pool at once, as a file with its own pool of that size used to run. + pool = S3ThreadPool(4) + call, state = overlapping_calls(2) + scope = _S3PoolScope(pool, 2) + try: + calls = [scope.submit(call) for _ in range(4)] + wait(calls, timeout=EVENT_TIMEOUT) + finally: + shutdown_bounded(scope) + pool.shutdown() + assert all(c.done() for c in calls) + assert state["peak"] == 2 + + def test_waiting_call_context(self): + # GH-1332: a call that waits for a place in the pool runs with the + # context of its submitter, so that its stop scope reaches its + # retries, although it enters the pool from the thread of the call + # that ends before it. + variable: contextvars.ContextVar[str] = contextvars.ContextVar("variable") + stop = threading.Event() + stop.set() + operation = ThrottledOperation() + release = threading.Event() + + def work(): + with pytest.raises(ClientError): + retry_api_call(operation, RetryConfig(attempt=3, multiplier=0)) + return variable.get() + + pool = S3ThreadPool(2) + scope = _S3PoolScope(pool, 1) + try: + blocking = scope.submit(release.wait, EVENT_TIMEOUT) + variable.set("submitter") + with stop_retries_on(stop): + waiting = scope.submit(work) + variable.set("other") + release.set() + assert waiting.result(EVENT_TIMEOUT) == "submitter" + assert blocking.result(EVENT_TIMEOUT) + finally: + release.set() + shutdown_bounded(scope) + pool.shutdown() + assert operation.calls == 1 + + def test_shutdown_waits_for_own_calls(self): + # GH-1332: shutting a scope down waits for its own calls, also those + # that wait for a place, and not for those of another scope; the pool + # keeps running for the other scopes. + pool = S3ThreadPool(2) + own_release, other_release = threading.Event(), threading.Event() + own, other = _S3PoolScope(pool, 1), _S3PoolScope(pool, 1) + events = [] + + def slow_append(event): + # Still running when a shutdown that does not wait returns. + time.sleep(0.2) + events.append(event) + + try: + other_call = other.submit(other_release.wait, EVENT_TIMEOUT) + own.submit(own_release.wait, EVENT_TIMEOUT) + own.submit(slow_append, "waiting call") + own_release.set() + shutdown_bounded(own) + assert events == ["waiting call"] + assert not other_call.done() + with pytest.raises(RuntimeError, match="after shutdown"): + own.submit(sum, [1, 2]) + # The pool, whose other thread the other scope holds, still runs calls. + assert pool.submit(sum, [1, 2]).result(EVENT_TIMEOUT) == 3 + finally: + own_release.set() + other_release.set() + shutdown_bounded(other) + pool.shutdown() + assert other_call.result(EVENT_TIMEOUT) + + def test_cancel_waiting_call(self): + # GH-1332: a call that waits for a place in the pool can be cancelled + # through its future, and then never runs. + pool = S3ThreadPool(2) + scope = _S3PoolScope(pool, 1) + release = threading.Event() + events = [] + try: + running = scope.submit(release.wait, EVENT_TIMEOUT) + waiting = scope.submit(events.append, "waiting call") + assert waiting.cancel() + release.set() + shutdown_bounded(scope) + finally: + release.set() + pool.shutdown() + assert running.result(EVENT_TIMEOUT) + assert events == [] + + def test_interrupted_submit(self): + # GH-1332: an interrupt before the pool takes a call cancels the call, + # so that the shutdown of the scope does not wait for it, and the + # scope runs the next calls. + pool = S3ThreadPool(1) + scope = _S3PoolScope(pool, 1) + try: + with ( + mock.patch.object(pool, "submit", side_effect=KeyboardInterrupt), + pytest.raises(KeyboardInterrupt), + ): + scope.submit(sum, [1, 2]) + assert scope.submit(sum, [3, 4]).result(EVENT_TIMEOUT) == 7 + shutdown_bounded(scope) + finally: + pool.shutdown() + + def test_interrupted_submit_after_acceptance(self): + # GH-1332: an interrupt after the pool accepted a call that has not + # started cancels the call, which then neither runs nor keeps its + # place when the pool gets to it, so that the next calls still run. + pool = S3ThreadPool(1) + scope = _S3PoolScope(pool, 1) + release = threading.Event() + events = [] + + class InterruptedSet(set): + def add(self, element): + super().add(element) + # Once, after the pool has accepted the call. + scope._entered = set(self) + raise KeyboardInterrupt + + try: + # Holds the only thread of the pool, so that the call does not start. + holding = pool.submit(release.wait, EVENT_TIMEOUT) + scope._entered = InterruptedSet() + with pytest.raises(KeyboardInterrupt): + scope.submit(events.append, "interrupted call") + following = scope.submit(sum, [1, 2]) + release.set() + assert holding.result(EVENT_TIMEOUT) + assert following.result(EVENT_TIMEOUT) == 3 + shutdown_bounded(scope) + finally: + release.set() + pool.shutdown() + assert events == [] + + def test_waiting_calls_run_after_pool_shutdown(self): + # GH-1332: once the pool takes no more calls, as when the interpreter + # exits, the calls that wait in a scope run in the thread of the call + # that ends before them, as the calls queued in a pool of their own + # used to drain, instead of failing. + pool = S3ThreadPool(1) + scope = _S3PoolScope(pool, 1) + release = threading.Event() + try: + running = scope.submit(release.wait, EVENT_TIMEOUT) + waiting = scope.submit(sum, [1, 2]) + pool.shutdown(wait=False) + release.set() + assert running.result(EVENT_TIMEOUT) + assert waiting.result(EVENT_TIMEOUT) == 3 + finally: + release.set() + shutdown_bounded(scope) + + def test_release_waits_for_waiting_calls(self): + # GH-1332: a release of the pool waits for the calls that wait in a + # scope, so that it ends the threads that run them. A waiting call + # used to enter the pool after its release had ended the threads, + # and to start threads that nothing released. + pool = S3ThreadPool(1) + scope = _S3PoolScope(pool, 1) + release = threading.Event() + threads = [] + + def record(): + threads.append(threading.current_thread()) + return release.wait(EVENT_TIMEOUT) + + try: + calls = [scope.submit(record) for _ in range(2)] + pool.release() + release.set() + assert all(c.result(EVENT_TIMEOUT) for c in calls) + for thread in threads: + thread.join(EVENT_TIMEOUT) + assert not thread.is_alive() + finally: + release.set() + shutdown_bounded(scope) + pool.shutdown() + + class TestS3AioExecutor: def test_loop_ignored(self): # GH-1332: the loop argument is ignored with a warning, and the calls diff --git a/tests/pyathena/pandas/test_cursor.py b/tests/pyathena/pandas/test_cursor.py index bb0fae569..fd70756da 100644 --- a/tests/pyathena/pandas/test_cursor.py +++ b/tests/pyathena/pandas/test_cursor.py @@ -1,5 +1,6 @@ import contextlib import csv +import inspect import io import math import random @@ -14,10 +15,12 @@ import pandas as pd import pytest +from pyathena.aio.pandas.cursor import AioPandasCursor from pyathena.error import DatabaseError, DataError, OperationalError, ProgrammingError from pyathena.filesystem.s3 import S3FileSystem from pyathena.formatter import DefaultParameterFormatter from pyathena.model import AthenaQueryExecution +from pyathena.pandas.async_cursor import AsyncPandasCursor from pyathena.pandas.converter import DefaultPandasTypeConverter from pyathena.pandas.cursor import PandasCursor from pyathena.pandas.result_set import AthenaPandasResultSet @@ -1927,3 +1930,12 @@ def column(name): assert column("col_json_not_null").tolist() == [9007199254740995, 9007199254740997] assert column("col_binary").tolist() == [b"\x01", None] assert column("col_time").tolist() == [datetime(2017, 1, 1, 1, 2, 3).time(), None] + + +@pytest.mark.parametrize("cursor_class", [PandasCursor, AsyncPandasCursor, AioPandasCursor]) +def test_s3_max_workers_default(cursor_class): + # GH-1332: the pandas cursors default s3_max_workers to None, so that the + # filesystem that reads the result files takes the connection pool of the + # S3 client; they used to default to cpu_count() * 5. + default = inspect.signature(cursor_class.__init__).parameters["s3_max_workers"].default + assert default is None diff --git a/tests/pyathena/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index 23a33cb18..d5b73e9e2 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -43,6 +43,7 @@ LocalBucketFileSystem, ThrottledOperation, assert_fetches_fail, + mock_connection, query_results_response, ) @@ -1548,6 +1549,20 @@ def _write_unload_tables(fs, tables, **write_options): return manifest +class TestAthenaPandasResultSetMetadataReader: + @pytest.mark.parametrize(("max_workers", "pool_size"), [(None, 7), (3, 3)]) + def test_filesystem_request_pool_size(self, max_workers, pool_size): + # GH-1332: without s3_max_workers, the filesystem that reads the result + # files sizes its request pool from the connection's S3 client, whose + # connections it shares; it used to take cpu_count() * 5 threads. An + # explicit value is used as given. + kwargs = {} if max_workers is None else {"max_workers": max_workers} + reader = AthenaPandasResultSetMetadataReader( + mock_connection(7), MagicMock(), RetryConfig(), **kwargs + ) + assert reader.filesystem.max_workers == pool_size + + class _InMemoryS3Object: """One S3 object served from memory to a boto3 client, with the ranges it reads. diff --git a/tests/pyathena/polars/test_cursor.py b/tests/pyathena/polars/test_cursor.py index 595136485..4950f70b9 100644 --- a/tests/pyathena/polars/test_cursor.py +++ b/tests/pyathena/polars/test_cursor.py @@ -6,6 +6,7 @@ # SPDX-License-Identifier: MIT import contextlib +import inspect import logging import random import string @@ -18,8 +19,10 @@ import polars as pl import pytest +from pyathena.aio.polars.cursor import AioPolarsCursor from pyathena.error import DatabaseError, DataError, OperationalError, ProgrammingError from pyathena.model import AthenaQueryExecution +from pyathena.polars.async_cursor import AsyncPolarsCursor from pyathena.polars.converter import DefaultPolarsTypeConverter from pyathena.polars.cursor import PolarsCursor from pyathena.polars.result_set import AthenaPolarsResultSet @@ -1109,3 +1112,12 @@ def test_read_options(self, execute_kwargs): expected = {**cursor_kwargs, **execute_kwargs} expected["max_workers"] = expected.pop("s3_max_workers") assert {key: kwargs[key] for key in expected} == expected + + +@pytest.mark.parametrize("cursor_class", [PolarsCursor, AsyncPolarsCursor, AioPolarsCursor]) +def test_s3_max_workers_default(cursor_class): + # GH-1332: the Polars cursors default s3_max_workers to None, so that the + # filesystem that reads the result files takes the connection pool of the + # S3 client; they used to default to cpu_count() * 5. + default = inspect.signature(cursor_class.__init__).parameters["s3_max_workers"].default + assert default is None diff --git a/tests/pyathena/s3fs/test_result_set.py b/tests/pyathena/s3fs/test_result_set.py index 9c7b6d412..9ae446165 100644 --- a/tests/pyathena/s3fs/test_result_set.py +++ b/tests/pyathena/s3fs/test_result_set.py @@ -20,6 +20,7 @@ from pyathena.s3fs.reader import AthenaCSVReader, EmptyStringAsNullCSVReader from pyathena.s3fs.result_set import AthenaS3FSResultSet from pyathena.util import RetryConfig +from tests.pyathena.util import mock_connection class _TupleCSVReader: @@ -53,7 +54,7 @@ def _result_set( filesystem = MagicMock() filesystem._open.return_value = stream filesystem_class = MagicMock(return_value=filesystem) - connection = MagicMock() + connection = mock_connection() connection.s3_client.head_object.return_value = {"ContentLength": len(data.encode())} connection.client.get_query_results.return_value = { "ResultSet": { diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index d967fe793..5871aa124 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -6,9 +6,12 @@ # SPDX-License-Identifier: MIT import asyncio +import dis +import gc import io import os import signal +import sys import threading import time import traceback @@ -213,6 +216,25 @@ async def cancel_when_started(task, started, release): release.set() +def poll_until(predicate, message, timeout=EVENT_TIMEOUT, collect=False): + """Wait, in the calling thread, until a condition holds. + + Args: + predicate: The condition, a callable without arguments. + message: The message of the assertion that fails on the timeout. + timeout: The longest wait, in seconds. + collect: Whether to run the cyclic garbage collector before each check. + """ + deadline = time.monotonic() + timeout + while True: + if collect: + gc.collect() + if predicate(): + return + assert time.monotonic() < deadline, message + time.sleep(0.001) + + async def wait_until(predicate, timeout=EVENT_TIMEOUT): """Wait, without blocking the event loop, until a condition holds. @@ -246,6 +268,255 @@ async def await_bounded(awaitable, timeout=EVENT_TIMEOUT): return task.result() +class OpcodeInterrupter: + """Calls a function with a ``KeyboardInterrupt`` raised at the n-th interrupt point of a target. + + The interrupt is raised in the calling thread, at the n-th point where + the target's code executes an instruction at which a signal can raise + it: + - the entry of the function; + - the instruction after a call returns; + - the target of a backward jump; + - a ``with`` statement, which models an interrupt of its blocking + ``acquire()``. + + CPython checks for signals only there, not between the instructions of a + ``with`` statement that has acquired its lock. Other code and other + threads are not interrupted. Python 3.12 and later use ``sys.monitoring``, since a trace + function gets no opcode events on the first call of a code object + there; Python 3.11 uses ``sys.settrace``. + + Args: + code: The code object of the target. + n: The interrupt point to raise the interrupt at, counted from 1. + action: A callable to run at that point instead of raising the + interrupt, in whichever thread reaches it, to interleave another + thread's call there. Only Python 3.12 and later support it. + """ + + def __init__(self, code, n, action=None): + self.code = code + self.n = n + self.action = action + self.count = 0 + self.fired = False + # The injected interrupt, kept with its traceback as a caller that + # handles it would keep it. + self.error = None + self._thread = threading.get_ident() + self._points = self._interrupt_points(code) + + def __call__(self, fn, *args, **kwargs): + """Call a function with the interrupt armed. + + Args: + fn: The function, which runs the target. + *args: Positional arguments of the function. + **kwargs: Keyword arguments of the function. + + Returns: + The return value of the function, or None if the interrupt fired. + + Raises: + KeyboardInterrupt: If the function raises one that was not injected. + """ + with self._armed(): + try: + return fn(*args, **kwargs) + except KeyboardInterrupt as e: + if not self.fired: + raise + self.error = e + return None + + @staticmethod + def _interrupt_points(code): + """Find the offsets of the instructions at which a signal can raise an interrupt. + + Args: + code: The code object. + + Returns: + The set of offsets. + """ + instructions = list(dis.get_instructions(code)) + entries = dis.Bytecode(code).exception_entries + + def handler(offset): + # The innermost exception handler that covers the offset. + covering = [e for e in entries if e.start <= offset < e.end] + return min(covering, key=lambda e: e.end - e.start).target if covering else None + + points = set() + for i, instruction in enumerate(instructions): + name = instruction.opname + if name == "RESUME" or name.startswith("BEFORE_WITH"): + points.add(instruction.offset) + elif name in ("CALL", "CALL_KW", "CALL_FUNCTION_EX") and i + 1 < len(instructions): + after = instructions[i + 1].offset + # A signal checked after the call is raised at the call, so + # the interrupt goes to the next instruction only when the + # same handler covers it; otherwise, as at a `return` that + # ends a `try`, it is raised before the call. + same = handler(after) == handler(instruction.offset) + points.add(after if same else instruction.offset) + elif name == "JUMP_BACKWARD": + # The target offset; Instruction.jump_target is new in 3.13. + points.add(instruction.argval) + return points + + def _instruction(self, offset): + """Count an interrupt point of the target in the calling thread, interrupting at the n-th. + + Args: + offset: The offset of the instruction about to run. + + Raises: + KeyboardInterrupt: At the n-th interrupt point. + """ + if self.fired or offset not in self._points: + return + if self.action is None and threading.get_ident() != self._thread: + return + self.count += 1 + if self.count == self.n: + self.fired = True + if self.action is not None: + self.action() + return + raise KeyboardInterrupt + + @contextmanager + def _armed(self): + monitoring = getattr(sys, "monitoring", None) + if monitoring is None: + previous = sys.gettrace() + sys.settrace(self._trace_call) + try: + yield + finally: + sys.settrace(previous) + return + tool = monitoring.DEBUGGER_ID + monitoring.use_tool_id(tool, "pyathena-tests") + try: + monitoring.register_callback( + tool, monitoring.events.INSTRUCTION, lambda code, offset: self._instruction(offset) + ) + monitoring.set_local_events(tool, self.code, monitoring.events.INSTRUCTION) + yield + finally: + monitoring.set_local_events(tool, self.code, 0) + monitoring.register_callback(tool, monitoring.events.INSTRUCTION, None) + monitoring.free_tool_id(tool) + + def _trace_call(self, frame, event, arg): + if frame.f_code is not self.code: + return None + frame.f_trace_opcodes = True + return self._trace_opcode + + def _trace_opcode(self, frame, event, arg): + if event == "opcode": + self._instruction(frame.f_lasti) + return self._trace_opcode + + +@contextmanager +def sigint_when(ready, timeout=EVENT_TIMEOUT): + """Deliver a real SIGINT to the main thread once a condition holds. + + The signal interrupts a blocking call of the main thread, such as the + acquisition of a lock or a wait for a future, as a user's Ctrl-C does. + Only the first one is raised as ``KeyboardInterrupt``; it is repeated + until handled, since a signal that arrives just before the main thread + blocks can miss waking it. A late signal is ignored after the block. + + Args: + ready: The condition, a callable without arguments, checked in a + thread of the helper. + timeout: The longest wait for the condition, in seconds. + + Yields: + A list that holds the handled signal once it was raised. + """ + handled = [] + stop = threading.Event() + + def handle_interrupt(signum, frame): + # The list is checked instead of an event, whose set() takes a lock + # that a nested call of this handler would wait for forever. + if not handled: + handled.append(signum) + raise KeyboardInterrupt + + def interrupt(): + deadline = time.monotonic() + timeout + while not ready(): + if stop.is_set() or time.monotonic() > deadline: + return + time.sleep(0.001) + while not handled and not stop.is_set(): + signal.pthread_kill(threading.main_thread().ident, signal.SIGINT) + time.sleep(0.05) + + thread = threading.Thread(target=interrupt, daemon=True) + previous_handler = signal.signal(signal.SIGINT, handle_interrupt) + try: + thread.start() + yield handled + finally: + signal.signal(signal.SIGINT, lambda signum, frame: None) + stop.set() + thread.join() + signal.signal(signal.SIGINT, previous_handler) + + +def interrupt_each_opcode(target, case, limit=2000): + """Run a case once for each opcode of a target, interrupting the target there. + + For n = 1, 2, and so on, ``case`` is called with an + :class:`OpcodeInterrupter` for the n-th instruction, which it uses to call the + code under test, and then checks the invariants that must hold after the + interrupt. The loop ends with the first n that the target never reaches. + + Args: + target: The function whose instructions are interrupted, or a method of it. + case: The callable that sets up, runs, and checks one attempt. + limit: The largest number of attempts, which a target that never + returns would exceed. + + Returns: + The number of opcodes that were interrupted. + """ + code = getattr(target, "__func__", target).__code__ + n = 0 + while True: + n += 1 + assert n <= limit, f"the target ran more than {limit} opcodes" + interrupter = OpcodeInterrupter(code, n) + case(interrupter) + if not interrupter.fired: + return n - 1 + + +def mock_connection(max_pool_connections=10): + """Return a mock connection whose S3 client has a connection pool of the given size. + + The filesystems size their request pool from the ``max_pool_connections`` + of the S3 client when ``max_workers`` is not given. + + Args: + max_pool_connections: The connection pool size of the S3 client. + + Returns: + The mock connection. + """ + connection = MagicMock() + connection.s3_client.meta.config.max_pool_connections = max_pool_connections + return connection + + def run_bounded(main, timeout=EVENT_TIMEOUT): """Run a coroutine with ``asyncio.run()``, failing if it does not end within the timeout. @@ -1363,3 +1634,6 @@ def open(self, path, *args, **kwargs): file = super().open(path, *args, **kwargs) self.opened.append((kwargs, file)) return file + + def close(self): + """Do nothing, as PyAthena's filesystems release their idle threads here."""