python · gpshead · May 29, 2017 · May 27, 2017 · May 27, 2017 · May 27, 2017
diff --git a/Doc/c-api/sys.rst b/Doc/c-api/sys.rst
@@ -49,9 +49,10 @@ Operating System Utilities
 
 .. 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.
+   Function to update internal interpreter state after a process fork.
+   This must be called from the child process after calling :c:func:`fork`,
+   or any similar function that clones the current process, if there is
+   any chance the process will call back into the Python interpreter.
    Only available on systems where :c:func:`fork` is defined.
 
    .. versionadded:: 3.7

diff --git a/Doc/library/os.rst b/Doc/library/os.rst
@@ -3280,16 +3280,22 @@ written in Python, such as a mail server's external command delivery program.
    subprocesses.
 
 
-.. function:: register_at_fork(func, when)
+.. function:: register_at_fork(*, before=None, after_in_parent=None, \
+                               after_in_child=None)
 
-   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:
+   Register callables to be executed when a new child process is forked
+   using :func:`os.fork` or similar process cloning APIs.
+   The parameters are optional and keyword-only.
+   Each specifies a different call point.
 
-   * *"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.
+   * *before* is a function called before forking a child process.
+   * *after_in_parent* is a function called from the parent process
+     after forking a child process.
+   * *after_in_child* is a function called from the child process.
+
+   These calls are only made if control is expected to return to the
+   Python interpreter.  A typical :mod:`subprocess` launch will not
+   trigger them as the child is not going to re-enter the interpreter.
 
    Functions registered for execution before forking are called in
    reverse registration order.  Functions registered for execution
@@ -3300,6 +3306,8 @@ written in Python, such as a mail server's external command delivery program.
    call those functions, unless it explicitly calls :c:func:`PyOS_BeforeFork`,
    :c:func:`PyOS_AfterFork_Parent` and :c:func:`PyOS_AfterFork_Child`.
 
+   There is no way to unregister a function.
+
    Availability: Unix.
 
    .. versionadded:: 3.7

diff --git a/Lib/random.py b/Lib/random.py
@@ -765,7 +765,7 @@ def _test(N=2000):
 getrandbits = _inst.getrandbits
 
 if hasattr(_os, "fork"):
-    _os.register_at_fork(_inst.seed, when='child')
+    _os.register_at_fork(after_in_child=_inst.seed)
 
 
 if __name__ == '__main__':

diff --git a/Lib/test/test_posix.py b/Lib/test/test_posix.py
@@ -189,19 +189,41 @@ def test_waitid(self):
             self.assertEqual(pid, res.si_pid)
 
     @unittest.skipUnless(hasattr(os, 'fork'), "test needs os.fork()")
-    def test_register_after_fork(self):
+    def test_register_at_fork(self):
+        with self.assertRaises(TypeError, msg="Positional args not allowed"):
+            os.register_at_fork(lambda: None)
+        with self.assertRaises(TypeError, msg="Args must be callable"):
+            os.register_at_fork(before=2)
+        with self.assertRaises(TypeError, msg="Args must be callable"):
+            os.register_at_fork(after_in_child="three")
+        with self.assertRaises(TypeError, msg="Args must be callable"):
+            os.register_at_fork(after_in_parent=b"Five")
+        with self.assertRaises(TypeError, msg="Args must not be None"):
+            os.register_at_fork(before=None)
+        with self.assertRaises(TypeError, msg="Args must not be None"):
+            os.register_at_fork(after_in_child=None)
+        with self.assertRaises(TypeError, msg="Args must not be None"):
+            os.register_at_fork(after_in_parent=None)
+        with self.assertRaises(TypeError, msg="Invalid arg was allowed"):
+            # Ensure a combination of valid and invalid is an error.
+            os.register_at_fork(before=None, after_in_parent=lambda: 3)
+        with self.assertRaises(TypeError, msg="Invalid arg was allowed"):
+            # Ensure a combination of valid and invalid is an error.
+            os.register_at_fork(before=lambda: None, after_in_child='')
+        # We test actual registrations in their own process so as not to
+        # pollute this one.  There is no way to unregister for cleanup.
         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')
