/*

 *  atexit - allow programmer to define multiple exit functions to be executed

 *  upon normal program termination.

 *

 *   Translated from atexit.py by Collin Winter.

 +   Copyright 2007 Python Software Foundation.

 */

#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

/* ===================================================================== */

/* Callback machinery. */

static inline struct atexit_state*

get_atexit_state(void)

{

    PyInterpreterState *interp = _PyInterpreterState_GET();

    return &interp->atexit;

}

int

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;

}

/* Clear all callbacks without calling them */

static void

atexit_cleanup(struct atexit_state *state)

{

    PyList_Clear(state->callbacks);

}

PyStatus

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

}

void

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

}

void

_PyAtExit_Call(PyInterpreterState *interp)

{

    struct atexit_state *state = &interp->atexit;

    atexit_callfuncs(state);

}

/* ===================================================================== */

/* Module methods. */

PyDoc_STRVAR(atexit_register__doc__,

"register($module, func, /, *args, **kwargs)\n\

--\n\

\n\

Register a function to be executed upon normal program termination\n\

\n\

    func - function to be called at exit\n\

    args - optional arguments to pass to func\n\

    kwargs - optional keyword arguments to pass to func\n\

\n\

    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__,

"_run_exitfuncs($module, /)\n\

--\n\

\n\

Run all registered exit functions.\n\

\n\

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($module, /)\n\

--\n\

\n\

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__,

"_ncallbacks($module, /)\n\

--\n\

\n\

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__,

"unregister($module, func, /)\n\

--\n\

\n\

Unregister an exit function which was previously registered using\n\

atexit.register\n\

\n\

    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 */

};

/* ===================================================================== */

/* Initialization function. */

PyDoc_STRVAR(atexit__doc__,

"allow programmer to define multiple exit functions to be executed\n\

upon normal program termination.\n\

\n\

Two public functions, register and unregister, are defined.\n\

");

static PyModuleDef_Slot atexitmodule_slots[] = {

    {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,

};

PyMODINIT_FUNC

PyInit_atexit(void)

{

    return PyModuleDef_Init(&atexitmodule);

}