#include "Python.h"

#include "pycore_pyerrors.h" // _PyErr_ClearExcState()

#include <stddef.h> // offsetof()

_Py_IDENTIFIER(__asyncio_running_event_loop__);

_Py_IDENTIFIER(_asyncio_future_blocking);

_Py_IDENTIFIER(add_done_callback);

_Py_IDENTIFIER(call_soon);

_Py_IDENTIFIER(cancel);

_Py_IDENTIFIER(get_event_loop);

_Py_IDENTIFIER(send);

_Py_IDENTIFIER(throw);

static PyObject *asyncio_mod;

static PyObject *traceback_extract_stack;

static PyObject *asyncio_get_event_loop_policy;

static PyObject *asyncio_future_repr_info_func;

static PyObject *asyncio_iscoroutine_func;

static PyObject *asyncio_task_get_stack_func;

static PyObject *asyncio_task_print_stack_func;

static PyObject *asyncio_task_repr_info_func;

static PyObject *asyncio_InvalidStateError;

static PyObject *asyncio_CancelledError;

static PyObject *context_kwname;

static int module_initialized;

static PyObject *cached_running_holder;

static volatile uint64_t cached_running_holder_tsid;

static uint64_t task_name_counter = 0;

static PyObject *all_tasks;

static PyObject *current_tasks;

static PyObject *iscoroutine_typecache;

typedef enum {

STATE_PENDING,

STATE_CANCELLED,

STATE_FINISHED

} fut_state;

#define FutureObj_HEAD(prefix) \

PyObject_HEAD \

PyObject *prefix##_loop; \

PyObject *prefix##_callback0; \

PyObject *prefix##_context0; \

PyObject *prefix##_callbacks; \

PyObject *prefix##_exception; \

PyObject *prefix##_result; \

PyObject *prefix##_source_tb; \

PyObject *prefix##_cancel_msg; \

fut_state prefix##_state; \

int prefix##_log_tb; \

int prefix##_blocking; \

PyObject *dict; \

PyObject *prefix##_weakreflist; \

_PyErr_StackItem prefix##_cancelled_exc_state;

typedef struct {

FutureObj_HEAD(fut)

} FutureObj;

typedef struct {

FutureObj_HEAD(task)

PyObject *task_fut_waiter;

PyObject *task_coro;

PyObject *task_name;

PyObject *task_context;

int task_must_cancel;

int task_log_destroy_pending;

} TaskObj;

typedef struct {

PyObject_HEAD

TaskObj *sw_task;

PyObject *sw_arg;

} TaskStepMethWrapper;

typedef struct {

PyObject_HEAD

TaskObj *ww_task;

} TaskWakeupMethWrapper;

typedef struct {

PyObject_HEAD

PyObject *rl_loop;

#if defined(HAVE_GETPID) && !defined(MS_WINDOWS)

pid_t rl_pid;

} PyRunningLoopHolder;

static PyTypeObject FutureType;

static PyTypeObject TaskType;

static PyTypeObject PyRunningLoopHolder_Type;

#define Future_CheckExact(obj) Py_IS_TYPE(obj, &FutureType)

#define Task_CheckExact(obj) Py_IS_TYPE(obj, &TaskType)

#define Future_Check(obj) PyObject_TypeCheck(obj, &FutureType)

#define Task_Check(obj) PyObject_TypeCheck(obj, &TaskType)

#include "clinic/_asynciomodule.c.h"

static PyObject * future_new_iter(PyObject *);

static PyRunningLoopHolder * new_running_loop_holder(PyObject *);

static int

_is_coroutine(PyObject *coro)

{

/* 'coro' is not a native coroutine, call asyncio.iscoroutine()

PyObject *res = PyObject_CallOneArg(asyncio_iscoroutine_func, coro);

if (res == NULL) {

return -1;

}

int is_res_true = PyObject_IsTrue(res);

Py_DECREF(res);

if (is_res_true <= 0) {

return is_res_true;

}

if (PySet_GET_SIZE(iscoroutine_typecache) < 100) {

/* Just in case we don't want to cache more than 100

if (PySet_Add(iscoroutine_typecache, (PyObject*) Py_TYPE(coro))) {

return -1;

}

}

return 1;

}

static inline int

is_coroutine(PyObject *coro)

{

if (PyCoro_CheckExact(coro)) {

return 1;

}

/* Check if `type(coro)` is in the cache.

int has_it = PySet_Contains(

iscoroutine_typecache, (PyObject*) Py_TYPE(coro));

if (has_it == 0) {

/* type(coro) is not in iscoroutine_typecache */

return _is_coroutine(coro);

}

