python · corona10 · Apr 17, 2020 · Apr 17, 2020 · Apr 17, 2020
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -267,6 +267,7 @@ PyInterpreterState_New(void)
         return NULL;
     }
 
+    interp->nexitmodule = 0;
     interp->tstate_next_unique_id = 0;
 
     interp->audit_hooks = NULL;
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,2 @@
		:mod:`atexit` now supports more than once loading per interpreter. Patch by
		Dong-hee Na.