[Python-checkins] GH-66285: fix forking in `asyncio` (#99539)

[kumaraditya303]

https://github.com/python/cpython/commit/0c1fbc17b48f56c6070d335ed291a24c91ed190a
commit: 0c1fbc17b48f56c6070d335ed291a24c91ed190a
branch: main
author: Kumar Aditya <59607654+kumaraditya303 at users.noreply.github.com>
committer: kumaraditya303 <59607654+kumaraditya303 at users.noreply.github.com>
date: 2022-11-24T09:10:27+05:30
summary:

GH-66285: fix forking in `asyncio` (#99539)

`asyncio` now does not shares event loop and signal wakeupfd in forked processes.

files:
A Misc/NEWS.d/next/Library/2022-11-17-10-56-47.gh-issue-66285.KvjlaB.rst
M Lib/asyncio/events.py
M Lib/test/test_asyncio/test_unix_events.py

diff --git a/Lib/asyncio/events.py b/Lib/asyncio/events.py
index a327ba54a323..39a01048b4c8 100644
--- a/Lib/asyncio/events.py
+++ b/Lib/asyncio/events.py
@@ -17,6 +17,7 @@
 import subprocess
 import sys
 import threading
+import signal
 
 from . import format_helpers
 
@@ -665,6 +666,14 @@ class _Local(threading.local):
 
     def __init__(self):
         self._local = self._Local()
+        if hasattr(os, 'fork'):
+            def on_fork():
+                # Reset the loop and wakeupfd in the forked child process.
+                self._local = self._Local()
+                signal.set_wakeup_fd(-1)
+
+            os.register_at_fork(after_in_child=on_fork)
+
 
     def get_event_loop(self):
         """Get the event loop for the current context.
diff --git a/Lib/test/test_asyncio/test_unix_events.py b/Lib/test/test_asyncio/test_unix_events.py
index 93e8611f184d..4e1dab2f86b4 100644
--- a/Lib/test/test_asyncio/test_unix_events.py
+++ b/Lib/test/test_asyncio/test_unix_events.py
@@ -11,10 +11,13 @@
 import sys
 import threading
 import unittest
+import time
 from unittest import mock
 import warnings
+import multiprocessing
 from test.support import os_helper
 from test.support import socket_helper
+from test.support import wait_process
 
 if sys.platform == 'win32':
     raise unittest.SkipTest('UNIX only')
@@ -1867,5 +1870,97 @@ async def runner():
             wsock.close()
 
 
+ at unittest.skipUnless(hasattr(os, 'fork'), 'requires os.fork()')
+class TestFork(unittest.IsolatedAsyncioTestCase):
+
+    async def test_fork_not_share_event_loop(self):
+        # The forked process should not share the event loop with the parent
+        loop = asyncio.get_running_loop()
+        r, w = os.pipe()
+        self.addCleanup(os.close, r)
+        self.addCleanup(os.close, w)
+        pid = os.fork()
+        if pid == 0:
+            # child
+            try:
+                loop = asyncio.get_event_loop_policy().get_event_loop()
+                os.write(w, str(id(loop)).encode())
+            finally:
+                os._exit(0)
+        else:
+            # parent
+            child_loop = int(os.read(r, 100).decode())
+            self.assertNotEqual(child_loop, id(loop))
+            wait_process(pid, exitcode=0)
+
+    def test_fork_signal_handling(self):
+        # Sending signal to the forked process should not affect the parent
+        # process
+        ctx = multiprocessing.get_context('fork')
+        manager = ctx.Manager()
+        self.addCleanup(manager.shutdown)
+        child_started = manager.Event()
+        child_handled = manager.Event()
+        parent_handled = manager.Event()
+
+        def child_main():
+            signal.signal(signal.SIGTERM, lambda *args: child_handled.set())
+            child_started.set()
+            time.sleep(1)
+
+        async def main():
+            loop = asyncio.get_running_loop()
+            loop.add_signal_handler(signal.SIGTERM, lambda *args: parent_handled.set())
+
+            process = ctx.Process(target=child_main)
+            process.start()
+            child_started.wait()
+            os.kill(process.pid, signal.SIGTERM)
+            process.join()
+
+            async def func():
+                await asyncio.sleep(0.1)
+                return 42
+
+            # Test parent's loop is still functional
+            self.assertEqual(await asyncio.create_task(func()), 42)
+
+        asyncio.run(main())
+
+        self.assertFalse(parent_handled.is_set())
+        self.assertTrue(child_handled.is_set())
+
+    def test_fork_asyncio_run(self):
+        ctx = multiprocessing.get_context('fork')
+        manager = ctx.Manager()
+        self.addCleanup(manager.shutdown)
+        result = manager.Value('i', 0)
+
+        async def child_main():
+            await asyncio.sleep(0.1)
+            result.value = 42
+
+        process = ctx.Process(target=lambda: asyncio.run(child_main()))
+        process.start()
+        process.join()
+
+        self.assertEqual(result.value, 42)
+
+    def test_fork_asyncio_subprocess(self):
+        ctx = multiprocessing.get_context('fork')
+        manager = ctx.Manager()
+        self.addCleanup(manager.shutdown)
+        result = manager.Value('i', 1)
+
+        async def child_main():
+            proc = await asyncio.create_subprocess_exec(sys.executable, '-c', 'pass')
+            result.value = await proc.wait()
+
+        process = ctx.Process(target=lambda: asyncio.run(child_main()))
+        process.start()
+        process.join()
+
+        self.assertEqual(result.value, 0)
+
 if __name__ == '__main__':
     unittest.main()
diff --git a/Misc/NEWS.d/next/Library/2022-11-17-10-56-47.gh-issue-66285.KvjlaB.rst b/Misc/NEWS.d/next/Library/2022-11-17-10-56-47.gh-issue-66285.KvjlaB.rst
new file mode 100644
index 000000000000..ebd821738827
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2022-11-17-10-56-47.gh-issue-66285.KvjlaB.rst
@@ -0,0 +1 @@
+Fix :mod:`asyncio` to not share event loop and signal wakeupfd in forked processes. Patch by Kumar Aditya.