/* either an error has occurred or

return has_it;

}

static PyObject *

get_future_loop(PyObject *fut)

{

/* Implementation of `asyncio.futures._get_loop` */

_Py_IDENTIFIER(get_loop);

_Py_IDENTIFIER(_loop);

PyObject *getloop;

if (Future_CheckExact(fut) || Task_CheckExact(fut)) {

PyObject *loop = ((FutureObj *)fut)->fut_loop;

Py_INCREF(loop);

return loop;

}

if (_PyObject_LookupAttrId(fut, &PyId_get_loop, &getloop) < 0) {

return NULL;

}

if (getloop != NULL) {

PyObject *res = PyObject_CallNoArgs(getloop);

Py_DECREF(getloop);

return res;

}

return _PyObject_GetAttrId(fut, &PyId__loop);

}

static int

get_running_loop(PyObject **loop)

{

PyObject *rl;

PyThreadState *ts = PyThreadState_Get();

uint64_t ts_id = PyThreadState_GetID(ts);

if (ts_id == cached_running_holder_tsid && cached_running_holder != NULL) {

// Fast path, check the cache.

rl = cached_running_holder; // borrowed

}

else {

PyObject *ts_dict = _PyThreadState_GetDict(ts); // borrowed

if (ts_dict == NULL) {

goto not_found;

}

rl = _PyDict_GetItemIdWithError(

ts_dict, &PyId___asyncio_running_event_loop__); // borrowed

if (rl == NULL) {

if (PyErr_Occurred()) {

goto error;

}

else {

goto not_found;

}

}

cached_running_holder = rl; // borrowed

cached_running_holder_tsid = ts_id;

}

assert(Py_IS_TYPE(rl, &PyRunningLoopHolder_Type));

PyObject *running_loop = ((PyRunningLoopHolder *)rl)->rl_loop;

if (running_loop == Py_None) {

goto not_found;

}

#if defined(HAVE_GETPID) && !defined(MS_WINDOWS)

/* On Windows there is no getpid, but there is also no os.fork(),

if (getpid() != ((PyRunningLoopHolder *)rl)->rl_pid) {

goto not_found;

}

Py_INCREF(running_loop);

*loop = running_loop;

return 0;

not_found:

*loop = NULL;

return 0;

error:

*loop = NULL;

return -1;

}

static int

set_running_loop(PyObject *loop)

{

PyObject *ts_dict = NULL;

PyThreadState *tstate = PyThreadState_Get();

if (tstate != NULL) {

ts_dict = _PyThreadState_GetDict(tstate); // borrowed

}

if (ts_dict == NULL) {

PyErr_SetString(

PyExc_RuntimeError, "thread-local storage is not available");

return -1;

}

PyRunningLoopHolder *rl = new_running_loop_holder(loop);

if (rl == NULL) {

return -1;

}

if (_PyDict_SetItemId(

ts_dict, &PyId___asyncio_running_event_loop__, (PyObject *)rl) < 0)

{

Py_DECREF(rl); // will cleanup loop & current_pid

return -1;

}

Py_DECREF(rl);

cached_running_holder = (PyObject *)rl;

cached_running_holder_tsid = PyThreadState_GetID(tstate);

return 0;

}

static PyObject *

get_event_loop(void)

{

PyObject *loop;

PyObject *policy;

if (get_running_loop(&loop)) {

return NULL;

}

if (loop != NULL) {

return loop;

}

policy = PyObject_CallNoArgs(asyncio_get_event_loop_policy);

if (policy == NULL) {

return NULL;

}

loop = _PyObject_CallMethodIdNoArgs(policy, &PyId_get_event_loop);

Py_DECREF(policy);

return loop;

}

static int

call_soon(PyObject *loop, PyObject *func, PyObject *arg, PyObject *ctx)

{

PyObject *handle;

PyObject *stack[3];

Py_ssize_t nargs;

if (ctx == NULL) {

handle = _PyObject_CallMethodIdObjArgs(

loop, &PyId_call_soon, func, arg, NULL);

}

else {

/* Use FASTCALL to pass a keyword-only argument to call_soon */

PyObject *callable = _PyObject_GetAttrId(loop, &PyId_call_soon);

if (callable == NULL) {

return -1;

}

/* All refs in 'stack' are borrowed. */

nargs = 1;

stack[0] = func;

if (arg != NULL) {

stack[1] = arg;

nargs++;

}

stack[nargs] = (PyObject *)ctx;

handle = PyObject_Vectorcall(callable, stack, nargs, context_kwname);

Py_DECREF(callable);

}

