gh-83856: Honor atexit for all multiprocessing start methods

Lib/multiprocessing/forkserver.py

@@ -1,3 +1,4 @@
+import atexit
 import errno
 import os
 import selectors
@@ -271,13 +272,16 @@ def sigchld_handler(*_unused):
                                 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)
                             except Exception:
                                 sys.excepthook(*sys.exc_info())
                                 sys.stderr.flush()
                             finally:
+                                atexit._run_exitfuncs()
                                 os._exit(code)
                         else:
                             # Send pid to client process

Lib/multiprocessing/popen_fork.py

@@ -1,3 +1,4 @@
+import atexit
 import os
 import signal
 
@@ -66,10 +67,13 @@ def _launch(self, process_obj):
         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)

Lib/multiprocessing/process.py

@@ -310,11 +310,8 @@ def _bootstrap(self, parent_sentinel=None):
                 # _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

Lib/test/test_atexit.py

@@ -1,9 +1,11 @@
 import atexit
 import os
+import sys
+import tempfile
 import textwrap
 import unittest
 from test import support
-from test.support import script_helper
+from test.support import script_helper, os_helper
 
 
 class GeneralTest(unittest.TestCase):
@@ -46,6 +48,43 @@ def test_atexit_instances(self):
         self.assertEqual(res.out.decode().splitlines(), ["atexit2", "atexit1"])
         self.assertFalse(res.err)
 
+    def test_multiprocessing(self):
+        if sys.platform == "win32":
+            contexts = ["spawn"]
+        elif sys.platform == "darwin":
+            contexts = ["forkserver", "spawn"]
+        else:
+            contexts = ["fork", "forkserver", "spawn"]
+
+        with os_helper.temp_dir() as temp_dir:
+            output_path = os.path.join(temp_dir, 'output.txt')
+            code = textwrap.dedent(rf"""
+                import multiprocessing
+                import os
+                import tempfile
+
+                def test(arg):
+                    import atexit
+                    def exit_handler():
+                        with open({repr(output_path)}, 'a') as f:
+                            f.write(f'|{{arg}}|\n')
+                    atexit.register(exit_handler)
+
+                if __name__ == '__main__':
+                    for context in {contexts}:
+                        p = multiprocessing.get_context(context).Process(target=test, args=(context,))
+                        p.start()
+                        p.join()
+            """)
+            file_name = script_helper.make_script(temp_dir, 'foo', code)
+            script_helper.run_test_script(file_name)
+
+            with open(output_path, "r") as f:
+                out = f.read()
+
+        for context in contexts:
+            self.assertIn(f"|{context}|", out)
+
 
 @support.cpython_only
 class SubinterpreterTest(unittest.TestCase):

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