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
This commit is contained in:
Antoine Pitrou 2017-05-27 17:50:54 +02:00 committed by GitHub
parent f931fd1c2a
commit 346cbd351e
15 changed files with 365 additions and 68 deletions

View file

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

View file

@ -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, ...)

View file

@ -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 */

View file

@ -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

View file

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

View file

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

View file

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

View file

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

View file

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

View file

@ -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.

View file

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

View file

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

View file

@ -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

View file

@ -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
}

View file

@ -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
}