diff --git a/Lib/test/test_atexit.py b/Lib/test/test_atexit.py --- a/Lib/test/test_atexit.py +++ b/Lib/test/test_atexit.py @@ -2,6 +2,7 @@ import sys import unittest import io import atexit +import _testcapi from test import support ### helpers @@ -23,7 +24,9 @@ def raise1(): def raise2(): raise SystemError -class TestCase(unittest.TestCase): + +class GeneralTest(unittest.TestCase): + def setUp(self): self.save_stdout = sys.stdout self.save_stderr = sys.stderr @@ -122,8 +125,43 @@ class TestCase(unittest.TestCase): self.assertEqual(l, [5]) +class SubinterpreterTest(unittest.TestCase): + + def test_callbacks_leak(self): + # This test shows a leak in refleak mode if atexit doesn't + # take care to free callbacks in its per-subinterpreter module + # state. + n = atexit._ncallbacks() + code = r"""if 1: + import atexit + def f(): + pass + atexit.register(f) + del atexit + """ + ret = _testcapi.run_in_subinterp(code) + self.assertEqual(ret, 0) + self.assertEqual(atexit._ncallbacks(), n) + + def test_callbacks_leak_refcycle(self): + # Similar to the above, but with a refcycle through the atexit + # module. + n = atexit._ncallbacks() + code = r"""if 1: + import atexit + def f(): + pass + atexit.register(f) + atexit.__atexit = atexit + """ + ret = _testcapi.run_in_subinterp(code) + self.assertEqual(ret, 0) + self.assertEqual(atexit._ncallbacks(), n) + + def test_main(): - support.run_unittest(TestCase) + support.run_unittest(__name__) + if __name__ == "__main__": test_main() diff --git a/Modules/atexitmodule.c b/Modules/atexitmodule.c --- a/Modules/atexitmodule.c +++ b/Modules/atexitmodule.c @@ -10,8 +10,6 @@ /* Forward declaration (for atexit_cleanup) */ static PyObject *atexit_clear(PyObject*, PyObject*); -/* Forward declaration (for atexit_callfuncs) */ -static void atexit_cleanup(PyObject*); /* Forward declaration of module object */ static struct PyModuleDef atexitmodule; @@ -33,6 +31,35 @@ typedef struct { #define GET_ATEXIT_STATE(mod) ((atexitmodule_state*)PyModule_GetState(mod)) +static void +atexit_delete_cb(atexitmodule_state *modstate, int i) +{ + atexit_callback *cb; + + cb = modstate->atexit_callbacks[i]; + modstate->atexit_callbacks[i] = NULL; + Py_DECREF(cb->func); + Py_DECREF(cb->args); + Py_XDECREF(cb->kwargs); + PyMem_Free(cb); +} + +/* Clear all callbacks without calling them */ +static void +atexit_cleanup(atexitmodule_state *modstate) +{ + atexit_callback *cb; + int i; + for (i = 0; i < modstate->ncallbacks; i++) { + cb = modstate->atexit_callbacks[i]; + if (cb == NULL) + continue; + + atexit_delete_cb(modstate, i); + } + modstate->ncallbacks = 0; +} + /* Installed into pythonrun.c's atexit mechanism */ static void @@ -78,34 +105,12 @@ atexit_callfuncs(void) } } - atexit_cleanup(module); + atexit_cleanup(modstate); if (exc_type) PyErr_Restore(exc_type, exc_value, exc_tb); } -static void -atexit_delete_cb(PyObject *self, int i) -{ - atexitmodule_state *modstate; - atexit_callback *cb; - - modstate = GET_ATEXIT_STATE(self); - cb = modstate->atexit_callbacks[i]; - modstate->atexit_callbacks[i] = NULL; - Py_DECREF(cb->func); - Py_DECREF(cb->args); - Py_XDECREF(cb->kwargs); - PyMem_Free(cb); -} - -static void -atexit_cleanup(PyObject *self) -{ - PyObject *r = atexit_clear(self, NULL); - Py_DECREF(r); -} - /* ===================================================================== */ /* Module methods. */ @@ -194,21 +199,50 @@ Clear the list of previously registered static PyObject * atexit_clear(PyObject *self, PyObject *unused) { + atexit_cleanup(GET_ATEXIT_STATE(self)); + Py_RETURN_NONE; +} + +PyDoc_STRVAR(atexit_ncallbacks__doc__, +"_ncallbacks() -> int\n\ +\n\ +Return the number of registered exit functions."); + +static PyObject * +atexit_ncallbacks(PyObject *self, PyObject *unused) +{ atexitmodule_state *modstate; - atexit_callback *cb; - int i; modstate = GET_ATEXIT_STATE(self); + return PyLong_FromSsize_t(modstate->ncallbacks); +} + +static int +atexit_m_traverse(PyObject *self, visitproc visit, void *arg) +{ + int i; + atexitmodule_state *modstate; + + modstate = GET_ATEXIT_STATE(self); for (i = 0; i < modstate->ncallbacks; i++) { - cb = modstate->atexit_callbacks[i]; + atexit_callback *cb = modstate->atexit_callbacks[i]; if (cb == NULL) continue; + Py_VISIT(cb->func); + Py_VISIT(cb->args); + Py_VISIT(cb->kwargs); + } + return 0; +} - atexit_delete_cb(self, i); - } - modstate->ncallbacks = 0; - Py_RETURN_NONE; +static int +atexit_m_clear(PyObject *self) +{ + atexitmodule_state *modstate; + modstate = GET_ATEXIT_STATE(self); + atexit_cleanup(modstate); + return 0; } static void @@ -216,6 +250,7 @@ atexit_free(PyObject *m) { atexitmodule_state *modstate; modstate = GET_ATEXIT_STATE(m); + atexit_cleanup(modstate); PyMem_Free(modstate->atexit_callbacks); } @@ -246,7 +281,7 @@ atexit_unregister(PyObject *self, PyObje if (eq < 0) return NULL; if (eq) - atexit_delete_cb(self, i); + atexit_delete_cb(modstate, i); } Py_RETURN_NONE; } @@ -260,6 +295,8 @@ static PyMethodDef atexit_methods[] = { atexit_unregister__doc__}, {"_run_exitfuncs", (PyCFunction) atexit_run_exitfuncs, METH_NOARGS, atexit_run_exitfuncs__doc__}, + {"_ncallbacks", (PyCFunction) atexit_ncallbacks, METH_NOARGS, + atexit_ncallbacks__doc__}, {NULL, NULL} /* sentinel */ }; @@ -275,15 +312,15 @@ Two public functions, register and unreg static struct PyModuleDef atexitmodule = { - PyModuleDef_HEAD_INIT, - "atexit", - atexit__doc__, - sizeof(atexitmodule_state), - atexit_methods, - NULL, - NULL, - NULL, - (freefunc)atexit_free + PyModuleDef_HEAD_INIT, + "atexit", + atexit__doc__, + sizeof(atexitmodule_state), + atexit_methods, + NULL, + atexit_m_traverse, + atexit_m_clear, + (freefunc)atexit_free }; PyMODINIT_FUNC