split setting task name during function declaration stage

Author: alex2211-put

ydb/_topic_common/common.py

@@ -29,11 +29,13 @@ def wrapper(rpc_state, response_pb, driver=None):
 
     return wrapper
 
-
-def wrap_create_asyncio_task(func: typing.Callable, task_name: str, *args, **kwargs):
-    if sys.hexversion < 0x03080000:
-        return asyncio.create_task(func(*args, **kwargs))
-    return asyncio.create_task(func(*args, **kwargs), name=task_name)
+if sys.hexversion < 0x03080000:
+    def wrap_set_name_for_asyncio_task(task: asyncio.Task, task_name: str) -> asyncio.Task:
+        task.set_name(task_name)
+        return task
+else:
+    def wrap_set_name_for_asyncio_task(task: asyncio.Task, task_name: str) -> asyncio.Task:
+        return task
 
 
 _shared_event_loop_lock = threading.Lock()

ydb/_topic_common/common_test.py

@@ -1,12 +1,13 @@
 import asyncio
+import sys
 import threading
 import time
 import typing
 
 import grpc
 import pytest
 
-from .common import CallFromSyncToAsync
+from .common import CallFromSyncToAsync, wrap_set_name_for_asyncio_task
 from .._grpc.grpcwrapper.common_utils import (
     GrpcWrapperAsyncIO,
     ServerStatus,
@@ -75,6 +76,21 @@ async def async_failed():
         with pytest.raises(TestError):
             await callback_from_asyncio(async_failed)
 
+    async def test_task_name_on_asyncio_task(self):
+        task_name = "asyncio task"
+        loop = asyncio.get_running_loop()
+
+        async def some_async_task():
+            await asyncio.sleep(0)
+            return 1
+
+        asyncio_task = loop.create_task(some_async_task())
+        wrap_set_name_for_asyncio_task(asyncio_task, task_name=task_name)
+
+        if sys.hexversion >= 0x03080000:
+            assert asyncio_task.get_name() == task_name
+
+
 
 @pytest.mark.asyncio
 class TestGrpcWrapperAsyncIO:

ydb/_topic_reader/topic_reader_asyncio.py

@@ -3,7 +3,6 @@
 import asyncio
 import concurrent.futures
 import gzip
-import sys
 import typing
 from asyncio import Task
 from collections import deque
@@ -89,10 +88,8 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
 
     def __del__(self):
         if not self._closed:
-            if sys.hexversion < 0x03080000:
-                self._loop.create_task(self.close(flush=False))
-            else:
-                self._loop.create_task(self.close(flush=False), name="close reader")
+            task = self._loop.create_task(self.close(flush=False))
+            topic_common.wrap_set_name_for_asyncio_task(task, task_name="close reader")
 
     async def wait_message(self):
         """
@@ -343,17 +340,29 @@ async def _start(self, stream: IGrpcWrapperAsyncIO, init_message: StreamReadMess
         self._update_token_event.set()
 
         self._background_tasks.add(
-            topic_common.wrap_create_asyncio_task(self._read_messages_loop, "read_messages_loop"),
+            topic_common.wrap_set_name_for_asyncio_task(
+                asyncio.create_task(self._read_messages_loop()),
+                task_name="read_messages_loop",
+            ),
         )
         self._background_tasks.add(
-            topic_common.wrap_create_asyncio_task(self._decode_batches_loop, "decode_batches"),
+            topic_common.wrap_set_name_for_asyncio_task(
+                asyncio.create_task(self._decode_batches_loop()),
+                task_name="decode_batches",
+            ),
         )
         if self._get_token_function:
             self._background_tasks.add(
-                topic_common.wrap_create_asyncio_task(self._update_token_loop, "update_token_loop"),
+                topic_common.wrap_set_name_for_asyncio_task(
+                    asyncio.create_task(self._update_token_loop()),
+                    task_name="update_token_loop",
+                ),
             )
         self._background_tasks.add(
-            topic_common.wrap_create_asyncio_task(self._handle_background_errors, "handle_background_errors"),
+            topic_common.wrap_set_name_for_asyncio_task(
+                asyncio.create_task(self._handle_background_errors()),
+                task_name="handle_background_errors",
+            ),
         )
 
     async def wait_error(self):

ydb/_topic_writer/topic_writer_asyncio.py

@@ -232,8 +232,14 @@ def __init__(self, driver: SupportedDriverType, settings: WriterSettings):
         self._new_messages = asyncio.Queue()
         self._stop_reason = self._loop.create_future()
         self._background_tasks = [
-            topic_common.wrap_create_asyncio_task(self._connection_loop, "connection_loop"),
-            topic_common.wrap_create_asyncio_task(self._encode_loop, "encode_loop"),
+            topic_common.wrap_set_name_for_asyncio_task(
+                asyncio.create_task(self._connection_loop()),
+                task_name="connection_loop",
+            ),
+            topic_common.wrap_set_name_for_asyncio_task(
+                asyncio.create_task(self._encode_loop()),
+                task_name="encode_loop",
+            ),
         ]
 
         self._state_changed = asyncio.Event()
@@ -367,11 +373,13 @@ async def _connection_loop(self):
 
                 self._stream_connected.set()
 
-                send_loop = topic_common.wrap_create_asyncio_task(self._send_loop, "writer send loop", stream_writer)
-                receive_loop = topic_common.wrap_create_asyncio_task(
-                    self._read_loop,
-                    "writer receive loop",
-                    stream_writer,
+                send_loop = topic_common.wrap_set_name_for_asyncio_task(
+                    asyncio.create_task(self._send_loop(stream_writer)),
+                    task_name="writer send loop",
+                )
+                receive_loop = topic_common.wrap_set_name_for_asyncio_task(
+                    asyncio.create_task(self._read_loop(stream_writer)),
+                    task_name="writer receive loop",
                 )
 
                 tasks = [send_loop, receive_loop]
@@ -658,9 +666,9 @@ async def _start(self, stream: IGrpcWrapperAsyncIO, init_message: StreamWriteMes
 
         if self._update_token_interval is not None:
             self._update_token_event.set()
-            self._update_token_task = topic_common.wrap_create_asyncio_task(
-                self._update_token_loop,
-                "update_token_loop",
+            self._update_token_task = topic_common.wrap_set_name_for_asyncio_task(
+                asyncio.create_task(self._update_token_loop()),
+                task_name="update_token_loop",
             )
 
     @staticmethod