Merge pull request #3200 from A5rocks/clear-after-fork

Clear worker thread cache after forking

Author: A5rocks

newsfragments/2764.bugfix.rst

@@ -0,0 +1 @@
+Clear Trio's cache of worker threads upon `os.fork`.

src/trio/_core/_tests/test_thread_cache.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import os
 import threading
 import time
 from contextlib import contextmanager
@@ -195,3 +196,47 @@ def deliver(_: object) -> NoReturn:
     err = capfd.readouterr().err
     assert "don't do this" in err
     assert "delivering result" in err
+
+
+@pytest.mark.skipif(not hasattr(os, "fork"), reason="os.fork isn't supported")
+def test_clear_thread_cache_after_fork() -> None:
+    assert hasattr(os, "fork")
+
+    def foo() -> None:
+        pass
+
+    # ensure the thread cache exists
+    done = threading.Event()
+    start_thread_soon(foo, lambda _: done.set())
+    done.wait()
+
+    child_pid = os.fork()
+
+    # try using it
+    done = threading.Event()
+    start_thread_soon(foo, lambda _: done.set())
+    done.wait()
+
+    if child_pid != 0:
+        # if this test fails, this will hang, triggering a timeout.
+        os.waitpid(child_pid, 0)
+    else:
+        # this is necessary because os._exit doesn't unwind the stack,
+        # so coverage doesn't get to automatically stop and save
+        # coverage information.
+        try:
+            import coverage
+
+            cov = coverage.Coverage.current()
+            # the following pragmas are necessary because if coverage:
+            #  - isn't running, then it can't record the branch not
+            #    taken
+            #  - isn't installed, then it can't record the ImportError
+
+            if cov:  # pragma: no branch
+                cov.stop()
+                cov.save()
+        except ImportError:  # pragma: no cover
+            pass
+
+        os._exit(0)  # pragma: no cover  # coverage was stopped above.

src/trio/_core/_thread_cache.py

@@ -2,6 +2,7 @@
 
 import ctypes
 import ctypes.util
+import os
 import sys
 import traceback
 from functools import partial
@@ -300,3 +301,15 @@ def start_thread_soon(
 
     """
     THREAD_CACHE.start_thread_soon(fn, deliver, name)
+
+
+def clear_worker_threads() -> None:
+    # This is OK because the child process does not actually have any
+    # worker threads. Additionally, while WorkerThread keeps a strong
+    # reference and so would get affected, the only place those are
+    # stored is here.
+    THREAD_CACHE._idle_workers.clear()
+
+
+if hasattr(os, "register_at_fork"):
+    os.register_at_fork(after_in_child=clear_worker_threads)