diff --git a/ydb/_topic_common/common.py b/ydb/_topic_common/common.py index 678f5ff61..7b7c6532a 100644 --- a/ydb/_topic_common/common.py +++ b/ydb/_topic_common/common.py @@ -1,7 +1,9 @@ from __future__ import annotations import asyncio +import atexit import concurrent.futures +import logging import threading import typing from typing import Optional @@ -11,6 +13,8 @@ TimeoutType = typing.Union[int, float, None] +logger = logging.getLogger(__name__) + def wrap_operation(rpc_state, response_pb, driver=None): return operation.Operation(rpc_state, response_pb, driver) @@ -33,9 +37,13 @@ def wrapper(rpc_state, response_pb, driver=None): _shared_event_loop_lock = threading.Lock() _shared_event_loop: Optional[asyncio.AbstractEventLoop] = None +_shared_event_loop_thread: Optional[threading.Thread] = None +_shared_event_loop_atexit_registered = False def _get_shared_event_loop() -> asyncio.AbstractEventLoop: + global _shared_event_loop_thread, _shared_event_loop_atexit_registered + if _shared_event_loop is not None: return _shared_event_loop @@ -56,7 +64,19 @@ def on_loop_started(): loop_ready.set() event_loop.call_soon(on_loop_started) - event_loop.run_forever() + try: + event_loop.run_forever() + finally: + try: + pending = asyncio.all_tasks(event_loop) + for task in pending: + task.cancel() + if pending: + event_loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True)) + except Exception: + logger.debug("Error while cancelling shared event loop tasks", exc_info=True) + finally: + event_loop.close() t = threading.Thread( target=start_event_loop, @@ -70,9 +90,49 @@ def on_loop_started(): if _shared_event_loop is None: raise RuntimeError("Event loop was not properly initialized") + _shared_event_loop_thread = t + + if not _shared_event_loop_atexit_registered: + atexit.register(_shutdown_shared_event_loop) + _shared_event_loop_atexit_registered = True + return _shared_event_loop +def _shutdown_shared_event_loop(timeout: TimeoutType = 30) -> None: + """ + Stop the shared topic asyncio loop thread and wait for it to exit. + + Sync topic/query clients run coroutines on a process-wide background loop. + Closing writers/readers cancels their tasks, but leaves the loop thread in + run_forever()/epoll. Joining it before interpreter finalization avoids a + TSan data race between Py_FinalizeEx and that worker thread. + """ + global _shared_event_loop, _shared_event_loop_thread + + with _shared_event_loop_lock: + loop = _shared_event_loop + thread = _shared_event_loop_thread + _shared_event_loop = None + _shared_event_loop_thread = None + + if loop is None: + return + + try: + if loop.is_running(): + loop.call_soon_threadsafe(loop.stop) + except RuntimeError: + logger.debug("Shared event loop already stopped", exc_info=True) + + # Join under the lock so a concurrent _get_shared_event_loop() cannot + # start a replacement thread while this one is still exiting. + if thread is not None and thread.is_alive(): + thread.join(timeout=timeout) + if thread.is_alive(): + logger.warning("Shared ydb topic event loop thread did not stop in time") + + class CallFromSyncToAsync: _loop: asyncio.AbstractEventLoop diff --git a/ydb/_topic_common/common_test.py b/ydb/_topic_common/common_test.py index b31f9af9a..93563db4d 100644 --- a/ydb/_topic_common/common_test.py +++ b/ydb/_topic_common/common_test.py @@ -6,7 +6,7 @@ import grpc import pytest -from .common import CallFromSyncToAsync +from .common import CallFromSyncToAsync, _get_shared_event_loop, _shutdown_shared_event_loop from .._grpc.grpcwrapper.common_utils import ( GrpcWrapperAsyncIO, ServerStatus, @@ -288,3 +288,190 @@ def callback(): with pytest.raises(TestError): caller.call_sync(callback) assert callback_eventloop is separate_loop + + +def _shared_loop_threads_alive() -> bool: + return any(t.name == "Common ydb topic event loop" and t.is_alive() for t in threading.enumerate()) + + +def _force_stop_real_shared_loop(): + """Ensure module globals do not retain a live shared loop between mocked tests.""" + import ydb._topic_common.common as common + + loop = common._shared_event_loop + thread = common._shared_event_loop_thread + common._shared_event_loop = None + common._shared_event_loop_thread = None + if loop is not None: + try: + if loop.is_running(): + loop.call_soon_threadsafe(loop.stop) + except RuntimeError: + pass + if thread is not None and thread.is_alive(): + thread.join(timeout=5) + + +class TestSharedEventLoop: + def teardown_method(self): + _shutdown_shared_event_loop() + _force_stop_real_shared_loop() + + def test_shutdown_joins_thread(self): + loop = _get_shared_event_loop() + assert _shared_loop_threads_alive() + + fut = asyncio.run_coroutine_threadsafe(asyncio.sleep(0.01), loop) + assert fut.result(1) is None + + _shutdown_shared_event_loop() + assert not _shared_loop_threads_alive() + + def test_shutdown_is_idempotent_when_unused(self): + _shutdown_shared_event_loop() + _shutdown_shared_event_loop() + assert not _shared_loop_threads_alive() + + def test_recreate_after_shutdown(self): + loop = _get_shared_event_loop() + fut = asyncio.run_coroutine_threadsafe(asyncio.sleep(0.01), loop) + assert fut.result(1) is None + + _shutdown_shared_event_loop() + assert not _shared_loop_threads_alive() + + loop2 = _get_shared_event_loop() + assert loop2 is not loop + assert _shared_loop_threads_alive() + + fut2 = asyncio.run_coroutine_threadsafe(asyncio.sleep(0.01), loop2) + assert fut2.result(1) is None + + _shutdown_shared_event_loop() + assert not _shared_loop_threads_alive() + + def test_get_returns_same_loop(self): + loop1 = _get_shared_event_loop() + loop2 = _get_shared_event_loop() + assert loop1 is loop2 + + def test_shutdown_cancels_pending_tasks(self): + loop = _get_shared_event_loop() + started = threading.Event() + cancelled = threading.Event() + + async def long_running(): + started.set() + try: + await asyncio.sleep(3600) + except asyncio.CancelledError: + cancelled.set() + raise + + asyncio.run_coroutine_threadsafe(long_running(), loop) + assert started.wait(timeout=1) + + _shutdown_shared_event_loop() + assert not _shared_loop_threads_alive() + assert cancelled.wait(timeout=1) + + def test_shutdown_skips_stop_when_loop_not_running(self): + import ydb._topic_common.common as common + + _force_stop_real_shared_loop() + + class FakeLoop: + def __init__(self): + self.stop_called = False + + def is_running(self): + return False + + def call_soon_threadsafe(self, callback): + self.stop_called = True + + loop = FakeLoop() + common._shared_event_loop = loop + common._shared_event_loop_thread = None + + _shutdown_shared_event_loop() + + assert not loop.stop_called + assert common._shared_event_loop is None + + def test_shutdown_handles_stop_runtime_error(self, caplog): + import logging + + import ydb._topic_common.common as common + + _force_stop_real_shared_loop() + + class FakeLoop: + def is_running(self): + return True + + def stop(self): + return None + + def call_soon_threadsafe(self, callback): + raise RuntimeError("loop is closed") + + common._shared_event_loop = FakeLoop() + common._shared_event_loop_thread = None + + with caplog.at_level(logging.DEBUG): + _shutdown_shared_event_loop() + + assert common._shared_event_loop is None + assert any("already stopped" in r.message for r in caplog.records) + + def test_shutdown_warns_when_thread_does_not_stop(self, caplog): + import logging + + import ydb._topic_common.common as common + + _force_stop_real_shared_loop() + + class FakeLoop: + def is_running(self): + return True + + def stop(self): + return None + + def call_soon_threadsafe(self, callback): + return None + + class StuckThread: + def is_alive(self): + return True + + def join(self, timeout=None): + return None + + common._shared_event_loop = FakeLoop() + common._shared_event_loop_thread = StuckThread() + + with caplog.at_level(logging.WARNING): + _shutdown_shared_event_loop(timeout=0.01) + + assert common._shared_event_loop is None + assert any("did not stop in time" in r.message for r in caplog.records) + + def test_cleanup_logs_when_cancelling_tasks_fails(self, monkeypatch, caplog): + import logging + + loop = _get_shared_event_loop() + fut = asyncio.run_coroutine_threadsafe(asyncio.sleep(0.01), loop) + assert fut.result(1) is None + + def boom(_event_loop): + raise RuntimeError("all_tasks failed") + + monkeypatch.setattr(asyncio, "all_tasks", boom) + + with caplog.at_level(logging.DEBUG): + _shutdown_shared_event_loop() + + assert not _shared_loop_threads_alive() + assert any("cancelling shared event loop tasks" in r.message for r in caplog.records)