From a0d908e1e2c7342148721ae1114ea1cc1f77041d Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 5 Nov 2026 01:29:49 +0900 Subject: [PATCH 01/22] Share one request pool per filesystem among its files and operations (GH-1332) S3FileSystem owns a request pool of max_workers threads, which defaults to the max_pool_connections of its S3 client; a client that it builds takes max_pool_connections from max_workers unless config_kwargs sets it. _create_executor() returns a scope on that pool: at most max_workers of its calls are in the pool at once, the others wait in the scope with the context of their submitter, and shutting it down waits for its own calls only, so files no longer own or shut down a pool. S3FileSystem gets close() and with support, which release the idle threads of the pool. AioS3FileSystem sizes its pool of blocking calls from the internal S3FileSystem, and runs the requests that it fans out (part copies, annotation copies, DeleteObjects batches) and the requests of its streamed files in the request pool, while the calls that wait for them (completions, commits, aborts) stay in its own pool. Its close() also releases the request pool. S3Executor gains run(), run_shielded(), and wait_shielded(), which used to be S3ThreadPool's. In a forked child process, a file that the parent opened raises RuntimeError from its reads, writes, flushes, commit, and discard, and closing it only marks it closed, since its upload belongs to the parent. Co-Authored-By: Claude Opus 5.5 --- benchmarks/README.md | 4 +- benchmarks/pyathena_bench/filesystem.py | 2 +- docs/aio.md | 25 +- docs/filesystem.md | 18 +- pyathena/filesystem/s3.py | 147 +++++++- pyathena/filesystem/s3_async.py | 218 +++++++----- pyathena/filesystem/s3_executor.py | 317 +++++++++++++----- tests/pyathena/filesystem/test_s3.py | 113 +++++++ tests/pyathena/filesystem/test_s3_async.py | 183 ++++++---- tests/pyathena/filesystem/test_s3_executor.py | 105 +++++- tests/pyathena/s3fs/test_result_set.py | 3 +- tests/pyathena/util.py | 17 + 12 files changed, 891 insertions(+), 261 deletions(-) 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..6f12aff51 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 diff --git a/docs/filesystem.md b/docs/filesystem.md index 307f4852b..2ae896519 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -100,6 +100,21 @@ 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. `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. + 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 +657,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/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index e79b2a663..b7e117b5b 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -36,7 +36,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 +170,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 +197,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 +224,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 +272,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 +303,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 +313,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: @@ -1102,16 +1137,23 @@ def _create_executor(self, max_workers: int) -> S3Executor: """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. @@ -2908,8 +2950,13 @@ 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 <= 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 @@ -2976,13 +3023,76 @@ 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. + """ + 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 +3106,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 +3299,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 +3334,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() diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 029054cab..eaf897e94 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -61,18 +61,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 +102,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 +119,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 +140,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 +153,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 +169,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 +350,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`` in the request pool of the internal + ``S3FileSystem``. Args: path: S3 path (s3://bucket/key) or list of paths to delete. @@ -379,7 +388,7 @@ async def _delete_objects(self, paths: list[str], **kwargs) -> None: batches, params = self._sync_fs._delete_batches(paths, **kwargs) results = await asyncio.gather( *[ - self._thread_pool.run(self._sync_fs._delete_batch, batch, **params) + self._sync_fs._request_pool.run(self._sync_fs._delete_batch, batch, **params) for batch in batches ], return_exceptions=True, @@ -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. @@ -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): @@ -758,7 +768,7 @@ async def _copy_annotation(i: int, name: str) -> None: # S3FileSystem does. return try: - await self._thread_pool.run(_copy_annotation_in_thread, i, name) + await self._sync_fs._request_pool.run(_copy_annotation_in_thread, i, name) except Exception: # The running copies make no further attempt. failed = True @@ -891,9 +901,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. @@ -1178,8 +1188,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 +1212,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 @@ -1240,8 +1254,8 @@ 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. 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 +1277,13 @@ 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. + self._requests = fs._sync_fs._create_executor(max_workers=max_workers) + 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 +1316,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 +1332,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 +1350,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. @@ -1364,12 +1405,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 +1440,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 +1464,23 @@ 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. + + 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) + # A commit that is still running waits for the requests on. + self._requests.shutdown(wait=False) async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: if self.closed: @@ -1445,11 +1499,14 @@ 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) + await self.fs._thread_pool.run_shielded(self._discard) finally: - self._pool.shutdown(wait=False) + self._requests.shutdown(wait=False) def _close_without_commit_blocking(self) -> None: """Close the file without committing it; blocks. @@ -1464,7 +1521,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 +1529,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: @@ -1577,7 +1634,7 @@ async def _commit(self) -> None: 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)) + commit = self.fs._thread_pool.submit(self._complete_upload, cast(BytesIO, self.buffer)) except BaseException: if commit is None: # Not queued, or queued by a submit() that an interrupt @@ -1586,7 +1643,7 @@ async def _commit(self) -> None: 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 +1698,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 +1709,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 +1725,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..326ae18d4 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -17,7 +17,9 @@ import warnings import weakref from abc import ABCMeta, abstractmethod +from collections import deque from collections.abc import Callable +from concurrent import futures from concurrent.futures import Future from concurrent.futures.thread import ThreadPoolExecutor from multiprocessing import cpu_count @@ -141,6 +143,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. @@ -278,93 +367,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. @@ -461,6 +463,147 @@ 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. 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. + + 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 futures of the calls that were submitted and have not ended. + self._futures: set[Future[Any]] = set() + # The calls that wait for a place in the pool, and the number of + # places that the calls of the scope hold. + self._waiting: deque[tuple[contextvars.Context, _S3WorkItem[Any]]] = deque() + self._running = 0 + 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. + + 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 has been shut down. + """ + future: Future[T] = Future() + item = _S3WorkItem(future, fn, args, kwargs) + # Captured here, since a waiting call enters the pool from the thread + # of the call that ends before it. + context = contextvars.copy_context() + with self._lock: + if self._shutdown: + raise RuntimeError("cannot schedule new futures after shutdown") + self._futures.add(future) + enter = self._running < self.max_workers + if enter: + self._running += 1 + else: + self._waiting.append((context, item)) + future.add_done_callback(self._untrack) + if enter: + self._enter(context, item) + 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 = list(self._futures) + if wait: + futures.wait(pending) + + def _enter(self, context: contextvars.Context, item: _S3WorkItem[Any]) -> None: + """Hand a call to the pool, in the place that the scope holds for it. + + Args: + context: The context of the submitter of the call. + item: The call. + """ + try: + entered = self._pool.submit(context.run, item.run) + except BaseException as e: + # Such as after the shutdown of the pool: the future of the call + # gets the error, and the place goes to the next call. + rejected: Future[None] = Future() + rejected.set_running_or_notify_cancel() + rejected.set_exception(e) + item.settle(rejected) + self._leave(rejected) + if not isinstance(e, Exception): + raise + return + entered.add_done_callback(item.settle) + entered.add_done_callback(self._leave) + + def _leave(self, entered: Future[Any]) -> None: + """Give the place of a call that has ended to the next waiting call. + + Args: + entered: The finished future of the call in the pool. + """ + with self._lock: + if self._waiting: + context, item = self._waiting.popleft() + else: + self._running -= 1 + return + self._enter(context, item) + + def _untrack(self, future: Future[Any]) -> None: + """Stop tracking a call that has ended or was cancelled. + + Args: + future: The future of the call. + """ + with self._lock: + self._futures.discard(future) + + class S3AioExecutor(S3ThreadPool): """A thread pool kept for compatibility, an :class:`S3ThreadPool`. diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 50038d249..1f947ebb5 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -55,7 +55,11 @@ MULTIPART_COPY_KWARGS, MULTIPART_COPY_SIZE, ThrottledOperation, + forking, + mock_connection, + overlapping_calls, retried, + run_in_forked_child, stub_multipart_copy, ) @@ -86,6 +90,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 +2685,114 @@ 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() + + @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" diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index f1bf51820..069d374e6 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,36 @@ 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._fetch_range(0, 1)): + 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_write_parts_in_flight(self, fs): # GH-1159: a write keeps at most max_workers parts in flight, and waits @@ -2242,15 +2289,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 +2330,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 +2443,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: @@ -2666,7 +2713,7 @@ async def 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) + await asyncio.to_thread(shutdown_bounded, f._requests) assert f.closed assert released == [True] @@ -2724,7 +2771,7 @@ async def 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) + await asyncio.to_thread(shutdown_bounded, f._requests) assert f.closed assert released == [True] @@ -2859,7 +2906,7 @@ def get_object(path, range_, **kwargs): 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) + await asyncio.to_thread(f._requests.shutdown, True) for _ in range(3): await asyncio.sleep(0) assert len(read) == 3 @@ -2895,7 +2942,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 +3015,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,7 +3061,7 @@ 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 diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index 2402d6a4e..274a86c7e 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -19,7 +19,7 @@ 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, @@ -27,6 +27,7 @@ forking, overlapping_calls, run_in_forked_child, + shutdown_bounded, ) @@ -310,6 +311,108 @@ def hold_lock(): assert code == 0 +class TestS3PoolScope: + 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 == [] + + 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/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..cffc07f86 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -246,6 +246,23 @@ async def await_bounded(awaitable, timeout=EVENT_TIMEOUT): return task.result() +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. From acfe321dca384b8246a34c1336e51fd7b07bab92 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 5 Nov 2026 01:43:18 +0900 Subject: [PATCH 02/22] Hand scope places on in the pool thread, bound aio rm, release on close (GH-1332) A call of a scope hands its place to the next waiting call in the pool thread, before it ends there, so that an interrupt of the submitter cannot keep a place, a release of the pool waits for the waiting calls, and a rejected call no longer recurses. An interrupt before a call reaches the pool cancels it and frees its place. Calls cancelled while they wait are settled without entering the pool. The DeleteObjects batches of AioS3FileSystem._rm() run through a scope of max_workers, as those of S3FileSystem.rm() do, so that a large deletion does not queue ahead of the requests of the files. The pandas and s3fs result sets release the threads of their filesystem when they close. Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3_async.py | 18 ++-- pyathena/filesystem/s3_executor.py | 85 +++++++++++++------ pyathena/pandas/result_set.py | 4 + pyathena/s3fs/result_set.py | 3 + tests/pyathena/filesystem/test_s3_async.py | 15 ++-- tests/pyathena/filesystem/test_s3_executor.py | 56 ++++++++++++ 6 files changed, 140 insertions(+), 41 deletions(-) diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index eaf897e94..76b882814 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -351,8 +351,8 @@ async def _rm( """Delete objects with DeleteObjects requests. See :meth:`S3FileSystem.rm`. The DeleteObjects requests run in - parallel with ``asyncio.gather`` in the request pool of the internal - ``S3FileSystem``. + 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. @@ -386,11 +386,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._sync_fs._request_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) @@ -1178,8 +1179,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``. """ diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index 326ae18d4..50f8f93b6 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -500,7 +500,7 @@ def __init__(self, pool: S3ThreadPool, max_workers: int) -> None: self._futures: set[Future[Any]] = set() # The calls that wait for a place in the pool, and the number of # places that the calls of the scope hold. - self._waiting: deque[tuple[contextvars.Context, _S3WorkItem[Any]]] = deque() + self._waiting: deque[tuple[contextvars.Context, _S3WorkItem[Any], Future[Any]]] = deque() self._running = 0 self._shutdown = False @@ -535,10 +535,26 @@ def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]: if enter: self._running += 1 else: - self._waiting.append((context, item)) - future.add_done_callback(self._untrack) - if enter: - self._enter(context, item) + self._waiting.append((context, item, future)) + entered = False + try: + future.add_done_callback(self._untrack) + if enter: + entered = self._enter(context, item) + if not entered: + # The pool rejected the call, whose future has its error. + self._next() + except BaseException: + # An interrupt before the call reached the pool: its place goes to + # the next call, and its future is cancelled, so that nothing + # waits for it. + if enter and not entered: + self._next() + if future.cancel(): + future.set_running_or_notify_cancel() + with self._lock: + self._futures.discard(future) + raise return future @override @@ -557,42 +573,61 @@ def shutdown(self, wait: bool = True) -> None: if wait: futures.wait(pending) - def _enter(self, context: contextvars.Context, item: _S3WorkItem[Any]) -> None: + def _enter(self, context: contextvars.Context, item: _S3WorkItem[Any]) -> bool: """Hand a call to the pool, in the place that the scope holds for it. Args: context: The context of the submitter of the call. item: The call. + + Returns: + Whether the pool accepted the call. If it rejected it, such as + after its shutdown, the future of the call has the error. """ try: - entered = self._pool.submit(context.run, item.run) - except BaseException as e: - # Such as after the shutdown of the pool: the future of the call - # gets the error, and the place goes to the next call. + self._pool.submit(self._run, context, item) + except Exception as e: rejected: Future[None] = Future() rejected.set_running_or_notify_cancel() rejected.set_exception(e) item.settle(rejected) - self._leave(rejected) - if not isinstance(e, Exception): - raise - return - entered.add_done_callback(item.settle) - entered.add_done_callback(self._leave) + return False + return True + + def _run(self, context: contextvars.Context, item: _S3WorkItem[Any]) -> None: + """Run a call in the pool, then give its place to the next waiting call. - def _leave(self, entered: Future[Any]) -> None: - """Give the place of a call that has ended to the next waiting call. + The place is handed on in the thread of the call, before the call + ends in the pool, so that an interrupt of the submitter cannot keep + it, and a release of the pool waits for the calls that wait here. Args: - entered: The finished future of the call in the pool. + context: The context of the submitter of the call. + item: The call. """ - with self._lock: - if self._waiting: - context, item = self._waiting.popleft() - else: - self._running -= 1 + try: + context.run(item.run) + finally: + self._next() + + def _next(self) -> None: + """Give a place of the scope to the next waiting call, or free it. + + Calls that were cancelled while they waited are settled without + entering the pool, and so are those that the pool rejects. + """ + while True: + with self._lock: + if not self._waiting: + self._running -= 1 + return + context, item, future = self._waiting.popleft() + if future.cancelled(): + # Marks the future as notified, without running the call. + item.run() + continue + if self._enter(context, item): return - self._enter(context, item) def _untrack(self, future: Future[Any]) -> None: """Stop tracking a call that has ended or was cancelled. diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index 34c183acb..77004d3da 100644 --- a/pyathena/pandas/result_set.py +++ b/pyathena/pandas/result_set.py @@ -2426,6 +2426,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/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_async.py b/tests/pyathena/filesystem/test_s3_async.py index 069d374e6..b17741514 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -2904,13 +2904,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. + # 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) - for _ in range(3): - await asyncio.sleep(0) 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 @@ -3065,10 +3064,10 @@ def get_object(path, range_, **kwargs): # 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 274a86c7e..5c62c56b4 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -412,6 +412,62 @@ def test_cancel_waiting_call(self): assert running.result(EVENT_TIMEOUT) assert events == [] + def test_interrupted_submit(self): + # GH-1332: an interrupt before a call reaches the pool cancels the + # call and gives its place to the next one, so that the scope neither + # stops running calls nor waits for the interrupted one. + pool = S3ThreadPool(1) + scope = _S3PoolScope(pool, 1) + add_done_callback = Future.add_done_callback + interrupts = [KeyboardInterrupt] + + def interrupt_once(future, fn): + # The first callback is the scope's, added after it reserved a + # place for the call. + if interrupts: + raise interrupts.pop() + add_done_callback(future, fn) + + try: + with ( + mock.patch.object( + Future, "add_done_callback", autospec=True, side_effect=interrupt_once + ), + 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_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): From 404490b7850b9fa83ac3d9630b1eaada97727d95 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 5 Nov 2026 01:45:44 +0900 Subject: [PATCH 03/22] Run the annotation copies through a scope and document raising parallelism (GH-1332) The annotation copies of AioS3FileSystem's multipart copy run through a scope of max_workers in the request pool, as the DeleteObjects batches do, instead of an asyncio semaphore around the pool; a failed copy keeps the waiting ones from starting through its stop event, as before. The docs say how to raise the parallel requests now that max_workers defaults to the client's connection pool. Co-Authored-By: Claude Opus 5.5 --- docs/filesystem.md | 3 ++- pyathena/filesystem/s3_async.py | 26 ++++++++++---------------- 2 files changed, 12 insertions(+), 17 deletions(-) diff --git a/docs/filesystem.md b/docs/filesystem.md index 2ae896519..e70fe39b3 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -108,7 +108,8 @@ filesystem, with `max_workers` threads. Each file or operation runs at most its 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. `close()`, or the end of a `with` block, releases the threads of the pool once they +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 diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 76b882814..480f885c2 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -721,10 +721,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 @@ -762,19 +762,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._sync_fs._request_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)], From fd7a6b24d975d3e1d2904158cc8616cb6f15821d Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 5 Nov 2026 02:04:00 +0900 Subject: [PATCH 04/22] Count a scope's call once the pool has accepted it (GH-1332) A call of a scope takes a place only once the pool has accepted it, and frees it when it ends, so that an interrupt of the submitter at any point leaves no place taken and none freed twice; an interrupt before the pool takes a call cancels the call. Once the pool takes no more calls, as when the interpreter exits, the thread of a call that ends runs the calls that wait in its scope. AioS3StreamedFile.read() rejects a file inherited across a fork also when it reads nothing. The queued-commit test of the streamed file holds the thread of the filesystem that runs the commit, and the commit tests wait for that pool. Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3_async.py | 18 ++ pyathena/filesystem/s3_executor.py | 177 ++++++++++-------- tests/pyathena/filesystem/test_s3_async.py | 26 ++- tests/pyathena/filesystem/test_s3_executor.py | 39 ++-- 4 files changed, 150 insertions(+), 110 deletions(-) diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 480f885c2..aa31e6aae 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -1391,6 +1391,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. diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index 50f8f93b6..a07fa1cc3 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -468,10 +468,15 @@ class _S3PoolScope(S3Executor): 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. 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. + 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: + at worst, one more call runs at once for a while. Args: pool: The shared pool. @@ -498,10 +503,10 @@ def __init__(self, pool: S3ThreadPool, max_workers: int) -> None: self._lock = threading.Lock() # The futures of the calls that were submitted and have not ended. self._futures: set[Future[Any]] = set() - # The calls that wait for a place in the pool, and the number of - # places that the calls of the scope hold. - self._waiting: deque[tuple[contextvars.Context, _S3WorkItem[Any], Future[Any]]] = deque() - self._running = 0 + # 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() + self._entered: set[_S3ScopeCall] = set() self._shutdown = False @override @@ -509,7 +514,8 @@ 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. + only before the call starts. An interrupt of the submission cancels + the call. Args: fn: The callable to run. @@ -520,40 +526,35 @@ def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]: The future of the call. Raises: - RuntimeError: If the scope has been shut down. + RuntimeError: If the scope or the pool has been shut down. """ future: Future[T] = Future() - item = _S3WorkItem(future, fn, args, kwargs) - # Captured here, since a waiting call enters the pool from the thread - # of the call that ends before it. - context = contextvars.copy_context() - with self._lock: - if self._shutdown: - raise RuntimeError("cannot schedule new futures after shutdown") - self._futures.add(future) - enter = self._running < self.max_workers - if enter: - self._running += 1 - else: - self._waiting.append((context, item, future)) - entered = False + # 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 + ) try: + with self._lock: + if self._shutdown: + raise RuntimeError("cannot schedule new futures after shutdown") + self._futures.add(future) + enter = 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) future.add_done_callback(self._untrack) - if enter: - entered = self._enter(context, item) - if not entered: - # The pool rejected the call, whose future has its error. - self._next() except BaseException: - # An interrupt before the call reached the pool: its place goes to - # the next call, and its future is cancelled, so that nothing - # waits for it. - if enter and not entered: - self._next() + # 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. if future.cancel(): future.set_running_or_notify_cancel() - with self._lock: - self._futures.discard(future) + with self._lock: + self._futures.discard(future) raise return future @@ -573,61 +574,51 @@ def shutdown(self, wait: bool = True) -> None: if wait: futures.wait(pending) - def _enter(self, context: contextvars.Context, item: _S3WorkItem[Any]) -> bool: - """Hand a call to the pool, in the place that the scope holds for it. + def _run(self, call: _S3ScopeCall) -> None: + """Run a call in the pool, then let the next waiting calls enter it. - Args: - context: The context of the submitter of the call. - item: The call. - - Returns: - Whether the pool accepted the call. If it rejected it, such as - after its shutdown, the future of the call has the error. - """ - try: - self._pool.submit(self._run, context, item) - except Exception as e: - rejected: Future[None] = Future() - rejected.set_running_or_notify_cancel() - rejected.set_exception(e) - item.settle(rejected) - return False - return True - - def _run(self, context: contextvars.Context, item: _S3WorkItem[Any]) -> None: - """Run a call in the pool, then give its place to the next waiting call. - - The place is handed on in the thread of the call, before the call - ends in the pool, so that an interrupt of the submitter cannot keep - it, and a release of the pool waits for the calls that wait here. + 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: - context: The context of the submitter of the call. - item: The call. + call: The call. """ - try: - context.run(item.run) - finally: - self._next() + # The work item settles the future also on a failure. + call.context.run(call.item.run) + self._next(call) - def _next(self) -> None: - """Give a place of the scope to the next waiting call, or free it. + 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, and so are those that the pool rejects. + entering the pool. + + Args: + ended: The call that has ended. """ + with self._lock: + self._entered.discard(ended) while True: with self._lock: - if not self._waiting: - self._running -= 1 + if not self._waiting or len(self._entered) >= self.max_workers: return - context, item, future = self._waiting.popleft() - if future.cancelled(): - # Marks the future as notified, without running the call. - item.run() - continue - if self._enter(context, item): - return + call = self._waiting.popleft() + if call.future.cancelled(): + run_inline = True + else: + run_inline = False + try: + self._pool.submit(self._run, call) + except RuntimeError: + # The pool takes no more calls. + run_inline = True + else: + self._entered.add(call) + if run_inline: + # A cancelled call is only marked as notified. + call.context.run(call.item.run) def _untrack(self, future: Future[Any]) -> None: """Stop tracking a call that has ended or was cancelled. @@ -639,6 +630,32 @@ def _untrack(self, future: Future[Any]) -> None: self._futures.discard(future) +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/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index b17741514..efbc98a6b 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -2258,7 +2258,7 @@ def test_file_inherited_across_fork(self, fs): asyncio.run(f.write(b"x")) async def use(): - for call in (lambda: f.write(b"y"), f.flush, lambda: f._fetch_range(0, 1)): + 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() @@ -2712,8 +2712,9 @@ async def flush_and_exit(): await await_bounded(close() if cancelled == "close" else flush_and_exit()) finally: release.set() - # Waits for the commit, and for an abort queued after it. - await asyncio.to_thread(shutdown_bounded, f._requests) + # Waits for the commit, and for an abort queued after it, in the + # threads of the filesystem. + await asyncio.to_thread(shutdown_bounded, fs._thread_pool) assert f.closed assert released == [True] @@ -2726,20 +2727,15 @@ async def test_flush_cancelled_while_queued_then_exit_with_exception(self, fs): # 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 @@ -2748,13 +2744,15 @@ def record_commit(pool, fn, *args, **kwargs): queued.set() return submit(pool, 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() @@ -2766,12 +2764,12 @@ async def flush_and_exit(): pytest.raises(asyncio.CancelledError), ): # Bounded, so that a regression whose exit waits for the - # held part fails here and releases it. + # held thread 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._requests) + await asyncio.to_thread(shutdown_bounded, fs._thread_pool) assert f.closed assert released == [True] diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index 5c62c56b4..226b355b3 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -413,26 +413,14 @@ def test_cancel_waiting_call(self): assert events == [] def test_interrupted_submit(self): - # GH-1332: an interrupt before a call reaches the pool cancels the - # call and gives its place to the next one, so that the scope neither - # stops running calls nor waits for the interrupted one. + # 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) - add_done_callback = Future.add_done_callback - interrupts = [KeyboardInterrupt] - - def interrupt_once(future, fn): - # The first callback is the scope's, added after it reserved a - # place for the call. - if interrupts: - raise interrupts.pop() - add_done_callback(future, fn) - try: with ( - mock.patch.object( - Future, "add_done_callback", autospec=True, side_effect=interrupt_once - ), + mock.patch.object(pool, "submit", side_effect=KeyboardInterrupt), pytest.raises(KeyboardInterrupt), ): scope.submit(sum, [1, 2]) @@ -441,6 +429,25 @@ def interrupt_once(future, fn): finally: pool.shutdown() + 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 From ddc302bd668b8990982656889251a631efc2c38a Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 5 Nov 2026 02:20:46 +0900 Subject: [PATCH 05/22] Settle a call interrupted after the pool accepted it through its work item (GH-1332) An interrupt after the pool accepted a call that had not started marked its future notified directly, so that the work item raised when the pool ran it, and the call kept its place in the scope for good. The work item now settles the cancelled future once, and the place is freed even if the call raises. The docstring states that each such interrupt can let one more call run at once. Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3_executor.py | 20 +++++++----- tests/pyathena/filesystem/test_s3_executor.py | 32 +++++++++++++++++++ 2 files changed, 44 insertions(+), 8 deletions(-) diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index a07fa1cc3..903081b95 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -475,8 +475,10 @@ class _S3PoolScope(S3Executor): 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: - at worst, one more call runs at once for a while. + 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. @@ -515,7 +517,7 @@ def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]: As with ``ThreadPoolExecutor``, the returned future can be cancelled only before the call starts. An interrupt of the submission cancels - the call. + the call unless it has started. Args: fn: The callable to run. @@ -550,9 +552,10 @@ def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]: 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. + # 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(): - future.set_running_or_notify_cancel() + call.item.run() with self._lock: self._futures.discard(future) raise @@ -585,9 +588,10 @@ def _run(self, call: _S3ScopeCall) -> None: Args: call: The call. """ - # The work item settles the future also on a failure. - call.context.run(call.item.run) - self._next(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. diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index 226b355b3..2f09646b1 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -429,6 +429,38 @@ def test_interrupted_submit(self): 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 From 4786e58bf4e5929b8200bba3eaac56a3f8c2531e Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 5 Nov 2026 16:03:16 +0900 Subject: [PATCH 06/22] Keep scope order under concurrent submitters, free cancelled calls, inject interrupts (GH-1332) Per the specification on #1332: - S3: a submitter queues behind any waiting call, so a call submitted while a place frees does not pass an older waiting call. - S8: a call cancelled while it waits leaves the queue at once, so that its arguments and context are not kept until a place frees. tests/pyathena/util.py gains interrupt_each_opcode(), which raises a KeyboardInterrupt at each point of a target where a signal can raise it: function entry, after a call returns, a backward jump target, and a with statement. It uses sys.monitoring on Python 3.12 and later, where a trace function gets no opcode events on the first call of a code object, and sys.settrace on 3.11. The submit() of the pool and the scope, and submit_to() through a scope, are interrupted at each of their points, after which nothing waits forever and the next call runs. Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3_executor.py | 17 +- tests/pyathena/filesystem/test_s3_executor.py | 145 ++++++++++++++++ tests/pyathena/util.py | 157 ++++++++++++++++++ 3 files changed, 318 insertions(+), 1 deletion(-) diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index 903081b95..fe60decf1 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -541,7 +541,9 @@ def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]: if self._shutdown: raise RuntimeError("cannot schedule new futures after shutdown") self._futures.add(future) - enter = len(self._entered) < self.max_workers + # 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 and len(self._entered) < self.max_workers if enter: # Counted only once the pool has accepted it. self._pool.submit(self._run, call) @@ -627,11 +629,24 @@ def _next(self, ended: _S3ScopeCall) -> None: 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. """ + dropped = None with self._lock: self._futures.discard(future) + if future.cancelled(): + for call in self._waiting: + if call.future is future: + self._waiting.remove(call) + dropped = call + break + if dropped is not None: + # Marks the future as notified, without running the call. + dropped.item.run() class _S3ScopeCall: diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index 2f09646b1..c0c0dc5ec 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -25,6 +25,7 @@ EVENT_TIMEOUT, ThrottledOperation, forking, + interrupt_each_opcode, overlapping_calls, run_in_forked_child, shutdown_bounded, @@ -32,6 +33,26 @@ 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) + events = [] + try: + future = interrupt(pool.submit, events.append, "call") + assert pool.submit(sum, [1, 2]).result(EVENT_TIMEOUT) == 3 + pool.release() + shutdown_bounded(pool) + if future is not None: + assert future.done() + assert events in ([], ["call"]) + finally: + 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: @@ -312,6 +333,130 @@ def hold_lock(): 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() + assert scope.submit(sum, [1, 2]).result(EVENT_TIMEOUT) == 3 + shutdown_bounded(scope) + # A call frees its place right after its future is settled. + deadline = time.monotonic() + EVENT_TIMEOUT + while scope._entered: + assert time.monotonic() < deadline, "a place was never freed" + time.sleep(0.001) + assert not scope._waiting + if future is not None: + assert future.done() + 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 and future.cancel(): + future.set_running_or_notify_cancel() + release.set() + assert not wait([future], timeout=EVENT_TIMEOUT).not_done + shutdown_bounded(scope) + deadline = time.monotonic() + EVENT_TIMEOUT + while scope._entered: + assert time.monotonic() < deadline, "a place was never freed" + time.sleep(0.001) + assert events == ([] if future.cancelled() else ["call"]) + finally: + release.set() + pool.shutdown() + + assert interrupt_each_opcode(S3Executor.submit_to, case) > 0 + + def test_concurrent_submitter_does_not_pass_waiting_call(self): + # GH-1332 (S3): a call submitted from another thread at the moment a + # place frees does not enter ahead of a call that waits for it. The + # place used to be freed and handed on under two acquisitions of the + # lock, between which a new submitter took it. + pool = S3ThreadPool(2) + scope = _S3PoolScope(pool, 1) + release = threading.Event() + order = [] + late: list[Future[None]] = [] + submitter = threading.Thread( + target=lambda: late.append(scope.submit(order.append, "late")), daemon=True + ) + + class FreeingSet(set): + def discard(self, element): + super().discard(element) + if not submitter.is_alive() and not late: + # The submitter takes the free place now unless the + # scope holds it off until the place is handed on. + submitter.start() + submitter.join(0.2) + + try: + scope.submit(release.wait, EVENT_TIMEOUT) + scope._entered = FreeingSet(scope._entered) + waiting = scope.submit(order.append, "waiting") + release.set() + assert waiting.result(EVENT_TIMEOUT) is None + submitter.join(EVENT_TIMEOUT) + late[0].result(EVENT_TIMEOUT) + shutdown_bounded(scope) + finally: + release.set() + pool.shutdown() + assert order == ["waiting", "late"] + + 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 + assert not scope._waiting + finally: + release.set() + shutdown_bounded(scope) + 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. diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index cffc07f86..05d2cd096 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -6,9 +6,12 @@ # SPDX-License-Identifier: MIT import asyncio +import dis import io +import itertools import os import signal +import sys import threading import time import traceback @@ -246,6 +249,160 @@ 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. + """ + + def __init__(self, code, n): + self.code = code + self.n = n + self.count = 0 + self.fired = False + 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: + if not self.fired: + raise + 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)) + 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): + points.add(instructions[i + 1].offset) + elif name == "JUMP_BACKWARD" and instruction.jump_target is not None: + points.add(instruction.jump_target) + 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 or threading.get_ident() != self._thread: + return + self.count += 1 + if self.count == self.n: + self.fired = True + 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 + + +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__ + for n in itertools.count(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 + raise AssertionError("unreachable") + + def mock_connection(max_pool_connections=10): """Return a mock connection whose S3 client has a connection pool of the given size. From 7d53dc9b0ae2ba87b3e3524715b6735e75ccc89c Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 5 Nov 2026 16:06:08 +0900 Subject: [PATCH 07/22] Wait for a streamed file's requests and commit when it closes (GH-1332) Per F2 of the specification on #1332, AioS3StreamedFile.close() and the close without commit of an async with block that raises wait, without blocking the event loop, for the requests of the file and a scheduled commit, also after a read, flush, or the close itself was cancelled. They used to return while a cancelled read's request or a queued commit still ran. _S3PoolScope.pending() returns the futures to wait for. Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3_async.py | 41 +++++++++-- pyathena/filesystem/s3_executor.py | 10 +++ tests/pyathena/filesystem/test_s3_async.py | 85 +++++++++++++++++----- 3 files changed, 112 insertions(+), 24 deletions(-) diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index aa31e6aae..1fc7e8368 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -32,7 +32,7 @@ from pyathena.aio.util import create_in_thread from pyathena.filesystem.s3 import CompressedBuffer, S3File, S3FileSystem, S3FileTarget from pyathena.filesystem.s3_core import S3Core -from pyathena.filesystem.s3_executor import S3ThreadPool +from pyathena.filesystem.s3_executor import S3ThreadPool, _S3PoolScope from pyathena.filesystem.s3_object import ( S3CompleteMultipartUpload, S3Metadata, @@ -1276,7 +1276,9 @@ def __init__( # 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. - self._requests = fs._sync_fs._create_executor(max_workers=max_workers) + self._requests = cast(_S3PoolScope, fs._sync_fs._create_executor(max_workers=max_workers)) + # The scheduled commit, which close() waits for. + self._commit_future: Future[None] | None = None self.max_workers = max_workers # The process that opened the file; a forked child rejects it. self._pid = os.getpid() @@ -1480,8 +1482,11 @@ async def flush(self, force: bool = False) -> None: async def close(self) -> None: """Close the file, committing the written data. - In a forked child process, a file that the parent opened is only - marked closed, without committing or aborting anything. + 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 @@ -1493,8 +1498,28 @@ async def close(self) -> None: await self.flush(force=True) finally: self.closed = True - # A commit that is still running waits for the requests on. - self._requests.shutdown(wait=False) + await self._drain() + + async def _drain(self) -> None: + """Wait, without blocking the event loop, for the requests and the commit of the file. + + A cancellation of the waiting task leaves them running, as the + cancellation of a commit does, and is re-raised. + """ + self._requests.shutdown(wait=False) + pending = self._requests.pending() + if self._commit_future is not None: + pending.append(self._commit_future) + if not pending: + return + waiters = [asyncio.wrap_future(future) for future in pending] + for waiter in waiters: + # The outcomes belong to the reads and the commit that awaited + # them; retrieved here, so that asyncio does not log them. + waiter.add_done_callback(lambda f: f.cancelled() or f.exception()) + # Unlike gather(), wait() neither raises their outcomes nor cancels + # them when this task is cancelled. + await asyncio.wait(waiters) async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: if self.closed: @@ -1520,7 +1545,8 @@ async def _close_without_commit(self) -> None: if self._begin_close_without_commit(): await self.fs._thread_pool.run_shielded(self._discard) finally: - self._requests.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. @@ -1649,6 +1675,7 @@ async def _commit(self) -> None: session.begin_complete() # Passed, as closing the file without committing it drops the buffer. commit = self.fs._thread_pool.submit(self._complete_upload, cast(BytesIO, self.buffer)) + self._commit_future = commit except BaseException: if commit is None: # Not queued, or queued by a submit() that an interrupt diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index fe60decf1..18030726a 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -579,6 +579,16 @@ def shutdown(self, wait: bool = True) -> None: if wait: futures.wait(pending) + def pending(self) -> list[Future[Any]]: + """Return the futures of the calls that were submitted and have not ended. + + Returns: + The futures, which an asyncio caller can wait for without + blocking the event loop, such as with :meth:`wait_shielded`. + """ + with self._lock: + return list(self._futures) + def _run(self, call: _S3ScopeCall) -> None: """Run a call in the pool, then let the next waiting calls enter it. diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index efbc98a6b..e6881bb27 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -2676,7 +2676,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() @@ -2707,20 +2709,68 @@ 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, in the - # threads of the filesystem. - await asyncio.to_thread(shutdown_bounded, fs._thread_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() + release.set() + await await_bounded(close) + finally: + release.set() + if close is not None: + close.cancel() + assert finished == [(0, block_size)] + @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 @@ -2758,18 +2808,19 @@ async def flush_and_exit(): 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 thread 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, fs._thread_pool) + with mock.patch.object(S3ThreadPool, "submit", 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] From 8bfe3b4945b0e19bdfac5ddb2e43a38eb930b545 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 5 Nov 2026 16:09:18 +0900 Subject: [PATCH 08/22] Free a parallel read's data on an interrupt at any point (GH-1332) Per F5 of the specification on #1332, the cleanup of a parallel read of S3File also covers the code after its wait: an interrupt there, such as while the results are joined, used to leave their futures in the frame that the traceback of the interrupt keeps. The interrupt injection of the tests raises an interrupt that a signal checked after a call would raise at the call itself before the call when the next instruction is outside the handler that covers the call, as at a return that ends a try block; the read is interrupted at each of its points. Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3.py | 31 +++++++++++---------- tests/pyathena/filesystem/test_s3.py | 40 ++++++++++++++++++++++++++++ tests/pyathena/util.py | 21 +++++++++++++-- 3 files changed, 74 insertions(+), 18 deletions(-) diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index b7e117b5b..130c5b4c8 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -3428,6 +3428,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: @@ -3440,30 +3441,28 @@ 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. + # 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; the frames of the wait, which keep the futures, are + # cleared, and clear_frames() skips this frame, which is + # executing. traceback.clear_frames(e.__traceback__) + done.clear() while futures: futures.pop().cancel() # 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/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index 1f947ebb5..d2276a532 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -56,6 +56,7 @@ MULTIPART_COPY_SIZE, ThrottledOperation, forking, + interrupt_each_opcode, mock_connection, overlapping_calls, retried, @@ -3608,6 +3609,45 @@ 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. + deadline = time.monotonic() + EVENT_TIMEOUT + while any(ref() is not None for ref in read): + assert time.monotonic() < deadline, "the data of a read was kept" + time.sleep(0.001) + finally: + f.close() + fs._request_pool.shutdown() + + assert interrupt_each_opcode(S3File._fetch_range, case) > 0 + 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 diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index 05d2cd096..12161c2e8 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -277,6 +277,9 @@ def __init__(self, code, n): self.n = n 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) @@ -297,9 +300,10 @@ def __call__(self, fn, *args, **kwargs): with self._armed(): try: return fn(*args, **kwargs) - except KeyboardInterrupt: + except KeyboardInterrupt as e: if not self.fired: raise + self.error = e return None @staticmethod @@ -313,13 +317,26 @@ def _interrupt_points(code): 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): - points.add(instructions[i + 1].offset) + 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" and instruction.jump_target is not None: points.add(instruction.jump_target) return points From c8b7300d9b36c1f168ccf459d9f3afa61633306c Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 5 Nov 2026 16:11:39 +0900 Subject: [PATCH 09/22] Default the cursors' s3_max_workers to the client's connection pool (GH-1332) Per decision 1 on #1332, the pandas and Polars cursors, their async and aio variants, and their result sets default s3_max_workers to None, so that the filesystem that reads the result files sizes its request pool from the max_pool_connections of the connection's S3 client (10 unless set with s3_config). They used to default to cpu_count() * 5, which ran more requests at once than the client keeps connections for. An explicit value is used as given. docs/usage.md also says that Polars uses it for UNLOAD results read in chunks through PyArrow. Co-Authored-By: Claude Opus 5.5 --- docs/usage.md | 7 +++---- pyathena/aio/pandas/cursor.py | 8 +++++--- pyathena/aio/polars/cursor.py | 8 +++++--- pyathena/pandas/async_cursor.py | 7 +++++-- pyathena/pandas/cursor.py | 8 +++++--- pyathena/pandas/result_set.py | 13 ++++++++----- pyathena/polars/async_cursor.py | 7 +++++-- pyathena/polars/cursor.py | 8 +++++--- pyathena/polars/result_set.py | 7 ++++--- tests/pyathena/pandas/test_cursor.py | 12 ++++++++++++ tests/pyathena/pandas/test_result_set.py | 15 +++++++++++++++ tests/pyathena/polars/test_cursor.py | 12 ++++++++++++ 12 files changed, 84 insertions(+), 28 deletions(-) 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/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 77004d3da..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, 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/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 From abde069bd5e3e0aeba4a96e9de21be3986e6633a Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 5 Nov 2026 16:19:13 +0900 Subject: [PATCH 10/22] Add barrier cases with a real SIGINT for waits that block (GH-1332) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per §6.2 of the specification on #1332, sigint_when() in the test helpers delivers a real SIGINT to the main thread once a condition holds, so that a blocking wait is interrupted inside its C call. - A shutdown of a scope interrupted while it waits for a running call leaves the call running, and a later shutdown waits for it (I1, S5). - A parallel read of S3File interrupted while it waits for its ranges cancels those that have not started and frees the data of the others (F5). 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 rather than when the reads end; the test collects before it checks. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/filesystem/test_s3.py | 50 +++++++++++++++++++ tests/pyathena/filesystem/test_s3_executor.py | 33 ++++++++++++ tests/pyathena/util.py | 50 +++++++++++++++++++ 3 files changed, 133 insertions(+) diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index d2276a532..c5c3fe5a6 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -61,6 +61,7 @@ overlapping_calls, retried, run_in_forked_child, + sigint_when, stub_multipart_copy, ) @@ -3648,6 +3649,55 @@ def get_object(path, range_, **kwargs): 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 + deadline = time.monotonic() + EVENT_TIMEOUT + while any(ref() is not None for ref in read): + assert time.monotonic() < deadline, "the data of a read was kept" + gc.collect() + time.sleep(0.001) + # 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 diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index c0c0dc5ec..b60d8ba2d 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -29,6 +29,7 @@ overlapping_calls, run_in_forked_child, shutdown_bounded, + sigint_when, ) @@ -457,6 +458,38 @@ class Body: 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: concurrent.futures + # installs its waiter on the future, inside the block. + with sigint_when(lambda: pending._waiters), 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_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. diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index 12161c2e8..92b770be3 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -393,6 +393,56 @@ def _trace_opcode(self, frame, event, arg): 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. From 7b77f375df9b06f1aad7665ce232fad59be24ef6 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 5 Nov 2026 16:34:40 +0900 Subject: [PATCH 11/22] Give the local test filesystem the close() of PyAthena's filesystems (GH-1332) The pandas result set closes the filesystem that it created, to release its idle threads. The tests that read a local result file through LocalBucketFileSystem in place of S3FileSystem need the same close(). Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/util.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index 92b770be3..ce0aefc76 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -1604,3 +1604,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.""" From cebbe077a9fb0a7e0caf9ee446d7fb03602a84ab Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 5 Nov 2026 17:14:24 +0900 Subject: [PATCH 12/22] Run a pool thread's scope call in that thread, and default files to the pool (GH-1332) From the review of #1359 against the specification on #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. S3ThreadPool marks its threads (owns_current_thread()). - S5: a waiting call that the pool fails to take for a reason other than its shutdown, such as a thread that cannot start, fails instead of being lost with its place. - S8: a call cancelled while it waits is dropped by its future in O(1), and its arguments and context are freed at once; the entry left in the queue is skipped. - F1: S3File and AioS3StreamedFile default max_workers to that of the filesystem, and S3File without an executor takes a scope of the filesystem's request pool instead of a pool of its own. _create_executor() is typed to return the scope. - F4: files tell that they crossed a fork by a fork count that an at-fork hook bumps, instead of calling os.getpid() on every read and write. The tests check the scope through what a caller observes; the mock filesystem of the TestS3File helpers gets a request pool. Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3.py | 39 ++++++---- pyathena/filesystem/s3_async.py | 18 ++--- pyathena/filesystem/s3_executor.py | 72 ++++++++++++++----- tests/pyathena/filesystem/test_s3.py | 31 +++++++- tests/pyathena/filesystem/test_s3_async.py | 10 +++ tests/pyathena/filesystem/test_s3_executor.py | 59 +++++++++++---- 6 files changed, 174 insertions(+), 55 deletions(-) diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 130c5b4c8..7b2d5ae6a 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -16,8 +16,7 @@ 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 typing import Any, BinaryIO, ClassVar, cast from urllib.parse import unquote_plus import botocore.exceptions @@ -1133,7 +1132,7 @@ 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 @@ -2883,6 +2882,9 @@ class S3File(AbstractBufferedFile): fs: S3FileSystem buffer: BytesIO | None + # The forks of this process so far, which a file compares with those at + # its opening to tell that it was inherited across a fork. + _forks: ClassVar[int] = 0 def __init__( self, @@ -2890,7 +2892,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", @@ -2919,13 +2921,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. @@ -2953,10 +2955,12 @@ def __init__( ``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() + # The forks at the opening; a forked child rejects the file. + self._forks_at_open = S3File._forks 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 @@ -2976,7 +2980,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": @@ -3071,6 +3075,11 @@ def flush(self, force: bool = False) -> None: self._check_process() super().flush(force=force) + @classmethod + def _count_fork(cls) -> None: + """Count a fork, in the child process.""" + cls._forks += 1 + def _inherited(self) -> bool: """Whether this process is a forked child of the process that opened the file. @@ -3078,7 +3087,7 @@ def _inherited(self) -> bool: True in a forked child process of the process that opened the file. """ - return getattr(self, "_pid", os.getpid()) != os.getpid() + return getattr(self, "_forks_at_open", S3File._forks) != S3File._forks def _check_process(self) -> None: """Reject the use of a file that the parent of a forked process opened. @@ -3485,3 +3494,7 @@ def _get_ranges( else: ranges.append((start, end)) return ranges + + +if hasattr(os, "register_at_fork"): + os.register_at_fork(after_in_child=S3File._count_fork) diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 1fc7e8368..4cab118a8 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 @@ -32,7 +31,7 @@ from pyathena.aio.util import create_in_thread from pyathena.filesystem.s3 import CompressedBuffer, S3File, S3FileSystem, S3FileTarget from pyathena.filesystem.s3_core import S3Core -from pyathena.filesystem.s3_executor import S3ThreadPool, _S3PoolScope +from pyathena.filesystem.s3_executor import S3ThreadPool from pyathena.filesystem.s3_object import ( S3CompleteMultipartUpload, S3Metadata, @@ -1234,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, @@ -1251,7 +1250,8 @@ def __init__( mode: The file mode: ``rb``, ``wb``, ``ab``, or ``xb``. version_id: The version ID to read. max_workers: The number of parallel requests of the file, in the - request pool of the filesystem. + 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 @@ -1276,12 +1276,14 @@ def __init__( # 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. - self._requests = cast(_S3PoolScope, fs._sync_fs._create_executor(max_workers=max_workers)) + 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, which close() waits for. self._commit_future: Future[None] | None = None self.max_workers = max_workers - # The process that opened the file; a forked child rejects it. - self._pid = os.getpid() + # The forks at the opening; a forked child rejects the file. + self._forks_at_open = S3File._forks target = S3FileTarget.resolve( fs._sync_fs, path, @@ -1358,7 +1360,7 @@ def _inherited(self) -> bool: True in a forked child process of the process that opened the file. """ - return getattr(self, "_pid", os.getpid()) != os.getpid() + return getattr(self, "_forks_at_open", S3File._forks) != S3File._forks def _check_process(self) -> None: """Reject the use of a file that the parent of a forked process opened. diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index 18030726a..40bfc8200 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -265,6 +265,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. @@ -426,6 +430,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. @@ -455,7 +468,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 @@ -508,6 +525,8 @@ def __init__(self, pool: S3ThreadPool, max_workers: int) -> None: # 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 @@ -536,6 +555,15 @@ def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]: 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). + with self._lock: + if self._shutdown: + raise RuntimeError("cannot schedule new futures after shutdown") + call.item.run() + return future try: with self._lock: if self._shutdown: @@ -543,13 +571,14 @@ def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]: self._futures.add(future) # 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 and len(self._entered) < self.max_workers + 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, @@ -621,20 +650,29 @@ def _next(self, ended: _S3ScopeCall) -> None: if not self._waiting or len(self._entered) >= self.max_workers: return call = self._waiting.popleft() - if call.future.cancelled(): + 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. run_inline = True + except Exception as e: + # Such as a failure to start a thread: the call fails. + run_inline, error = False, e else: run_inline = False - try: - self._pool.submit(self._run, call) - except RuntimeError: - # The pool takes no more calls. - run_inline = True - else: - self._entered.add(call) + self._entered.add(call) if run_inline: - # A cancelled call is only marked as notified. 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. @@ -645,18 +683,14 @@ def _untrack(self, future: Future[Any]) -> None: Args: future: The future of the call. """ - dropped = None with self._lock: self._futures.discard(future) - if future.cancelled(): - for call in self._waiting: - if call.future is future: - self._waiting.remove(call) - dropped = call - break + dropped = self._waiting_calls.pop(future, None) if future.cancelled() else None if dropped is not None: - # Marks the future as notified, without running the call. + # 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: diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index c5c3fe5a6..edd8e8e74 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -42,7 +42,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 @@ -2748,6 +2748,29 @@ def read(_): 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 @@ -6440,6 +6463,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 e6881bb27..15d0b13a2 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -2278,6 +2278,16 @@ def child(): 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 diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index b60d8ba2d..964835cef 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -350,14 +350,9 @@ def case(interrupt): 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) - # A call frees its place right after its future is settled. - deadline = time.monotonic() + EVENT_TIMEOUT - while scope._entered: - assert time.monotonic() < deadline, "a place was never freed" - time.sleep(0.001) - assert not scope._waiting if future is not None: assert future.done() assert events in ([], ["call"]) @@ -385,11 +380,8 @@ def case(interrupt): future.set_running_or_notify_cancel() release.set() assert not wait([future], timeout=EVENT_TIMEOUT).not_done + assert scope.submit(sum, [1, 2]).result(EVENT_TIMEOUT) == 3 shutdown_bounded(scope) - deadline = time.monotonic() + EVENT_TIMEOUT - while scope._entered: - assert time.monotonic() < deadline, "a place was never freed" - time.sleep(0.001) assert events == ([] if future.cancelled() else ["call"]) finally: release.set() @@ -399,9 +391,9 @@ def case(interrupt): def test_concurrent_submitter_does_not_pass_waiting_call(self): # GH-1332 (S3): a call submitted from another thread at the moment a - # place frees does not enter ahead of a call that waits for it. The - # place used to be freed and handed on under two acquisitions of the - # lock, between which a new submitter took it. + # place frees does not enter ahead of a call that waits for it. A + # barrier case (spec section 6.2): the submitter is started from the + # internal point where the place is freed, before it is handed on. pool = S3ThreadPool(2) scope = _S3PoolScope(pool, 1) release = threading.Event() @@ -452,7 +444,6 @@ class Body: assert waiting.cancel() gc.collect() assert body_ref() is None - assert not scope._waiting finally: release.set() shutdown_bounded(scope) @@ -490,6 +481,46 @@ def call(): 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. From 80ad4f2b4641abc948978faa01ff01f0df11316c Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 5 Nov 2026 17:16:54 +0900 Subject: [PATCH 13/22] Share the polling and SIGINT test helpers (GH-1332) The tests that wait for the data of reads to be freed use poll_until(), and the interrupted-creation test of the multipart copy uses sigint_when() instead of its own copy of the SIGINT handling. interrupt_each_opcode() loops without unreachable code. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/filesystem/test_s3.py | 78 ++++++++-------------------- tests/pyathena/util.py | 26 ++++++++-- 2 files changed, 45 insertions(+), 59 deletions(-) diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index edd8e8e74..bdf2f5382 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 @@ -59,6 +58,7 @@ interrupt_each_opcode, mock_connection, overlapping_calls, + poll_until, retried, run_in_forked_child, sigint_when, @@ -3662,10 +3662,9 @@ def get_object(path, range_, **kwargs): assert interrupt.error.__traceback__ is not None # The threads leave the frames of the reads right after the # close has waited for them. - deadline = time.monotonic() + EVENT_TIMEOUT - while any(ref() is not None for ref in read): - assert time.monotonic() < deadline, "the data of a read was kept" - time.sleep(0.001) + 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() @@ -3709,11 +3708,11 @@ def get_object(path, range_, **kwargs): release.set() f.close() assert info.value.__traceback__ is not None - deadline = time.monotonic() + EVENT_TIMEOUT - while any(ref() is not None for ref in read): - assert time.monotonic() < deadline, "the data of a read was kept" - gc.collect() - time.sleep(0.001) + 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: @@ -3815,12 +3814,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) @@ -3836,52 +3834,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 diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index ce0aefc76..ce31ca002 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -7,8 +7,8 @@ import asyncio import dis +import gc import io -import itertools import os import signal import sys @@ -216,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. @@ -461,13 +480,14 @@ def interrupt_each_opcode(target, case, limit=2000): The number of opcodes that were interrupted. """ code = getattr(target, "__func__", target).__code__ - for n in itertools.count(1): + 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 - raise AssertionError("unreachable") def mock_connection(max_pool_connections=10): From dee243f7788458781944ae458588a607cc04362b Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 5 Nov 2026 17:49:21 +0900 Subject: [PATCH 14/22] Track every call of a file and scope through cancellation (GH-1332) From the independent review against the specification: - A streamed file's close drains its requests, commit, and abort also when its task is cancelled, repeatedly, and raises the cancellation after they end. The commit and abort are tracked before they are queued, so an interrupt of the queuing or a cancellation of the abort no longer leaves them untracked (F2, I1). - A call that a pool thread runs itself is tracked by the scope, whose shutdown waits for it, and runs in the captured context (S4, S5). - A place that frees passes to the oldest waiting call under the lock that submitters take (S3); the entries of cancelled waiting calls are compacted so that they do not accumulate behind a long call (S8). - A failed or interrupted parallel read drops its futures in a finally, so a second interrupt during the cleanup cannot keep their data (F5). - Files compare the process ID again: an at-fork hook registered before this module's could close an inherited file before a counter updates. - The opcode injector reads jump targets from argval, which Python 3.12 supports, and the scope and pool tests assert the bounds of S9 and P6 without notifying futures themselves or relying on a timed barrier. Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3.py | 47 ++++----- pyathena/filesystem/s3_async.py | 61 +++++++----- pyathena/filesystem/s3_executor.py | 19 +++- tests/pyathena/filesystem/test_s3_async.py | 68 ++++++++++--- tests/pyathena/filesystem/test_s3_executor.py | 98 +++++++++++++------ tests/pyathena/util.py | 5 +- 6 files changed, 201 insertions(+), 97 deletions(-) diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 7b2d5ae6a..001a84787 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -16,7 +16,7 @@ from dataclasses import dataclass, field from datetime import datetime from io import BytesIO -from typing import Any, BinaryIO, ClassVar, cast +from typing import Any, BinaryIO, cast from urllib.parse import unquote_plus import botocore.exceptions @@ -2882,9 +2882,6 @@ class S3File(AbstractBufferedFile): fs: S3FileSystem buffer: BytesIO | None - # The forks of this process so far, which a file compares with those at - # its opening to tell that it was inherited across a fork. - _forks: ClassVar[int] = 0 def __init__( self, @@ -2959,8 +2956,8 @@ def __init__( max_workers = fs.max_workers if max_workers <= 0: raise ValueError("max_workers must be greater than 0") - # The forks at the opening; a forked child rejects the file. - self._forks_at_open = S3File._forks + # 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 @@ -3075,11 +3072,6 @@ def flush(self, force: bool = False) -> None: self._check_process() super().flush(force=force) - @classmethod - def _count_fork(cls) -> None: - """Count a fork, in the child process.""" - cls._forks += 1 - def _inherited(self) -> bool: """Whether this process is a forked child of the process that opened the file. @@ -3087,7 +3079,9 @@ def _inherited(self) -> bool: True in a forked child process of the process that opened the file. """ - return getattr(self, "_forks_at_open", S3File._forks) != S3File._forks + # 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. @@ -3461,16 +3455,21 @@ def _fetch_range(self, start: int, end: int) -> bytes: # 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; the frames of the wait, which keep the futures, are - # cleared, and clear_frames() skips this frame, which is - # executing. - traceback.clear_frames(e.__traceback__) - done.clear() - while futures: - futures.pop().cancel() - # After the cancellation, which a thread that it frees could - # otherwise precede by taking a queued read. - stop.set() + # 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() + # After the cancellation, which a thread that it frees + # could otherwise precede by taking a queued read. + stop.set() + traceback.clear_frames(e.__traceback__) + finally: + done.clear() + futures.clear() raise return self.fs.core.get_object(path, ranges[0], **request_kwargs) @@ -3494,7 +3493,3 @@ def _get_ranges( else: ranges.append((start, end)) return ranges - - -if hasattr(os, "register_at_fork"): - os.register_at_fork(after_in_child=S3File._count_fork) diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 4cab118a8..8d9073dcd 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -1279,11 +1279,12 @@ def __init__( 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, which close() waits for. + # The scheduled commit and abort, which close() waits for. self._commit_future: Future[None] | None = None + self._discard_future: Future[None] | None = None self.max_workers = max_workers - # The forks at the opening; a forked child rejects the file. - self._forks_at_open = S3File._forks + # The process that opened the file; a forked child rejects it. + self._pid = os.getpid() target = S3FileTarget.resolve( fs._sync_fs, path, @@ -1360,7 +1361,7 @@ def _inherited(self) -> bool: True in a forked child process of the process that opened the file. """ - return getattr(self, "_forks_at_open", S3File._forks) != S3File._forks + 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. @@ -1503,25 +1504,33 @@ async def close(self) -> None: await self._drain() async def _drain(self) -> None: - """Wait, without blocking the event loop, for the requests and the commit of the file. + """Wait, without blocking the event loop, for the requests, commit, and abort of the file. - A cancellation of the waiting task leaves them running, as the - cancellation of a commit does, and is re-raised. + 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). """ self._requests.shutdown(wait=False) pending = self._requests.pending() - if self._commit_future is not None: - pending.append(self._commit_future) + pending.extend(f for f in (self._commit_future, self._discard_future) if f is not None) if not pending: return waiters = [asyncio.wrap_future(future) for future in pending] for waiter in waiters: - # The outcomes belong to the reads and the commit that awaited - # them; retrieved here, so that asyncio does not log them. + # 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()) - # Unlike gather(), wait() neither raises their outcomes nor cancels - # them when this task is cancelled. - await asyncio.wait(waiters) + cancelled = False + while True: + try: + # Unlike gather(), wait() neither raises their outcomes nor + # cancels them when this task is cancelled. + await asyncio.wait(waiters) + break + 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: @@ -1545,7 +1554,12 @@ async def _close_without_commit(self) -> None: return try: if self._begin_close_without_commit(): - await self.fs._thread_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) + await self.fs._thread_pool.wait_shielded(discard) finally: # A commit that was scheduled before runs on, and is waited for. await self._drain() @@ -1672,17 +1686,20 @@ 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.fs._thread_pool.submit(self._complete_upload, cast(BytesIO, self.buffer)) - self._commit_future = commit + self.fs._thread_pool.submit_to( + commit, self._complete_upload, cast(BytesIO, self.buffer) + ) 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 diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index 40bfc8200..3531fa5a6 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -558,11 +558,14 @@ def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]: 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). + # 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") - call.item.run() + self._futures.add(future) + future.add_done_callback(self._untrack) + call.context.run(call.item.run) return future try: with self._lock: @@ -643,10 +646,14 @@ def _next(self, ended: _S3ScopeCall) -> None: Args: ended: The call that has ended. """ - with self._lock: - self._entered.discard(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() @@ -686,6 +693,10 @@ def _untrack(self, future: Future[Any]) -> None: with self._lock: self._futures.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. diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index 15d0b13a2..ae5b83ef8 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -2560,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: @@ -2588,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() @@ -2773,13 +2773,18 @@ def get_object(path, range_, **kwargs): 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() - await await_bounded(close) + with pytest.raises(asyncio.CancelledError): + await await_bounded(close) + assert finished == [(0, block_size)] finally: release.set() if close is not None: close.cancel() - assert finished == [(0, block_size)] @pytest.mark.asyncio async def test_flush_cancelled_while_queued_then_exit_with_exception(self, fs): @@ -2797,12 +2802,12 @@ def hold(): released.append(release.wait(EVENT_TIMEOUT)) 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) fs._thread_pool = S3ThreadPool(1) f = await fs.open_async("s3://bucket/key", "wb", block_size=block_size) @@ -2818,7 +2823,7 @@ async def flush_and_exit(): flush.cancel() await flush - with mock.patch.object(S3ThreadPool, "submit", autospec=True, side_effect=record_commit): + 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)) @@ -2843,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): @@ -2859,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: @@ -2917,6 +2922,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): diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index 964835cef..c3fbbfc67 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -45,6 +45,11 @@ def case(interrupt): future = interrupt(pool.submit, events.append, "call") assert pool.submit(sum, [1, 2]).result(EVENT_TIMEOUT) == 3 pool.release() + # At most one call more than max_workers runs at once (I2). + call, state = overlapping_calls(2, timeout=0.5, grace=0.1) + for running in [pool.submit(call) for _ in range(3)]: + running.result(EVENT_TIMEOUT) + assert state["peak"] <= 2 shutdown_bounded(pool) if future is not None: assert future.done() @@ -355,7 +360,11 @@ def case(interrupt): shutdown_bounded(scope) if future is not None: assert future.done() - assert events in ([], ["call"]) + 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() @@ -376,10 +385,10 @@ def case(interrupt): try: scope.submit(release.wait, EVENT_TIMEOUT) interrupt(scope.submit_to, future, events.append, "call") - if interrupt.fired and future.cancel(): - future.set_running_or_notify_cancel() + if interrupt.fired: + # As the write session does with its tracked request. + future.cancel() release.set() - assert not wait([future], timeout=EVENT_TIMEOUT).not_done assert scope.submit(sum, [1, 2]).result(EVENT_TIMEOUT) == 3 shutdown_bounded(scope) assert events == ([] if future.cancelled() else ["call"]) @@ -389,42 +398,73 @@ def case(interrupt): assert interrupt_each_opcode(S3Executor.submit_to, case) > 0 - def test_concurrent_submitter_does_not_pass_waiting_call(self): - # GH-1332 (S3): a call submitted from another thread at the moment a - # place frees does not enter ahead of a call that waits for it. A - # barrier case (spec section 6.2): the submitter is started from the - # internal point where the place is freed, before it is handed on. + 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. A place that frees passes to the oldest + # waiting call under the lock that new submitters take, so a call + # submitted meanwhile, from any thread, queues behind it. pool = S3ThreadPool(2) scope = _S3PoolScope(pool, 1) release = threading.Event() order = [] - late: list[Future[None]] = [] - submitter = threading.Thread( - target=lambda: late.append(scope.submit(order.append, "late")), daemon=True - ) - - class FreeingSet(set): - def discard(self, element): - super().discard(element) - if not submitter.is_alive() and not late: - # The submitter takes the free place now unless the - # scope holds it off until the place is handed on. - submitter.start() - submitter.join(0.2) - try: scope.submit(release.wait, EVENT_TIMEOUT) - scope._entered = FreeingSet(scope._entered) - waiting = scope.submit(order.append, "waiting") + waiting = [scope.submit(order.append, name) for name in ("first", "second")] release.set() - assert waiting.result(EVENT_TIMEOUT) is None - submitter.join(EVENT_TIMEOUT) - late[0].result(EVENT_TIMEOUT) + for future in waiting: + future.result(EVENT_TIMEOUT) shutdown_bounded(scope) finally: release.set() pool.shutdown() - assert order == ["waiting", "late"] + 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 = threading.Event(), threading.Event() + + def call(): + started.set() + return release.wait(EVENT_TIMEOUT) + + try: + nested = pool.submit(lambda: scope.submit(call).result(EVENT_TIMEOUT)) + assert started.wait(EVENT_TIMEOUT) + shut_down = threading.Event() + waiter = threading.Thread( + target=lambda: (scope.shutdown(wait=True), shut_down.set()), daemon=True + ) + waiter.start() + assert not shut_down.wait(0.2) + release.set() + assert shut_down.wait(EVENT_TIMEOUT) + assert 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 diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index ce31ca002..83649fc3e 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -356,8 +356,9 @@ def handler(offset): # 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" and instruction.jump_target is not None: - points.add(instruction.jump_target) + 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): From db9c9e0ff40fac97c0820256bc817b45607d7a9f Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 5 Nov 2026 18:17:50 +0900 Subject: [PATCH 15/22] Cancel an abort that is never queued, and clear read frames first (GH-1332) From the follow-up review of the repairs: - An abort of a streamed file whose queuing fails or is interrupted is cancelled, so the drain of the close no longer waits for it forever. - The cleanup of a failed or interrupted parallel read clears the frames of the error first, in a finally, so a second interrupt during the cancellation, whose error keeps the first as its context, cannot keep the futures of the wait. - Tests: an interrupt after the commit has started, an abort whose queuing is interrupted, a second interrupt during the read cleanup, an interrupted pool call that runs while others overlap it, and a shutdown that has begun before the inline call ends. Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3.py | 17 ++-- pyathena/filesystem/s3_async.py | 9 ++- tests/pyathena/filesystem/test_s3.py | 45 +++++++++++ tests/pyathena/filesystem/test_s3_async.py | 77 +++++++++++++++++++ tests/pyathena/filesystem/test_s3_executor.py | 48 ++++++++---- 5 files changed, 176 insertions(+), 20 deletions(-) diff --git a/pyathena/filesystem/s3.py b/pyathena/filesystem/s3.py index 001a84787..480f377b9 100644 --- a/pyathena/filesystem/s3.py +++ b/pyathena/filesystem/s3.py @@ -3463,13 +3463,18 @@ def _fetch_range(self, start: int, end: int) -> bytes: # keeps a future. while futures: futures.pop().cancel() - # After the cancellation, which a thread that it frees - # could otherwise precede by taking a queued read. - stop.set() - traceback.clear_frames(e.__traceback__) finally: - done.clear() - futures.clear() + # Also if an interrupt stops the cancellation, whose error + # keeps this one as its context; the frames first, as + # they keep the futures of the wait. + try: + traceback.clear_frames(e.__traceback__) + finally: + done.clear() + futures.clear() + # After the cancellation, which a thread that it frees + # could otherwise precede by taking a queued read. + stop.set() raise return self.fs.core.get_object(path, ranges[0], **request_kwargs) diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 8d9073dcd..924caf245 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -1558,7 +1558,14 @@ async def _close_without_commit(self) -> None: # 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) + try: + self.fs._thread_pool.submit_to(discard, self._discard) + except BaseException: + # Not queued, or queued by a submission that an interrupt + # stopped and that now finds it cancelled: the drain does + # not wait for an abort that never runs. + discard.cancel() + raise await self.fs._thread_pool.wait_shielded(discard) finally: # A commit that was scheduled before runs on, and is waited for. diff --git a/tests/pyathena/filesystem/test_s3.py b/tests/pyathena/filesystem/test_s3.py index bdf2f5382..3808fcc11 100644 --- a/tests/pyathena/filesystem/test_s3.py +++ b/tests/pyathena/filesystem/test_s3.py @@ -3671,6 +3671,51 @@ def get_object(path, range_, **kwargs): 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.", diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index ae5b83ef8..b4b263d29 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -2880,6 +2880,83 @@ 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 = threading.Event() + events = [] + + def complete_multipart_upload(*args, **kwargs): + started.set() + time.sleep(0.2) + 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) + raise KeyboardInterrupt + + 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 diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index c3fbbfc67..32526fcde 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -40,21 +40,40 @@ def test_submit_interrupted_at_each_opcode(self): # returns. def case(interrupt): pool = S3ThreadPool(1) + hold = threading.Event() + lock = threading.Lock() + state = {"active": 0, "peak": 0} events = [] + + def call(name, timeout): + with lock: + state["active"] += 1 + state["peak"] = max(state["peak"], state["active"]) + try: + hold.wait(timeout) + finally: + with lock: + state["active"] -= 1 + events.append(name) + try: - future = interrupt(pool.submit, events.append, "call") - assert pool.submit(sum, [1, 2]).result(EVENT_TIMEOUT) == 3 + # 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.05) for _ in range(3)] + time.sleep(0.1) + hold.set() + for other in others: + other.result(EVENT_TIMEOUT) # At most one call more than max_workers runs at once (I2). - call, state = overlapping_calls(2, timeout=0.5, grace=0.1) - for running in [pool.submit(call) for _ in range(3)]: - running.result(EVENT_TIMEOUT) assert state["peak"] <= 2 shutdown_bounded(pool) if future is not None: assert future.done() - assert events in ([], ["call"]) + assert events.count("other") == 3 + assert events.count("call") <= 1 finally: + hold.set() pool.shutdown() assert interrupt_each_opcode(S3ThreadPool.submit, case) > 0 @@ -400,9 +419,7 @@ def case(interrupt): 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. A place that frees passes to the oldest - # waiting call under the lock that new submitters take, so a call - # submitted meanwhile, from any thread, queues behind it. + # order they were submitted. pool = S3ThreadPool(2) scope = _S3PoolScope(pool, 1) release = threading.Event() @@ -433,11 +450,16 @@ def call(): try: nested = pool.submit(lambda: scope.submit(call).result(EVENT_TIMEOUT)) assert started.wait(EVENT_TIMEOUT) - shut_down = threading.Event() - waiter = threading.Thread( - target=lambda: (scope.shutdown(wait=True), shut_down.set()), daemon=True - ) + shutting_down, shut_down = threading.Event(), threading.Event() + + def shutdown(): + shutting_down.set() + scope.shutdown(wait=True) + shut_down.set() + + waiter = threading.Thread(target=shutdown, daemon=True) waiter.start() + assert shutting_down.wait(EVENT_TIMEOUT) assert not shut_down.wait(0.2) release.set() assert shut_down.wait(EVENT_TIMEOUT) From 770541b790b50924c093859c90b0deb3815860d2 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 5 Nov 2026 18:27:11 +0900 Subject: [PATCH 16/22] Make the interrupt and shutdown tests independent of timing (GH-1332) The pool's interrupt test waits for a third overlapping call, which only a pool exceeding its bound starts; the scope's shutdown test releases the inline call from another thread while the test thread shuts the scope down; and the commit test holds the commit until the interrupt has been raised. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/filesystem/test_s3_async.py | 11 +++-- tests/pyathena/filesystem/test_s3_executor.py | 44 ++++++++++--------- 2 files changed, 31 insertions(+), 24 deletions(-) diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index b4b263d29..f6f5eaee8 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -2886,12 +2886,14 @@ def test_close_interrupted_after_queuing_commit(self, fs): # 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 = threading.Event() + started, interrupted = threading.Event(), threading.Event() events = [] def complete_multipart_upload(*args, **kwargs): started.set() - time.sleep(0.2) + # 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) @@ -2900,7 +2902,10 @@ 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) - raise KeyboardInterrupt + try: + raise KeyboardInterrupt + finally: + interrupted.set() async def write_and_close(): async with asyncio.timeout(EVENT_TIMEOUT): diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index 32526fcde..44467d283 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -45,14 +45,22 @@ def case(interrupt): state = {"active": 0, "peak": 0} events = [] + overlap = threading.Condition(lock) + def call(name, timeout): - with lock: + 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: - hold.wait(timeout) + if name == "call": + hold.wait(timeout) finally: - with lock: + with overlap: state["active"] -= 1 events.append(name) @@ -60,8 +68,10 @@ def call(name, timeout): # 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.05) for _ in range(3)] - time.sleep(0.1) + 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) @@ -441,29 +451,21 @@ def test_shutdown_waits_for_call_from_pool_thread(self): # thread of its pool runs itself (T1). pool = S3ThreadPool(2) scope = _S3PoolScope(pool, 1) - started, release = threading.Event(), threading.Event() + started, release, ended = threading.Event(), threading.Event(), threading.Event() def call(): started.set() - return release.wait(EVENT_TIMEOUT) + release.wait(EVENT_TIMEOUT) + ended.set() try: nested = pool.submit(lambda: scope.submit(call).result(EVENT_TIMEOUT)) assert started.wait(EVENT_TIMEOUT) - shutting_down, shut_down = threading.Event(), threading.Event() - - def shutdown(): - shutting_down.set() - scope.shutdown(wait=True) - shut_down.set() - - waiter = threading.Thread(target=shutdown, daemon=True) - waiter.start() - assert shutting_down.wait(EVENT_TIMEOUT) - assert not shut_down.wait(0.2) - release.set() - assert shut_down.wait(EVENT_TIMEOUT) - assert nested.result(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() From 9898f176930dd3206e7c106344fc2c50fd0769cf Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 5 Nov 2026 21:41:15 +0900 Subject: [PATCH 17/22] Wait only for a commit or abort that was queued or started (GH-1332) A streamed file marks its commit and abort as queued once their queuing returns, and its drain cancels one that is neither queued nor started, so the close does not wait forever for one that was never queued, also when a second interrupt stops the cleanup of the first. The abort no longer needs a cancellation of its own when its queuing fails. The docs state the interrupt guarantees: the cleanup of one interrupt at a time cancels the requests that have not started and releases the data of the others; under further interrupts during that cleanup, a call still does not wait forever, but the data can stay referenced until the garbage collector frees it and a multipart upload can be left incomplete. Co-Authored-By: Claude Opus 5.5 --- docs/aio.md | 2 ++ docs/filesystem.md | 11 +++++++++++ pyathena/filesystem/s3_async.py | 24 +++++++++++++----------- 3 files changed, 26 insertions(+), 11 deletions(-) diff --git a/docs/aio.md b/docs/aio.md index 6f12aff51..806c844f1 100644 --- a/docs/aio.md +++ b/docs/aio.md @@ -394,6 +394,8 @@ async with await fs.open_async("s3://my-bucket/data/file.bin", "rb") as f: 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. +- 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 e70fe39b3..37d08fff9 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -116,6 +116,17 @@ In a forked child process, create a new filesystem. A file that the parent opene `RuntimeError` there from its reads, writes, and flushes, `commit()`, and `discard()`, and closing it only marks it closed, without uploading or aborting anything. +An interrupt, such as Ctrl+C, while a file or operation submits or waits for its +requests cancels the requests that have not started. The requests that are running +finish, and their data is released once they end. Closing the file waits for them. 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 diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 924caf245..5106bfd1c 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -1279,9 +1279,11 @@ def __init__( 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. + # 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() @@ -1508,11 +1510,16 @@ async def _drain(self) -> None: 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). + 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() - pending.extend(f for f in (self._commit_future, self._discard_future) if f is not None) + 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] @@ -1558,14 +1565,8 @@ async def _close_without_commit(self) -> None: # for it also after a cancellation of this task. discard: Future[None] = Future() self._discard_future = discard - try: - self.fs._thread_pool.submit_to(discard, self._discard) - except BaseException: - # Not queued, or queued by a submission that an interrupt - # stopped and that now finds it cancelled: the drain does - # not wait for an abort that never runs. - discard.cancel() - raise + self.fs._thread_pool.submit_to(discard, self._discard) + self._queued.add(discard) await self.fs._thread_pool.wait_shielded(discard) finally: # A commit that was scheduled before runs on, and is waited for. @@ -1703,6 +1704,7 @@ async def _commit(self) -> None: self.fs._thread_pool.submit_to( commit, self._complete_upload, cast(BytesIO, self.buffer) ) + self._queued.add(commit) except BaseException: if commit.cancel(): # Not started, and now never: a commit that was queued by a From b9763f80ada7a336a22fd465d2def5dfdf35fb8d Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 5 Nov 2026 21:50:50 +0900 Subject: [PATCH 18/22] Let a scope wait only for calls that can still run (GH-1332) A submission that the pool rejected and that an interrupt then stopped before its cleanup cancelled it left a pending future that a shutdown of the scope, or the drain of a streamed file, waited for forever. The scope now waits only for calls that have entered the pool, wait for a place, run in a thread of the pool, or have started, and cancels any other, so that no placement of interrupts makes it wait forever (I1). The docs narrow the cancellation on interrupt to parallel range reads, which is where it applies, and state the limit for an interrupt during the cleanup of an interrupted commit. Co-Authored-By: Claude Opus 5.5 --- docs/aio.md | 3 +- docs/filesystem.md | 11 +++--- pyathena/filesystem/s3_executor.py | 36 ++++++++++++++----- tests/pyathena/filesystem/test_s3_executor.py | 15 ++++++++ 4 files changed, 50 insertions(+), 15 deletions(-) diff --git a/docs/aio.md b/docs/aio.md index 806c844f1..5bf72b50b 100644 --- a/docs/aio.md +++ b/docs/aio.md @@ -393,7 +393,8 @@ 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` diff --git a/docs/filesystem.md b/docs/filesystem.md index 37d08fff9..e9f6eaa91 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -116,12 +116,11 @@ In a forked child process, create a new filesystem. A file that the parent opene `RuntimeError` there from its reads, writes, and flushes, `commit()`, and `discard()`, and closing it only marks it closed, without uploading or aborting anything. -An interrupt, such as Ctrl+C, while a file or operation submits or waits for its -requests cancels the requests that have not started. The requests that are running -finish, and their data is released once they end. Closing the file waits for them. 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 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 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 diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index 3531fa5a6..30bd83104 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -520,8 +520,10 @@ def __init__(self, pool: S3ThreadPool, max_workers: int) -> None: self.max_workers = max_workers self._pool = pool self._lock = threading.Lock() - # The futures of the calls that were submitted and have not ended. - self._futures: set[Future[Any]] = set() + # 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() @@ -563,7 +565,8 @@ def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]: with self._lock: if self._shutdown: raise RuntimeError("cannot schedule new futures after shutdown") - self._futures.add(future) + self._futures[future] = call + self._inline.add(future) future.add_done_callback(self._untrack) call.context.run(call.item.run) return future @@ -571,7 +574,7 @@ def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]: with self._lock: if self._shutdown: raise RuntimeError("cannot schedule new futures after shutdown") - self._futures.add(future) + 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 @@ -591,7 +594,7 @@ def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]: if future.cancel(): call.item.run() with self._lock: - self._futures.discard(future) + self._futures.pop(future, None) raise return future @@ -607,19 +610,35 @@ def shutdown(self, wait: bool = True) -> None: """ with self._lock: self._shutdown = True - pending = list(self._futures) + pending = self.pending() if wait: futures.wait(pending) 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, 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: - return list(self._futures) + 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 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. @@ -691,7 +710,8 @@ def _untrack(self, future: Future[Any]) -> None: future: The future of the call. """ with self._lock: - self._futures.discard(future) + 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 diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index 44467d283..b45dedebc 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -6,6 +6,7 @@ # SPDX-License-Identifier: MIT import asyncio +import contextlib import contextvars import gc import threading @@ -427,6 +428,20 @@ def case(interrupt): 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 + 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. From 6a6e73718a0f6b4b1328954273d14c94f62c3dad Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 5 Nov 2026 21:59:32 +0900 Subject: [PATCH 19/22] Keep a waiting call that a pool thread runs after shutdown (GH-1332) A waiting call that the ending call's thread runs because the pool takes no more calls is tracked as an inline call, so that a concurrent shutdown of the scope waits for it instead of cancelling it (S7). The scope also leaves cancelled calls out of the calls it waits for, since they never run, also if a second interrupt stopped their notification (I1). OpcodeInterrupter can run an action at an interrupt point instead of raising, which the test uses to interleave a shutdown at each point of the hand-on. Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3_executor.py | 11 ++-- tests/pyathena/filesystem/test_s3_executor.py | 51 +++++++++++++++++++ tests/pyathena/util.py | 13 ++++- 3 files changed, 70 insertions(+), 5 deletions(-) diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index 30bd83104..da49e861f 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -619,8 +619,8 @@ def pending(self) -> list[Future[Any]]: 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, so that nothing waits forever - for a call that never runs. + 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 @@ -632,6 +632,9 @@ def pending(self) -> list[Future[Any]]: 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: @@ -684,8 +687,10 @@ def _next(self, ended: _S3ScopeCall) -> None: self._pool.submit(self._run, call) except RuntimeError: # The pool takes no more calls, as when the interpreter - # exits: this thread runs the call. + # 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 diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index b45dedebc..1622e779a 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -9,6 +9,7 @@ import contextlib import contextvars import gc +import sys import threading import time import traceback @@ -24,6 +25,7 @@ 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, @@ -442,6 +444,55 @@ def case(interrupt): 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_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. diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index 83649fc3e..5871aa124 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -289,11 +289,15 @@ class OpcodeInterrupter: 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): + 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 @@ -370,11 +374,16 @@ def _instruction(self, offset): Raises: KeyboardInterrupt: At the n-th interrupt point. """ - if self.fired or offset not in self._points or threading.get_ident() != self._thread: + 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 From 48fef9af9cd74c4b9bfa5e86f243c6a9888c86f5 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 5 Nov 2026 22:10:56 +0900 Subject: [PATCH 20/22] Wait for a scope's calls through done callbacks (GH-1332) The shutdown of a scope waited with concurrent.futures.wait(), which counts a cancelled future as done only once the pool reaches its call and marks it as notified. A cancelled call kept the shutdown waiting for unrelated calls ahead of it in the pool, and forever if a further interrupt stopped that notification. The shutdown now waits through done callbacks, which a cancellation runs at once (S5, I1). Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3_executor.py | 20 +++++++-- tests/pyathena/filesystem/test_s3_executor.py | 42 +++++++++++++++++-- 2 files changed, 56 insertions(+), 6 deletions(-) diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index da49e861f..be565cbba 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -19,7 +19,6 @@ from abc import ABCMeta, abstractmethod from collections import deque from collections.abc import Callable -from concurrent import futures from concurrent.futures import Future from concurrent.futures.thread import ThreadPoolExecutor from multiprocessing import cpu_count @@ -611,8 +610,23 @@ def shutdown(self, wait: bool = True) -> None: with self._lock: self._shutdown = True pending = self.pending() - if wait: - futures.wait(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. + remaining = [len(pending)] + lock = threading.Lock() + ended = threading.Event() + + def count(_: Future[Any]) -> None: + with lock: + remaining[0] -= 1 + if remaining[0] == 0: + ended.set() + + for future in pending: + future.add_done_callback(count) + ended.wait() def pending(self) -> list[Future[Any]]: """Return the futures of the calls that were submitted and have not ended. diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index 1622e779a..bbdc744de 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -493,6 +493,34 @@ def end_first_call(): 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_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. @@ -599,9 +627,17 @@ def call(): running = scope.submit(call) assert started.wait(EVENT_TIMEOUT) (pending,) = scope.pending() - # Sent once the shutdown waits for the call: concurrent.futures - # installs its waiter on the future, inside the block. - with sigint_when(lambda: pending._waiters), pytest.raises(KeyboardInterrupt): + # 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() From 524b22b02e836e2266d42307a86b7a8329ae1940 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 5 Nov 2026 22:17:38 +0900 Subject: [PATCH 21/22] Recheck future states while waiting for a scope or a drain (GH-1332) An interrupt of Future.cancel() after it changes the state and before it runs the callbacks skips them, so a wait that relies only on done callbacks never ends. The shutdown of a scope and the drain of a streamed file now also check the states of the futures every second, as fsspec's sync() does, and end once all of them are done (I1, S5). Co-Authored-By: Claude Opus 5.5 --- pyathena/filesystem/s3_async.py | 7 ++-- pyathena/filesystem/s3_executor.py | 19 ++++----- tests/pyathena/filesystem/test_s3_executor.py | 39 +++++++++++++++++++ 3 files changed, 50 insertions(+), 15 deletions(-) diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 5106bfd1c..8a20ce546 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -1528,12 +1528,13 @@ async def _drain(self) -> None: # awaited them; retrieved here, so that asyncio does not log them. waiter.add_done_callback(lambda f: f.cancelled() or f.exception()) cancelled = False - while True: + # 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) - break + await asyncio.wait(waiters, timeout=1.0) except asyncio.CancelledError: cancelled = True if cancelled: diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index be565cbba..a5f43b9ef 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -613,20 +613,15 @@ def shutdown(self, wait: bool = True) -> None: 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. - remaining = [len(pending)] - lock = threading.Lock() + # 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() - - def count(_: Future[Any]) -> None: - with lock: - remaining[0] -= 1 - if remaining[0] == 0: - ended.set() - for future in pending: - future.add_done_callback(count) - ended.wait() + 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. diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index bbdc744de..e774599c9 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -521,6 +521,45 @@ def test_shutdown_returns_once_calls_are_cancelled(self): 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. From f18ee1068c8ed6df9063dc18aec9412afea3ab56 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Thu, 5 Nov 2026 22:25:04 +0900 Subject: [PATCH 22/22] Document that an interrupt of a close stops its wait (GH-1332) Co-Authored-By: Claude Opus 5.5 --- docs/filesystem.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/docs/filesystem.md b/docs/filesystem.md index e9f6eaa91..15f58f551 100644 --- a/docs/filesystem.md +++ b/docs/filesystem.md @@ -118,7 +118,8 @@ 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 while a +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: