Merge pull request #801 from LincolnPuzey/fix_gh_791

Fix #791: Fix infinite loop when async actor raises `TimeoutError` exception

Author: LincolnPuzey

dramatiq/asyncio.py

@@ -146,8 +146,27 @@ async def wrapped_coro() -> R:
                     # Use a timeout to be able to catch asynchronously
                     # raised dramatiq exceptions (Interrupt).
                     return future.result(timeout=self.interrupt_check_ival)
-                except concurrent.futures.TimeoutError:
-                    continue
+                except (
+                    # TODO replace with built-in TimeoutError once 3.10 support dropped.
+                    concurrent.futures.TimeoutError
+                ):
+                    # NOTE: TimeoutError caught here could be from future.result() timing out (i.e. future not done yet),
+                    # or a TimeoutError raised inside the future itself (future is done).
+                    if not future.done():
+                        # future not done, so .result() must've timed out. continue to wait again.
+                        continue
+
+                # If execution reaches here, it means a TimeoutError was caught above, and the future is done.
+                # There are 3 possibilities here:
+                # 1. TimeoutError was raised inside the future. This will re-raise it.
+                # 2. First .result() call timed out, but the future completed by the time .done() was called.
+                #     a) This will return the future's result, or
+                #     b) raise the Exception that happened in the future.
+                return future.result(timeout=0)
+                # This is outside the 'except' block to avoid any
+                # "During handling of the above exception, another exception occurred" messages.
+                # zero timeout used because future is now done.
+
         except Interrupt as e:
             # Asynchronously raised from another thread: cancel the
             # future.

tests/middleware/test_asyncio.py

@@ -82,6 +82,24 @@ async def raise_error():
     assert e.traceback[-1].name == "raise_actual_error"
 
 
+def test_event_loop_thread_run_coroutine_timeout_exception(started_thread: EventLoopThread):
+    """Test that TimeoutError in coroutine doesn't lead to infinite loop.
+
+    Regression test for https://github.com/Bogdanp/dramatiq/issues/791
+    """
+
+    async def raise_actual_error():
+        raise TimeoutError("something took too long")
+
+    async def raise_error():
+        await raise_actual_error()
+
+    coro = raise_error()
+
+    with pytest.raises(TimeoutError, match="something took too long"):
+        started_thread.run_coroutine(coro)
+
+
 @pytest.mark.skipif(
     threading.current_platform not in threading.supported_platforms,
     reason="Threading not supported on this platform.",