#include "Python.h"

#include "pycore_atexit.h" // export _Py_AtExit()

#include "pycore_initconfig.h" // _PyStatus_NO_MEMORY

#include "pycore_interp.h" // PyInterpreterState.atexit

#include "pycore_pystate.h" // _PyInterpreterState_GET

static inline struct atexit_state*

get_atexit_state(void)

{

PyInterpreterState *interp = _PyInterpreterState_GET();

return &interp->atexit;

}

PyUnstable_AtExit(PyInterpreterState *interp,

atexit_datacallbackfunc func, void *data)

{

PyThreadState *tstate = _PyThreadState_GET();

_Py_EnsureTstateNotNULL(tstate);

assert(tstate->interp == interp);

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;

_PyAtExit_LockCallbacks(state);

atexit_callback *top = state->ll_callbacks;

if (top == NULL) {

state->ll_callbacks = callback;

}

else {

callback->next = top;

state->ll_callbacks = callback;

}

_PyAtExit_UnlockCallbacks(state);

return 0;

}

static void

atexit_cleanup(struct atexit_state *state)

{

PyList_Clear(state->callbacks);

}

_PyAtExit_Init(PyInterpreterState *interp)

{

struct atexit_state *state = &interp->atexit;

// _PyAtExit_Init() must only be called once

assert(state->callbacks == NULL);

state->callbacks = PyList_New(0);

if (state->callbacks == NULL) {

return _PyStatus_NO_MEMORY();

}

return _PyStatus_OK();

}

_PyAtExit_Fini(PyInterpreterState *interp)

{

// In theory, there shouldn't be any threads left by now, so we

// won't lock this.

struct atexit_state *state = &interp->atexit;

atexit_cleanup(state);

Py_CLEAR(state->callbacks);

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);

}

}

static void

atexit_callfuncs(struct atexit_state *state)

{

assert(!PyErr_Occurred());

assert(state->callbacks != NULL);

assert(PyList_CheckExact(state->callbacks));

// Create a copy of the list for thread safety

PyObject *copy = PyList_GetSlice(state->callbacks, 0, PyList_GET_SIZE(state->callbacks));

if (copy == NULL)

{

PyErr_FormatUnraisable("Exception ignored while "

"copying atexit callbacks");

atexit_cleanup(state);

return;

}

for (Py_ssize_t i = 0; i < PyList_GET_SIZE(copy); ++i) {

// We don't have to worry about evil borrowed references, because

// no other threads can access this list.

PyObject *tuple = PyList_GET_ITEM(copy, i);

assert(PyTuple_CheckExact(tuple));

PyObject *func = PyTuple_GET_ITEM(tuple, 0);

PyObject *args = PyTuple_GET_ITEM(tuple, 1);

PyObject *kwargs = PyTuple_GET_ITEM(tuple, 2);

PyObject *res = PyObject_Call(func,

args,

kwargs == Py_None ? NULL : kwargs);

if (res == NULL) {

PyErr_FormatUnraisable(

"Exception ignored in atexit callback %R", func);

}

else {

Py_DECREF(res);

}

}

Py_DECREF(copy);

atexit_cleanup(state);

assert(!PyErr_Occurred());

}

_PyAtExit_Call(PyInterpreterState *interp)

{

struct atexit_state *state = &interp->atexit;

atexit_callfuncs(state);

}