+            os.register_at_fork(before=lambda: os.write(w, b'A'))
+            os.register_at_fork(after_in_parent=lambda: os.write(w, b'C'))
+            os.register_at_fork(after_in_child=lambda: os.write(w, b'E'))
+            os.register_at_fork(before=lambda: os.write(w, b'B'),
+                                after_in_parent=lambda: os.write(w, b'D'),
+                                after_in_child=lambda: os.write(w, b'F'))
 
             pid = os.fork()
             if pid == 0:

diff --git a/Lib/threading.py b/Lib/threading.py
@@ -1359,5 +1359,5 @@ def _after_fork():
         assert len(_active) == 1
 
 
-if hasattr(_os, "fork"):
-    _os.register_at_fork(_after_fork, when="child")
+if hasattr(_os, "register_at_fork"):
+    _os.register_at_fork(after_in_child=_after_fork)
diff --git a/Modules/_posixsubprocess.c b/Modules/_posixsubprocess.c
@@ -651,14 +651,6 @@ subprocess_fork_exec(PyObject* self, PyObject *args)
             goto cleanup;
     }
 
-    if (preexec_fn != Py_None) {
-        preexec_fn_args_tuple = PyTuple_New(0);
-        if (!preexec_fn_args_tuple)
-            goto cleanup;
-        PyOS_BeforeFork();
-        need_after_fork = 1;
-    }
-
     if (cwd_obj != Py_None) {
         if (PyUnicode_FSConverter(cwd_obj, &cwd_obj2) == 0)
             goto cleanup;
@@ -668,6 +660,17 @@ subprocess_fork_exec(PyObject* self, PyObject *args)
         cwd_obj2 = NULL;
     }
 
+    /* This must be the last thing done before fork() because we do not
+     * want to call PyOS_BeforeFork() if there is any chance of another
+     * error leading to the cleanup: code without calling fork(). */
+    if (preexec_fn != Py_None) {
+        preexec_fn_args_tuple = PyTuple_New(0);
+        if (!preexec_fn_args_tuple)
+            goto cleanup;
+        PyOS_BeforeFork();
+        need_after_fork = 1;
+    }
+
     pid = fork();
     if (pid == 0) {
         /* Child process */
@@ -722,8 +725,6 @@ subprocess_fork_exec(PyObject* self, PyObject *args)
     return PyLong_FromPid(pid);
 
 cleanup:
-    if (need_after_fork)
-        PyOS_AfterFork_Parent();
     if (envp)
         _Py_FreeCharPArray(envp);
     if (argv)

diff --git a/Modules/clinic/posixmodule.c.h b/Modules/clinic/posixmodule.c.h
@@ -1828,40 +1828,44 @@ os_spawnve(PyObject *module, PyObject **args, Py_ssize_t nargs, PyObject *kwname
 #if defined(HAVE_FORK)
 
 PyDoc_STRVAR(os_register_at_fork__doc__,
-"register_at_fork($module, func, /, when)\n"
+"register_at_fork($module, /, *, before=None, after_in_child=None,\n"
+"                 after_in_parent=None)\n"
 "--\n"
 "\n"
-"Register a callable object to be called when forking.\n"
+"Register callables to be called when forking a new process.\n"
 "\n"
-"  func\n"
-"    Function or callable\n"
-"  when\n"
-"    \'before\', \'child\' or \'parent\'\n"
+"  before\n"
+"    A callable to be called in the parent before the fork() syscall.\n"
+"  after_in_child\n"
+"    A callable to be called in the child after fork().\n"
+"  after_in_parent\n"
+"    A callable to be called in the parent after fork().\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.");
+"\'before\' callbacks are called in reverse order.\n"
+"\'after_in_child\' and \'after_in_parent\' callbacks are called in order.");
 
 #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);
+os_register_at_fork_impl(PyObject *module, PyObject *before,
+                         PyObject *after_in_child, PyObject *after_in_parent);
 
 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;
+    static const char * const _keywords[] = {"before", "after_in_child", "after_in_parent", NULL};
+    static _PyArg_Parser _parser = {"|$OOO:register_at_fork", _keywords, 0};
+    PyObject *before = NULL;
+    PyObject *after_in_child = NULL;
+    PyObject *after_in_parent = NULL;
 
     if (!_PyArg_ParseStackAndKeywords(args, nargs, kwnames, &_parser,
-        &func, &when)) {
+        &before, &after_in_child, &after_in_parent)) {
         goto exit;
     }
