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

Author: pitrou

Doc/c-api/sys.rst

@@ -26,13 +26,52 @@ Operating System Utilities
    one of the strings ``'<stdin>'`` or ``'???'``.
 
 
+.. c:function:: void PyOS_BeforeFork()
+
+   Function to prepare some internal state before a process fork.  This
+   should be called before calling :c:func:`fork` or any similar function
+   that clones the current process.
+   Only available on systems where :c:func:`fork` is defined.
+
+   .. versionadded:: 3.7
+
+
+.. c:function:: void PyOS_AfterFork_Parent()
+
+   Function to update some internal state after a process fork.  This
+   should be called from the parent process after calling :c:func:`fork`
+   or any similar function that clones the current process, regardless
+   of whether process cloning was successful.
+   Only available on systems where :c:func:`fork` is defined.
+
+   .. versionadded:: 3.7
+
+
+.. c:function:: void PyOS_AfterFork_Child()
+
+   Function to update some internal state after a process fork.  This
+   should be called from the child process after calling :c:func:`fork`
+   or any similar function that clones the current process.
+   Only available on systems where :c:func:`fork` is defined.
+
+   .. versionadded:: 3.7
+
+   .. seealso::
+      :func:`os.register_at_fork` allows registering custom Python functions
+      to be called by :c:func:`PyOS_BeforeFork()`,
+      :c:func:`PyOS_AfterFork_Parent` and  :c:func:`PyOS_AfterFork_Child`.
+
+
 .. c:function:: void PyOS_AfterFork()
 
    Function to update some internal state after a process fork; this should be
    called in the new process if the Python interpreter will continue to be used.
    If a new executable is loaded into the new process, this function does not need
    to be called.
 
+   .. deprecated:: 3.7
+      This function is superseded by :c:func:`PyOS_AfterFork_Child()`.
+
 
 .. c:function:: int PyOS_CheckStack()
 

Doc/library/os.rst

@@ -3280,6 +3280,31 @@ written in Python, such as a mail server's external command delivery program.
    subprocesses.
 
 
+.. function:: register_at_fork(func, when)
+
+   Register *func* as a function to be executed when a new child process
+   is forked.  *when* is a string specifying at which point the function is
+   called and can take the following values:
+
+   * *"before"* means the function is called before forking a child process;
+   * *"parent"* means the function is called from the parent process after
+     forking a child process;
+   * *"child"* means the function is called from the child process.
+
+   Functions registered for execution before forking are called in
+   reverse registration order.  Functions registered for execution
+   after forking (either in the parent or in the child) are called
+   in registration order.
+
+   Note that :c:func:`fork` calls made by third-party C code may not
+   call those functions, unless it explicitly calls :c:func:`PyOS_BeforeFork`,
+   :c:func:`PyOS_AfterFork_Parent` and :c:func:`PyOS_AfterFork_Child`.
+
+   Availability: Unix.
+
+   .. versionadded:: 3.7
+
+
 .. function:: spawnl(mode, path, ...)
               spawnle(mode, path, ..., env)
               spawnlp(mode, file, ...)

Include/intrcheck.h

@@ -7,10 +7,17 @@ extern "C" {
 
 PyAPI_FUNC(int) PyOS_InterruptOccurred(void);
 PyAPI_FUNC(void) PyOS_InitInterrupts(void);
-PyAPI_FUNC(void) PyOS_AfterFork(void);
+#ifdef HAVE_FORK
+PyAPI_FUNC(void) PyOS_BeforeFork(void);
+PyAPI_FUNC(void) PyOS_AfterFork_Parent(void);
+PyAPI_FUNC(void) PyOS_AfterFork_Child(void);
+#endif
+/* Deprecated, please use PyOS_AfterFork_Child() instead */
+PyAPI_FUNC(void) PyOS_AfterFork(void) Py_DEPRECATED(3.7);
 
 #ifndef Py_LIMITED_API
 PyAPI_FUNC(int) _PyOS_IsMainThread(void);
+PyAPI_FUNC(void) _PySignal_AfterFork(void);
 
 #ifdef MS_WINDOWS
 /* windows.h is not included by Python.h so use void* instead of HANDLE */

Include/pystate.h

@@ -74,6 +74,11 @@ typedef struct _is {
     PyObject *import_func;
     /* Initialized to PyEval_EvalFrameDefault(). */
     _PyFrameEvalFunction eval_frame;
+#ifdef HAVE_FORK
+    PyObject *before_forkers;
+    PyObject *after_forkers_parent;
+    PyObject *after_forkers_child;
+#endif
 } PyInterpreterState;
 #endif
 

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)
 

Lib/multiprocessing/popen_fork.py

@@ -68,9 +68,6 @@ def _launch(self, process_obj):
         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)

Lib/random.py

@@ -46,6 +46,7 @@
 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 @@ def _test(N=2000):
 setstate = _inst.setstate
 getrandbits = _inst.getrandbits
 
+if hasattr(_os, "fork"):
+    _os.register_at_fork(_inst.seed, when='child')
+
+
 if __name__ == '__main__':
     _test()

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 @@ def test_waitid(self):
             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)

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 @@ def __init__(self, newarg=None):
                 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()

Misc/NEWS

@@ -341,6 +341,8 @@ Extension Modules
 Library
 -------
 
+- bpo-16500: Allow registering at-fork handlers.
+
 - bpo-30470: Deprecate invalid ctypes call protection on Windows.  Patch by
   Mariatta Wijaya.