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

Doc/c-api/sys.rst

@@ -26,6 +26,42 @@ 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
@@ -33,6 +69,9 @@ Operating System Utilities
    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 @@ 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)

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()

Lib/test/test_posix.py

 "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)

Lib/test/test_random.py

 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()

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 @@ exit:
 
 #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 @@ exit:
     #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 @@ exit:
 #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]*/

Modules/posixmodule.c

@@ -25,6 +25,7 @@
 #define PY_SSIZE_T_CLEAN
 
 #include "Python.h"
+#include "pythread.h"
 #include "structmember.h"
 #ifndef MS_WINDOWS
 #include "posixmodule.h"
@@ -394,6 +395,95 @@ static int win32_can_symlink = 0;
 #define MODNAME "posix"
 #endif
 
+
+#ifdef HAVE_FORK
+static void
+run_at_forkers(PyObject *lst, int reverse)
+{
+    Py_ssize_t i;
+    PyObject *cpy;
+
+    if (lst != NULL) {
+        assert(PyList_CheckExact(lst));
+
+        /* Use a list copy in case register_at_fork() is called from
+         * one of the callbacks.
+         */
+        cpy = PyList_GetSlice(lst, 0, PyList_GET_SIZE(lst));
+        if (cpy == NULL)
+            PyErr_WriteUnraisable(lst);
+        else {
+            if (reverse)
+                PyList_Reverse(cpy);
+            for (i = 0; i < PyList_GET_SIZE(cpy); i++) {
+                PyObject *func, *res;
+                func = PyList_GET_ITEM(cpy, i);
+                res = PyObject_CallObject(func, NULL);
+                if (res == NULL)
+                    PyErr_WriteUnraisable(func);
+                else
+                    Py_DECREF(res);
+            }
+            Py_DECREF(cpy);
+        }
+    }
+}
+
+void
+PyOS_BeforeFork(void)
+{
+    run_at_forkers(PyThreadState_Get()->interp->before_forkers, 1);
+
+    _PyImport_AcquireLock();
+}
+
+void
+PyOS_AfterFork_Parent(void)
+{
+    if (_PyImport_ReleaseLock() <= 0)
+        Py_FatalError("failed releasing import lock after fork");
+
+    run_at_forkers(PyThreadState_Get()->interp->after_forkers_parent, 0);
+}
+
+void
+PyOS_AfterFork_Child(void)
+{
+#ifdef WITH_THREAD
+    /* PyThread_ReInitTLS() must be called early, to make sure that the TLS API
+     * can be called safely. */
+    PyThread_ReInitTLS();
+    _PyGILState_Reinit();
+    PyEval_ReInitThreads();
+    _PyImport_ReInitLock();
+#endif
+    _PySignal_AfterFork();
+
+    run_at_forkers(PyThreadState_Get()->interp->after_forkers_child, 0);
+}
+
+static int
+register_at_forker(PyObject **lst, PyObject *func)
+{
+    if (*lst == NULL) {
+        *lst = PyList_New(0);
+        if (*lst == NULL)
+            return -1;
+    }
+    return PyList_Append(*lst, func);
+}
+#endif
+
+/* Legacy wrapper */
+void
+PyOS_AfterFork(void)
+{
+#ifdef HAVE_FORK
+    PyOS_AfterFork_Child();
+#endif
+}
+
+
 #ifdef MS_WINDOWS
 /* defined in fileutils.c */
 PyAPI_FUNC(void) _Py_time_t_to_FILE_TIME(time_t, int, FILETIME *);
