bpo-31901: atexit callbacks should be run at subinterpreter shutdown (#4611)

Change atexit behavior and PEP-489 multiphase init support.

Author: Marcel Plch

Doc/library/atexit.rst

@@ -20,6 +20,9 @@ at interpreter termination time they will be run in the order ``C``, ``B``,
 program is killed by a signal not handled by Python, when a Python fatal
 internal error is detected, or when :func:`os._exit` is called.
 
+.. versionchanged:: 3.7
+    When used with C-API subinterpreters, registered functions
+    are local to the interpreter they were registered in.
 
 .. function:: register(func, *args, **kwargs)
 

Include/internal/pystate.h

@@ -90,7 +90,6 @@ typedef struct pyruntimestate {
 #define NEXITFUNCS 32
     void (*exitfuncs[NEXITFUNCS])(void);
     int nexitfuncs;
-    void (*pyexitfunc)(void);
 
     struct _gc_runtime_state gc;
     struct _warnings_runtime_state warnings;

Include/pylifecycle.h

@@ -92,7 +92,7 @@ PyAPI_FUNC(void) Py_EndInterpreter(PyThreadState *);
  * exit functions.
  */
 #ifndef Py_LIMITED_API
-PyAPI_FUNC(void) _Py_PyAtExit(void (*func)(void));
+PyAPI_FUNC(void) _Py_PyAtExit(void (*func)(PyObject *), PyObject *);
 #endif
 PyAPI_FUNC(int) Py_AtExit(void (*func)(void));
 

Include/pystate.h

@@ -131,6 +131,9 @@ typedef struct _is {
     PyObject *after_forkers_parent;
     PyObject *after_forkers_child;
 #endif
+    /* AtExit module */
+    void (*pyexitfunc)(PyObject *);
+    PyObject *pyexitmodule;
 } PyInterpreterState;
 #endif   /* !Py_LIMITED_API */
 

Lib/test/test_atexit.py

@@ -2,6 +2,7 @@
 import unittest
 import io
 import atexit
+import os
 from test import support
 from test.support import script_helper
 
@@ -203,6 +204,24 @@ def f():
         self.assertEqual(ret, 0)
         self.assertEqual(atexit._ncallbacks(), n)
 
+    def test_callback_on_subinterpreter_teardown(self):
+        # This tests if a callback is called on
+        # subinterpreter teardown.
+        expected = b"The test has passed!"
+        r, w = os.pipe()
+
+        code = r"""if 1:
+            import os
+            import atexit
+            def callback():
+                os.write({:d}, b"The test has passed!")
+            atexit.register(callback)
+        """.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()

Misc/NEWS.d/next/Core and Builtins/2017-11-28-15-04-14.bpo-31901.mDeCLK.rst

@@ -0,0 +1 @@
+The `atexit` module now has its callback stored per interpreter.

Modules/atexitmodule.c

@@ -63,15 +63,13 @@ atexit_cleanup(atexitmodule_state *modstate)
 /* Installed into pylifecycle.c's atexit mechanism */
 
 static void
-atexit_callfuncs(void)
+atexit_callfuncs(PyObject *module)
 {
     PyObject *exc_type = NULL, *exc_value, *exc_tb, *r;
     atexit_callback *cb;
-    PyObject *module;
     atexitmodule_state *modstate;
     int i;
 
-    module = PyState_FindModule(&atexitmodule);
     if (module == NULL)
         return;
     modstate = GET_ATEXIT_STATE(module);
@@ -185,7 +183,7 @@ Run all registered exit functions.");
 static PyObject *
 atexit_run_exitfuncs(PyObject *self, PyObject *unused)
 {
-    atexit_callfuncs();
+    atexit_callfuncs(self);
     if (PyErr_Occurred())
         return NULL;
     Py_RETURN_NONE;
@@ -225,13 +223,15 @@ atexit_m_traverse(PyObject *self, visitproc visit, void *arg)
     atexitmodule_state *modstate;
 
     modstate = GET_ATEXIT_STATE(self);
-    for (i = 0; i < modstate->ncallbacks; i++) {
-        atexit_callback *cb = modstate->atexit_callbacks[i];
-        if (cb == NULL)
-            continue;
-        Py_VISIT(cb->func);
-        Py_VISIT(cb->args);
-        Py_VISIT(cb->kwargs);
+    if (modstate != NULL) {
+        for (i = 0; i < modstate->ncallbacks; 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;
 }
@@ -241,7 +241,9 @@ atexit_m_clear(PyObject *self)
 {
     atexitmodule_state *modstate;
     modstate = GET_ATEXIT_STATE(self);
-    atexit_cleanup(modstate);
+    if (modstate != NULL) {
+        atexit_cleanup(modstate);
+    }
     return 0;
 }
 
@@ -250,8 +252,10 @@ atexit_free(PyObject *m)
 {
     atexitmodule_state *modstate;
     modstate = GET_ATEXIT_STATE(m);
-    atexit_cleanup(modstate);
-    PyMem_Free(modstate->atexit_callbacks);
+    if (modstate != NULL) {
+        atexit_cleanup(modstate);
+        PyMem_Free(modstate->atexit_callbacks);
+    }
 }
 
 PyDoc_STRVAR(atexit_unregister__doc__,
@@ -310,14 +314,34 @@ upon normal program termination.\n\
 Two public functions, register and unregister, are defined.\n\
 ");
 
+static int
+atexit_exec(PyObject *m) {
+    atexitmodule_state *modstate;
+
+    modstate = GET_ATEXIT_STATE(m);
+    modstate->callback_len = 32;
+    modstate->ncallbacks = 0;
+    modstate->atexit_callbacks = PyMem_New(atexit_callback*,
+                                           modstate->callback_len);
+    if (modstate->atexit_callbacks == NULL)
+        return -1;
+
+    _Py_PyAtExit(atexit_callfuncs, m);
+    return 0;
+}
+
+static PyModuleDef_Slot atexit_slots[] = {
+    {Py_mod_exec, atexit_exec},
+    {0, NULL}
+};
 
 static struct PyModuleDef atexitmodule = {
     PyModuleDef_HEAD_INIT,
     "atexit",
     atexit__doc__,
     sizeof(atexitmodule_state),
     atexit_methods,
-    NULL,
+    atexit_slots,
     atexit_m_traverse,
     atexit_m_clear,
     (freefunc)atexit_free
@@ -326,21 +350,5 @@ static struct PyModuleDef atexitmodule = {
 PyMODINIT_FUNC
 PyInit_atexit(void)
 {
-    PyObject *m;
-    atexitmodule_state *modstate;
-
-    m = PyModule_Create(&atexitmodule);
-    if (m == NULL)
-        return NULL;
-
-    modstate = GET_ATEXIT_STATE(m);
-    modstate->callback_len = 32;
-    modstate->ncallbacks = 0;
-    modstate->atexit_callbacks = PyMem_New(atexit_callback*,
-                                           modstate->callback_len);
-    if (modstate->atexit_callbacks == NULL)
-        return NULL;
-
-    _Py_PyAtExit(atexit_callfuncs);
-    return m;
+    return PyModuleDef_Init(&atexitmodule);
 }

Programs/_testembed.c

@@ -126,7 +126,6 @@ static int test_forced_io_encoding(void)
     return 0;
 }
 
-
 /*********************************************************
  * Test parts of the C-API that work before initialization
  *********************************************************/

Python/pylifecycle.c

@@ -56,7 +56,7 @@ static _PyInitError initfsencoding(PyInterpreterState *interp);
 static _PyInitError initsite(void);
 static _PyInitError init_sys_streams(PyInterpreterState *interp);
 static _PyInitError initsigs(void);
-static void call_py_exitfuncs(void);
+static void call_py_exitfuncs(PyInterpreterState *);
 static void wait_for_thread_shutdown(void);
 static void call_ll_exitfuncs(void);
 extern int _PyUnicode_Init(void);
@@ -1006,6 +1006,10 @@ Py_FinalizeEx(void)
 
     wait_for_thread_shutdown();
 
+    /* Get current thread state and interpreter pointer */
+    tstate = PyThreadState_GET();
+    interp = tstate->interp;
+
     /* The interpreter is still entirely intact at this point, and the
      * exit funcs may be relying on that.  In particular, if some thread
      * or exit func is still waiting to do an import, the import machinery
@@ -1015,11 +1019,8 @@ Py_FinalizeEx(void)
      * threads created thru it, so this also protects pending imports in
      * the threads created via Threading.
      */
-    call_py_exitfuncs();
 
-    /* Get current thread state and interpreter pointer */
-    tstate = PyThreadState_GET();
-    interp = tstate->interp;
+    call_py_exitfuncs(interp);
 
     /* Copy the core config, PyInterpreterState_Delete() free
        the core config memory */
@@ -1412,6 +1413,8 @@ Py_EndInterpreter(PyThreadState *tstate)
 
     wait_for_thread_shutdown();
 
+    call_py_exitfuncs(interp);
+
     if (tstate != interp->tstate_head || tstate->next != NULL)
         Py_FatalError("Py_EndInterpreter: not the last thread");
 
@@ -2023,20 +2026,28 @@ _Py_FatalInitError(_PyInitError err)
 #  include "pythread.h"
 
 /* For the atexit module. */
-void _Py_PyAtExit(void (*func)(void))
+void _Py_PyAtExit(void (*func)(PyObject *), PyObject *module)
 {
+    PyThreadState *ts;
+    PyInterpreterState *is;
+    
+    ts = PyThreadState_GET();
+    is = ts->interp;
+
     /* Guard against API misuse (see bpo-17852) */
-    assert(_PyRuntime.pyexitfunc == NULL || _PyRuntime.pyexitfunc == func);
-    _PyRuntime.pyexitfunc = func;
+    assert(is->pyexitfunc == NULL || is->pyexitfunc == func);
+
+    is->pyexitfunc = func;
+    is->pyexitmodule = module;
 }
 
 static void
-call_py_exitfuncs(void)
+call_py_exitfuncs(PyInterpreterState *istate)
 {
-    if (_PyRuntime.pyexitfunc == NULL)
+    if (istate->pyexitfunc == NULL)
         return;
 
-    (*_PyRuntime.pyexitfunc)();
+    (*istate->pyexitfunc)(istate->pyexitmodule);
     PyErr_Clear();
 }
 

Python/pystate.c

@@ -153,6 +153,8 @@ PyInterpreterState_New(void)
     interp->after_forkers_parent = NULL;
     interp->after_forkers_child = NULL;
 #endif
+    interp->pyexitfunc = NULL;
+    interp->pyexitmodule = NULL;
 
     HEAD_LOCK();
     interp->next = _PyRuntime.interpreters.head;