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

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

6 files changed, 117 insertions, 15 deletions

diff --git a/Lib/asyncio/base_events.py b/Lib/asyncio/base_events.py
index 78fe2a719f..ee13d1a780 100644
--- a/Lib/asyncio/base_events.py
+++ b/Lib/asyncio/base_events.py
@@ -384,18 +384,20 @@ class BaseEventLoop(events.AbstractEventLoop):
         """Create a Future object attached to the loop."""
         return futures.Future(loop=self)
 
-    def create_task(self, coro):
+    def create_task(self, coro, *, name=None):
         """Schedule a coroutine object.
 
         Return a task object.
         """
         self._check_closed()
         if self._task_factory is None:
-            task = tasks.Task(coro, loop=self)
+            task = tasks.Task(coro, loop=self, name=name)
             if task._source_traceback:
                 del task._source_traceback[-1]
         else:
             task = self._task_factory(self, coro)
+            tasks._set_task_name(task, name)
+
         return task
 
     def set_task_factory(self, factory):
diff --git a/Lib/asyncio/base_tasks.py b/Lib/asyncio/base_tasks.py
index 3ce51f6a98..e2da462fde 100644
--- a/Lib/asyncio/base_tasks.py
+++ b/Lib/asyncio/base_tasks.py
@@ -12,11 +12,13 @@ def _task_repr_info(task):
         # replace status
         info[0] = 'cancelling'
 
+    info.insert(1, 'name=%r' % task.get_name())
+
     coro = coroutines._format_coroutine(task._coro)
-    info.insert(1, f'coro=<{coro}>')
+    info.insert(2, f'coro=<{coro}>')
 
     if task._fut_waiter is not None:
-        info.insert(2, f'wait_for={task._fut_waiter!r}')
+        info.insert(3, f'wait_for={task._fut_waiter!r}')
     return info
 
 
diff --git a/Lib/asyncio/events.py b/Lib/asyncio/events.py
index e4e632206a..58a60a0b16 100644
--- a/Lib/asyncio/events.py
+++ b/Lib/asyncio/events.py
@@ -277,7 +277,7 @@ class AbstractEventLoop:
 
     # Method scheduling a coroutine object: create a task.
 
-    def create_task(self, coro):
+    def create_task(self, coro, *, name=None):
         raise NotImplementedError
 
     # Methods for interacting with threads.
diff --git a/Lib/asyncio/tasks.py b/Lib/asyncio/tasks.py
index 72792a25cf..03d71d37f0 100644
--- a/Lib/asyncio/tasks.py
+++ b/Lib/asyncio/tasks.py
@@ -13,6 +13,7 @@ import concurrent.futures
 import contextvars
 import functools
 import inspect
+import itertools
 import types
 import warnings
 import weakref
@@ -23,6 +24,11 @@ from . import events
 from . import futures
 from .coroutines import coroutine
 
+# Helper to generate new task names
+# 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__
+
 
 def current_task(loop=None):
     """Return a currently executed task."""
@@ -48,6 +54,16 @@ def _all_tasks_compat(loop=None):
     return {t for t in _all_tasks if futures._get_loop(t) is loop}
 
 
+def _set_task_name(task, name):
+    if name is not None:
+        try:
+            set_name = task.set_name
+        except AttributeError:
+            pass
+        else:
+            set_name(name)
+
+
 class Task(futures._PyFuture):  # Inherit Python Task implementation
                                 # from a Python Future implementation.
 
@@ -94,7 +110,7 @@ class Task(futures._PyFuture):  # Inherit Python Task implementation
                       stacklevel=2)
         return _all_tasks_compat(loop)
 
-    def __init__(self, coro, *, loop=None):
+    def __init__(self, coro, *, loop=None, name=None):
         super().__init__(loop=loop)
         if self._source_traceback:
             del self._source_traceback[-1]
@@ -104,6 +120,11 @@ class Task(futures._PyFuture):  # Inherit Python Task implementation
             self._log_destroy_pending = False
             raise TypeError(f"a coroutine was expected, got {coro!r}")
 
+        if name is None:
+            self._name = f'Task-{_task_name_counter()}'
+        else:
+            self._name = str(name)
+
         self._must_cancel = False
         self._fut_waiter = None
         self._coro = coro
@@ -126,6 +147,12 @@ class Task(futures._PyFuture):  # Inherit Python Task implementation
     def _repr_info(self):
         return base_tasks._task_repr_info(self)
 
+    def get_name(self):
+        return self._name
+
+    def set_name(self, value):
+        self._name = str(value)
+
     def set_result(self, result):
         raise RuntimeError('Task does not support set_result operation')
 
@@ -312,13 +339,15 @@ else:
     Task = _CTask = _asyncio.Task
 
 
-def create_task(coro):
+def create_task(coro, *, name=None):
     """Schedule the execution of a coroutine object in a spawn task.
 
     Return a Task object.
     """
     loop = events.get_running_loop()
-    return loop.create_task(coro)
+    task = loop.create_task(coro)
+    _set_task_name(task, name)
+    return task
 
 
 # wait() and as_completed() similar to those in PEP 3148.
diff --git a/Lib/test/test_asyncio/test_base_events.py b/Lib/test/test_asyncio/test_base_events.py
index f3ae140448..e108637953 100644
--- a/Lib/test/test_asyncio/test_base_events.py
+++ b/Lib/test/test_asyncio/test_base_events.py
@@ -825,6 +825,34 @@ class BaseEventLoopTests(test_utils.TestCase):
         task._log_destroy_pending = False
         coro.close()
 
