diff --git a/mode/threads.py b/mode/threads.py index 1cff5fe..5c1573a 100644 --- a/mode/threads.py +++ b/mode/threads.py @@ -20,6 +20,7 @@ from .services import Service from .utils.futures import ( + all_tasks, maybe_async, maybe_set_exception, maybe_set_result, @@ -109,6 +110,10 @@ def __init__( raise NotImplementedError("executor argument no longer supported") self.parent_loop = loop or get_event_loop() self.thread_loop = thread_loop or asyncio.new_event_loop() + # Only close the loop on shutdown if we created it ourselves -- + # a caller who passed their own thread_loop= may expect to keep + # managing its lifecycle. + self._close_thread_loop_on_stop = thread_loop is None self._thread_started = Event(loop=self.parent_loop) if Worker is not None: self.Worker = Worker @@ -209,6 +214,9 @@ def _start_thread(self) -> None: # shutdown here, since _shutdown_thread will not execute. self.set_shutdown() raise + finally: + if self._close_thread_loop_on_stop: + self.thread_loop.close() async def stop(self) -> None: if self._started.is_set(): @@ -241,6 +249,30 @@ async def _shutdown_thread(self) -> None: await self._default_stop_futures() await self._default_stop_exit_stacks() + async def _gather_remaining_tasks(self) -> None: + # _default_stop_futures() above only cancels/awaits futures that + # were registered with *this* Service via add_future()/@Service.task. + # A driver library running inside this thread (e.g. a Kafka client's + # internal consumer poll/heartbeat loop) may schedule tasks directly + # on thread_loop, bypassing that tracking entirely. Left unswept, + # those tasks are simply abandoned the moment this coroutine returns + # and thread_loop stops running -- producing "coroutine ... was + # never awaited" / "Task was destroyed but it is pending" warnings + # at some later, unpredictable point (e.g. interpreter shutdown). + # Mirrors Worker._gather_all(), which does the same sweep for the + # main loop. + current = asyncio.current_task() + tasks = { + task + for task in all_tasks(loop=self.thread_loop) + if task is not current and not task.done() + } + if not tasks: + return + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + async def _serve(self) -> None: try: # start the service @@ -260,6 +292,7 @@ async def _serve(self) -> None: raise finally: await self._shutdown_thread() + await self._gather_remaining_tasks() @Service.task async def _thread_keepalive(self) -> None: diff --git a/tests/functional/test_threads.py b/tests/functional/test_threads.py new file mode 100644 index 0000000..6c0a2ea --- /dev/null +++ b/tests/functional/test_threads.py @@ -0,0 +1,65 @@ +import asyncio + +import pytest + +from mode.threads import ServiceThread + + +class UntrackedTaskThread(ServiceThread): + """Mimics a driver library (e.g. a Kafka client) scheduling a task + directly on ``thread_loop``, bypassing Service's own + add_future()/@Service.task tracking entirely. + """ + + debug_task: asyncio.Task + + async def on_thread_started(self) -> None: + self.debug_task = asyncio.ensure_future(self._untracked_sleep()) + + async def _untracked_sleep(self) -> None: + await asyncio.sleep(5.0) + + +@pytest.mark.asyncio +async def test_ServiceThread__sweeps_untracked_tasks_on_stop(): + # Regression test for #54: a task scheduled on thread_loop outside of + # Service's own future-tracking must still be cancelled and awaited to + # completion when the thread stops, instead of being silently + # abandoned -- which later triggers "coroutine ... was never awaited" / + # "Task was destroyed but it is pending" warnings at some + # unpredictable later point (e.g. interpreter shutdown), whenever the + # GC happens to finalize the orphaned task/coroutine object. Asserting + # directly on the task's own state (rather than trying to catch that + # GC-timing-dependent warning) is what actually proves the sweep ran. + thread = UntrackedTaskThread() + await thread.start() + await asyncio.sleep(0.2) # let on_thread_started's task actually begin + assert not thread.debug_task.done() + + await thread.stop() + + assert thread.debug_task.done() + assert thread.debug_task.cancelled() + + +@pytest.mark.asyncio +async def test_ServiceThread__closes_self_created_thread_loop(): + thread = ServiceThread() + await thread.start() + loop = thread.thread_loop + await thread.stop() + + assert loop.is_closed() + + +@pytest.mark.asyncio +async def test_ServiceThread__does_not_close_externally_provided_thread_loop(): + external_loop = asyncio.new_event_loop() + try: + thread = ServiceThread(thread_loop=external_loop) + await thread.start() + await thread.stop() + + assert not external_loop.is_closed() + finally: + external_loop.close()