Sweep untracked tasks and close thread_loop on ServiceThread shutdown

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:

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()