Skip to content

Commit c84538d

Browse files
authored
Removing loop name from asyncio.create_task (#489)
* Update topic_writer_asyncio.py * Update topic_reader_asyncio.py * Update topic_writer_asyncio.py * Update topic_reader_asyncio.py fix linters * Update topic_reader_asyncio.py * Update topic_reader_asyncio.py * Update topic_reader_asyncio.py * Update topic_reader_asyncio.py * add wrapper for asyncio.create_task * fix linters * fix linters * fix linters * fix linters * fix linters * fix linters * fix linters * split setting task name during function declaration stage * fix tests * fix tests * fix linters * fix linters * fix linters * fix tests
1 parent b0cce07 commit c84538d

File tree

4 files changed

+73
-11
lines changed

4 files changed

+73
-11
lines changed

ydb/_topic_common/common.py

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import asyncio
22
import concurrent.futures
3+
import sys
34
import threading
45
import typing
56
from typing import Optional
@@ -29,6 +30,18 @@ def wrapper(rpc_state, response_pb, driver=None):
2930
return wrapper
3031

3132

33+
if sys.hexversion < 0x03080000:
34+
35+
def wrap_set_name_for_asyncio_task(task: asyncio.Task, task_name: str) -> asyncio.Task:
36+
return task
37+
38+
else:
39+
40+
def wrap_set_name_for_asyncio_task(task: asyncio.Task, task_name: str) -> asyncio.Task:
41+
task.set_name(task_name)
42+
return task
43+
44+
3245
_shared_event_loop_lock = threading.Lock()
3346
_shared_event_loop: Optional[asyncio.AbstractEventLoop] = None
3447

ydb/_topic_common/common_test.py

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
import grpc
77
import pytest
88

9-
from .common import CallFromSyncToAsync
9+
from .common import CallFromSyncToAsync, wrap_set_name_for_asyncio_task
1010
from .._grpc.grpcwrapper.common_utils import (
1111
GrpcWrapperAsyncIO,
1212
ServerStatus,
@@ -75,6 +75,19 @@ async def async_failed():
7575
with pytest.raises(TestError):
7676
await callback_from_asyncio(async_failed)
7777

78+
async def test_task_name_on_asyncio_task(self):
79+
task_name = "asyncio task"
80+
loop = asyncio.get_running_loop()
81+
82+
async def some_async_task():
83+
await asyncio.sleep(0)
84+
return 1
85+
86+
asyncio_task = loop.create_task(some_async_task())
87+
wrap_set_name_for_asyncio_task(asyncio_task, task_name=task_name)
88+
89+
assert asyncio_task.get_name() == task_name
90+
7891

7992
@pytest.mark.asyncio
8093
class TestGrpcWrapperAsyncIO:

ydb/_topic_reader/topic_reader_asyncio.py

Lines changed: 25 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010

1111
import ydb
1212
from .. import _apis, issues
13+
from .._topic_common import common as topic_common
1314
from .._utilities import AtomicCounter
1415
from ..aio import Driver
1516
from ..issues import Error as YdbError, _process_response
@@ -87,7 +88,8 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
8788

8889
def __del__(self):
8990
if not self._closed:
90-
self._loop.create_task(self.close(flush=False), name="close reader")
91+
task = self._loop.create_task(self.close(flush=False))
92+
topic_common.wrap_set_name_for_asyncio_task(task, task_name="close reader")
9193

9294
async def wait_message(self):
9395
"""
@@ -337,12 +339,30 @@ async def _start(self, stream: IGrpcWrapperAsyncIO, init_message: StreamReadMess
337339

338340
self._update_token_event.set()
339341

340-
self._background_tasks.add(asyncio.create_task(self._read_messages_loop(), name="read_messages_loop"))
341-
self._background_tasks.add(asyncio.create_task(self._decode_batches_loop(), name="decode_batches"))
342+
self._background_tasks.add(
343+
topic_common.wrap_set_name_for_asyncio_task(
344+
asyncio.create_task(self._read_messages_loop()),
345+
task_name="read_messages_loop",
346+
),
347+
)
348+
self._background_tasks.add(
349+
topic_common.wrap_set_name_for_asyncio_task(
350+
asyncio.create_task(self._decode_batches_loop()),
351+
task_name="decode_batches",
352+
),
353+
)
342354
if self._get_token_function:
343-
self._background_tasks.add(asyncio.create_task(self._update_token_loop(), name="update_token_loop"))
355+
self._background_tasks.add(
356+
topic_common.wrap_set_name_for_asyncio_task(
357+
asyncio.create_task(self._update_token_loop()),
358+
task_name="update_token_loop",
359+
),
360+
)
344361
self._background_tasks.add(
345-
asyncio.create_task(self._handle_background_errors(), name="handle_background_errors")
362+
topic_common.wrap_set_name_for_asyncio_task(
363+
asyncio.create_task(self._handle_background_errors()),
364+
task_name="handle_background_errors",
365+
),
346366
)
347367

348368
async def wait_error(self):

ydb/_topic_writer/topic_writer_asyncio.py

Lines changed: 21 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
issues,
2929
)
3030
from .._errors import check_retriable_error
31+
from .._topic_common import common as topic_common
3132
from ..retries import RetrySettings
3233
from .._grpc.grpcwrapper.ydb_topic_public_types import PublicCodec
3334
from .._grpc.grpcwrapper.ydb_topic import (
@@ -231,8 +232,14 @@ def __init__(self, driver: SupportedDriverType, settings: WriterSettings):
231232
self._new_messages = asyncio.Queue()
232233
self._stop_reason = self._loop.create_future()
233234
self._background_tasks = [
234-
asyncio.create_task(self._connection_loop(), name="connection_loop"),
235-
asyncio.create_task(self._encode_loop(), name="encode_loop"),
235+
topic_common.wrap_set_name_for_asyncio_task(
236+
asyncio.create_task(self._connection_loop()),
237+
task_name="connection_loop",
238+
),
239+
topic_common.wrap_set_name_for_asyncio_task(
240+
asyncio.create_task(self._encode_loop()),
241+
task_name="encode_loop",
242+
),
236243
]
237244

238245
self._state_changed = asyncio.Event()
@@ -366,8 +373,14 @@ async def _connection_loop(self):
366373

367374
self._stream_connected.set()
368375

369-
send_loop = asyncio.create_task(self._send_loop(stream_writer), name="writer send loop")
370-
receive_loop = asyncio.create_task(self._read_loop(stream_writer), name="writer receive loop")
376+
send_loop = topic_common.wrap_set_name_for_asyncio_task(
377+
asyncio.create_task(self._send_loop(stream_writer)),
378+
task_name="writer send loop",
379+
)
380+
receive_loop = topic_common.wrap_set_name_for_asyncio_task(
381+
asyncio.create_task(self._read_loop(stream_writer)),
382+
task_name="writer receive loop",
383+
)
371384

372385
tasks = [send_loop, receive_loop]
373386
done, _ = await asyncio.wait([send_loop, receive_loop], return_when=asyncio.FIRST_COMPLETED)
@@ -653,7 +666,10 @@ async def _start(self, stream: IGrpcWrapperAsyncIO, init_message: StreamWriteMes
653666

654667
if self._update_token_interval is not None:
655668
self._update_token_event.set()
656-
self._update_token_task = asyncio.create_task(self._update_token_loop(), name="update_token_loop")
669+
self._update_token_task = topic_common.wrap_set_name_for_asyncio_task(
670+
asyncio.create_task(self._update_token_loop()),
671+
task_name="update_token_loop",
672+
)
657673

658674
@staticmethod
659675
def _ensure_ok(message: WriterMessagesFromServerToClient):

0 commit comments

Comments
 (0)