+    def test_create_named_task_with_default_factory(self):
+        async def test():
+            pass
+
+        loop = asyncio.new_event_loop()
+        task = loop.create_task(test(), name='test_task')
+        try:
+            self.assertEqual(task.get_name(), 'test_task')
+        finally:
+            loop.run_until_complete(task)
+            loop.close()
+
+    def test_create_named_task_with_custom_factory(self):
+        def task_factory(loop, coro):
+            return asyncio.Task(coro, loop=loop)
+
+        async def test():
+            pass
+
+        loop = asyncio.new_event_loop()
+        loop.set_task_factory(task_factory)
+        task = loop.create_task(test(), name='test_task')
+        try:
+            self.assertEqual(task.get_name(), 'test_task')
+        finally:
+            loop.run_until_complete(task)
+            loop.close()
+
     def test_run_forever_keyboard_interrupt(self):
         # Python issue #22601: ensure that the temporary task created by
         # run_forever() consumes the KeyboardInterrupt and so don't log
diff --git a/Lib/test/test_asyncio/test_tasks.py b/Lib/test/test_asyncio/test_tasks.py
index a5442f5fdf..c9305936de 100644
--- a/Lib/test/test_asyncio/test_tasks.py
+++ b/Lib/test/test_asyncio/test_tasks.py
@@ -87,8 +87,8 @@ class BaseTaskTests:
     Task = None
     Future = None
 
-    def new_task(self, loop, coro):
-        return self.__class__.Task(coro, loop=loop)
+    def new_task(self, loop, coro, name='TestTask'):
+        return self.__class__.Task(coro, loop=loop, name=name)
 
     def new_future(self, loop):
         return self.__class__.Future(loop=loop)
@@ -295,12 +295,12 @@ class BaseTaskTests:
         coro = format_coroutine(coro_qualname, 'running', src,
                                 t._source_traceback, generator=True)
         self.assertEqual(repr(t),
-                         '<Task pending %s cb=[<Dummy>()]>' % coro)
+                         "<Task pending name='TestTask' %s cb=[<Dummy>()]>" % coro)
 
         # test cancelling Task
         t.cancel()  # Does not take immediate effect!
         self.assertEqual(repr(t),
-                         '<Task cancelling %s cb=[<Dummy>()]>' % coro)
+                         "<Task cancelling name='TestTask' %s cb=[<Dummy>()]>" % coro)
 
         # test cancelled Task
         self.assertRaises(asyncio.CancelledError,
@@ -308,7 +308,7 @@ class BaseTaskTests:
         coro = format_coroutine(coro_qualname, 'done', src,
                                 t._source_traceback)
         self.assertEqual(repr(t),
-                         '<Task cancelled %s>' % coro)
+                         "<Task cancelled name='TestTask' %s>" % coro)
 
         # test finished Task
         t = self.new_task(self.loop, notmuch())
@@ -316,7 +316,36 @@ class BaseTaskTests:
         coro = format_coroutine(coro_qualname, 'done', src,
                                 t._source_traceback)
         self.assertEqual(repr(t),
-                         "<Task finished %s result='abc'>" % coro)
+                         "<Task finished name='TestTask' %s result='abc'>" % coro)
+
+    def test_task_repr_autogenerated(self):
+        @asyncio.coroutine
+        def notmuch():
+            return 123
+
+        t1 = self.new_task(self.loop, notmuch(), None)
+        t2 = self.new_task(self.loop, notmuch(), None)
+        self.assertNotEqual(repr(t1), repr(t2))
+
+        match1 = re.match("^<Task pending name='Task-(\d+)'", repr(t1))
+        self.assertIsNotNone(match1)
+        match2 = re.match("^<Task pending name='Task-(\d+)'", repr(t2))
+        self.assertIsNotNone(match2)
+
+        # Autogenerated task names should have monotonically increasing numbers
+        self.assertLess(int(match1.group(1)), int(match2.group(1)))
+        self.loop.run_until_complete(t1)
+        self.loop.run_until_complete(t2)
+
+    def test_task_repr_name_not_str(self):
+        @asyncio.coroutine
+        def notmuch():
+            return 123
+
+        t = self.new_task(self.loop, notmuch())
+        t.set_name({6})
+        self.assertEqual(t.get_name(), '{6}')
+        self.loop.run_until_complete(t)
 
     def test_task_repr_coro_decorator(self):
         self.loop.set_debug(False)
@@ -376,7 +405,7 @@ class BaseTaskTests:
                                 t._source_traceback,
                                 generator=not coroutines._DEBUG)
         self.assertEqual(repr(t),
-                         '<Task pending %s cb=[<Dummy>()]>' % coro)
+                         "<Task pending name='TestTask' %s cb=[<Dummy>()]>" % coro)
         self.loop.run_until_complete(t)
 
     def test_task_repr_wait_for(self):
@@ -2260,6 +2289,18 @@ class BaseTaskTests:
 
         self.loop.run_until_complete(coro())
 
+    def test_bare_create_named_task(self):
+
+        async def coro_noop():
+            pass
+
+        async def coro():
+            task = asyncio.create_task(coro_noop(), name='No-op')
+            self.assertEqual(task.get_name(), 'No-op')
+            await task
+
+        self.loop.run_until_complete(coro())
+
     def test_context_1(self):
         cvar = contextvars.ContextVar('cvar', default='nope')