bpo-16500: Allow registering at-fork handlers (#1715)

* bpo-16500: Allow registering at-fork handlers

* Address Serhiy's comments

* Add doc for new C API

* Add doc for new Python-facing function

* Add NEWS entry + doc nit

5 files changed, 64 insertions, 8 deletions

diff --git a/Lib/multiprocessing/forkserver.py b/Lib/multiprocessing/forkserver.py
index 6e09539993..8156dae3b7 100644
--- a/Lib/multiprocessing/forkserver.py
+++ b/Lib/multiprocessing/forkserver.py
@@ -210,11 +210,6 @@ def _serve_one(s, listener, alive_r, handlers):
     # send pid to client processes
     write_unsigned(child_w, os.getpid())
 
-    # reseed random number generator
-    if 'random' in sys.modules:
-        import random
-        random.seed()
-
     # run process object received over pipe
     code = spawn._main(child_r)
 
diff --git a/Lib/multiprocessing/popen_fork.py b/Lib/multiprocessing/popen_fork.py
index d2ebd7cfbe..683b52d227 100644
--- a/Lib/multiprocessing/popen_fork.py
+++ b/Lib/multiprocessing/popen_fork.py
@@ -68,9 +68,6 @@ class Popen(object):
         if self.pid == 0:
             try:
                 os.close(parent_r)
-                if 'random' in sys.modules:
-                    import random
-                    random.seed()
                 code = process_obj._bootstrap()
             finally:
                 os._exit(code)
diff --git a/Lib/random.py b/Lib/random.py
index ad1c9167b0..52df7d8f74 100644
--- a/Lib/random.py
+++ b/Lib/random.py
@@ -46,6 +46,7 @@ from _collections_abc import Set as _Set, Sequence as _Sequence
 from hashlib import sha512 as _sha512
 import itertools as _itertools
 import bisect as _bisect
+import os as _os
 
 __all__ = ["Random","seed","random","uniform","randint","choice","sample",
            "randrange","shuffle","normalvariate","lognormvariate",
@@ -763,5 +764,9 @@ getstate = _inst.getstate
 setstate = _inst.setstate
 getrandbits = _inst.getrandbits
 
+if hasattr(_os, "fork"):
+    _os.register_at_fork(_inst.seed, when='child')
+
+
 if __name__ == '__main__':
     _test()
diff --git a/Lib/test/test_posix.py b/Lib/test/test_posix.py
index 029d0815e9..a72f83c8dc 100644
--- a/Lib/test/test_posix.py
+++ b/Lib/test/test_posix.py
@@ -1,6 +1,7 @@
 "Test posix functions"
 
 from test import support
+from test.support.script_helper import assert_python_ok
 android_not_root = support.android_not_root
 
 # Skip these tests if there is no posix module.
@@ -187,6 +188,45 @@ class PosixTester(unittest.TestCase):
             res = posix.waitid(posix.P_PID, pid, posix.WEXITED)
             self.assertEqual(pid, res.si_pid)
 
+    @unittest.skipUnless(hasattr(os, 'fork'), "test needs os.fork()")
+    def test_register_after_fork(self):
+        code = """if 1:
+            import os
+
+            r, w = os.pipe()
+            fin_r, fin_w = os.pipe()
+
+            os.register_at_fork(lambda: os.write(w, b'A'), when='before')
+            os.register_at_fork(lambda: os.write(w, b'B'), when='before')
+            os.register_at_fork(lambda: os.write(w, b'C'), when='parent')
+            os.register_at_fork(lambda: os.write(w, b'D'), when='parent')
+            os.register_at_fork(lambda: os.write(w, b'E'), when='child')
+            os.register_at_fork(lambda: os.write(w, b'F'), when='child')
+
+            pid = os.fork()
+            if pid == 0:
+                # At this point, after-forkers have already been executed
+                os.close(w)
+                # Wait for parent to tell us to exit
+                os.read(fin_r, 1)
+                os._exit(0)
+            else:
+                try:
+                    os.close(w)
+                    with open(r, "rb") as f:
+                        data = f.read()
+                        assert len(data) == 6, data
+                        # Check before-fork callbacks
+                        assert data[:2] == b'BA', data
+                        # Check after-fork callbacks
+                        assert sorted(data[2:]) == list(b'CDEF'), data
+                        assert data.index(b'C') < data.index(b'D'), data
+                        assert data.index(b'E') < data.index(b'F'), data
+                finally:
+                    os.write(fin_w, b'!')
+            """
+        assert_python_ok('-c', code)
+
     @unittest.skipUnless(hasattr(posix, 'lockf'), "test needs posix.lockf()")
     def test_lockf(self):
         fd = os.open(support.TESTFN, os.O_WRONLY | os.O_CREAT)
diff --git a/Lib/test/test_random.py b/Lib/test/test_random.py
index 45468c7ce4..f657b46b3a 100644
--- a/Lib/test/test_random.py
+++ b/Lib/test/test_random.py
@@ -1,6 +1,7 @@
 import unittest
 import unittest.mock
 import random
+import os
 import time
 import pickle
 import warnings
@@ -902,6 +903,24 @@ class TestModule(unittest.TestCase):
                 random.Random.__init__(self)
         Subclass(newarg=1)
 
+    @unittest.skipUnless(hasattr(os, "fork"), "fork() required")
+    def test_after_fork(self):
+        # Test the global Random instance gets reseeded in child
+        r, w = os.pipe()
+        if os.fork() == 0:
+            try:
+                val = random.getrandbits(128)
+                with open(w, "w") as f:
+                    f.write(str(val))
+            finally:
+                os._exit(0)
+        else:
+            os.close(w)
+            val = random.getrandbits(128)
+            with open(r, "r") as f:
+                child_val = eval(f.read())
+            self.assertNotEqual(val, child_val)
+
 
 if __name__ == "__main__":
     unittest.main()