Stop shared topic event loop thread on process shutdown

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
 

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)