8000 bpo-16500: Allow registering at-fork handlers (#1715) · python/cpython@346cbd3 · GitHub
[go: up one dir, main page]
Skip to content

Commit 346cbd3

Browse files
pitroupitrou
authored
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
1 parent f931fd1 commit 346cbd3

File tree

15 files changed

+365
-68
lines changed

15 files changed

+365
-68
lines changed

‎Doc/c-api/sys.rst

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,13 +26,52 @@ Operating System Utilities
2626
one of the strings ``'<stdin>'`` or ``'???'``.
2727
2828
29+
.. c:function:: void PyOS_BeforeFork()
30+
31+
Function to prepare some internal state before a process fork. This
32+
should be called before calling :c:func:`fork` or any similar function
33+
that clones the current process.
34+
Only available on systems where :c:func:`fork` is defined.
35+
36+
.. versionadded:: 3.7
37+
38+
39+
.. c:function:: void PyOS_AfterFork_Parent()
40+
41+
Function to update some internal state after a process fork. This
42+
should be called from the parent process after calling :c:func:`fork`
43+
or any similar function that clones the current process, regardless
44+
of whether process cloning was successful.
45+
Only available on systems where :c:func:`fork` is defined.
46+
47+
.. versionadded:: 3.7
48+
49+
50+
.. c:function:: void PyOS_AfterFork_Child()
51+
52+
Function to update some internal state after a process fork. This
53+
should be called from the child process after calling :c:func:`fork`
54+
or any similar function that clones the current process.
55+
Only available on systems where :c:func:`fork` is defined.
56+
57+
.. versionadded:: 3.7
58+
59+
.. seealso::
60+
:func:`os.register_at_fork` allows registering custom Python functions
61+
to be called by :c:func:`PyOS_BeforeFork()`,
62+
:c:func:`PyOS_AfterFork_Parent` and :c:func:`PyOS_AfterFork_Child`.
63+
64+
2965
.. c:function:: void PyOS_AfterFork()
3066
3167
Function to update some internal state after a process fork; this should be
3268
called in the new process if the Python interpreter will continue to be used.
3369
If a new executable is loaded into the new process, this function does not need
3470
to be called.
3571
72+
.. deprecated:: 3.7
73+
This function is superseded by :c:func:`PyOS_AfterFork_Child()`.
74+
3675
3776
.. c:function:: int PyOS_CheckStack()
3877

‎Doc/library/os.rst

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3280,6 +3280,31 @@ written in Python, such as a mail server's external command delivery program.
32803280
subprocesses.
32813281

32823282

3283+
.. function:: register_at_fork(func, when)
3284+
3285+
Register *func* as a function to be executed when a new child process
3286+
is forked. *when* is a string specifying at which point the function is
3287+
called and can take the following values:
3288+
3289+
* *"before"* means the function is called before forking a child process;
3290+
* *"parent"* means the function is called from the parent process after
3291+
forking a child process;
3292+
* *"child"* means the function is called from the child process.
3293+
3294+
Functions registered for execution before forking are called in
3295+
reverse registration order. Functions registered for execution
3296+
after forking (either in the parent or in the child) are called
3297+
in registration order.
3298+
3299+
Note that :c:func:`fork` calls made by third-party C code may not
3300+
call those functions, unless it explicitly calls :c:func:`PyOS_BeforeFork`,
3301+
:c:func:`PyOS_AfterFork_Parent` and :c:func:`PyOS_AfterFork_Child`.
3302+
3303+
Availability: Unix.
3304+
3305+
.. versionadded:: 3.7
3306+
3307+
32833308
.. function:: spawnl(mode, path, ...)
32843309
spawnle(mode, path, ..., env)
32853310
spawnlp(mode, file, ...)

‎Include/intrcheck.h

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,10 +7,17 @@ extern "C" {
77

88
PyAPI_FUNC(int) PyOS_InterruptOccurred(void);
99
PyAPI_FUNC(void) PyOS_InitInterrupts(void);
10-
PyAPI_FUNC(void) PyOS_AfterFork(void);
10+
#ifdef HAVE_FORK
11+
PyAPI_FUNC(void) PyOS_BeforeFork(void);
12+
PyAPI_FUNC(void) PyOS_AfterFork_Parent(void);
13+
PyAPI_FUNC(void) PyOS_AfterFork_Child(void);
14+
#endif
15+
/* Deprecated, please use PyOS_AfterFork_Child() instead */
16+
PyAPI_FUNC(void) PyOS_AfterFork(void) Py_DEPRECATED(3.7);
1117

1218
#ifndef Py_LIMITED_API
1319
PyAPI_FUNC(int) _PyOS_IsMainThread(void);
20+
PyAPI_FUNC(void) _PySignal_AfterFork(void);
1421

1522
#ifdef MS_WINDOWS
1623
/* windows.h is not included by Python.h so use void* instead of HANDLE */

‎Include/pystate.h

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,11 @@ typedef struct _is {
7474
PyObject *import_func;
7575
/* Initialized to PyEval_EvalFrameDefault(). */
7676
_PyFrameEvalFunction eval_frame;
77+
#ifdef HAVE_FORK
78+
PyObject *before_forkers;
79+
PyObject *after_forkers_parent;
80+
PyObject *after_forkers_child;
81+
#endif
7782
} PyInterpreterState;
7883
#endif
7984

‎Lib/multiprocessing/forkserver.py

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -210,11 +210,6 @@ def _serve_one(s, listener, alive_r, handlers):
210210
# send pid to client processes
211211
write_unsigned(child_w, os.getpid())
212212

213-
# reseed random number generator
214-
if 'random' in sys.modules:
215-
import random
216-
random.seed()
217-
218213
# run process object received over pipe
219214
code = spawn._main(child_r)
220215

‎Lib/multiprocessing/popen_fork.py

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -68,9 +68,6 @@ def _launch(self, process_obj):
6868
if self.pid == 0:
6969
try:
7070
os.close(parent_r)
71-
if 'random' in sys.modules:
72-
import random
73-
random.seed()
7471
code = process_obj._bootstrap()
7572
finally:
7673
os._exit(code)

‎Lib/random.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646
from hashlib import sha512 as _sha512
4747
import itertools as _itertools
4848
import bisect as _bisect
49+
import os as _os
4950

5051
__all__ = ["Random","seed","random","uniform","randint","choice","sample",
5152
"randrange","shuffle","normalvariate","lognormvariate",
@@ -763,5 +764,9 @@ def _test(N=2000):
763764
setstate = _inst.setstate
764765
getrandbits = _inst.getrandbits
765766

767+
if hasattr(_os, "fork"):
768+
_os.register_at_fork(_inst.seed, when='child')
769+
770+
766771
if __name__ == '__main__':
767772
_test()

‎Lib/test/test_posix.py

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
"Test posix functions"
22

33
from test import support
4+
from test.support.script_helper import assert_python_ok
45
android_not_root = support.android_not_root
56

67
# Skip these tests if there is no posix module.
@@ -187,6 +188,45 @@ def test_waitid(self):
187188
res = posix.waitid(posix.P_PID, pid, posix.WEXITED)
188189
self.assertEqual(pid, res.si_pid)
189190

191+
@unittest.skipUnless(hasattr(os, 'fork'), "test needs os.fork()")
192+
def test_register_after_fork(self):
193+
code = """if 1:
194+
import os
195+
196+
r, w = os.pipe()
197+
fin_r, fin_w = os.pipe()
198+
199+
os.register_at_fork(lambda: os.write(w, b'A'), when='before')
200+
os.register_at_fork(lambda: os.write(w, b'B'), when='before')
201+
os.register_at_fork(lambda: os.write(w, b'C'), when='parent')
202+
os.register_at_fork(lambda: os.write(w, b'D'), when='parent')
203+
os.register_at_fork(lambda: os.write(w, b'E'), when='child')
204+
os.register_at_fork(lambda: os.write(w, b'F'), when='child')
205+
206+
pid = os.fork()
207+
if pid == 0:
208+
# At this point, after-forkers have already been executed
209+
os.close(w)
210+
# Wait for parent to tell us to exit
211+
os.read(fin_r, 1)
212+
os._exit(0)
213+
else:
214+
try:
215+
os.close(w)
216+
with open(r, "rb") as f:
217+
data = f.read()
218+
assert len(data) == 6, data
219+
# Check before-fork callbacks
220+
assert data[:2] == b'BA', data
221+
# Check after-fork callbacks
222+
assert sorted(data[2:]) == list(b'CDEF'), data
223+
assert data.index(b'C') < data.index(b'D'), data
224+
assert data.index(b'E') < data.index(b'F'), data
225+
finally:
226+
os.write(fin_w, b'!')
227+
"""
228+
assert_python_ok('-c', code)
229+
190230
@unittest.skipUnless(hasattr(posix, 'lockf'), "test needs posix.lockf()")
191231
def test_lockf(self):
192232
fd = os.open(support.TESTFN, os.O_WRONLY | os.O_CREAT)

‎Lib/test/test_random.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import unittest
22
import unittest.mock
33
import random
4+
import os
45
import time
56
import pickle
67
import warnings
@@ -902,6 +903,24 @@ def __init__(self, newarg=None):
902903
random.Random.__init__(self)
903904
Subclass(newarg=1)
904905

906+
@unittest.skipUnless(hasattr(os, "fork"), "fork() required")
907+
def test_after_fork(self):
908+
# Test the global Random instance gets reseeded in child
909+
r, w = os.pipe()
910+
if os.fork() == 0:
911+
try:
912+
val = random.getrandbits(128)
913+
with open(w, "w") as f:
914+
f.write(str(val))
915+
finally:
916+
os._exit(0)
917+
else:
918+
os.close(w)
919+
val = random.getrandbits(128)
920+
with open(r, "r") as f:
921+
child_val = eval(f.read())
922+
self.assertNotEqual(val, child_val)
923+
905924

906925
if __name__ == "__main__":
907926
unittest.main()
< 10000 /div>

‎Misc/NEWS

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -341,6 +341,8 @@ Extension Modules
341341
Library
342342
-------
343343

344+
- bpo-16500: Allow registering at-fork handlers.
345+
344346
- bpo-30470: Deprecate invalid ctypes call protection on Windows. Patch by
345347
Mariatta Wijaya.
346348

‎Modules/_posixsubprocess.c

Lines changed: 8 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -559,9 +559,7 @@ subprocess_fork_exec(PyObject* self, PyObject *args)
559559
int need_to_reenable_gc = 0;
560560
char *const *exec_array, *const *argv = NULL, *const *envp = NULL;
561561
Py_ssize_t arg_num;
562-
#ifdef WITH_THREAD
563-
int import_lock_held = 0;
564-
#endif
562+
int need_after_fork = 0;
565563

566564
if (!PyArg_ParseTuple(
567565
args, "OOpO!OOiiiiiiiiiiO:fork_exec",
@@ -657,10 +655,8 @@ subprocess_fork_exec(PyObject* self, PyObject *args)
657655
preexec_fn_args_tuple = PyTuple_New(0);
658656
if (!preexec_fn_args_tuple)
659657
goto cleanup;
660-
#ifdef WITH_THREAD
661-
_PyImport_AcquireLock();
662-
import_lock_held = 1;
663-
#endif
658+
PyOS_BeforeFork();
659+
need_after_fork = 1;
664660
}
665661

666662
if (cwd_obj != Py_None) {
@@ -686,7 +682,7 @@ subprocess_fork_exec(PyObject* self, PyObject *args)
686682
* This call may not be async-signal-safe but neither is calling
687683
* back into Python. The user asked us to use hope as a strategy
688684
* to avoid deadlock... */
689-
PyOS_AfterFork();
685+
PyOS_AfterFork_Child();
690686
}
691687

692688
child_exec(exec_array, argv, envp, cwd,
@@ -703,17 +699,10 @@ subprocess_fork_exec(PyObject* self, PyObject *args)
703699
/* Capture the errno exception before errno can be clobbered. */
704700
PyErr_SetFromErrno(PyExc_OSError);
705701
}
706-
#ifdef WITH_THREAD
707-
if (preexec_fn != Py_None
708-
&& _PyImport_ReleaseLock() < 0 && !PyErr_Occurred()) {
709-
PyErr_SetString(PyExc_RuntimeError,
710-
"not holding the import lock");
711-
pid = -1;
712-
}
713-
import_lock_held = 0;
714-
#endif
715702

716703
/* Parent process */
704+
if (need_after_fork)
705+
PyOS_AfterFork_Parent();
717706
if (envp)
718707
_Py_FreeCharPArray(envp);
719708
if (argv)
@@ -733,10 +722,8 @@ subprocess_fork_exec(PyObject* self, PyObject *args)
733722
return PyLong_FromPid(pid);
734723

735724
cleanup:
736-
#ifdef WITH_THREAD
737-
if (import_lock_held)
738-
_PyImport_ReleaseLock();
739-
#endif
725+
if (need_after_fork)
726+
PyOS_AfterFork_Parent();
740727
if (envp)
741728
_Py_FreeCharPArray(envp);
742729
if (argv)

‎Modules/clinic/posixmodule.c.h

Lines changed: 49 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1825,6 +1825,50 @@ os_spawnve(PyObject *module, PyObject **args, Py_ssize_t nargs, PyObject *kwname
18251825

18261826
#endif /* (defined(HAVE_SPAWNV) || defined(HAVE_WSPAWNV)) */
18271827

1828+
#if defined(HAVE_FORK)
1829+
1830+
PyDoc_STRVAR(os_register_at_fork__doc__,
1831+
"register_at_fork($module, func, /, when)\n"
1832+
"--\n"
1833+
"\n"
1834+
"Register a callable object to be called when forking.\n"
1835+
"\n"
1836+
" func\n"
1837+
" Function or callable\n"
1838+
" when\n"
1839+
" \'before\', \'child\' or \'parent\'\n"
1840+
"\n"
1841+
"\'before\' callbacks are called in reverse order before forking.\n"
1842+
"\'child\' callbacks are called in order after forking, in the child process.\n"
1843+
"\'parent\' callbacks are called in order after forking, in the parent process.");
1844+
1845+
#define OS_REGISTER_AT_FORK_METHODDEF \
1846+
{"register_at_fork", (PyCFunction)os_register_at_fork, METH_FASTCALL, os_register_at_fork__doc__},
1847+
1848+
static PyObject *
1849+
os_register_at_fork_impl(PyObject *module, PyObject *func, const char *when);
1850+
1851+
static PyObject *
1852+
os_register_at_fork(PyObject *module, PyObject **args, Py_ssize_t nargs, PyObject *kwnames)
1853+
{
1854+
PyObject *return_value = NULL;
1855+
static const char * const _keywords[] = {"", "when", NULL};
1856+
static _PyArg_Parser _parser = {"Os:register_at_fork", _keywords, 0};
1857+
PyObject *func;
1858+
const char *when;
1859+
1860+
if (!_PyArg_ParseStackAndKeywords(args, nargs, kwnames, &_parser,
1861+
&func, &when)) {
1862+
goto exit;
1863+
}
1864+
return_value = os_register_at_fork_impl(module, func, when);
1865+
1866+
exit:
1867+
return return_value;
1868+
}
1869+
1870+
#endif /* defined(HAVE_FORK) */
1871+
18281872
#if defined(HAVE_FORK1)
18291873

18301874
PyDoc_STRVAR(os_fork1__doc__,
@@ -6122,6 +6166,10 @@ os_getrandom(PyObject *module, PyObject **args, Py_ssize_t nargs, PyObject *kwna
61226166
#define OS_SPAWNVE_METHODDEF
61236167
#endif /* !defined(OS_SPAWNVE_METHODDEF) */
61246168

6169+
#ifndef OS_REGISTER_AT_FORK_METHODDEF
6170+
#define OS_REGISTER_AT_FORK_METHODDEF
6171+
#endif /* !defined(OS_REGISTER_AT_FORK_METHODDEF) */
6172+
61256173
#ifndef OS_FORK1_METHODDEF
61266174
#define OS_FORK1_METHODDEF
61276175
#endif /* !defined(OS_FORK1_METHODDEF) */
@@ -6493,4 +6541,4 @@ os_getrandom(PyObject *module, PyObject **args, Py_ssize_t nargs, PyObject *kwna
64936541
#ifndef OS_GETRANDOM_METHODDEF
64946542
#define OS_GETRANDOM_METHODDEF
64956543
#endif /* !defined(OS_GETRANDOM_METHODDEF) */
6496-
/*[clinic end generated code: output=5529857101c08b49 input=a9049054013a1b77]*/
6544+
/*[clinic end generated code: output=699e11c5579a104e input=a9049054013a1b77]*/

0 commit comments

Comments
 (0)
0