Add asyncio task benchmark counters

Author: KRRT7

Lib/asyncio/base_events.py

@@ -468,6 +468,7 @@ def create_task(self, coro, **kwargs):
 
         Return a task object.
         """
+        tasks._benchmark_count('py_loop_create_task_calls')
         self._check_closed()
         if self._task_factory is not None:
             return self._task_factory(self, coro, **kwargs)

Lib/asyncio/tasks.py

@@ -32,6 +32,33 @@
 # This uses itertools.count() instead of a "+= 1" operation because the latter
 # is not thread safe. See bpo-11866 for a longer explanation.
 _task_name_counter = itertools.count(1).__next__
+_benchmark_counters = {
+    'py_create_task_calls': 0,
+    'py_loop_create_task_calls': 0,
+    'py_gather_calls': 0,
+    'py_gather_input_count': 0,
+    'py_gather_done_futures': 0,
+    'py_gather_callback_registrations': 0,
+}
+
+
+def _benchmark_count(name, delta=1):
+    _benchmark_counters[name] = _benchmark_counters.get(name, 0) + delta
+
+
+def _reset_benchmark_counters():
+    for key in list(_benchmark_counters):
+        _benchmark_counters[key] = 0
+    if _c_reset_benchmark_counters is not None:
+        _c_reset_benchmark_counters()
+
+
+def _get_benchmark_counters():
+    counters = dict(_benchmark_counters)
+    if _c_get_benchmark_counters is not None:
+        for key, value in _c_get_benchmark_counters().items():
+            counters[key] = value
+    return counters
 
 
 def current_task(loop=None):
@@ -384,13 +411,19 @@ def __wakeup(self, future):
 else:
     # _CTask is needed for tests.
     Task = _CTask = _asyncio.Task
+    _c_get_benchmark_counters = getattr(_asyncio, '_get_benchmark_counters', None)
+    _c_reset_benchmark_counters = getattr(_asyncio, '_reset_benchmark_counters', None)
+if '_c_get_benchmark_counters' not in globals():
+    _c_get_benchmark_counters = None
+    _c_reset_benchmark_counters = None
 
 
 def create_task(coro, **kwargs):
     """Schedule the execution of a coroutine object in a spawn task.
 
     Return a Task object.
     """
+    _benchmark_count('py_create_task_calls')
     loop = events.get_running_loop()
     return loop.create_task(coro, **kwargs)
 
@@ -798,6 +831,7 @@ def gather(*coros_or_futures, return_exceptions=False):
     after catching an exception (raised by one of the awaitables) from
     gather won't cancel any other awaitables.
     """
+    _benchmark_count('py_gather_calls')
     if not coros_or_futures:
         loop = events.get_event_loop()
         outer = loop.create_future()
@@ -876,6 +910,7 @@ def _done_callback(fut, cur_task=cur_task):
     outer = None  # bpo-46672
     for arg in coros_or_futures:
         if arg not in arg_to_fut:
+            _benchmark_count('py_gather_input_count')
             fut = ensure_future(arg, loop=loop)
             if loop is None:
                 loop = futures._get_loop(fut)
@@ -888,11 +923,13 @@ def _done_callback(fut, cur_task=cur_task):
             nfuts += 1
             arg_to_fut[arg] = fut
             if fut.done():
+                _benchmark_count('py_gather_done_futures')
                 done_futs.append(fut)
             else:
                 if cur_task is not None:
                     futures.future_add_to_awaited_by(fut, cur_task)
                 fut.add_done_callback(_done_callback)
+                _benchmark_count('py_gather_callback_registrations')
 
         else:
             # There's a duplicate Future object in coros_or_futures.

Modules/_asynciomodule.c