if (handle == NULL) {

return -1;

}

Py_DECREF(handle);

return 0;

}

static inline int

future_is_alive(FutureObj *fut)

{

return fut->fut_loop != NULL;

}

static inline int

future_ensure_alive(FutureObj *fut)

{

if (!future_is_alive(fut)) {

PyErr_SetString(PyExc_RuntimeError,

"Future object is not initialized.");

return -1;

}

return 0;

}

#define ENSURE_FUTURE_ALIVE(fut) \

do { \

assert(Future_Check(fut) || Task_Check(fut)); \

if (future_ensure_alive((FutureObj*)fut)) { \

return NULL; \

} \

} while(0);

static int

future_schedule_callbacks(FutureObj *fut)

{

Py_ssize_t len;

Py_ssize_t i;

if (fut->fut_callback0 != NULL) {

/* There's a 1st callback */

int ret = call_soon(

fut->fut_loop, fut->fut_callback0,

(PyObject *)fut, fut->fut_context0);

Py_CLEAR(fut->fut_callback0);

Py_CLEAR(fut->fut_context0);

if (ret) {

/* If an error occurs in pure-Python implementation,

Py_CLEAR(fut->fut_callbacks);

return ret;

}

/* we called the first callback, now try calling

}

if (fut->fut_callbacks == NULL) {

/* No more callbacks, return. */

return 0;

}

len = PyList_GET_SIZE(fut->fut_callbacks);

if (len == 0) {

/* The list of callbacks was empty; clear it and return. */

Py_CLEAR(fut->fut_callbacks);

return 0;

}

for (i = 0; i < len; i++) {

PyObject *cb_tup = PyList_GET_ITEM(fut->fut_callbacks, i);

PyObject *cb = PyTuple_GET_ITEM(cb_tup, 0);

PyObject *ctx = PyTuple_GET_ITEM(cb_tup, 1);

if (call_soon(fut->fut_loop, cb, (PyObject *)fut, ctx)) {

/* If an error occurs in pure-Python implementation,

Py_CLEAR(fut->fut_callbacks);

return -1;

}

}

Py_CLEAR(fut->fut_callbacks);

return 0;

}

static int

future_init(FutureObj *fut, PyObject *loop)

{

PyObject *res;

int is_true;

_Py_IDENTIFIER(get_debug);

// Same to FutureObj_clear() but not clearing fut->dict

Py_CLEAR(fut->fut_loop);

Py_CLEAR(fut->fut_callback0);

Py_CLEAR(fut->fut_context0);

Py_CLEAR(fut->fut_callbacks);

Py_CLEAR(fut->fut_result);

Py_CLEAR(fut->fut_exception);

Py_CLEAR(fut->fut_source_tb);

Py_CLEAR(fut->fut_cancel_msg);

_PyErr_ClearExcState(&fut->fut_cancelled_exc_state);

fut->fut_state = STATE_PENDING;

fut->fut_log_tb = 0;

fut->fut_blocking = 0;

if (loop == Py_None) {

loop = get_event_loop();

if (loop == NULL) {

return -1;

}

}

else {

Py_INCREF(loop);

}

fut->fut_loop = loop;

res = _PyObject_CallMethodIdNoArgs(fut->fut_loop, &PyId_get_debug);

if (res == NULL) {

return -1;

}

is_true = PyObject_IsTrue(res);

Py_DECREF(res);

if (is_true < 0) {

return -1;

}

if (is_true && !_Py_IsFinalizing()) {

/* Only try to capture the traceback if the interpreter is not being

fut->fut_source_tb = PyObject_CallNoArgs(traceback_extract_stack);

if (fut->fut_source_tb == NULL) {

return -1;

}

}

return 0;

}

static PyObject *

future_set_result(FutureObj *fut, PyObject *res)

{

if (future_ensure_alive(fut)) {

return NULL;

}

if (fut->fut_state != STATE_PENDING) {

PyErr_SetString(asyncio_InvalidStateError, "invalid state");

return NULL;

}

assert(!fut->fut_result);

Py_INCREF(res);

fut->fut_result = res;

fut->fut_state = STATE_FINISHED;

if (future_schedule_callbacks(fut) == -1) {

return NULL;

}

Py_RETURN_NONE;

}

static PyObject *

future_set_exception(FutureObj *fut, PyObject *exc)

{

PyObject *exc_val = NULL;

if (fut->fut_state != STATE_PENDING) {

PyErr_SetString(asyncio_InvalidStateError, "invalid state");

return NULL;

}

if (PyExceptionClass_Check(exc)) {

exc_val = PyObject_CallNoArgs(exc);

if (exc_val == NULL) {

return NULL;

}

if (fut->fut_state != STATE_PENDING) {

Py_DECREF(exc_val);

PyErr_SetString(asyncio_InvalidStateError, "invalid state");

return NULL;

}

}

else {

exc_val = exc;

Py_INCREF(exc_val);

}

if (!PyExceptionInstance_Check(exc_val)) {

Py_DECREF(exc_val);

PyErr_SetString(PyExc_TypeError, "invalid exception object");

return NULL;

}

if (Py_IS_TYPE(exc_val, (PyTypeObject *)PyExc_StopIteration)) {

Py_DECREF(exc_val);

PyErr_SetString(PyExc_TypeError,

"StopIteration interacts badly with generators "

"and cannot be raised into a Future");

return NULL;

}

assert(!fut->fut_exception);

fut->fut_exception = exc_val;

fut->fut_state = STATE_FINISHED;

if (future_schedule_callbacks(fut) == -1) {

return NULL;

}

fut->fut_log_tb = 1;

Py_RETURN_NONE;

}

static PyObject *

create_cancelled_error(PyObject *msg)

{

PyObject *exc;

if (msg == NULL || msg == Py_None) {

exc = PyObject_CallNoArgs(asyncio_CancelledError);

} else {

exc = PyObject_CallOneArg(asyncio_CancelledError, msg);

}

return exc;

}

static void

future_set_cancelled_error(FutureObj *fut)

{

PyObject *exc = create_cancelled_error(fut->fut_cancel_msg);

PyErr_SetObject(asyncio_CancelledError, exc);

Py_DECREF(exc);

_PyErr_ChainStackItem(&fut->fut_cancelled_exc_state);

}

static int

future_get_result(FutureObj *fut, PyObject **result)

{

if (fut->fut_state == STATE_CANCELLED) {

future_set_cancelled_error(fut);

return -1;

}

if (fut->fut_state != STATE_FINISHED) {

PyErr_SetString(asyncio_InvalidStateError, "Result is not set.");

return -1;

}

fut->fut_log_tb = 0;

if (fut->fut_exception != NULL) {

Py_INCREF(fut->fut_exception);

*result = fut->fut_exception;

return 1;

}

Py_INCREF(fut->fut_result);

*result = fut->fut_result;

return 0;

}

static PyObject *

future_add_done_callback(FutureObj *fut, PyObject *arg, PyObject *ctx)

{

if (!future_is_alive(fut)) {

PyErr_SetString(PyExc_RuntimeError, "uninitialized Future object");

return NULL;

}

if (fut->fut_state != STATE_PENDING) {

/* The future is done/cancelled, so schedule the callback

if (call_soon(fut->fut_loop, arg, (PyObject*) fut, ctx)) {

return NULL;

}

}

else {

/* The future is pending, add a callback.

if (fut->fut_callbacks == NULL && fut->fut_callback0 == NULL) {

Py_INCREF(arg);

fut->fut_callback0 = arg;

Py_INCREF(ctx);

fut->fut_context0 = ctx;

}

else {

PyObject *tup = PyTuple_New(2);

if (tup == NULL) {

return NULL;

}

Py_INCREF(arg);

PyTuple_SET_ITEM(tup, 0, arg);

Py_INCREF(ctx);

PyTuple_SET_ITEM(tup, 1, (PyObject *)ctx);

if (fut->fut_callbacks != NULL) {

int err = PyList_Append(fut->fut_callbacks, tup);

if (err) {

Py_DECREF(tup);

return NULL;

}

Py_DECREF(tup);

}

else {

fut->fut_callbacks = PyList_New(1);

if (fut->fut_callbacks == NULL) {

return NULL;

}

PyList_SET_ITEM(fut->fut_callbacks, 0, tup); /* borrow */

}

}

}

Py_RETURN_NONE;

}

