diff --git a/docs/aio.md b/docs/aio.md index 1a6016068..ab83aea8a 100644 --- a/docs/aio.md +++ b/docs/aio.md @@ -334,7 +334,8 @@ multipart uploads. `S3FileSystem` and `S3File` use a pluggable executor abstraction (`S3Executor`) for parallel operations. Two implementations are provided: -- `S3ThreadPoolExecutor` — wraps `ThreadPoolExecutor` (default for sync usage) +- `S3ThreadPoolExecutor` — runs the calls in a pool of `max_workers` threads (default + for sync usage); in a forked child process, the pool creates new threads - `S3AioExecutor` — dispatches work via `asyncio.run_coroutine_threadsafe` to an `S3AioThreadPool` of its own @@ -343,6 +344,11 @@ so multipart uploads and parallel range reads are dispatched through the event l a pool of `max_workers` threads of each file. An instance created with `asynchronous=True` has no event loop of its own, so its file handles use `S3ThreadPoolExecutor`. +In a forked child process, the event loop of the parent does not run. +An `AioS3FileSystem` created in the parent without `asynchronous=True` raises +`RuntimeError` from `open()` there, as fsspec's synchronous methods do, and so does +`S3AioExecutor.submit()`; a filesystem created in the child works. + ### Async file handles `open_async()` returns an `AioS3StreamedFile`, whose `read()`, `write()`, and `close()` diff --git a/pyathena/filesystem/s3_async.py b/pyathena/filesystem/s3_async.py index 4d93376f6..57203a7d0 100644 --- a/pyathena/filesystem/s3_async.py +++ b/pyathena/filesystem/s3_async.py @@ -848,10 +848,15 @@ def _create_executor(self, max_workers: int) -> S3Executor: Returns: An ``S3AioExecutor`` on the event loop of this filesystem, or an ``S3ThreadPoolExecutor`` if it has none. + + Raises: + RuntimeError: If this instance has an event loop and the call is + made in a forked child process of the process that created + it, where the loop does not run. """ if self._loop is None: return S3ThreadPoolExecutor(max_workers=max_workers) - return S3AioExecutor(loop=self._loop, max_workers=max_workers) + return S3AioExecutor(loop=self.loop, max_workers=max_workers) def _open( self, diff --git a/pyathena/filesystem/s3_executor.py b/pyathena/filesystem/s3_executor.py index 9b22cdbc0..1e9b3f2e6 100644 --- a/pyathena/filesystem/s3_executor.py +++ b/pyathena/filesystem/s3_executor.py @@ -178,28 +178,32 @@ def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None: class S3ThreadPoolExecutor(S3Executor): - """Executor that delegates to a ``ThreadPoolExecutor``. + """Executor that runs the submitted calls in a pool of threads. This is the default executor used by ``S3File`` and ``S3FileSystem`` - for synchronous parallel operations. + for synchronous parallel operations. The calls run in an + :class:`S3AioThreadPool`, so that a forked child process, whose copy of + the executor has no threads, creates new ones. """ def __init__(self, max_workers: int) -> None: - """Initialize the executor with a new ``ThreadPoolExecutor``. + """Initialize the executor without starting any thread. Args: max_workers: The maximum number of threads of the thread pool. + + Raises: + ValueError: If ``max_workers`` is not positive. """ - self._executor = ThreadPoolExecutor(max_workers=max_workers) + self._pool = S3AioThreadPool(max_workers) @override def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]: - context = contextvars.copy_context() - return self._executor.submit(context.run, fn, *args, **kwargs) + return self._pool.submit(fn, *args, **kwargs) @override def shutdown(self, wait: bool = True) -> None: - self._executor.shutdown(wait=wait) + self._pool.shutdown(wait=wait) class S3AioThreadPool(S3Executor): @@ -498,9 +502,15 @@ class S3AioExecutor(S3Executor): max_workers: The maximum number of submitted functions that run at once. Raises: - RuntimeError: If the event loop is not running when ``submit`` is called. + RuntimeError: If the event loop is not running when ``submit`` is + called, or ``submit`` is called in a forked child process, where + the event loop of the parent does not run. """ + # The executors of this process, which a forked child detaches from the + # event loop of its parent. + _instances: ClassVar[weakref.WeakSet[S3AioExecutor]] = weakref.WeakSet() + def __init__( self, loop: asyncio.AbstractEventLoop | None = None, @@ -522,9 +532,13 @@ def __init__( # Orders the dispatches with the shutdown. self._lock = threading.Lock() self._shutdown = False + self._forked = False + S3AioExecutor._instances.add(self) @override def submit(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> Future[T]: + if self._forked: + raise RuntimeError("S3AioExecutor cannot be used in a forked child process.") if self._loop is not None and self._loop.is_running(): # The future of run_coroutine_threadsafe can be cancelled while # the function keeps running in its thread, so the returned future @@ -596,3 +610,21 @@ def shutdown(self, wait: bool = True) -> None: if not loop.is_running() or sys.is_finalizing(): break self._pool.shutdown(wait=True) + + @classmethod + def _reset_after_fork(cls) -> None: + """Detach the executors that a forked child copied from the event loop of its parent. + + The copy of the loop reports running, but nothing runs it, and a lock + that a thread of the parent held stays held. Each executor then + raises from :meth:`submit`, and :meth:`shutdown` only shuts down its + pool. + """ + for executor in list(cls._instances): + executor._loop = None + executor._lock = threading.Lock() + executor._forked = True + + +if hasattr(os, "register_at_fork"): + os.register_at_fork(after_in_child=S3AioExecutor._reset_after_fork) diff --git a/tests/pyathena/filesystem/test_s3_async.py b/tests/pyathena/filesystem/test_s3_async.py index b04dd7e32..ff98fd09b 100644 --- a/tests/pyathena/filesystem/test_s3_async.py +++ b/tests/pyathena/filesystem/test_s3_async.py @@ -44,9 +44,11 @@ ThrottledOperation, await_bounded, cancel_when_started, + forking, overlapping_calls, retried, run_bounded, + run_in_forked_child, shutdown_bounded, stub_multipart_copy, wait_until, @@ -127,6 +129,23 @@ async def _multipart_copy(release, fs=None, **kwargs): await asyncio.to_thread(shutdown_bounded, fs._thread_pool) stubber.assert_no_pending_responses() + @forking + def test_open_in_forked_child(self): + # GH-1336: in a forked child, a filesystem created in the parent + # raises on open(), as fsspec's synchronous methods do there. It used + # to return a file whose reads waited forever for the event loop of + # the parent. + fs = AioS3FileSystem( + key="dummy", secret="dummy", region_name="us-east-1", skip_instance_cache=True + ) + + def child(): + with pytest.raises(RuntimeError, match="not fork-safe"): + fs.open("s3://bucket/key", "rb") + return True + + assert run_in_forked_child(child) == 0 + @pytest.mark.parametrize("size", [0, 10]) @pytest.mark.asyncio async def test_copy_object_with_multipart_upload_small_head_object_size(self, size): diff --git a/tests/pyathena/filesystem/test_s3_executor.py b/tests/pyathena/filesystem/test_s3_executor.py index ef4de1767..d6a7bdc51 100644 --- a/tests/pyathena/filesystem/test_s3_executor.py +++ b/tests/pyathena/filesystem/test_s3_executor.py @@ -8,8 +8,6 @@ import asyncio import contextvars import gc -import os -import signal import threading import time import traceback @@ -23,7 +21,14 @@ from pyathena.filesystem.s3_executor import S3AioExecutor, S3AioThreadPool, S3ThreadPoolExecutor from pyathena.util import RetryConfig, retry_api_call, stop_retries_on -from tests.pyathena.util import EVENT_TIMEOUT, ThrottledOperation, overlapping_calls +from tests.pyathena.util import ( + EVENT_TIMEOUT, + ThrottledOperation, + forking, + overlapping_calls, + run_in_forked_child, + shutdown_bounded, +) class TestS3AioThreadPool: @@ -275,8 +280,7 @@ def test_release_after_interrupted_submit(self): assert not threads[0].is_alive() pool.shutdown() - @pytest.mark.skipif(not hasattr(os, "fork"), reason="os.fork() is not available.") - @pytest.mark.filterwarnings("ignore:.*fork\\(\\) may lead to deadlocks:DeprecationWarning") + @forking def test_fork(self): # A forked child, whose copy of the pool has no threads, creates new # ones, also if a thread of the parent held the lock of the pool. @@ -289,38 +293,22 @@ def test_fork(self): def hold_lock(): with pool._lock: held.set() - release.wait(EVENT_TIMEOUT) + release.wait() - holder = threading.Thread(target=hold_lock) + holder = threading.Thread(target=hold_lock, daemon=True) holder.start() try: assert held.wait(EVENT_TIMEOUT) - pid = os.fork() - finally: - release.set() - if pid == 0: - # The child exits here, whatever happens, instead of running pytest. - code = 1 - try: - if pool.submit(sum, [3, 4]).result(timeout=EVENT_TIMEOUT) == 7: - code = 0 - finally: - os._exit(code) - status = (0, 0) - try: # A child whose pool waits for the threads of the parent would # not exit. - deadline = time.monotonic() + 2 * EVENT_TIMEOUT - while (status := os.waitpid(pid, os.WNOHANG)) == (0, 0): - assert time.monotonic() < deadline, "the child did not exit" - time.sleep(0.05) + code = run_in_forked_child( + lambda: pool.submit(sum, [3, 4]).result(timeout=EVENT_TIMEOUT) == 7 + ) finally: - if status == (0, 0): - os.kill(pid, signal.SIGKILL) - os.waitpid(pid, 0) + release.set() holder.join(EVENT_TIMEOUT) pool.shutdown() - assert os.waitstatus_to_exitcode(status[1]) == 0 + assert code == 0 class TestS3AioExecutor: @@ -502,6 +490,52 @@ def block(): loop.run_until_complete(asyncio.sleep(0)) loop.close() + @forking + def test_fork(self): + # GH-1336: in a forked child, where the event loop of the parent does + # not run, submit() raises and shutdown() returns, also if a thread + # of the parent held the lock of the executor. Both used to wait + # forever for the copy of the loop, which reports running. + loop = asyncio.new_event_loop() + thread = threading.Thread(target=loop.run_forever, daemon=True) + thread.start() + executor = S3AioExecutor(loop=loop) + held, release = threading.Event(), threading.Event() + + def hold_lock(): + with executor._lock: + held.set() + release.wait() + + def child(): + with pytest.raises(RuntimeError, match="forked child"): + executor.submit(sum, [3, 4]) + executor.shutdown(wait=True) + return True + + holder = threading.Thread(target=hold_lock, daemon=True) + # The executor submits only to a loop that runs. + running = threading.Event() + loop.call_soon_threadsafe(running.set) + try: + assert running.wait(EVENT_TIMEOUT) + assert executor.submit(sum, [1, 2]).result(timeout=EVENT_TIMEOUT) == 3 + holder.start() + assert held.wait(EVENT_TIMEOUT) + code = run_in_forked_child(child) + finally: + release.set() + if holder.ident is not None: + holder.join(EVENT_TIMEOUT) + shutdown_bounded(executor) + loop.call_soon_threadsafe(loop.stop) + thread.join(EVENT_TIMEOUT) + assert not thread.is_alive() + # Runs the dispatch tasks that the stop left pending. + loop.run_until_complete(asyncio.sleep(0)) + loop.close() + assert code == 0 + def test_submit_without_running_loop(self): with pytest.raises(RuntimeError, match="requires a running event loop"): S3AioExecutor().submit(sum, [1, 2]) @@ -545,6 +579,23 @@ async def main(): assert pending.cancelled() +class TestS3ThreadPoolExecutor: + @forking + def test_fork(self): + # GH-1336: a forked child, whose copy of the executor has no threads, + # runs the calls submitted there. After the executor had run a call + # in the parent, they used to wait forever. + executor = S3ThreadPoolExecutor(max_workers=1) + try: + assert executor.submit(sum, [1, 2]).result(timeout=EVENT_TIMEOUT) == 3 + code = run_in_forked_child( + lambda: executor.submit(sum, [3, 4]).result(timeout=EVENT_TIMEOUT) == 7 + ) + finally: + executor.shutdown() + assert code == 0 + + class TestS3Executor: @pytest.mark.parametrize("executor_class", [S3ThreadPoolExecutor, S3AioThreadPool]) def test_submit_context(self, executor_class): diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index 89ac327ce..277b020fe 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -7,8 +7,11 @@ import asyncio import io +import os +import signal import threading import time +import traceback from concurrent.futures import Future, wait from contextlib import contextmanager from datetime import UTC, datetime, timedelta, timezone @@ -277,6 +280,64 @@ def run(): return outcome.get("result") +def forking(test): + """Mark a test that forks a child process. + + The test is skipped where ``os.fork()`` is not available, and the warning + of Python 3.12 and later about forking a multi-threaded process is + ignored. + + Args: + test: The test function. + + Returns: + The marked test function. + """ + test = pytest.mark.filterwarnings( + "ignore:.*fork\\(\\) may lead to deadlocks:DeprecationWarning" + )(test) + return pytest.mark.skipif(not hasattr(os, "fork"), reason="os.fork() is not available.")(test) + + +def run_in_forked_child(child, timeout=2 * EVENT_TIMEOUT): + """Run a check in a forked child process, failing if the child does not exit in time. + + The child exits with the status 0 if the check returns a true value, and + 1 otherwise or if it raises, after printing the traceback, instead of + returning to pytest. A child that + a regression blocks is killed, so that the test fails instead of hanging. + + Args: + child: The check, called without arguments in the child. + timeout: The longest run of the child, in seconds. + + Returns: + The exit code of the child. + """ + pid = os.fork() + if pid == 0: + code = 1 + try: + if child(): + code = 0 + except BaseException: + # Shown in the captured output of the failed test. + traceback.print_exc() + finally: + os._exit(code) + status = (0, 0) + try: + deadline = time.monotonic() + timeout + while (status := os.waitpid(pid, os.WNOHANG)) == (0, 0): + assert time.monotonic() < deadline, f"the child did not exit within {timeout} s" + time.sleep(0.05) + finally: + if status == (0, 0): + os.kill(pid, signal.SIGKILL) + os.waitpid(pid, 0) + return os.waitstatus_to_exitcode(status[1]) + + def shutdown_bounded(pool, timeout=EVENT_TIMEOUT): """Shut down an executor and wait for its calls, failing if they run past the timeout.