1 files changed, 48 insertions, 9 deletions

diff --git a/Modules/atexitmodule.c b/Modules/atexitmodule.c
index a1c511e..47afd7f 100644
--- a/Modules/atexitmodule.c
+++ b/Modules/atexitmodule.c
@@ -7,6 +7,7 @@
  */
 
 #include "Python.h"
+#include "pycore_atexit.h"
 #include "pycore_initconfig.h"    // _PyStatus_NO_MEMORY
 #include "pycore_interp.h"        // PyInterpreterState.atexit
 #include "pycore_pystate.h"       // _PyInterpreterState_GET
@@ -22,10 +23,36 @@ get_atexit_state(void)
 }
 
 
+int
+_Py_AtExit(PyInterpreterState *interp,
+           atexit_datacallbackfunc func, void *data)
+{
+    assert(interp == _PyInterpreterState_GET());
+    atexit_callback *callback = PyMem_Malloc(sizeof(atexit_callback));
+    if (callback == NULL) {
+        PyErr_NoMemory();
+        return -1;
+    }
+    callback->func = func;
+    callback->data = data;
+    callback->next = NULL;
+
+    struct atexit_state *state = &interp->atexit;
+    if (state->ll_callbacks == NULL) {
+        state->ll_callbacks = callback;
+        state->last_ll_callback = callback;
+    }
+    else {
+        state->last_ll_callback->next = callback;
+    }
+    return 0;
+}
+
+
 static void
 atexit_delete_cb(struct atexit_state *state, int i)
 {
-    atexit_callback *cb = state->callbacks[i];
+    atexit_py_callback *cb = state->callbacks[i];
     state->callbacks[i] = NULL;
 
     Py_DECREF(cb->func);
@@ -39,7 +66,7 @@ atexit_delete_cb(struct atexit_state *state, int i)
 static void
 atexit_cleanup(struct atexit_state *state)
 {
-    atexit_callback *cb;
+    atexit_py_callback *cb;
     for (int i = 0; i < state->ncallbacks; i++) {
         cb = state->callbacks[i];
         if (cb == NULL)
@@ -60,7 +87,7 @@ _PyAtExit_Init(PyInterpreterState *interp)
 
     state->callback_len = 32;
     state->ncallbacks = 0;
-    state->callbacks = PyMem_New(atexit_callback*, state->callback_len);
+    state->callbacks = PyMem_New(atexit_py_callback*, state->callback_len);
     if (state->callbacks == NULL) {
         return _PyStatus_NO_MEMORY();
     }
@@ -75,6 +102,18 @@ _PyAtExit_Fini(PyInterpreterState *interp)
     atexit_cleanup(state);
     PyMem_Free(state->callbacks);
     state->callbacks = NULL;
+
+    atexit_callback *next = state->ll_callbacks;
+    state->ll_callbacks = NULL;
+    while (next != NULL) {
+        atexit_callback *callback = next;
+        next = callback->next;
+        atexit_datacallbackfunc exitfunc = callback->func;
+        void *data = callback->data;
+        // It was allocated in _PyAtExit_AddCallback().
+        PyMem_Free(callback);
+        exitfunc(data);
+    }
 }
 
 
@@ -88,7 +127,7 @@ atexit_callfuncs(struct atexit_state *state)
     }
 
     for (int i = state->ncallbacks - 1; i >= 0; i--) {
-        atexit_callback *cb = state->callbacks[i];
+        atexit_py_callback *cb = state->callbacks[i];
         if (cb == NULL) {
             continue;
         }
@@ -152,17 +191,17 @@ atexit_register(PyObject *module, PyObject *args, PyObject *kwargs)
 
     struct atexit_state *state = get_atexit_state();
     if (state->ncallbacks >= state->callback_len) {
-        atexit_callback **r;
+        atexit_py_callback **r;
         state->callback_len += 16;
-        size_t size = sizeof(atexit_callback*) * (size_t)state->callback_len;
-        r = (atexit_callback**)PyMem_Realloc(state->callbacks, size);
+        size_t size = sizeof(atexit_py_callback*) * (size_t)state->callback_len;
+        r = (atexit_py_callback**)PyMem_Realloc(state->callbacks, size);
         if (r == NULL) {
             return PyErr_NoMemory();
         }
         state->callbacks = r;
     }
 
-    atexit_callback *callback = PyMem_Malloc(sizeof(atexit_callback));
+    atexit_py_callback *callback = PyMem_Malloc(sizeof(atexit_py_callback));
     if (callback == NULL) {
         return PyErr_NoMemory();
     }
@@ -233,7 +272,7 @@ atexit_unregister(PyObject *module, PyObject *func)
     struct atexit_state *state = get_atexit_state();
     for (int i = 0; i < state->ncallbacks; i++)
     {
-        atexit_callback *cb = state->callbacks[i];
+        atexit_py_callback *cb = state->callbacks[i];
         if (cb == NULL) {
             continue;
         }