Handle the loop closing between ThreadsafeProxy's is_closed() check and dispatch

bellows/thread.py

@@ -2,6 +2,7 @@
 from concurrent.futures import ThreadPoolExecutor
 import contextlib
 import functools
+import inspect
 import logging
 
 LOGGER = logging.getLogger(__name__)
@@ -13,10 +14,25 @@ class EventLoopThread:
     def __init__(self):
         self.loop = None
         self.thread_complete = None
+        # Published by `force_stop()` before it schedules anything. `loop.is_closed()`
+        # only becomes `True` once the loop has actually been closed, so it cannot answer
+        # "will this loop still run what I hand it?" during the shutdown window: between
+        # `force_stop()` and `loop.close()` the loop accepts work it will never run.
+        self.stopping = False
 
     def run_coroutine_threadsafe(self, coroutine):
         current_loop = asyncio.get_event_loop()
-        future = asyncio.run_coroutine_threadsafe(coroutine, self.loop)
+        # Snapshot: the worker thread publishes `None` when it exits
+        loop = self.loop
+        if loop is None or self.stopping:
+            coroutine.close()
+            raise RuntimeError("Event loop is not running")
+        try:
+            future = asyncio.run_coroutine_threadsafe(coroutine, loop)
+        except RuntimeError:
+            # The worker thread may close the loop after our None check
+            coroutine.close()
+            raise
         return asyncio.wrap_future(future, loop=current_loop)
 
     def _thread_main(self, init_task):
@@ -36,6 +52,9 @@ def _thread_main(self, init_task):
     async def start(self):
         current_loop = asyncio.get_event_loop()
         if self.loop is not None and not self.loop.is_closed():
+            # Note this returns a thread that is still stopping, if one is: reusing a loop
+            # that is winding down is not safe, and spawning a second thread while the
+            # first is in its `finally` would have it close the new loop out from under us
             return
 
         executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix=__name__)
@@ -45,6 +64,8 @@ async def start(self):
         async def init_task():
             current_loop.call_soon_threadsafe(thread_started_future.set_result, None)
 
