7 files changed, 101 insertions, 113 deletions

diff --git a/Include/internal/pycore_interp.h b/Include/internal/pycore_interp.h
index 12416c2..9ec5358 100644
--- a/Include/internal/pycore_interp.h
+++ b/Include/internal/pycore_interp.h
@@ -159,6 +159,20 @@ struct _Py_exc_state {
 };
 
 
+// atexit state
+typedef struct {
+    PyObject *func;
+    PyObject *args;
+    PyObject *kwargs;
+} atexit_callback;
+
+struct atexit_state {
+    atexit_callback **callbacks;
+    int ncallbacks;
+    int callback_len;
+};
+
+
 /* interpreter state */
 
 #define _PY_NSMALLPOSINTS           257
@@ -234,12 +248,10 @@ struct _is {
     PyObject *after_forkers_child;
 #endif
 
-    /* AtExit module */
-    PyObject *atexit_module;
-
     uint64_t tstate_next_unique_id;
 
     struct _warnings_runtime_state warnings;
+    struct atexit_state atexit;
 
     PyObject *audit_hooks;
 
diff --git a/Include/internal/pycore_pylifecycle.h b/Include/internal/pycore_pylifecycle.h
index e6e4187..d1c23c8 100644
--- a/Include/internal/pycore_pylifecycle.h
+++ b/Include/internal/pycore_pylifecycle.h
@@ -55,6 +55,7 @@ extern PyStatus _PyTypes_Init(void);
 extern PyStatus _PyTypes_InitSlotDefs(void);
 extern PyStatus _PyImportZip_Init(PyThreadState *tstate);
 extern PyStatus _PyGC_Init(PyThreadState *tstate);
+extern PyStatus _PyAtExit_Init(PyThreadState *tstate);
 
 
 /* Various internal finalizers */
@@ -85,6 +86,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);
@@ -109,7 +111,7 @@ PyAPI_FUNC(void) _PyErr_Display(PyObject *file, PyObject *exception,
 
 PyAPI_FUNC(void) _PyThreadState_DeleteCurrent(PyThreadState *tstate);
 
-extern void _PyAtExit_Call(PyObject *module);
+extern void _PyAtExit_Call(PyThreadState *tstate);
 
 #ifdef __cplusplus
 }
diff --git a/Lib/test/test_atexit.py b/Lib/test/test_atexit.py
index 906f96d..29faaaf 100644
--- a/Lib/test/test_atexit.py
+++ b/Lib/test/test_atexit.py
@@ -170,6 +170,24 @@ class GeneralTest(unittest.TestCase):
         self.assertEqual(res.out.decode().splitlines(), ["two", "one"])
         self.assertFalse(res.err)
 
+    def test_atexit_instances(self):
+        # bpo-42639: It is safe to have more than one atexit instance.
+        code = textwrap.dedent("""
+            import sys
+            import atexit as atexit1
+            del sys.modules['atexit']
+            import atexit as atexit2
+            del sys.modules['atexit']
+
+            assert atexit2 is not atexit1
+
+            atexit1.register(print, "atexit1")
+            atexit2.register(print, "atexit2")
+        """)
+        res = script_helper.assert_python_ok("-c", code)
+        self.assertEqual(res.out.decode().splitlines(), ["atexit2", "atexit1"])
+        self.assertFalse(res.err)
+
 
 @support.cpython_only
 class SubinterpreterTest(unittest.TestCase):
diff --git a/Misc/NEWS.d/next/Core and Builtins/2020-12-09-01-55-10.bpo-42639.5pI5HG.rst b/Misc/NEWS.d/next/Core and Builtins/2020-12-09-01-55-10.bpo-42639.5pI5HG.rst
new file mode 100644
index 0000000..7999ee9
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2020-12-09-01-55-10.bpo-42639.5pI5HG.rst
@@ -0,0 +1,3 @@
+Make the :mod:`atexit` module state per-interpreter. It is now safe have more
+than one :mod:`atexit` module instance.
+Patch by Dong-hee Na and Victor Stinner.
diff --git a/Modules/atexitmodule.c b/Modules/atexitmodule.c
index 90ed7d5..ae13eae 100644
--- a/Modules/atexitmodule.c
+++ b/Modules/atexitmodule.c
@@ -7,38 +7,26 @@
  */
 
 #include "Python.h"
