gh-101659: Add _Py_AtExit() (gh-103298)

The function is like Py_AtExit() but for a single interpreter.  This is a companion to the atexit module's register() function, taking a C callback instead of a Python one.

We also update the _xxinterpchannels module to use _Py_AtExit(), which is the motivating case.  (This is inspired by pain points felt while working on gh-101660.)

4 files changed, 164 insertions, 32 deletions

diff --git a/Modules/_testcapimodule.c b/Modules/_testcapimodule.c
index 3d9a2ae..557a6d4 100644
--- a/Modules/_testcapimodule.c
+++ b/Modules/_testcapimodule.c
@@ -3381,6 +3381,37 @@ test_gc_visit_objects_exit_early(PyObject *Py_UNUSED(self),
 }
 
 
+struct atexit_data {
+    int called;
+};
+
+static void
+callback(void *data)
+{
+    ((struct atexit_data *)data)->called += 1;
+}
+
+static PyObject *
+test_atexit(PyObject *self, PyObject *Py_UNUSED(args))
+{
+    PyThreadState *oldts = PyThreadState_Swap(NULL);
+    PyThreadState *tstate = Py_NewInterpreter();
+
+    struct atexit_data data = {0};
+    int res = _Py_AtExit(tstate->interp, callback, (void *)&data);
+    Py_EndInterpreter(tstate);
+    PyThreadState_Swap(oldts);
+    if (res < 0) {
+        return NULL;
+    }
+    if (data.called == 0) {
+        PyErr_SetString(PyExc_RuntimeError, "atexit callback not called");
+        return NULL;
+    }
+    Py_RETURN_NONE;
+}
+
+
 static PyObject *test_buildvalue_issue38913(PyObject *, PyObject *);
 
 static PyMethodDef TestMethods[] = {
@@ -3525,6 +3556,7 @@ static PyMethodDef TestMethods[] = {
     {"function_set_kw_defaults", function_set_kw_defaults, METH_VARARGS, NULL},
     {"test_gc_visit_objects_basic", test_gc_visit_objects_basic, METH_NOARGS, NULL},
     {"test_gc_visit_objects_exit_early", test_gc_visit_objects_exit_early, METH_NOARGS, NULL},
+    {"test_atexit", test_atexit, METH_NOARGS},
     {NULL, NULL} /* sentinel */
 };
 
diff --git a/Modules/_xxinterpchannelsmodule.c b/Modules/_xxinterpchannelsmodule.c
index ef1cdca..13b005e 100644
--- a/Modules/_xxinterpchannelsmodule.c
+++ b/Modules/_xxinterpchannelsmodule.c
@@ -174,19 +174,7 @@ _release_xid_data(_PyCrossInterpreterData *data, int ignoreexc)
     }
     int res = _PyCrossInterpreterData_Release(data);
     if (res < 0) {
-        // XXX Fix this!
-        /* The owning interpreter is already destroyed.
-         * Ideally, this shouldn't ever happen.  When an interpreter is
-         * about to be destroyed, we should clear out all of its objects
-         * from every channel associated with that interpreter.
-         * For now we hack around that to resolve refleaks, by decref'ing
-         * the released object here, even if its the wrong interpreter.
-         * The owning interpreter has already been destroyed
-         * so we should be okay, especially since the currently
-         * shareable types are all very basic, with no GC.
-         * That said, it becomes much messier once interpreters
-         * no longer share a GIL, so this needs to be fixed before then. */
-        _PyCrossInterpreterData_Clear(NULL, data);
+        /* The owning interpreter is already destroyed. */
         if (ignoreexc) {
             // XXX Emit a warning?
             PyErr_Clear();
@@ -489,6 +477,30 @@ _channelqueue_get(_channelqueue *queue)
     return _channelitem_popped(item);
 }
 
+static void
+_channelqueue_drop_interpreter(_channelqueue *queue, int64_t interp)
+{
+    _channelitem *prev = NULL;
+    _channelitem *next = queue->first;
+    while (next != NULL) {
+        _channelitem *item = next;
+        next = item->next;
+        if (item->data->interp == interp) {
+            if (prev == NULL) {
+                queue->first = item->next;
+            }
+            else {
+                prev->next = item->next;
+            }
+            _channelitem_free(item);
+            queue->count -= 1;
+        }
+        else {
+            prev = item;
+        }
+    }
+}
+
 /* channel-interpreter associations */
 
 struct _channelend;
@@ -694,6 +706,20 @@ _channelends_close_interpreter(_channelends *ends, int64_t interp, int which)
 }
 
 static void