PyDoc_STRVAR(atexit_register__doc__,

func is returned to facilitate usage as a decorator.");

static PyObject *

atexit_register(PyObject *module, PyObject *args, PyObject *kwargs)

{

if (PyTuple_GET_SIZE(args) == 0) {

PyErr_SetString(PyExc_TypeError,

"register() takes at least 1 argument (0 given)");

return NULL;

}

PyObject *func = PyTuple_GET_ITEM(args, 0);

if (!PyCallable_Check(func)) {

PyErr_SetString(PyExc_TypeError,

"the first argument must be callable");

return NULL;

}

PyObject *func_args = PyTuple_GetSlice(args, 1, PyTuple_GET_SIZE(args));

if (func_args == NULL) {

return NULL;

}

PyObject *func_kwargs = kwargs;

if (func_kwargs == NULL)

{

func_kwargs = Py_None;

}

PyObject *callback = PyTuple_Pack(3, func, func_args, func_kwargs);

Py_DECREF(func_args);

if (callback == NULL)

{

return NULL;

}

struct atexit_state *state = get_atexit_state();

// atexit callbacks go in a LIFO order

if (PyList_Insert(state->callbacks, 0, callback) < 0)

{

Py_DECREF(callback);

return NULL;

}

Py_DECREF(callback);

return Py_NewRef(func);

}

PyDoc_STRVAR(atexit_run_exitfuncs__doc__,

If a callback raises an exception, it is logged with sys.unraisablehook.");

static PyObject *

atexit_run_exitfuncs(PyObject *module, PyObject *Py_UNUSED(dummy))

{

struct atexit_state *state = get_atexit_state();

atexit_callfuncs(state);

Py_RETURN_NONE;

}

PyDoc_STRVAR(atexit_clear__doc__,

Clear the list of previously registered exit functions.");

static PyObject *

atexit_clear(PyObject *module, PyObject *Py_UNUSED(dummy))

{

atexit_cleanup(get_atexit_state());

Py_RETURN_NONE;

}

PyDoc_STRVAR(atexit_ncallbacks__doc__,

Return the number of registered exit functions.");

static PyObject *

atexit_ncallbacks(PyObject *module, PyObject *Py_UNUSED(dummy))

{

struct atexit_state *state = get_atexit_state();

assert(state->callbacks != NULL);

assert(PyList_CheckExact(state->callbacks));

return PyLong_FromSsize_t(PyList_GET_SIZE(state->callbacks));

}

static int

atexit_unregister_locked(PyObject *callbacks, PyObject *func)

{

for (Py_ssize_t i = PyList_GET_SIZE(callbacks) - 1; i >= 0; --i) {

PyObject *tuple = Py_NewRef(PyList_GET_ITEM(callbacks, i));

assert(PyTuple_CheckExact(tuple));

PyObject *to_compare = PyTuple_GET_ITEM(tuple, 0);

int cmp = PyObject_RichCompareBool(func, to_compare, Py_EQ);

if (cmp < 0) {

Py_DECREF(tuple);

return -1;

}

if (cmp == 1) {

// We found a callback!

// But its index could have changed if it or other callbacks were

// unregistered during the comparison.

Py_ssize_t j = PyList_GET_SIZE(callbacks) - 1;

j = Py_MIN(j, i);

for (; j >= 0; --j) {

if (PyList_GET_ITEM(callbacks, j) == tuple) {

// We found the callback index! For real!

if (PyList_SetSlice(callbacks, j, j + 1, NULL) < 0) {

Py_DECREF(tuple);

return -1;

}

i = j;

break;

}

}

}

Py_DECREF(tuple);

if (i >= PyList_GET_SIZE(callbacks)) {

i = PyList_GET_SIZE(callbacks);

}

}

return 0;

}

PyDoc_STRVAR(atexit_unregister__doc__,

func - function to be unregistered");

static PyObject *

atexit_unregister(PyObject *module, PyObject *func)

{

struct atexit_state *state = get_atexit_state();

int result;

Py_BEGIN_CRITICAL_SECTION(state->callbacks);

result = atexit_unregister_locked(state->callbacks, func);

Py_END_CRITICAL_SECTION();

return result < 0 ? NULL : Py_None;

}

static PyMethodDef atexit_methods[] = {

{"register", _PyCFunction_CAST(atexit_register), METH_VARARGS|METH_KEYWORDS,

atexit_register__doc__},

{"_clear", atexit_clear, METH_NOARGS, atexit_clear__doc__},

{"unregister", atexit_unregister, METH_O, atexit_unregister__doc__},

{"_run_exitfuncs", atexit_run_exitfuncs, METH_NOARGS,

atexit_run_exitfuncs__doc__},

{"_ncallbacks", atexit_ncallbacks, METH_NOARGS,

atexit_ncallbacks__doc__},

{NULL, NULL} /* sentinel */

};

PyDoc_STRVAR(atexit__doc__,

");

static PyModuleDef_Slot atexitmodule_slots[] = {

_Py_ABI_SLOT,

{Py_mod_multiple_interpreters, Py_MOD_PER_INTERPRETER_GIL_SUPPORTED},

{Py_mod_gil, Py_MOD_GIL_NOT_USED},

{0, NULL}

};

static struct PyModuleDef atexitmodule = {

PyModuleDef_HEAD_INIT,

.m_name = "atexit",

.m_doc = atexit__doc__,

.m_size = 0,

.m_methods = atexit_methods,

.m_slots = atexitmodule_slots,

};

PyInit_atexit(void)

{

return PyModuleDef_Init(&atexitmodule);

}