static PyObject *

future_cancel(FutureObj *fut, PyObject *msg)

{

fut->fut_log_tb = 0;

if (fut->fut_state != STATE_PENDING) {

Py_RETURN_FALSE;

}

fut->fut_state = STATE_CANCELLED;

Py_XINCREF(msg);

Py_XSETREF(fut->fut_cancel_msg, msg);

if (future_schedule_callbacks(fut) == -1) {

return NULL;

}

Py_RETURN_TRUE;

}

static int

_asyncio_Future___init___impl(FutureObj *self, PyObject *loop)

{

return future_init(self, loop);

}

static int

FutureObj_clear(FutureObj *fut)

{

Py_CLEAR(fut->fut_loop);

Py_CLEAR(fut->fut_callback0);

Py_CLEAR(fut->fut_context0);

Py_CLEAR(fut->fut_callbacks);

Py_CLEAR(fut->fut_result);

Py_CLEAR(fut->fut_exception);

Py_CLEAR(fut->fut_source_tb);

Py_CLEAR(fut->fut_cancel_msg);

_PyErr_ClearExcState(&fut->fut_cancelled_exc_state);

Py_CLEAR(fut->dict);

return 0;

}

static int

FutureObj_traverse(FutureObj *fut, visitproc visit, void *arg)

