bpo-34270: Make it possible to name asyncio tasks (GH-8547)

Co-authored-by: Antti Haapala <antti.haapala@anttipatterns.com>

Modules/_asynciomodule.c

@@ -37,6 +37,8 @@ static PyObject *context_kwname;
 static PyObject *cached_running_holder;
 static volatile uint64_t cached_running_holder_tsid;
 
+/* Counter for autogenerated Task names */
+static uint64_t task_name_counter = 0;
 
 /* WeakSet containing all alive tasks. */
 static PyObject *all_tasks;
@@ -78,6 +80,7 @@ typedef struct {
     FutureObj_HEAD(task)
     PyObject *task_fut_waiter;
     PyObject *task_coro;
+    PyObject *task_name;
     PyContext *task_context;
     int task_must_cancel;
     int task_log_destroy_pending;
@@ -1934,13 +1937,15 @@ _asyncio.Task.__init__
     coro: object
     *
     loop: object = None
+    name: object = None
 
 A coroutine wrapped in a Future.
 [clinic start generated code]*/
 
 static int
-_asyncio_Task___init___impl(TaskObj *self, PyObject *coro, PyObject *loop)
-/*[clinic end generated code: output=9f24774c2287fc2f input=8d132974b049593e]*/
+_asyncio_Task___init___impl(TaskObj *self, PyObject *coro, PyObject *loop,
+                            PyObject *name)
+/*[clinic end generated code: output=88b12b83d570df50 input=352a3137fe60091d]*/
 {
     if (future_init((FutureObj*)self, loop)) {
         return -1;
@@ -1969,6 +1974,18 @@ _asyncio_Task___init___impl(TaskObj *self, PyObject *coro, PyObject *loop)
     Py_INCREF(coro);
     Py_XSETREF(self->task_coro, coro);
 
+    if (name == Py_None) {
+        name = PyUnicode_FromFormat("Task-%" PRIu64, ++task_name_counter);
+    } else if (!PyUnicode_Check(name)) {
+        name = PyObject_Str(name);
+    } else {
+        Py_INCREF(name);
+    }
+    Py_XSETREF(self->task_name, name);
+    if (self->task_name == NULL) {
+        return -1;
+    }
+
     if (task_call_step_soon(self, NULL)) {
         return -1;
     }
@@ -1981,6 +1998,7 @@ TaskObj_clear(TaskObj *task)
     (void)FutureObj_clear((FutureObj*) task);
     Py_CLEAR(task->task_context);
     Py_CLEAR(task->task_coro);
+    Py_CLEAR(task->task_name);
     Py_CLEAR(task->task_fut_waiter);
     return 0;
 }
@@ -1990,6 +2008,7 @@ TaskObj_traverse(TaskObj *task, visitproc visit, void *arg)
 {
     Py_VISIT(task->task_context);
     Py_VISIT(task->task_coro);
+    Py_VISIT(task->task_name);
     Py_VISIT(task->task_fut_waiter);
     (void)FutureObj_traverse((FutureObj*) task, visit, arg);
     return 0;
@@ -2297,6 +2316,41 @@ _asyncio_Task_set_exception(TaskObj *self, PyObject *exception)
     return NULL;
 }
 
+/*[clinic input]
+_asyncio.Task.get_name
+[clinic start generated code]*/
+
+static PyObject *
+_asyncio_Task_get_name_impl(TaskObj *self)
+/*[clinic end generated code: output=0ecf1570c3b37a8f input=a4a6595d12f4f0f8]*/
+{
+    if (self->task_name) {
+        Py_INCREF(self->task_name);
+        return self->task_name;
+    }
+
+    Py_RETURN_NONE;
+}
+
+/*[clinic input]
+_asyncio.Task.set_name
+
+    value: object
+    /
+[clinic start generated code]*/
+
+static PyObject *
+_asyncio_Task_set_name(TaskObj *self, PyObject *value)
+/*[clinic end generated code: output=138a8d51e32057d6 input=a8359b6e65f8fd31]*/
+{
+    PyObject *name = PyObject_Str(value);
+    if (name == NULL) {
+        return NULL;
+    }
+
+    Py_XSETREF(self->task_name, name);
+    Py_RETURN_NONE;
+}
 
 static void
 TaskObj_finalize(TaskObj *task)
@@ -2382,6 +2436,8 @@ static PyMethodDef TaskType_methods[] = {
     _ASYNCIO_TASK_GET_STACK_METHODDEF
     _ASYNCIO_TASK_PRINT_STACK_METHODDEF
     _ASYNCIO_TASK__REPR_INFO_METHODDEF
+    _ASYNCIO_TASK_GET_NAME_METHODDEF
+    _ASYNCIO_TASK_SET_NAME_METHODDEF
     {NULL, NULL}        /* Sentinel */
 };
 

Modules/clinic/_asynciomodule.c.h

@@ -253,28 +253,30 @@ _asyncio_Future__repr_info(FutureObj *self, PyObject *Py_UNUSED(ignored))
 }
 
 PyDoc_STRVAR(_asyncio_Task___init____doc__,
-"Task(coro, *, loop=None)\n"
+"Task(coro, *, loop=None, name=None)\n"
 "--\n"
 "\n"
 "A coroutine wrapped in a Future.");
 
 static int
-_asyncio_Task___init___impl(TaskObj *self, PyObject *coro, PyObject *loop);
+_asyncio_Task___init___impl(TaskObj *self, PyObject *coro, PyObject *loop,
+                            PyObject *name);
 
 static int
 _asyncio_Task___init__(PyObject *self, PyObject *args, PyObject *kwargs)
 {
     int return_value = -1;
-    static const char * const _keywords[] = {"coro", "loop", NULL};
-    static _PyArg_Parser _parser = {"O|$O:Task", _keywords, 0};
+    static const char * const _keywords[] = {"coro", "loop", "name", NULL};
+    static _PyArg_Parser _parser = {"O|$OO:Task", _keywords, 0};
     PyObject *coro;
     PyObject *loop = Py_None;
+    PyObject *name = Py_None;
 
     if (!_PyArg_ParseTupleAndKeywordsFast(args, kwargs, &_parser,
-        &coro, &loop)) {
+        &coro, &loop, &name)) {
         goto exit;
     }
-    return_value = _asyncio_Task___init___impl((TaskObj *)self, coro, loop);
+    return_value = _asyncio_Task___init___impl((TaskObj *)self, coro, loop, name);
 
 exit:
     return return_value;
@@ -500,6 +502,31 @@ PyDoc_STRVAR(_asyncio_Task_set_exception__doc__,
 #define _ASYNCIO_TASK_SET_EXCEPTION_METHODDEF    \
     {"set_exception", (PyCFunction)_asyncio_Task_set_exception, METH_O, _asyncio_Task_set_exception__doc__},
 
+PyDoc_STRVAR(_asyncio_Task_get_name__doc__,
+"get_name($self, /)\n"
+"--\n"
+"\n");
+
+#define _ASYNCIO_TASK_GET_NAME_METHODDEF    \
+    {"get_name", (PyCFunction)_asyncio_Task_get_name, METH_NOARGS, _asyncio_Task_get_name__doc__},
+
+static PyObject *
+_asyncio_Task_get_name_impl(TaskObj *self);
+
+static PyObject *
+_asyncio_Task_get_name(TaskObj *self, PyObject *Py_UNUSED(ignored))
+{
+    return _asyncio_Task_get_name_impl(self);
+}
+
+PyDoc_STRVAR(_asyncio_Task_set_name__doc__,
+"set_name($self, value, /)\n"
+"--\n"
+"\n");
+
+#define _ASYNCIO_TASK_SET_NAME_METHODDEF    \
+    {"set_name", (PyCFunction)_asyncio_Task_set_name, METH_O, _asyncio_Task_set_name__doc__},
+
 PyDoc_STRVAR(_asyncio__get_running_loop__doc__,
 "_get_running_loop($module, /)\n"
 "--\n"
@@ -711,4 +738,4 @@ _asyncio__leave_task(PyObject *module, PyObject *const *args, Py_ssize_t nargs,
 exit:
     return return_value;
 }
-/*[clinic end generated code: output=b6148b0134e7a819 input=a9049054013a1b77]*/
+/*[clinic end generated code: output=67da879c9f841505 input=a9049054013a1b77]*/