-    return_value = os_register_at_fork_impl(module, func, when);
+    return_value = os_register_at_fork_impl(module, before, after_in_child, after_in_parent);
 
 exit:
     return return_value;
@@ -6541,4 +6545,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=699e11c5579a104e input=a9049054013a1b77]*/
+/*[clinic end generated code: output=dce741f527ddbfa4 input=a9049054013a1b77]*/
diff --git a/Modules/posixmodule.c b/Modules/posixmodule.c
@@ -465,6 +465,8 @@ PyOS_AfterFork_Child(void)
 static int
 register_at_forker(PyObject **lst, PyObject *func)
 {
+    if (func == NULL)  /* nothing to register? do nothing. */
+        return 0;
     if (*lst == NULL) {
         *lst = PyList_New(0);
         if (*lst == NULL)
@@ -5309,52 +5311,67 @@ os_spawnve_impl(PyObject *module, int mode, path_t *path, PyObject *argv,
 
 
 #ifdef HAVE_FORK
+
+/* Helper function to validate arguments.
+   Returns 0 on success.  non-zero on failure with a TypeError raised.
+   If obj is non-NULL it must be callable.  */
+static int
+check_null_or_callable(PyObject *obj, const char* obj_name)
+{
+    if (obj && !PyCallable_Check(obj)) {
+        PyErr_Format(PyExc_TypeError, "'%s' must be callable, not %s",
+                     obj_name, Py_TYPE(obj)->tp_name);
+        return -1;
+    }
+    return 0;
+}
+
 /*[clinic input]
 os.register_at_fork
 
-    func: object
-        Function or callable
-    /
-    when: str
-        'before', 'child' or 'parent'
+    *
+    before: object=NULL
+        A callable to be called in the parent before the fork() syscall.
+    after_in_child: object=NULL
+        A callable to be called in the child after fork().
+    after_in_parent: object=NULL
+        A callable to be called in the parent after fork().
 
-Register a callable object to be called when forking.
+Register callables to be called when forking a new process.
 
-'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.
+'before' callbacks are called in reverse order.
+'after_in_child' and 'after_in_parent' callbacks are called in order.
 
 [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]*/
+os_register_at_fork_impl(PyObject *module, PyObject *before,
+                         PyObject *after_in_child, PyObject *after_in_parent)
+/*[clinic end generated code: output=5398ac75e8e97625 input=cd1187aa85d2312e]*/
 {
     PyInterpreterState *interp;
-    PyObject **lst;
 
-    if (!PyCallable_Check(func)) {
-        PyErr_Format(PyExc_TypeError,
-                     "expected callable object, got %R", Py_TYPE(func));
+    if (!before && !after_in_child && !after_in_parent) {
+        PyErr_SetString(PyExc_TypeError, "At least one argument is required.");
+        return NULL;
+    }
+    if (check_null_or_callable(before, "before") ||
+        check_null_or_callable(after_in_child, "after_in_child") ||
+        check_null_or_callable(after_in_parent, "after_in_parent")) {
         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);
+    if (register_at_forker(&interp->before_forkers, before)) {
         return NULL;
     }
-    if (register_at_forker(lst, func))
+    if (register_at_forker(&interp->after_forkers_child, after_in_child)) {
         return NULL;
-    else
-        Py_RETURN_NONE;
+    }
+    if (register_at_forker(&interp->after_forkers_parent, after_in_parent)) {
+        return NULL;
+    }
+    Py_RETURN_NONE;
 }
 #endif /* HAVE_FORK */