Add names to all asyncio tasks (#152)

* Attach names to all tasks

* Name task in wait_for

* Cover the kick event too

* Fix copy paste error

* Only run on python 312 for task naming check

Author: AlanCoding

dispatcherd/brokers/socket.py

@@ -6,6 +6,7 @@
 from typing import Any, AsyncGenerator, Callable, Coroutine, Iterator, Optional, Union
 
 from ..protocols import Broker as BrokerProtocol
+from ..service.asyncio_tasks import named_wait
 
 logger = logging.getLogger(__name__)
 
@@ -77,6 +78,9 @@ async def _add_client(self, reader: asyncio.StreamReader, writer: asyncio.Stream
         client = Client(self.client_ct, reader, writer)
         self.clients[self.client_ct] = client
         self.client_ct += 1
+        current_task = asyncio.current_task()
+        if current_task is not None:
+            current_task.set_name(f'socket_client_task_{client.client_id}')
         logger.info(f'Socket client_id={client.client_id} is connected')
 
         try:
@@ -89,7 +93,7 @@ async def _add_client(self, reader: asyncio.StreamReader, writer: asyncio.Stream
                 await self.incoming_queue.put((client.client_id, message))
                 # Wait for caller to potentially fill a reply queue
                 # this should realistically never take more than a trivial amount of time
-                await asyncio.wait_for(client.yield_clear.wait(), timeout=2)
+                await asyncio.wait_for(named_wait(client.yield_clear, f'internal_wait_for_client_{client.client_id}'), timeout=2)
                 client.yield_clear.clear()
                 await client.send_replies()
         except asyncio.TimeoutError:

dispatcherd/producers/base.py

@@ -10,6 +10,7 @@ def __init__(self) -> None:
 
 
 class BaseProducer(ProducerProtocol):
+    can_recycle: bool = False
 
     def __init__(self) -> None:
         self.events = ProducerEvents()

dispatcherd/producers/brokered.py

@@ -9,6 +9,8 @@
 
 
 class BrokeredProducer(BaseProducer):
+    can_recycle = True
+
     def __init__(self, broker: Broker) -> None:
         self.production_task: Optional[asyncio.Task] = None
         self.broker = broker

dispatcherd/protocols.py

@@ -75,6 +75,7 @@ class Producer(Protocol):
     """
 
     events: ProducerEvents
+    can_recycle: bool
 
     async def start_producing(self, dispatcher: 'DispatcherMain') -> None:
         """Starts tasks which will eventually call DispatcherMain.process_message - how tasks originate in the service"""

dispatcherd/service/asyncio_tasks.py

@@ -36,12 +36,15 @@ def ensure_fatal(task: asyncio.Task, exit_event: Optional[asyncio.Event] = None)
     return task  # nicety so this can be used as a wrapper
 
 
-async def wait_for_any(events: Iterable[asyncio.Event]) -> int:
+async def wait_for_any(events: Iterable[asyncio.Event], names: Optional[Iterable[str]] = None) -> int:
     """
     Wait for a list of events. If any of the events gets set, this function
     will return
     """
-    tasks = [asyncio.create_task(event.wait()) for event in events]
+    if names:
+        tasks = [asyncio.create_task(event.wait(), name=task_name) for (event, task_name) in zip(events, names)]
+    else:
+        tasks = [asyncio.create_task(event.wait()) for event in events]
     done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
     for task in pending:
         task.cancel()
@@ -53,3 +56,12 @@ async def wait_for_any(events: Iterable[asyncio.Event]) -> int:
             return i
 
     raise RuntimeError('Internal error - could done find any tasks that are done')
+
+
+async def named_wait(event: asyncio.Event, name: str) -> None:
+    """Add a name to waiting task so it is visible via debugging commands"""
+    current_task = asyncio.current_task()
+    if current_task:
+        current_task.set_name(name)
+
+    await event.wait()

dispatcherd/service/main.py

@@ -242,6 +242,8 @@ async def cancel_tasks(self) -> None:
     async def recycle_broker_producers(self) -> None:
         """For any producer in a broken state (likely due to external factors beyond our control) recycle it"""
         for producer in self.producers:
+            if not producer.can_recycle:
+                continue
             if producer.events.recycle_event.is_set():
                 await producer.recycle()
                 for task in producer.all_tasks():
@@ -251,12 +253,20 @@ async def recycle_broker_producers(self) -> None:
     async def main_loop_wait(self) -> None:
         """Wait for an event that requires some kind of action by the main loop"""
         events = [self.events.exit_event]
+        names = ['exit_event_wait']
         for producer in self.producers:
+            if not producer.can_recycle:
+                continue
             events.append(producer.events.recycle_event)
+            names.append(f'{str(producer)}_recycle_event_wait')
 
-        await wait_for_any(events)
+        await wait_for_any(events, names=names)
 
     async def main(self) -> None:
+        current_task = asyncio.current_task()
+        if current_task is not None:
+            current_task.set_name('dispatcherd_service_main')
+
         await self.connect_signals()
 
         try:

dispatcherd/service/next_wakeup_runner.py

@@ -4,7 +4,7 @@
 from abc import abstractmethod
 from typing import Any, Callable, Coroutine, Iterable, Optional
 
-from .asyncio_tasks import ensure_fatal
+from .asyncio_tasks import ensure_fatal, named_wait
 
 logger = logging.getLogger(__name__)
 
@@ -92,7 +92,7 @@ async def background_task(self) -> None:
                 delta = 0.1
 
             try:
-                await asyncio.wait_for(self.kick_event.wait(), timeout=delta)
+                await asyncio.wait_for(named_wait(self.kick_event, f'{self.name}_kick_event_wait'), timeout=delta)
             except asyncio.TimeoutError:
                 pass  # intended mechanism to hit the next schedule
             except asyncio.CancelledError:

tests/integration/conftest.py

@@ -0,0 +1,9 @@
+import sys
+
+import pytest
+
+
+@pytest.fixture
+def python312():
+    if sys.version_info < (3, 12):
+        pytest.skip("test requires python 3.12")

tests/integration/test_main.py

@@ -6,6 +6,7 @@
 import pytest
 
 from dispatcherd.config import temporary_settings
+from dispatcherd.service.control_tasks import _stack_from_task
 from tests.data import methods as test_methods
 
 SLEEP_METHOD = 'lambda: __import__("time").sleep(0.1)'
@@ -293,3 +294,18 @@ async def test_scale_up(apg_dispatcher, test_settings):
             break
     else:
         assert f'Never scaled up to expected 6 workers, have: {apg_dispatcher.pool.workers}'
+
+
+@pytest.mark.asyncio
+async def test_tasks_are_named(apg_dispatcher, python312):
+    wait_task = asyncio.create_task(apg_dispatcher.main_loop_wait(), name='this_is_for_test')
+
+    current_task = asyncio.current_task()
+    for task in asyncio.all_tasks():
+        if task is current_task:
+            continue
+        task_name = task.get_name()
+        assert not task_name.startswith('Task-'), _stack_from_task(task)
+
+    apg_dispatcher.events.exit_event.set()
+    await wait_task

tests/integration/test_socket_use.py

@@ -10,6 +10,7 @@
 from dispatcherd.control import Control
 from dispatcherd.factories import from_settings, get_control_from_settings, get_publisher_from_settings
 from dispatcherd.protocols import DispatcherMain
+from dispatcherd.service.control_tasks import _stack_from_task
 
 logger = logging.getLogger(__name__)
 
@@ -84,3 +85,22 @@ def alive_cmd():
     data = alive[0]
 
     assert data['node_id'] == 'socket-test-server'
+
+
+@pytest.mark.asyncio
+async def test_socket_tasks_are_named(asock_dispatcher, sock_control, python312):
+    loop = asyncio.get_event_loop()
+
+    def aio_tasks_cmd():
+        return sock_control.control_with_reply('aio_tasks')
+
+    aio_tasks = await loop.run_in_executor(None, aio_tasks_cmd)
+
+    current_task_name = asyncio.current_task().get_name()
+
+    assert len(aio_tasks) == 1
+    data = aio_tasks[0]
+    for task_name, task_stuff in data.items():
+        if task_name == current_task_name:
+            continue
+        assert not task_name.startswith('Task-'), task_stuff['stack']