+_channelends_drop_interpreter(_channelends *ends, int64_t interp)
+{
+    _channelend *end;
+    end = _channelend_find(ends->send, interp, NULL);
+    if (end != NULL) {
+        _channelends_close_end(ends, end, 1);
+    }
+    end = _channelend_find(ends->recv, interp, NULL);
+    if (end != NULL) {
+        _channelends_close_end(ends, end, 0);
+    }
+}
+
+static void
 _channelends_close_all(_channelends *ends, int which, int force)
 {
     // XXX Handle the ends.
@@ -841,6 +867,18 @@ done:
     return res;
 }
 
+static void
+_channel_drop_interpreter(_PyChannelState *chan, int64_t interp)
+{
+    PyThread_acquire_lock(chan->mutex, WAIT_LOCK);
+
+    _channelqueue_drop_interpreter(chan->queue, interp);
+    _channelends_drop_interpreter(chan->ends, interp);
+    chan->open = _channelends_is_open(chan->ends);
+
+    PyThread_release_lock(chan->mutex);
+}
+
 static int
 _channel_close_all(_PyChannelState *chan, int end, int force)
 {
@@ -1213,6 +1251,21 @@ done:
     return cids;
 }
 
+static void
+_channels_drop_interpreter(_channels *channels, int64_t interp)
+{
+    PyThread_acquire_lock(channels->mutex, WAIT_LOCK);
+
+    _channelref *ref = channels->head;
+    for (; ref != NULL; ref = ref->next) {
+        if (ref->chan != NULL) {
+            _channel_drop_interpreter(ref->chan, interp);
+        }
+    }
+
+    PyThread_release_lock(channels->mutex);
+}
+
 /* support for closing non-empty channels */
 
 struct _channel_closing {
@@ -1932,6 +1985,19 @@ _global_channels(void) {
 }
 
 
+static void
+clear_interpreter(void *data)
+{
+    if (_globals.module_count == 0) {
+        return;
+    }
+    PyInterpreterState *interp = (PyInterpreterState *)data;
+    assert(interp == _get_current_interp());
+    int64_t id = PyInterpreterState_GetID(interp);
+    _channels_drop_interpreter(&_globals.channels, id);
+}
+
+
 static PyObject *
 channel_create(PyObject *self, PyObject *Py_UNUSED(ignored))
 {
@@ -2339,6 +2405,10 @@ module_exec(PyObject *mod)
         goto error;
     }
 
+    // Make sure chnnels drop objects owned by this interpreter
+    PyInterpreterState *interp = _get_current_interp();
+    _Py_AtExit(interp, clear_interpreter, (void *)interp);
+
     return 0;
 
 error:
diff --git a/Modules/_xxsubinterpretersmodule.c b/Modules/_xxsubinterpretersmodule.c
index 1116467..884fb0d 100644
--- a/Modules/_xxsubinterpretersmodule.c
+++ b/Modules/_xxsubinterpretersmodule.c
@@ -67,16 +67,7 @@ _release_xid_data(_PyCrossInterpreterData *data, int ignoreexc)
     }
     int res = _PyCrossInterpreterData_Release(data);
     if (res < 0) {
-        // XXX Fix this!
-        /* The owning interpreter is already destroyed.
-         * Ideally, this shouldn't ever happen.  (It's highly unlikely.)
-         * For now we hack around that to resolve refleaks, by decref'ing
-         * the released object here, even if its the wrong interpreter.
-         * The owning interpreter has already been destroyed
-         * so we should be okay, especially since the currently
-         * shareable types are all very basic, with no GC.
-         * That said, it becomes much messier once interpreters
-         * no longer share a GIL, so this needs to be fixed before then. */
+        /* The owning interpreter is already destroyed. */
         _PyCrossInterpreterData_Clear(NULL, data);
         if (ignoreexc) {
             // XXX Emit a warning?
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;
         }