{

Py_VISIT(fut->fut_loop);

Py_VISIT(fut->fut_callback0);

Py_VISIT(fut->fut_context0);

Py_VISIT(fut->fut_callbacks);

Py_VISIT(fut->fut_result);

Py_VISIT(fut->fut_exception);

Py_VISIT(fut->fut_source_tb);

Py_VISIT(fut->fut_cancel_msg);

Py_VISIT(fut->dict);

_PyErr_StackItem *exc_state = &fut->fut_cancelled_exc_state;

Py_VISIT(exc_state->exc_type);

Py_VISIT(exc_state->exc_value);

Py_VISIT(exc_state->exc_traceback);

return 0;

}

static PyObject *

_asyncio_Future_result_impl(FutureObj *self)

{

PyObject *result;

if (!future_is_alive(self)) {

PyErr_SetString(asyncio_InvalidStateError,

"Future object is not initialized.");

return NULL;

}

int res = future_get_result(self, &result);

if (res == -1) {

return NULL;

}

if (res == 0) {

return result;

}

assert(res == 1);

PyErr_SetObject(PyExceptionInstance_Class(result), result);

Py_DECREF(result);

return NULL;

}

static PyObject *

_asyncio_Future_exception_impl(FutureObj *self)

{

if (!future_is_alive(self)) {

PyErr_SetString(asyncio_InvalidStateError,

"Future object is not initialized.");

return NULL;

}

if (self->fut_state == STATE_CANCELLED) {

future_set_cancelled_error(self);

return NULL;

}

if (self->fut_state != STATE_FINISHED) {

PyErr_SetString(asyncio_InvalidStateError, "Exception is not set.");

return NULL;

}

if (self->fut_exception != NULL) {

self->fut_log_tb = 0;

Py_INCREF(self->fut_exception);

return self->fut_exception;

}

Py_RETURN_NONE;

}

static PyObject *

_asyncio_Future_set_result(FutureObj *self, PyObject *result)

{

ENSURE_FUTURE_ALIVE(self)

return future_set_result(self, result);

}

static PyObject *

_asyncio_Future_set_exception(FutureObj *self, PyObject *exception)

{

ENSURE_FUTURE_ALIVE(self)

return future_set_exception(self, exception);

}

static PyObject *

_asyncio_Future_add_done_callback_impl(FutureObj *self, PyObject *fn,

PyObject *context)

{

if (context == NULL) {

context = PyContext_CopyCurrent();

if (context == NULL) {

return NULL;

}

PyObject *res = future_add_done_callback(self, fn, context);

Py_DECREF(context);

return res;

}

return future_add_done_callback(self, fn, context);

}

static PyObject *

_asyncio_Future_remove_done_callback(FutureObj *self, PyObject *fn)

{

PyObject *newlist;

Py_ssize_t len, i, j=0;

Py_ssize_t cleared_callback0 = 0;

ENSURE_FUTURE_ALIVE(self)

if (self->fut_callback0 != NULL) {

int cmp = PyObject_RichCompareBool(self->fut_callback0, fn, Py_EQ);

if (cmp == -1) {

return NULL;

}

if (cmp == 1) {

/* callback0 == fn */

Py_CLEAR(self->fut_callback0);

Py_CLEAR(self->fut_context0);

cleared_callback0 = 1;

}

}

if (self->fut_callbacks == NULL) {

return PyLong_FromSsize_t(cleared_callback0);

}