Make files usable after fork with either filesystem (GH-1336)

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()`

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,

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)

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):

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):