Fix multiprocessing deadlock after using prefect_test_harness on Linux (#19116)

Author: zzstoatzz

src/prefect/_internal/concurrency/services.py

@@ -5,8 +5,10 @@
 import concurrent.futures
 import contextlib
 import logging
+import os
 import queue
 import threading
+import weakref
 from collections.abc import AsyncGenerator, Awaitable, Coroutine, Generator, Hashable
 from typing import TYPE_CHECKING, Any, Generic, NoReturn, Optional, Union, cast
 
@@ -22,6 +24,31 @@
 Ts = TypeVarTuple("Ts")
 R = TypeVar("R", infer_variance=True)
 
+# Track all active services for fork handling
+_active_services: weakref.WeakSet[_QueueServiceBase[Any]] = weakref.WeakSet()
+
+
+def _reset_services_after_fork():
+    """
+    Reset service state after fork() to prevent multiprocessing deadlocks on Linux.
+
+    Called by os.register_at_fork() in the child process after fork().
+    """
+    for service in list(_active_services):
+        service.reset_for_fork()
+
+    # Reset the class-level instance tracking
+    _QueueServiceBase.reset_instances_for_fork()
+
+
+# Register fork handler if supported (POSIX systems)
+if hasattr(os, "register_at_fork"):
+    try:
+        os.register_at_fork(after_in_child=_reset_services_after_fork)
+    except RuntimeError:
+        # Might fail in certain contexts (e.g., if already in a child process)
+        pass
+
 
 class _QueueServiceBase(abc.ABC, Generic[T]):
     _instances: dict[int, Self] = {}
@@ -44,6 +71,25 @@ def __init__(self, *args: Hashable) -> None:
         )
         self._logger = logging.getLogger(f"{type(self).__name__}")
 
+        # Track this instance for fork handling
+        _active_services.add(self)
+
+    def reset_for_fork(self) -> None:
+        """Reset instance state after fork() to prevent deadlocks in child process."""
+        self._stopped = True
+        self._started = False
+        self._loop = None
+        self._done_event = None
+        self._task = None
+        self._queue = queue.Queue()
+        self._lock = threading.Lock()
+
+    @classmethod
+    def reset_instances_for_fork(cls) -> None:
+        """Reset class-level state after fork() to prevent deadlocks in child process."""
+        cls._instances.clear()
+        cls._instance_lock = threading.Lock()
+
     def start(self) -> None:
         logger.debug("Starting service %r", self)
         loop_thread = get_global_loop()

src/prefect/_internal/concurrency/threads.py

@@ -8,8 +8,10 @@
 import atexit
 import concurrent.futures
 import itertools
+import os
 import queue
 import threading
+import weakref
 from typing import Any, Optional
 
 from typing_extensions import TypeVar
@@ -22,6 +24,32 @@
 
 T = TypeVar("T", infer_variance=True)
 
+# Track all active instances for fork handling
+_active_instances: weakref.WeakSet[WorkerThread | EventLoopThread] = weakref.WeakSet()
+
+
+def _reset_after_fork_in_child():
+    """
+    Reset thread state after fork() to prevent multiprocessing deadlocks on Linux.
+
+    When fork() is called, the child process inherits all thread state and locks
+    from the parent, but only the calling thread continues. This leaves other threads'
+    locks in inconsistent states causing deadlocks.
+
+    This handler is called by os.register_at_fork() in the child process after fork().
+    """
+    for instance in list(_active_instances):
+        instance.reset_for_fork()
+
+
+# Register fork handler if supported (POSIX systems)
+if hasattr(os, "register_at_fork"):
+    try:
+        os.register_at_fork(after_in_child=_reset_after_fork_in_child)
+    except RuntimeError:
+        # Might fail in certain contexts (e.g., if already in a child process)
+        pass
+
 
 class WorkerThread(Portal):
     """
@@ -45,9 +73,19 @@ def __init__(
         self._submitted_count: int = 0
         self._lock = threading.Lock()
 
+        # Track this instance for fork handling
+        _active_instances.add(self)
+
         if not daemon:
             atexit.register(self.shutdown)
 
+    def reset_for_fork(self) -> None:
+        """Reset state after fork() to prevent deadlocks in child process."""
+        self._started = False
+        self._queue = queue.Queue()
+        self._lock = threading.Lock()
+        self._submitted_count = 0
+
     def start(self) -> None:
         """
         Start the worker thread.
@@ -145,9 +183,21 @@ def __init__(
         self._on_shutdown: list[Call[Any]] = []
         self._lock = threading.Lock()
 
+        # Track this instance for fork handling
+        _active_instances.add(self)
+
         if not daemon:
             atexit.register(self.shutdown)
 
+    def reset_for_fork(self) -> None:
+        """Reset state after fork() to prevent deadlocks in child process."""
+        self._loop = None
+        self._ready_future = concurrent.futures.Future()
+        self._shutdown_event = Event()
+        self._lock = threading.Lock()
+        self._submitted_count = 0
+        self._on_shutdown = []
+
     def start(self):
         """
         Start the worker thread; raises any exceptions encountered during startup.

tests/testing/test_utilites.py

@@ -1,3 +1,4 @@
+import sys
 import uuid
 import warnings
 from unittest.mock import MagicMock
@@ -15,6 +16,20 @@
 from prefect.testing.utilities import assert_does_not_warn, prefect_test_harness
 
 
+def _multiprocessing_worker():
+    """
+    Worker function for multiprocessing test. Must be at module level for pickling.
+    """
+    import os
+
+    # os._exit() is required here despite the underscore prefix. On Linux with fork(),
+    # the child process inherits Prefect's logging/event state. Normal exit (return or
+    # sys.exit) triggers Python cleanup that fails with this inherited state, causing
+    # exitcode=1. os._exit() bypasses cleanup and is documented for use "in the child
+    # process after os.fork()" - which is exactly this scenario.
+    os._exit(0)
+
+
 def test_assert_does_not_warn_no_warning():
     with assert_does_not_warn():
         pass
@@ -101,3 +116,35 @@ def test_prefect_test_harness_timeout(monkeypatch):
         server().start.assert_called_once_with(
             timeout=PREFECT_SERVER_EPHEMERAL_STARTUP_TIMEOUT_SECONDS.value()
         )
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="fork() not available on Windows")
+def test_multiprocessing_after_test_harness():
+    """
+    Test that multiprocessing works after using prefect_test_harness.
+
+    Regression test for issue #19112 - on Linux, multiprocessing.Process() would
+    deadlock after using the test harness because fork() inherited locked thread
+    state from background threads.
+    """
+    import multiprocessing
+
+    @task
+    def test_task():
+        return 1
+
+    @flow
+    def test_flow():
+        return test_task.submit()
+
+    # Use test harness which starts background threads
+    with prefect_test_harness():
+        test_flow()
+
+    # This should not deadlock
+    process = multiprocessing.Process(target=_multiprocessing_worker)
+    process.start()
+    process.join(timeout=5)
+
+    # Verify process completed successfully
+    assert process.exitcode == 0, "Process should complete without deadlock"