[3.13] gh-127791: Fix, document, and test PyUnstable_AtExit (GH-127793) (#127819)

* Fix merge conflicts.

* [3.13] gh-127791: Fix, document, and test `PyUnstable_AtExit` (GH-127793)
(cherry picked from commit d5d84c3)

Co-authored-by: Peter Bierma <[email protected]>

Author: ZeroIntensity

Doc/c-api/init.rst

@@ -557,6 +557,15 @@ Initializing and finalizing the interpreter
    customized Python that always runs in isolated mode using
    :c:func:`Py_RunMain`.
 
+.. c:function:: int PyUnstable_AtExit(PyInterpreterState *interp, void (*func)(void *), void *data)
+
+   Register an :mod:`atexit` callback for the target interpreter *interp*.
+   This is similar to :c:func:`Py_AtExit`, but takes an explicit interpreter and
+   data pointer for the callback.
+
+   The :term:`GIL` must be held for *interp*.
+
+   .. versionadded:: 3.13
 
 Process-wide parameters
 =======================

Doc/c-api/sys.rst

@@ -426,3 +426,7 @@ Process Control
    function registered last is called first. Each cleanup function will be called
    at most once.  Since Python's internal finalization will have completed before
    the cleanup function, no Python APIs should be called by *func*.
+
+   .. seealso::
+
+      :c:func:`PyUnstable_AtExit` for passing a ``void *data`` argument.

Include/internal/pycore_atexit.h

@@ -44,6 +44,7 @@ typedef struct {
 
 struct atexit_state {
     atexit_callback *ll_callbacks;
+    // Kept for ABI compatibility--do not use! (See GH-127791.)
     atexit_callback *last_ll_callback;
 
     // XXX The rest of the state could be moved to the atexit module state

Misc/NEWS.d/next/C API/2024-12-10-14-25-22.gh-issue-127791.YRw4GU.rst

@@ -0,0 +1,2 @@
+Fix loss of callbacks after more than one call to
+:c:func:`PyUnstable_AtExit`.

Modules/_testcapimodule.c

@@ -3338,6 +3338,54 @@ pyeval_getlocals(PyObject *module, PyObject *Py_UNUSED(args))
     return Py_XNewRef(PyEval_GetLocals());
 }
 
+struct atexit_data {
+    int called;
+    PyThreadState *tstate;
+    PyInterpreterState *interp;
+};
+
+static void
+atexit_callback(void *data)
+{
+    struct atexit_data *at_data = (struct atexit_data *)data;
+    // Ensure that the callback is from the same interpreter
+    assert(PyThreadState_Get() == at_data->tstate);
+    assert(PyInterpreterState_Get() == at_data->interp);
+    ++at_data->called;
+}
+
+static PyObject *
+test_atexit(PyObject *self, PyObject *Py_UNUSED(args))
+{
+    PyThreadState *oldts = PyThreadState_Swap(NULL);
+    PyThreadState *tstate = Py_NewInterpreter();
+
+    struct atexit_data data = {0};
+    data.tstate = PyThreadState_Get();
+    data.interp = PyInterpreterState_Get();
+
+    int amount = 10;
+    for (int i = 0; i < amount; ++i)
+    {
+        int res = PyUnstable_AtExit(tstate->interp, atexit_callback, (void *)&data);
+        if (res < 0) {
+            Py_EndInterpreter(tstate);
+            PyThreadState_Swap(oldts);
+            PyErr_SetString(PyExc_RuntimeError, "atexit callback failed");
+            return NULL;
+        }
+    }
+
+    Py_EndInterpreter(tstate);
+    PyThreadState_Swap(oldts);
+
+    if (data.called != amount) {
+        PyErr_SetString(PyExc_RuntimeError, "atexit callback not called");
+        return NULL;
+    }
+    Py_RETURN_NONE;
+}
+
 static PyMethodDef TestMethods[] = {
     {"set_errno",               set_errno,                       METH_VARARGS},
     {"test_config",             test_config,                     METH_NOARGS},
@@ -3483,6 +3531,7 @@ static PyMethodDef TestMethods[] = {
     {"function_set_warning", function_set_warning, METH_NOARGS},
     {"test_critical_sections", test_critical_sections, METH_NOARGS},
     {"pyeval_getlocals", pyeval_getlocals, METH_NOARGS},
+    {"test_atexit", test_atexit, METH_NOARGS},
     {NULL, NULL} /* sentinel */
 };
 

Modules/_testinternalcapi.c

@@ -1234,39 +1234,6 @@ unicode_transformdecimalandspacetoascii(PyObject *self, PyObject *arg)
     return _PyUnicode_TransformDecimalAndSpaceToASCII(arg);
 }
 
-
-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 = PyUnstable_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_pyobject_is_freed(const char *test_name, PyObject *op)
 {
@@ -2065,7 +2032,6 @@ static PyMethodDef module_functions[] = {
     {"_PyTraceMalloc_GetTraceback", tracemalloc_get_traceback, METH_VARARGS},
     {"test_tstate_capi", test_tstate_capi, METH_NOARGS, NULL},
     {"_PyUnicode_TransformDecimalAndSpaceToASCII", unicode_transformdecimalandspacetoascii, METH_O},
-    {"test_atexit", test_atexit, METH_NOARGS},
     {"check_pyobject_forbidden_bytes_is_freed",
                             check_pyobject_forbidden_bytes_is_freed, METH_NOARGS},
     {"check_pyobject_freed_is_freed", check_pyobject_freed_is_freed, METH_NOARGS},

Modules/atexitmodule.c

@@ -27,7 +27,10 @@ int
 PyUnstable_AtExit(PyInterpreterState *interp,
                   atexit_datacallbackfunc func, void *data)
 {
-    assert(interp == _PyInterpreterState_GET());
+    PyThreadState *tstate = _PyThreadState_GET();
+    _Py_EnsureTstateNotNULL(tstate);
+    assert(tstate->interp == interp);
+
     atexit_callback *callback = PyMem_Malloc(sizeof(atexit_callback));
     if (callback == NULL) {
         PyErr_NoMemory();
@@ -38,12 +41,13 @@ PyUnstable_AtExit(PyInterpreterState *interp,
     callback->next = NULL;
 
     struct atexit_state *state = &interp->atexit;
-    if (state->ll_callbacks == NULL) {
+    atexit_callback *top = state->ll_callbacks;
+    if (top == NULL) {
         state->ll_callbacks = callback;
-        state->last_ll_callback = callback;
     }
     else {
-        state->last_ll_callback->next = callback;
+        callback->next = top;
+        state->ll_callbacks = callback;
     }
     return 0;
 }