diff --git a/Include/cpython/pylifecycle.h b/Include/cpython/pylifecycle.h index a01e9c94f12d75f..16e8c0bceec08c2 100644 --- a/Include/cpython/pylifecycle.h +++ b/Include/cpython/pylifecycle.h @@ -42,7 +42,7 @@ PyAPI_FUNC(void) _Py_NO_RETURN Py_ExitStatusException(PyStatus err); /* Py_PyAtExit is for the atexit module, Py_AtExit is for low-level * exit functions. */ -PyAPI_FUNC(void) _Py_PyAtExit(void (*func)(PyObject *), PyObject *); +PyAPI_FUNC(int) _Py_PyAtExit(void (*func)(PyObject *), PyObject *); /* Restore signals that the interpreter has called SIG_IGN on to SIG_DFL. */ PyAPI_FUNC(void) _Py_RestoreSignals(void); diff --git a/Include/internal/pycore_interp.h b/Include/internal/pycore_interp.h index c6fc6aff5ab2553..40643f5d41463ab 100644 --- a/Include/internal/pycore_interp.h +++ b/Include/internal/pycore_interp.h @@ -121,8 +121,11 @@ struct _is { PyObject *after_forkers_child; #endif /* AtExit module */ - void (*pyexitfunc)(PyObject *); - PyObject *pyexitmodule; + +#define NEXITMODULE 32 + void (*pyexitfunc[NEXITMODULE])(PyObject *); + PyObject *pyexitmodule[NEXITMODULE]; + int nexitmodule; uint64_t tstate_next_unique_id; diff --git a/Lib/test/test_atexit.py b/Lib/test/test_atexit.py index 3105f6c378193ef..c52934dcb4dcdc7 100644 --- a/Lib/test/test_atexit.py +++ b/Lib/test/test_atexit.py @@ -222,6 +222,29 @@ def callback(): self.assertEqual(os.read(r, len(expected)), expected) os.close(r) + def test_multiple_loading_on_subinterpreter(self): + expected = b"atexit2 callback\natexit1 callback\n" + r, w = os.pipe() + + code = r"""if 1: + import sys + import os + def callback(msg): + os.write({:d}, msg) + atexit1 = sys.modules.pop('atexit', None) + if atexit1 is None: + import atexit as atexit1 + del sys.modules['atexit'] + + import atexit as atexit2 + atexit1.register(callback, b"atexit1 callback\n") + atexit2.register(callback, b"atexit2 callback\n") + """.format(w) + ret = support.run_in_subinterp(code) + os.close(w) + self.assertEqual(os.read(r, len(expected)), expected) + os.close(r) + if __name__ == "__main__": unittest.main() diff --git a/Misc/NEWS.d/next/Core and Builtins/2020-04-17-10-02-36.bpo-40288.twb5Dx.rst b/Misc/NEWS.d/next/Core and Builtins/2020-04-17-10-02-36.bpo-40288.twb5Dx.rst new file mode 100644 index 000000000000000..7ed0e545eaab979 --- /dev/null +++ b/Misc/NEWS.d/next/Core and Builtins/2020-04-17-10-02-36.bpo-40288.twb5Dx.rst @@ -0,0 +1,2 @@ +:mod:`atexit` now supports more than once loading per interpreter. Patch by +Dong-hee Na. diff --git a/Modules/atexitmodule.c b/Modules/atexitmodule.c index 8cef64ceb9b6bff..89b3b31fbce15be 100644 --- a/Modules/atexitmodule.c +++ b/Modules/atexitmodule.c @@ -327,7 +327,9 @@ atexit_exec(PyObject *m) { if (modstate->atexit_callbacks == NULL) return -1; - _Py_PyAtExit(atexit_callfuncs, m); + if (_Py_PyAtExit(atexit_callfuncs, m) < 0) { + return -1; + } return 0; } diff --git a/Python/pylifecycle.c b/Python/pylifecycle.c index 7909cdbf5b77204..29c9270b13f7a99 100644 --- a/Python/pylifecycle.c +++ b/Python/pylifecycle.c @@ -2340,25 +2340,33 @@ Py_ExitStatusException(PyStatus status) /* Clean up and exit */ /* For the atexit module. */ -void _Py_PyAtExit(void (*func)(PyObject *), PyObject *module) +int _Py_PyAtExit(void (*func)(PyObject *), PyObject *module) { PyInterpreterState *is = _PyInterpreterState_GET(); - +#define NEXITMODULE 32 + if (is->nexitmodule >= NEXITMODULE) { + PyErr_SetString(PyExc_ImportError, "atexit module can not be loaded more than 32"); + return -1; + } + int n = is->nexitmodule++; /* Guard against API misuse (see bpo-17852) */ - assert(is->pyexitfunc == NULL || is->pyexitfunc == func); + assert(is->pyexitfunc[n] == NULL || is->pyexitfunc[n] == func); - is->pyexitfunc = func; - is->pyexitmodule = module; + is->pyexitfunc[n] = func; + is->pyexitmodule[n] = module; + return is->nexitmodule; } static void call_py_exitfuncs(PyThreadState *tstate) { PyInterpreterState *interp = tstate->interp; - if (interp->pyexitfunc == NULL) - return; - - (*interp->pyexitfunc)(interp->pyexitmodule); + for (int i = interp->nexitmodule -1; i >= 0; i--) { + if (interp->pyexitfunc[i] == NULL) { + continue; + } + (*interp->pyexitfunc[i])(interp->pyexitmodule[i]); + } _PyErr_Clear(tstate); } diff --git a/Python/pystate.c b/Python/pystate.c index 84a694b32e5d0f7..97a6edd0e49b8b5 100644 --- a/Python/pystate.c +++ b/Python/pystate.c @@ -267,6 +267,7 @@ PyInterpreterState_New(void) return NULL; } + interp->nexitmodule = 0; interp->tstate_next_unique_id = 0; interp->audit_hooks = NULL;