+        self.stopping = False
+
         # Use current loop so current loop has a reference to the long-running thread
         # as one of its tasks
         thread_complete = current_loop.run_in_executor(
@@ -56,6 +77,10 @@ async def init_task():
         return thread_complete
 
     def force_stop(self):
+        # Published before anything is scheduled below: from here on the loop may stop at
+        # any moment, so work handed to it may never run.
+        self.stopping = True
+
         loop = self.loop
         if loop is None or loop.is_closed():
             return
@@ -80,9 +105,12 @@ class ThreadsafeProxy:
     using that object's methods are done on a particular event loop
     """
 
-    def __init__(self, obj, obj_loop):
+    def __init__(self, obj, obj_loop, loop_thread=None):
         self._obj = obj
         self._obj_loop = obj_loop
+        # The `EventLoopThread` running `obj_loop`, when there is one. Only it knows that a
+        # stop has been requested; the loop object itself still looks perfectly usable.
+        self._loop_thread = loop_thread
 
     def __getattr__(self, name):
         func = getattr(self._obj, name)
@@ -99,12 +127,34 @@ def func_wrapper(*args, **kwargs):
             call = functools.partial(func, *args, **kwargs)
             if loop == curr_loop:
                 return call()
+
+            def disconnected_result(message="Attempted to use a closed event loop"):
+                # Disconnected: sync calls are dropped, async calls resolve to None
+                LOGGER.warning(message)
+                if not inspect.iscoroutinefunction(func):
+                    return None
+                future = curr_loop.create_future()
+                future.set_result(None)
+                return future
+
             if loop.is_closed():
-                # Disconnected
-                LOGGER.warning("Attempted to use a closed event loop")
-                return
-            if asyncio.iscoroutinefunction(func):
-                future = asyncio.run_coroutine_threadsafe(call(), loop)
+                return disconnected_result()
+            if self._loop_thread is not None and self._loop_thread.stopping:
+                # The loop is still open, but it has been asked to stop: anything handed
+                # to it from here on may never run, and `run_coroutine_threadsafe()` will
+                # not complain. Without this branch the caller waits on a future that is
+                # never resolved -- a silent, unbounded hang.
+                return disconnected_result(
+                    "Attempted to use an event loop that is shutting down"
+                )
+            if inspect.iscoroutinefunction(func):
+                coro = call()
+                try:
+                    future = asyncio.run_coroutine_threadsafe(coro, loop)
+                except RuntimeError:
+                    # The worker thread may close the loop after our is_closed() check
+                    coro.close()
+                    return disconnected_result()
                 return asyncio.wrap_future(future, loop=curr_loop)
             else:
 
@@ -118,6 +168,10 @@ def check_result_wrapper():
                             ).format(self._obj.__class__.__name__, name)
                         )
 
-                loop.call_soon_threadsafe(check_result_wrapper)
+                try:
+                    loop.call_soon_threadsafe(check_result_wrapper)
+                except RuntimeError:
+                    # The worker thread may close the loop after our is_closed() check
+                    return disconnected_result()
 
         return func_wrapper

bellows/uart.py

@@ -105,7 +105,7 @@ async def reset(self):
             return await self._reset_future
 
 
-async def _connect(config, api):
+async def _connect(config, api, thread=None):
     loop = asyncio.get_event_loop()
 
     connection_done_future = loop.create_future()
@@ -129,7 +129,7 @@ async def _connect(config, api):
 
     await gateway.wait_until_connected()
 
-    thread_safe_protocol = ThreadsafeProxy(gateway, loop)
+    thread_safe_protocol = ThreadsafeProxy(gateway, loop, thread)
     return thread_safe_protocol, connection_done_future
 
 
@@ -140,7 +140,7 @@ async def connect(config, api, use_thread=True):
         await thread.start()
         try:
             protocol, connection_done = await thread.run_coroutine_threadsafe(
-                _connect(config, api)
+                _connect(config, api, thread)
             )
         except Exception:
             thread.force_stop()

tests/test_thread.py

@@ -1,6 +1,8 @@
 import asyncio
 from asyncio import timeout as asyncio_timeout
+import inspect
 import threading
+import time
 from unittest import mock
 
 import pytest
@@ -192,6 +194,94 @@ async def test_proxy_loop_closed():
     assert obj.test.call_count == 0
 
 
+async def test_proxy_loop_closed_async():
+    """An async call through a proxy to a closed loop is awaitable and resolves to None."""
+    loop = asyncio.new_event_loop()
+    obj = mock.MagicMock()
+    call_count = 0
+
+    async def magic():
+        nonlocal call_count
+        call_count += 1
+
+    obj.test = magic
+    proxy = ThreadsafeProxy(obj, loop)
+    loop.close()
+
+    assert await proxy.test() is None
+    assert call_count == 0
+
+
+@pytest.mark.filterwarnings("error::RuntimeWarning")
+async def test_proxy_loop_closed_during_async_dispatch(caplog):
+    """The loop closing between the `is_closed()` check and dispatch is handled."""
+    loop = asyncio.new_event_loop()
+    try:
+        obj = mock.MagicMock()
+        call_count = 0
+
+        async def magic():
+            nonlocal call_count
+            call_count += 1
+
+        obj.test = magic
+        proxy = ThreadsafeProxy(obj, loop)
+        loop.call_soon_threadsafe = mock.Mock(
+            side_effect=RuntimeError("Event loop is closed")
+        )
+
+        assert await proxy.test() is None
+
+        assert call_count == 0
+        assert "Attempted to use a closed event loop" in caplog.text
+    finally:
+        loop.close()
+
+
+async def test_proxy_loop_closed_during_sync_dispatch(caplog):
+    """The loop closing between the `is_closed()` check and dispatch is handled."""
+    loop = asyncio.new_event_loop()
+    try:
+        obj = mock.MagicMock()
+        obj.test.return_value = None
+        proxy = ThreadsafeProxy(obj, loop)
+        loop.call_soon_threadsafe = mock.Mock(
+            side_effect=RuntimeError("Event loop is closed")
+        )
+
+        proxy.test()
+
+        assert obj.test.call_count == 0
+        assert "Attempted to use a closed event loop" in caplog.text
+    finally:
+        loop.close()
+
+
+@pytest.mark.filterwarnings("error::RuntimeWarning")
+async def test_thread_run_coroutine_threadsafe_loop_not_running():
+    """A `RuntimeError` (not `AttributeError`) is raised when the loop is gone."""
+    thread = EventLoopThread()
+    assert thread.loop is None
+
+    with pytest.raises(RuntimeError):
+        # The coroutine is closed internally: no "never awaited" RuntimeWarning
+        thread.run_coroutine_threadsafe(asyncio.sleep(0))
+
+
+@pytest.mark.filterwarnings("error::RuntimeWarning")
+async def test_thread_run_coroutine_threadsafe_loop_closed_mid_dispatch():
+    """The coroutine is closed when the loop closes between snapshot and dispatch."""
+    thread = EventLoopThread()
+    thread.loop = asyncio.new_event_loop()
+    thread.loop.close()
+
+    coro = asyncio.sleep(0)
+    with pytest.raises(RuntimeError):
+        thread.run_coroutine_threadsafe(coro)
+
+    assert inspect.getcoroutinestate(coro) == inspect.CORO_CLOSED
+
+
 async def test_thread_task_cancellation_after_stop(thread):
     loop = asyncio.get_event_loop()
     obj = mock.MagicMock()
@@ -211,3 +301,114 @@ async def wait_forever():
         # This will stall forever without the patch
         async with asyncio_timeout(1):
             await proxy.wait_forever()
+
+
+@pytest.mark.filterwarnings("error::RuntimeWarning")
+async def test_proxy_loop_stopping_async(thread, caplog):
+    """An async call dispatched after `force_stop()` resolves instead of hanging."""
+    obj = mock.MagicMock()
+    call_count = 0
+    worker_loop = thread.loop
+
+    async def magic():
+        nonlocal call_count
+        call_count += 1
+        # Like a real `Gateway.disconnect()`, this cannot finish within the window: the
+        # loop stops with the task still pending, and closing drops it
+        await worker_loop.create_future()
+
+    obj.test = magic
+    proxy = ThreadsafeProxy(obj, worker_loop, thread)
+
+    # Wedge the worker thread so the stop cannot complete while we dispatch: the loop is
+    # then reliably in the window this pins, stopping but not yet closed
+    worker_loop.call_soon_threadsafe(lambda: time.sleep(0.5))
+    thread.force_stop()
+
+    async with asyncio_timeout(1):
+        # `run_coroutine_threadsafe()` accepts work for a stopping loop without raising,
+        # so without the patch this awaits a future that is never resolved
+        assert await proxy.test() is None
+
+    assert call_count == 0
+    assert "Attempted to use an event loop that is shutting down" in caplog.text
+
+
+async def test_proxy_loop_stopping_sync(caplog):
+    """A sync call dispatched to a stopping loop is dropped, not silently queued."""
+    loop = asyncio.new_event_loop()
+    try:
+        thread = EventLoopThread()
+        thread.loop = loop
+        thread.stopping = True
+
+        obj = mock.MagicMock()
+        obj.test.return_value = None
+        proxy = ThreadsafeProxy(obj, loop, thread)
+
+        proxy.test()
+
+        assert obj.test.call_count == 0
+        assert "Attempted to use an event loop that is shutting down" in caplog.text
+    finally:
+        loop.close()
+
+
+async def test_proxy_stopping_ignored_without_thread(thread):
+    """A proxy with no `EventLoopThread` behaves exactly as before."""
+    obj = mock.MagicMock()
+    obj.test.return_value = None
+    proxy = ThreadsafeProxy(obj, thread.loop)
+
+    # A proxy that was not given the thread has no way to see this, and must not guess
+    thread.stopping = True
+    proxy.test()
+    thread.stopping = False
+
+    await yield_other_thread(thread)
+
+    assert obj.test.call_count == 1
+
+
+@pytest.mark.filterwarnings("error::RuntimeWarning")
+async def test_thread_run_coroutine_threadsafe_stopping():
+    """Handing a coroutine to a stopping loop raises instead of hanging."""
+    thread = EventLoopThread()
+    thread.loop = asyncio.new_event_loop()
+    try:
+        thread.force_stop()
+        assert thread.stopping
+        assert not thread.loop.is_closed()
+
+        coro = asyncio.sleep(0)
+        with pytest.raises(RuntimeError):
+            thread.run_coroutine_threadsafe(coro)
+
+        # The coroutine is closed internally: no "never awaited" RuntimeWarning
+        assert inspect.getcoroutinestate(coro) == inspect.CORO_CLOSED
+    finally:
+        thread.loop.close()
+
+
+async def test_thread_start_clears_stopping():
+    """A restarted thread is usable again: `start()` clears the stopping flag."""
+    thread = EventLoopThread()
+    thread_complete = await thread.start()
+    thread.force_stop()
+    assert thread.stopping
+
+    async with asyncio_timeout(1):
+        await thread_complete
+
+    await thread.start()
+    assert not thread.stopping
+    result = await thread.run_coroutine_threadsafe(
+        asyncio.sleep(0, mock.sentinel.result)
+    )
+    assert result is mock.sentinel.result
+
+    thread.force_stop()
+    async with asyncio_timeout(1):
+        await thread.thread_complete
+    [t.join(1) for t in threading.enumerate() if "bellows" in t.name]
+    assert [t for t in threading.enumerate() if "bellows" in t.name] == []