bpo-16500: Allow registering at-fork handlers

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.
 

Modules/_posixsubprocess.c

@@ -559,9 +559,7 @@ subprocess_fork_exec(PyObject* self, PyObject *args)
     int need_to_reenable_gc = 0;
     char *const *exec_array, *const *argv = NULL, *const *envp = NULL;
     Py_ssize_t arg_num;
-#ifdef WITH_THREAD
-    int import_lock_held = 0;
-#endif
+    int need_after_fork = 0;
 
     if (!PyArg_ParseTuple(
             args, "OOpO!OOiiiiiiiiiiO:fork_exec",
@@ -657,10 +655,8 @@ subprocess_fork_exec(PyObject* self, PyObject *args)
         preexec_fn_args_tuple = PyTuple_New(0);
         if (!preexec_fn_args_tuple)
             goto cleanup;
-#ifdef WITH_THREAD
-        _PyImport_AcquireLock();
-        import_lock_held = 1;
-#endif
+        PyOS_BeforeFork();
+        need_after_fork = 1;
     }
 
     if (cwd_obj != Py_None) {
@@ -686,7 +682,7 @@ subprocess_fork_exec(PyObject* self, PyObject *args)
              * This call may not be async-signal-safe but neither is calling
              * back into Python.  The user asked us to use hope as a strategy
              * to avoid deadlock... */
-            PyOS_AfterFork();
+            PyOS_AfterFork_Child();
         }
 
         child_exec(exec_array, argv, envp, cwd,
@@ -703,17 +699,10 @@ subprocess_fork_exec(PyObject* self, PyObject *args)
         /* Capture the errno exception before errno can be clobbered. */
         PyErr_SetFromErrno(PyExc_OSError);
     }
-#ifdef WITH_THREAD
-    if (preexec_fn != Py_None
-        && _PyImport_ReleaseLock() < 0 && !PyErr_Occurred()) {
-        PyErr_SetString(PyExc_RuntimeError,
-                        "not holding the import lock");
-        pid = -1;
-    }
-    import_lock_held = 0;
-#endif
 
     /* Parent process */
+    if (need_after_fork)
+        PyOS_AfterFork_Parent();
     if (envp)
         _Py_FreeCharPArray(envp);
     if (argv)
@@ -733,10 +722,8 @@ subprocess_fork_exec(PyObject* self, PyObject *args)
     return PyLong_FromPid(pid);
 
 cleanup:
-#ifdef WITH_THREAD
-    if (import_lock_held)
-        _PyImport_ReleaseLock();
-#endif
+    if (need_after_fork)
+        PyOS_AfterFork_Parent();
     if (envp)
         _Py_FreeCharPArray(envp);
     if (argv)

Modules/clinic/posixmodule.c.h

@@ -1825,6 +1825,50 @@ os_spawnve(PyObject *module, PyObject **args, Py_ssize_t nargs, PyObject *kwname
 
 #endif /* (defined(HAVE_SPAWNV) || defined(HAVE_WSPAWNV)) */
 
+#if defined(HAVE_FORK)
+
+PyDoc_STRVAR(os_register_at_fork__doc__,
+"register_at_fork($module, func, /, when)\n"
+"--\n"
+"\n"
+"Register a callable object to be called when forking.\n"
+"\n"
+"  func\n"
+"    Function or callable\n"
+"  when\n"
+"    \'before\', \'child\' or \'parent\'\n"
+"\n"
+"\'before\' callbacks are called in reverse order before forking.\n"
+"\'child\' callbacks are called in order after forking, in the child process.\n"
+"\'parent\' callbacks are called in order after forking, in the parent process.");
+
+#define OS_REGISTER_AT_FORK_METHODDEF    \
+    {"register_at_fork", (PyCFunction)os_register_at_fork, METH_FASTCALL, os_register_at_fork__doc__},
+
+static PyObject *
+os_register_at_fork_impl(PyObject *module, PyObject *func, const char *when);
+
+static PyObject *
+os_register_at_fork(PyObject *module, PyObject **args, Py_ssize_t nargs, PyObject *kwnames)
+{
+    PyObject *return_value = NULL;
+    static const char * const _keywords[] = {"", "when", NULL};
+    static _PyArg_Parser _parser = {"Os:register_at_fork", _keywords, 0};
+    PyObject *func;
+    const char *when;
+
+    if (!_PyArg_ParseStackAndKeywords(args, nargs, kwnames, &_parser,
+        &func, &when)) {
+        goto exit;
+    }
+    return_value = os_register_at_fork_impl(module, func, when);
+
+exit:
+    return return_value;
+}
+
+#endif /* defined(HAVE_FORK) */
+
 #if defined(HAVE_FORK1)
 
 PyDoc_STRVAR(os_fork1__doc__,
@@ -6122,6 +6166,10 @@ os_getrandom(PyObject *module, PyObject **args, Py_ssize_t nargs, PyObject *kwna
     #define OS_SPAWNVE_METHODDEF
 #endif /* !defined(OS_SPAWNVE_METHODDEF) */
 
+#ifndef OS_REGISTER_AT_FORK_METHODDEF
+    #define OS_REGISTER_AT_FORK_METHODDEF
+#endif /* !defined(OS_REGISTER_AT_FORK_METHODDEF) */
+
 #ifndef OS_FORK1_METHODDEF
     #define OS_FORK1_METHODDEF
 #endif /* !defined(OS_FORK1_METHODDEF) */
@@ -6493,4 +6541,4 @@ os_getrandom(PyObject *module, PyObject **args, Py_ssize_t nargs, PyObject *kwna
 #ifndef OS_GETRANDOM_METHODDEF
     #define OS_GETRANDOM_METHODDEF
 #endif /* !defined(OS_GETRANDOM_METHODDEF) */
-/*[clinic end generated code: output=5529857101c08b49 input=a9049054013a1b77]*/
+/*[clinic end generated code: output=699e11c5579a104e input=a9049054013a1b77]*/