-#include "pycore_interp.h"        // PyInterpreterState.atexit_func
+#include "pycore_initconfig.h"    // _PyStatus_NO_MEMORY
+#include "pycore_interp.h"        // PyInterpreterState.atexit
 #include "pycore_pystate.h"       // _PyInterpreterState_GET
 
 /* ===================================================================== */
 /* Callback machinery. */
 
-typedef struct {
-    PyObject *func;
-    PyObject *args;
-    PyObject *kwargs;
-} atexit_callback;
-
-struct atexit_state {
-    atexit_callback **atexit_callbacks;
-    int ncallbacks;
-    int callback_len;
-};
-
 static inline struct atexit_state*
-get_atexit_state(PyObject *module)
+get_atexit_state(void)
 {
-    void *state = PyModule_GetState(module);
-    assert(state != NULL);
-    return (struct atexit_state *)state;
+    PyInterpreterState *interp = _PyInterpreterState_GET();
+    return &interp->atexit;
 }
 
 
 static void
 atexit_delete_cb(struct atexit_state *state, int i)
 {
-    atexit_callback *cb = state->atexit_callbacks[i];
-    state->atexit_callbacks[i] = NULL;
+    atexit_callback *cb = state->callbacks[i];
+    state->callbacks[i] = NULL;
 
     Py_DECREF(cb->func);
     Py_DECREF(cb->args);
@@ -53,7 +41,7 @@ atexit_cleanup(struct atexit_state *state)
 {
     atexit_callback *cb;
     for (int i = 0; i < state->ncallbacks; i++) {
-        cb = state->atexit_callbacks[i];
+        cb = state->callbacks[i];
         if (cb == NULL)
             continue;
 
@@ -62,25 +50,45 @@ atexit_cleanup(struct atexit_state *state)
     state->ncallbacks = 0;
 }
 
-/* Installed into pylifecycle.c's atexit mechanism */
 
-static void
-atexit_callfuncs(PyObject *module, int ignore_exc)
+PyStatus
+_PyAtExit_Init(PyThreadState *tstate)
 {
-    assert(!PyErr_Occurred());
+    struct atexit_state *state = &tstate->interp->atexit;
+    // _PyAtExit_Init() must only be called once
+    assert(state->callbacks == NULL);
 
-    if (module == NULL) {
-        return;
+    state->callback_len = 32;
+    state->ncallbacks = 0;
+    state->callbacks = PyMem_New(atexit_callback*, state->callback_len);
+    if (state->callbacks == NULL) {
+        return _PyStatus_NO_MEMORY();
     }
+    return _PyStatus_OK();
+}
+
+
+void
+_PyAtExit_Fini(PyInterpreterState *interp)
+{
+    struct atexit_state *state = &interp->atexit;
+    atexit_cleanup(state);
+    PyMem_Free(state->callbacks);
+}
+
+
+static void
+atexit_callfuncs(struct atexit_state *state, int ignore_exc)
+{
+    assert(!PyErr_Occurred());
 
-    struct atexit_state *state = get_atexit_state(module);
     if (state->ncallbacks == 0) {
         return;
     }
 
     PyObject *exc_type = NULL, *exc_value, *exc_tb;
     for (int i = state->ncallbacks - 1; i >= 0; i--) {
-        atexit_callback *cb = state->atexit_callbacks[i];
+        atexit_callback *cb = state->callbacks[i];
         if (cb == NULL) {
             continue;
         }
@@ -125,15 +133,17 @@ atexit_callfuncs(PyObject *module, int ignore_exc)
 
 
 void
-_PyAtExit_Call(PyObject *module)
+_PyAtExit_Call(PyThreadState *tstate)
 {
-    atexit_callfuncs(module, 1);
+    struct atexit_state *state = &tstate->interp->atexit;
+    atexit_callfuncs(state, 1);
 }
 
 
 /* ===================================================================== */
 /* Module methods. */
 
+
 PyDoc_STRVAR(atexit_register__doc__,
 "register(func, *args, **kwargs) -> func\n\
 \n\
@@ -161,15 +171,15 @@ atexit_register(PyObject *module, PyObject *args, PyObject *kwargs)
         return NULL;
     }
 
-    struct atexit_state *state = get_atexit_state(module);
+    struct atexit_state *state = get_atexit_state();
     if (state->ncallbacks >= state->callback_len) {
         atexit_callback **r;
         state->callback_len += 16;
-        r = (atexit_callback**)PyMem_Realloc(state->atexit_callbacks,
-                                      sizeof(atexit_callback*) * state->callback_len);
+        size_t size = sizeof(atexit_callback*) * (size_t)state->callback_len;
+        r = (atexit_callback**)PyMem_Realloc(state->callbacks, size);
         if (r == NULL)
             return PyErr_NoMemory();
-        state->atexit_callbacks = r;
+        state->callbacks = r;
     }
 
     atexit_callback *callback = PyMem_Malloc(sizeof(atexit_callback));
@@ -185,7 +195,7 @@ atexit_register(PyObject *module, PyObject *args, PyObject *kwargs)
     callback->func = Py_NewRef(func);
     callback->kwargs = Py_XNewRef(kwargs);
 
-    state->atexit_callbacks[state->ncallbacks++] = callback;
+    state->callbacks[state->ncallbacks++] = callback;
 
     return Py_NewRef(func);
 }
@@ -198,7 +208,8 @@ Run all registered exit functions.");
 static PyObject *
 atexit_run_exitfuncs(PyObject *module, PyObject *unused)
 {
-    atexit_callfuncs(module, 0);
+    struct atexit_state *state = get_atexit_state();
+    atexit_callfuncs(state, 0);
     if (PyErr_Occurred()) {
         return NULL;
     }
@@ -213,7 +224,7 @@ Clear the list of previously registered exit functions.");
 static PyObject *
 atexit_clear(PyObject *module, PyObject *unused)
 {
-    atexit_cleanup(get_atexit_state(module));
+    atexit_cleanup(get_atexit_state());
     Py_RETURN_NONE;
 }
 
@@ -225,41 +236,10 @@ Return the number of registered exit functions.");
 static PyObject *
 atexit_ncallbacks(PyObject *module, PyObject *unused)
 {
-    struct atexit_state *state = get_atexit_state(module);
+    struct atexit_state *state = get_atexit_state();
     return PyLong_FromSsize_t(state->ncallbacks);
 }
 
-static int
-atexit_m_traverse(PyObject *module, visitproc visit, void *arg)
-{
-    struct atexit_state *state = (struct atexit_state *)PyModule_GetState(module);
-    for (int i = 0; i < state->ncallbacks; i++) {
-        atexit_callback *cb = state->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 *module)
-{
-    struct atexit_state *state = (struct atexit_state *)PyModule_GetState(module);
-    atexit_cleanup(state);
-    return 0;
-}
-
-static void
-atexit_free(PyObject *module)
-{
-    struct atexit_state *state = (struct atexit_state *)PyModule_GetState(module);
-    atexit_cleanup(state);
-    PyMem_Free(state->atexit_callbacks);
-}
-
 PyDoc_STRVAR(atexit_unregister__doc__,
 "unregister(func) -> None\n\
 \n\
@@ -271,10 +251,10 @@ atexit.register\n\
 static PyObject *
 atexit_unregister(PyObject *module, PyObject *func)
 {
-    struct atexit_state *state = get_atexit_state(module);
+    struct atexit_state *state = get_atexit_state();
     for (int i = 0; i < state->ncallbacks; i++)
     {
-        atexit_callback *cb = state->atexit_callbacks[i];
+        atexit_callback *cb = state->callbacks[i];
         if (cb == NULL) {
             continue;
         }
@@ -305,6 +285,7 @@ static PyMethodDef atexit_methods[] = {
     {NULL, NULL}        /* sentinel */
 };
 
+
 /* ===================================================================== */
 /* Initialization function. */
 
@@ -315,37 +296,12 @@ upon normal program termination.\n\
 Two public functions, register and unregister, are defined.\n\
 ");
 
-static int
-atexit_exec(PyObject *module)
-{
-    struct atexit_state *state = get_atexit_state(module);
-    state->callback_len = 32;
-    state->ncallbacks = 0;
-    state->atexit_callbacks = PyMem_New(atexit_callback*, state->callback_len);
-    if (state->atexit_callbacks == NULL) {
-        return -1;
-    }
-
-    PyInterpreterState *interp = _PyInterpreterState_GET();
-    interp->atexit_module = module;
-    return 0;
-}
-
-static PyModuleDef_Slot atexit_slots[] = {
-    {Py_mod_exec, atexit_exec},
-    {0, NULL}
-};
-
 static struct PyModuleDef atexitmodule = {
     PyModuleDef_HEAD_INIT,
     .m_name = "atexit",
     .m_doc = atexit__doc__,
-    .m_size = sizeof(struct atexit_state),
+    .m_size = 0,
     .m_methods = atexit_methods,
-    .m_slots = atexit_slots,
-    .m_traverse = atexit_m_traverse,
-    .m_clear = atexit_m_clear,
-    .m_free = (freefunc)atexit_free
 };
 
 PyMODINIT_FUNC
diff --git a/Python/pylifecycle.c b/Python/pylifecycle.c
index 7b80d01..8d744c7 100644
--- a/Python/pylifecycle.c
+++ b/Python/pylifecycle.c
@@ -55,7 +55,6 @@ static PyStatus add_main_module(PyInterpreterState *interp);
 static PyStatus init_import_site(void);
 static PyStatus init_set_builtins_open(void);
 static PyStatus init_sys_streams(PyThreadState *tstate);
-static void call_py_exitfuncs(PyThreadState *tstate);
 static void wait_for_thread_shutdown(PyThreadState *tstate);
 static void call_ll_exitfuncs(_PyRuntimeState *runtime);
 
@@ -690,6 +689,12 @@ pycore_init_types(PyThreadState *tstate)
     if (_PyWarnings_InitState(tstate) < 0) {
         return _PyStatus_ERR("can't initialize warnings");
     }
+
+    status = _PyAtExit_Init(tstate);
+    if (_PyStatus_EXCEPTION(status)) {
+        return status;
+    }
+
     return _PyStatus_OK();
 }
 
@@ -1655,7 +1660,7 @@ Py_FinalizeEx(void)
      * the threads created via Threading.
      */
 
-    call_py_exitfuncs(tstate);
+    _PyAtExit_Call(tstate);
 
     /* Copy the core config, PyInterpreterState_Delete() free
        the core config memory */
@@ -1946,7 +1951,7 @@ Py_EndInterpreter(PyThreadState *tstate)
     // Wrap up existing "threading"-module-created, non-daemon threads.
     wait_for_thread_shutdown(tstate);
 
-    call_py_exitfuncs(tstate);
+    _PyAtExit_Call(tstate);
 
     if (tstate != interp->tstate_head || tstate->next != NULL) {
         Py_FatalError("not the last thread");
@@ -2633,15 +2638,6 @@ Py_ExitStatusException(PyStatus status)
 }
 
 
-/* Clean up and exit */
-
-static void
-call_py_exitfuncs(PyThreadState *tstate)
-{
-    _PyAtExit_Call(tstate->interp->atexit_module);
-}
-
-
 /* Wait until threading._shutdown completes, provided
    the threading module was imported in the first place.
    The shutdown routine will wait until all non-daemon
diff --git a/Python/pystate.c b/Python/pystate.c
index 8da583f..ead25b0 100644
--- a/Python/pystate.c
+++ b/Python/pystate.c
@@ -303,6 +303,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