bpo-40600: atexit only loads single time per interpreter

Include/internal/pycore_atexit.h

@@ -0,0 +1,26 @@
+#ifndef Py_INTERNAL_ATEXIT_H
+#define Py_INTERNAL_ATEXIT_H
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#ifndef Py_BUILD_CORE
+#  error "this header requires Py_BUILD_CORE define"
+#endif
+
+typedef struct {
+    PyObject *func;
+    PyObject *args;
+    PyObject *kwargs;
+} _Py_atexit_callback;
+
+struct _atexit_runtime_state {
+    _Py_atexit_callback **atexit_callbacks;
+    int ncallbacks;
+    int callback_len;
+};
+
+#ifdef __cplusplus
+}
+#endif
+#endif /* !Py_INTERNAL_ATEXIT_H */

Include/internal/pycore_interp.h

@@ -13,6 +13,7 @@ extern "C" {
 #include "pycore_gil.h"           // struct _gil_runtime_state
 #include "pycore_gc.h"            // struct _gc_runtime_state
 #include "pycore_warnings.h"      // struct _warnings_runtime_state
+#include "pycore_atexit.h"        // struct _atexit_runtime_state
 
 struct _pending_calls {
     PyThread_type_lock lock;
@@ -240,6 +241,7 @@ struct _is {
     uint64_t tstate_next_unique_id;
 
     struct _warnings_runtime_state warnings;
+    struct _atexit_runtime_state atexits;
 
     PyObject *audit_hooks;
 

Include/internal/pycore_pylifecycle.h

@@ -85,6 +85,7 @@ extern void _PyHash_Fini(void);
 extern void _PyTraceMalloc_Fini(void);
 extern void _PyWarnings_Fini(PyInterpreterState *interp);
 extern void _PyAST_Fini(PyInterpreterState *interp);
+extern void _PyAtExit_Fini(PyInterpreterState *interp);
 
 extern PyStatus _PyGILState_Init(PyThreadState *tstate);
 extern void _PyGILState_Fini(PyThreadState *tstate);

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"atexit1 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']
+            try:
+                import atexit as atexit2
+            except ImportError:
+                pass
+            atexit1.register(callback, b"atexit1 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()

Misc/NEWS.d/next/Core and Builtins/2020-12-09-01-55-10.bpo-40600.5pI5HG.rst

@@ -0,0 +1,2 @@
+:mod:`atexit` supports only once loading per interpreter. Patch by Dong-hee
+Na.

Modules/atexitmodule.c

@@ -7,6 +7,11 @@
  */
 
 #include "Python.h"
+#include "pycore_interp.h"
+#include "pycore_pystate.h"
+
+typedef struct _atexit_runtime_state atexitmodule_state;
+typedef _Py_atexit_callback atexit_callback;
 
 /* Forward declaration (for atexit_cleanup) */
 static PyObject *atexit_clear(PyObject*, PyObject*);
@@ -16,24 +21,13 @@ static struct PyModuleDef atexitmodule;
 /* ===================================================================== */
 /* Callback machinery. */
 
-typedef struct {
-    PyObject *func;
-    PyObject *args;
-    PyObject *kwargs;
-} atexit_callback;
-
-typedef struct {
-    atexit_callback **atexit_callbacks;
-    int ncallbacks;
-    int callback_len;
-} atexitmodule_state;
 
 static inline atexitmodule_state*
-get_atexit_state(PyObject *module)
+get_atexit_state()
 {
-    void *state = PyModule_GetState(module);
-    assert(state != NULL);
-    return (atexitmodule_state *)state;
+    PyInterpreterState *interp = _PyInterpreterState_GET();
+    atexitmodule_state *modstate = &interp->atexits;
+    return modstate;
 }
 
 
@@ -55,8 +49,7 @@ static void
 atexit_cleanup(atexitmodule_state *modstate)
 {
     atexit_callback *cb;
-    int i;
-    for (i = 0; i < modstate->ncallbacks; i++) {
+    for (int i = 0; i < modstate->ncallbacks; i++) {
         cb = modstate->atexit_callbacks[i];
         if (cb == NULL)
             continue;
@@ -78,7 +71,7 @@ atexit_callfuncs(PyObject *module)
 
     if (module == NULL)
         return;
-    modstate = get_atexit_state(module);
+    modstate = get_atexit_state();
 
     if (modstate->ncallbacks == 0)
         return;
@@ -136,7 +129,7 @@ atexit_register(PyObject *self, PyObject *args, PyObject *kwargs)
     atexit_callback *new_callback;
     PyObject *func = NULL;
 
-    modstate = get_atexit_state(self);
+    modstate = get_atexit_state();
 
     if (modstate->ncallbacks >= modstate->callback_len) {
         atexit_callback **r;
@@ -203,7 +196,7 @@ Clear the list of previously registered exit functions.");
 static PyObject *
 atexit_clear(PyObject *self, PyObject *unused)
 {
-    atexit_cleanup(get_atexit_state(self));
+    atexit_cleanup(get_atexit_state());
     Py_RETURN_NONE;
 }
 
@@ -217,48 +210,11 @@ atexit_ncallbacks(PyObject *self, PyObject *unused)
 {
     atexitmodule_state *modstate;
 
-    modstate = get_atexit_state(self);
+    modstate = get_atexit_state();
 
     return PyLong_FromSsize_t(modstate->ncallbacks);
 }
 
-static int
-atexit_m_traverse(PyObject *self, visitproc visit, void *arg)
-{
-    int i;
-    atexitmodule_state *modstate;
-
-    modstate = (atexitmodule_state *)PyModule_GetState(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);
-    }
-    return 0;
-}
-
-static int
-atexit_m_clear(PyObject *self)
-{
-    atexitmodule_state *modstate;
-    modstate = (atexitmodule_state *)PyModule_GetState(self);
-    atexit_cleanup(modstate);
-    return 0;
-}
-
-static void
-atexit_free(PyObject *m)
-{
-    atexitmodule_state *modstate;
-    modstate = (atexitmodule_state *)PyModule_GetState(m);
-    atexit_cleanup(modstate);
-    PyMem_Free(modstate->atexit_callbacks);
-}
-
 PyDoc_STRVAR(atexit_unregister__doc__,
 "unregister(func) -> None\n\
 \n\
@@ -274,7 +230,7 @@ atexit_unregister(PyObject *self, PyObject *func)
     atexit_callback *cb;
     int i, eq;
 
-    modstate = get_atexit_state(self);
+    modstate = get_atexit_state();
 
     for (i = 0; i < modstate->ncallbacks; i++)
     {
@@ -316,16 +272,21 @@ Two public functions, register and unregister, are defined.\n\
 ");
 
 static int
-atexit_exec(PyObject *m) {
-    atexitmodule_state *modstate;
-
-    modstate = get_atexit_state(m);
+atexit_exec(PyObject *m)
+{
+    atexitmodule_state *modstate = get_atexit_state();
+    if (modstate->atexit_callbacks != NULL) {
+        PyErr_SetString(PyExc_ImportError,
+                        "cannot load module more than once per interpreter");
+        return -1;
+    }
     modstate->callback_len = 32;
     modstate->ncallbacks = 0;
     modstate->atexit_callbacks = PyMem_New(atexit_callback*,
                                            modstate->callback_len);
-    if (modstate->atexit_callbacks == NULL)
+    if (modstate->atexit_callbacks == NULL) {
         return -1;
+    }
 
     _Py_PyAtExit(atexit_callfuncs, m);
     return 0;
@@ -338,16 +299,20 @@ static PyModuleDef_Slot atexit_slots[] = {
 
 static struct PyModuleDef atexitmodule = {
     PyModuleDef_HEAD_INIT,
-    "atexit",
-    atexit__doc__,
-    sizeof(atexitmodule_state),
-    atexit_methods,
-    atexit_slots,
-    atexit_m_traverse,
-    atexit_m_clear,
-    (freefunc)atexit_free
+    .m_name = "atexit",
+    .m_doc = atexit__doc__,
+    .m_methods = atexit_methods,
+    .m_slots = atexit_slots,
 };
 
+void
+_PyAtExit_Fini(PyInterpreterState *interp)
+{
+    atexitmodule_state *modstate = &interp->atexits;
+    atexit_cleanup(modstate);
+    PyMem_Free(modstate->atexit_callbacks);
+}
+
 PyMODINIT_FUNC
 PyInit_atexit(void)
 {

Python/pystate.c

@@ -9,6 +9,7 @@
 #include "pycore_pymem.h"         // _PyMem_SetDefaultAllocator()
 #include "pycore_pystate.h"       // _PyThreadState_GET()
 #include "pycore_sysmodule.h"
+#include "pycore_atexit.h"
 
 /* --------------------------------------------------------------------------
 CAUTION
@@ -303,6 +304,7 @@ interpreter_clear(PyInterpreterState *interp, PyThreadState *tstate)
 
     _PyAST_Fini(interp);
     _PyWarnings_Fini(interp);
+    _PyAtExit_Fini(interp);
 
     // All Python types must be destroyed before the last GC collection. Python
     // types create a reference cycle to themselves in their in their