@@ -186,6 +186,16 @@ typedef struct {
     /* Counter for autogenerated Task names */
     uint64_t task_name_counter;
 
+    /* Internal benchmarking counters for asyncio task hot paths. */
+    uint64_t bench_call_soon_calls;
+    uint64_t bench_task_init_calls;
+    uint64_t bench_task_call_step_soon_calls;
+    uint64_t bench_task_call_step_soon_with_arg_calls;
+    uint64_t bench_register_task_calls;
+    uint64_t bench_future_schedule_callbacks_calls;
+    uint64_t bench_future_schedule_callback0_calls;
+    uint64_t bench_future_schedule_callback_list_items;
+
     /* Pointer to the asyncio debug offset to avoid it to be optimized away
        by the compiler */
     void *debug_offsets;
@@ -359,6 +369,7 @@ static int
 call_soon(asyncio_state *state, PyObject *loop, PyObject *func, PyObject *arg,
           PyObject *ctx)
 {
+    state->bench_call_soon_calls++;
     PyObject *handle;
 
     if (ctx == NULL) {
@@ -423,6 +434,7 @@ static int
 future_schedule_callbacks(asyncio_state *state, FutureObj *fut)
 {
     _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(fut);
+    state->bench_future_schedule_callbacks_calls++;
 
     assert(fut->fut_state != STATE_PENDING);
 
@@ -434,6 +446,7 @@ future_schedule_callbacks(asyncio_state *state, FutureObj *fut)
     }
 
     if (fut->fut_callback0 != NULL) {
+        state->bench_future_schedule_callback0_calls++;
         /* There's a 1st callback */
 
         // Beware: An evil call_soon could alter fut_callback0 or fut_context0.
@@ -472,6 +485,7 @@ future_schedule_callbacks(asyncio_state *state, FutureObj *fut)
     PyObject *callbacks = fut->fut_callbacks;
     fut->fut_callbacks = NULL;
     Py_ssize_t n = PyList_GET_SIZE(callbacks);
+    state->bench_future_schedule_callback_list_items += (uint64_t)n;
     for (Py_ssize_t i = 0; i < n; i++) {
         assert(PyList_GET_SIZE(callbacks) == n);
         PyObject *cb_tup = PyList_GET_ITEM(callbacks, i);
@@ -2183,6 +2197,8 @@ register_task(_PyThreadStateImpl *ts, TaskObj *task)
         return;
     }
     struct llist_node *head = &ts->asyncio_tasks_head;
+    asyncio_state *state = get_asyncio_state_by_def((PyObject *)task);
+    state->bench_register_task_calls++;
     llist_insert_tail(head, &task->task_node);
 }
 
@@ -2308,6 +2324,7 @@ _asyncio_Task___init___impl(TaskObj *self, PyObject *coro, PyObject *loop,
     self->task_is_task = 1;
 
     asyncio_state *state = get_asyncio_state_by_def((PyObject *)self);
+    state->bench_task_init_calls++;
     int is_coro = is_coroutine(state, coro);
     if (is_coro == -1) {
         return -1;
@@ -2983,6 +3000,10 @@ TaskObj_dealloc(PyObject *self)
 static int
 task_call_step_soon(asyncio_state *state, TaskObj *task, PyObject *arg)
 {
+    state->bench_task_call_step_soon_calls++;
+    if (arg != NULL) {
+        state->bench_task_call_step_soon_with_arg_calls++;
+    }
     PyObject *cb = TaskStepMethWrapper_new(task, arg);
     if (cb == NULL) {
         return -1;
@@ -4327,6 +4348,58 @@ module_init(asyncio_state *state)
 
 PyDoc_STRVAR(module_doc, "Accelerator module for asyncio");
 
+static PyObject *
+_asyncio__get_benchmark_counters(PyObject *module, PyObject *Py_UNUSED(ignored))
+{
+    asyncio_state *state = get_asyncio_state(module);
+    PyObject *counters = PyDict_New();
+    if (counters == NULL) {
+        return NULL;
+    }
+
+#define SET_COUNTER(name) \
+    do { \
+        PyObject *value = PyLong_FromUnsignedLongLong(state->name); \
+        if (value == NULL) { \
+            Py_DECREF(counters); \
+            return NULL; \
+        } \
+        if (PyDict_SetItemString(counters, #name, value) < 0) { \
+            Py_DECREF(value); \
+            Py_DECREF(counters); \
+            return NULL; \
+        } \
+        Py_DECREF(value); \
+    } while (0)
+
+    SET_COUNTER(bench_call_soon_calls);
+    SET_COUNTER(bench_task_init_calls);
+    SET_COUNTER(bench_task_call_step_soon_calls);
+    SET_COUNTER(bench_task_call_step_soon_with_arg_calls);
+    SET_COUNTER(bench_register_task_calls);
+    SET_COUNTER(bench_future_schedule_callbacks_calls);
+    SET_COUNTER(bench_future_schedule_callback0_calls);
+    SET_COUNTER(bench_future_schedule_callback_list_items);
+
+#undef SET_COUNTER
+    return counters;
+}
+
+static PyObject *
+_asyncio__reset_benchmark_counters(PyObject *module, PyObject *Py_UNUSED(ignored))
+{
+    asyncio_state *state = get_asyncio_state(module);
+    state->bench_call_soon_calls = 0;
+    state->bench_task_init_calls = 0;
+    state->bench_task_call_step_soon_calls = 0;
+    state->bench_task_call_step_soon_with_arg_calls = 0;
+    state->bench_register_task_calls = 0;
+    state->bench_future_schedule_callbacks_calls = 0;
+    state->bench_future_schedule_callback0_calls = 0;
+    state->bench_future_schedule_callback_list_items = 0;
+    Py_RETURN_NONE;
+}
+
 static PyMethodDef asyncio_methods[] = {
     _ASYNCIO_CURRENT_TASK_METHODDEF
     _ASYNCIO_GET_EVENT_LOOP_METHODDEF
@@ -4343,6 +4416,8 @@ static PyMethodDef asyncio_methods[] = {
     _ASYNCIO_ALL_TASKS_METHODDEF
     _ASYNCIO_FUTURE_ADD_TO_AWAITED_BY_METHODDEF
     _ASYNCIO_FUTURE_DISCARD_FROM_AWAITED_BY_METHODDEF
+    {"_get_benchmark_counters", _asyncio__get_benchmark_counters, METH_NOARGS, NULL},
+    {"_reset_benchmark_counters", _asyncio__reset_benchmark_counters, METH_NOARGS, NULL},
     {NULL, NULL}
 };