@@ -5218,6 +5308,57 @@ os_spawnve_impl(PyObject *module, int mode, path_t *path, PyObject *argv,
 #endif /* HAVE_SPAWNV */
 
 
+#ifdef HAVE_FORK
+/*[clinic input]
+os.register_at_fork
+
+    func: object
+        Function or callable
+    /
+    when: str
+        'before', 'child' or 'parent'
+
+Register a callable object to be called when forking.
+
+'before' callbacks are called in reverse order before forking.
+'child' callbacks are called in order after forking, in the child process.
+'parent' callbacks are called in order after forking, in the parent process.
+
+[clinic start generated code]*/
+
+static PyObject *
+os_register_at_fork_impl(PyObject *module, PyObject *func, const char *when)
+/*[clinic end generated code: output=8943be81a644750c input=5fc05efa4d42eb84]*/
+{
+    PyInterpreterState *interp;
+    PyObject **lst;
+
+    if (!PyCallable_Check(func)) {
+        PyErr_Format(PyExc_TypeError,
+                     "expected callable object, got %R", Py_TYPE(func));
+        return NULL;
+    }
+    interp = PyThreadState_Get()->interp;
+
+    if (!strcmp(when, "before"))
+        lst = &interp->before_forkers;
+    else if (!strcmp(when, "child"))
+        lst = &interp->after_forkers_child;
+    else if (!strcmp(when, "parent"))
+        lst = &interp->after_forkers_parent;
+    else {
+        PyErr_Format(PyExc_ValueError, "unexpected value for `when`: '%s'",
+                     when);
+        return NULL;
+    }
+    if (register_at_forker(lst, func))
+        return NULL;
+    else
+        Py_RETURN_NONE;
+}
+#endif /* HAVE_FORK */
+
+
 #ifdef HAVE_FORK1
 /*[clinic input]
 os.fork1
@@ -5232,24 +5373,18 @@ os_fork1_impl(PyObject *module)
 /*[clinic end generated code: output=0de8e67ce2a310bc input=12db02167893926e]*/
 {
     pid_t pid;
-    int result = 0;
-    _PyImport_AcquireLock();
+
+    PyOS_BeforeFork();
     pid = fork1();
     if (pid == 0) {
         /* child: this clobbers and resets the import lock. */
-        PyOS_AfterFork();
+        PyOS_AfterFork_Child();
     } else {
         /* parent: release the import lock. */
-        result = _PyImport_ReleaseLock();
+        PyOS_AfterFork_Parent();
     }
     if (pid == -1)
         return posix_error();
-    if (result < 0) {
-        /* Don't clobber the OSError if the fork failed. */
-        PyErr_SetString(PyExc_RuntimeError,
-                        "not holding the import lock");
-        return NULL;
-    }
     return PyLong_FromPid(pid);
 }
 #endif /* HAVE_FORK1 */
@@ -5269,24 +5404,18 @@ os_fork_impl(PyObject *module)
 /*[clinic end generated code: output=3626c81f98985d49 input=13c956413110eeaa]*/
 {
     pid_t pid;
-    int result = 0;
-    _PyImport_AcquireLock();
+
+    PyOS_BeforeFork();
     pid = fork();
     if (pid == 0) {
         /* child: this clobbers and resets the import lock. */
-        PyOS_AfterFork();
+        PyOS_AfterFork_Child();
     } else {
         /* parent: release the import lock. */
-        result = _PyImport_ReleaseLock();
+        PyOS_AfterFork_Parent();
     }
     if (pid == -1)
         return posix_error();
-    if (result < 0) {
-        /* Don't clobber the OSError if the fork failed. */
-        PyErr_SetString(PyExc_RuntimeError,
-                        "not holding the import lock");
-        return NULL;
-    }
     return PyLong_FromPid(pid);
 }
 #endif /* HAVE_FORK */
@@ -5868,26 +5997,20 @@ static PyObject *
 os_forkpty_impl(PyObject *module)
 /*[clinic end generated code: output=60d0a5c7512e4087 input=f1f7f4bae3966010]*/
 {
-    int master_fd = -1, result = 0;
+    int master_fd = -1;
     pid_t pid;
 
-    _PyImport_AcquireLock();
+    PyOS_BeforeFork();
     pid = forkpty(&master_fd, NULL, NULL, NULL);
     if (pid == 0) {
         /* child: this clobbers and resets the import lock. */
-        PyOS_AfterFork();
+        PyOS_AfterFork_Child();
     } else {
         /* parent: release the import lock. */
-        result = _PyImport_ReleaseLock();
+        PyOS_AfterFork_Parent();
     }
     if (pid == -1)
         return posix_error();
-    if (result < 0) {
-        /* Don't clobber the OSError if the fork failed. */
-        PyErr_SetString(PyExc_RuntimeError,
-                        "not holding the import lock");
-        return NULL;
-    }
     return Py_BuildValue("(Ni)", PyLong_FromPid(pid), master_fd);
 }
 #endif /* HAVE_FORKPTY */
@@ -12265,6 +12388,7 @@ static PyMethodDef posix_methods[] = {
     OS_SPAWNVE_METHODDEF
     OS_FORK1_METHODDEF
     OS_FORK_METHODDEF
+    OS_REGISTER_AT_FORK_METHODDEF
     OS_SCHED_GET_PRIORITY_MAX_METHODDEF
     OS_SCHED_GET_PRIORITY_MIN_METHODDEF
     OS_SCHED_GETPARAM_METHODDEF

Modules/signalmodule.c

@@ -1618,21 +1618,15 @@ _clear_pending_signals(void)
 }
 
 void
-PyOS_AfterFork(void)
+_PySignal_AfterFork(void)
 {
     /* Clear the signal flags after forking so that they aren't handled
      * in both processes if they came in just before the fork() but before
      * the interpreter had an opportunity to call the handlers.  issue9535. */
     _clear_pending_signals();
 #ifdef WITH_THREAD
-    /* PyThread_ReInitTLS() must be called early, to make sure that the TLS API
-     * can be called safely. */
-    PyThread_ReInitTLS();
-    _PyGILState_Reinit();
-    PyEval_ReInitThreads();
     main_thread = PyThread_get_thread_ident();
     main_pid = getpid();
-    _PyImport_ReInitLock();
 #endif
 }
 

Python/pystate.c

@@ -117,6 +117,11 @@ PyInterpreterState_New(void)
 #else
         interp->dlopenflags = RTLD_LAZY;
 #endif
+#endif
+#ifdef HAVE_FORK
+        interp->before_forkers = NULL;
+        interp->after_forkers_parent = NULL;
+        interp->after_forkers_child = NULL;
 #endif
 
         HEAD_LOCK();
@@ -159,6 +164,11 @@ PyInterpreterState_Clear(PyInterpreterState *interp)
     Py_CLEAR(interp->builtins_copy);
     Py_CLEAR(interp->importlib);
     Py_CLEAR(interp->import_func);
+#ifdef HAVE_FORK
+    Py_CLEAR(interp->before_forkers);
+    Py_CLEAR(interp->after_forkers_parent);
+    Py_CLEAR(interp->after_forkers_child);
+#endif
 }