gh-83856: Honor atexit for all multiprocessing start methods (GH-114279)

Use atexit for all multiprocessing start methods to cleanup.
See the GH-114279 PR discussion and related issue for details as to why.

5 files changed, 34 insertions, 5 deletions

diff --git a/Lib/multiprocessing/forkserver.py b/Lib/multiprocessing/forkserver.py
index 4642707..53b8c49 100644
--- a/Lib/multiprocessing/forkserver.py
+++ b/Lib/multiprocessing/forkserver.py
@@ -1,3 +1,4 @@
+import atexit
 import errno
 import os
 import selectors
@@ -271,6 +272,8 @@ def main(listener_fd, alive_r, preload, main_path=None, sys_path=None):
                                 selector.close()
                                 unused_fds = [alive_r, child_w, sig_r, sig_w]
                                 unused_fds.extend(pid_to_fd.values())
+                                atexit._clear()
+                                atexit.register(util._exit_function)
                                 code = _serve_one(child_r, fds,
                                                   unused_fds,
                                                   old_handlers)
@@ -278,6 +281,7 @@ def main(listener_fd, alive_r, preload, main_path=None, sys_path=None):
                                 sys.excepthook(*sys.exc_info())
                                 sys.stderr.flush()
                             finally:
+                                atexit._run_exitfuncs()
                                 os._exit(code)
                         else:
                             # Send pid to client process
diff --git a/Lib/multiprocessing/popen_fork.py b/Lib/multiprocessing/popen_fork.py
index 625981c..a57ef6b 100644
--- a/Lib/multiprocessing/popen_fork.py
+++ b/Lib/multiprocessing/popen_fork.py
@@ -1,3 +1,4 @@
+import atexit
 import os
 import signal
 
@@ -66,10 +67,13 @@ class Popen(object):
         self.pid = os.fork()
         if self.pid == 0:
             try:
+                atexit._clear()
+                atexit.register(util._exit_function)
                 os.close(parent_r)
                 os.close(parent_w)
                 code = process_obj._bootstrap(parent_sentinel=child_r)
             finally:
+                atexit._run_exitfuncs()
                 os._exit(code)
         else:
             os.close(child_w)
diff --git a/Lib/multiprocessing/process.py b/Lib/multiprocessing/process.py
index 271ba3f..b45f7df 100644
--- a/Lib/multiprocessing/process.py
+++ b/Lib/multiprocessing/process.py
@@ -310,11 +310,8 @@ class BaseProcess(object):
                 # _run_after_forkers() is executed
                 del old_process
             util.info('child process calling self.run()')
-            try:
-                self.run()
-                exitcode = 0
-            finally:
-                util._exit_function()
+            self.run()
+            exitcode = 0
         except SystemExit as e:
             if e.code is None:
                 exitcode = 0
diff --git a/Lib/test/_test_multiprocessing.py b/Lib/test/_test_multiprocessing.py
index 5fc4181..46afdfc 100644
--- a/Lib/test/_test_multiprocessing.py
+++ b/Lib/test/_test_multiprocessing.py
@@ -6161,6 +6161,29 @@ class TestNamedResource(unittest.TestCase):
         self.assertFalse(err, msg=err.decode('utf-8'))
 
 
+class _TestAtExit(BaseTestCase):
+
+    ALLOWED_TYPES = ('processes',)
+
+    @classmethod
+    def _write_file_at_exit(self, output_path):
+        import atexit
+        def exit_handler():
+            with open(output_path, 'w') as f:
+                f.write("deadbeef")
+        atexit.register(exit_handler)
+
+    def test_atexit(self):
+        # gh-83856
+        with os_helper.temp_dir() as temp_dir:
+            output_path = os.path.join(temp_dir, 'output.txt')
+            p = self.Process(target=self._write_file_at_exit, args=(output_path,))
+            p.start()
+            p.join()
+            with open(output_path) as f:
+                self.assertEqual(f.read(), 'deadbeef')
+
+
 class MiscTestCase(unittest.TestCase):
     def test__all__(self):
         # Just make sure names in not_exported are excluded
diff --git a/Misc/NEWS.d/next/Library/2024-01-19-05-40-46.gh-issue-83856.jN5M80.rst b/Misc/NEWS.d/next/Library/2024-01-19-05-40-46.gh-issue-83856.jN5M80.rst
new file mode 100644
index 0000000..b2889f2
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2024-01-19-05-40-46.gh-issue-83856.jN5M80.rst
@@ -0,0 +1 @@
+Honor :mod:`atexit` for all :mod:`multiprocessing` start methods