fix: coordinate concurrent async runner operations

Submit concurrent work without holding completion locks, reject loop-thread synchronous reentry, and define close-during-startup cancellation.

Author: marandaneto

posthog/_async_utils.py

@@ -57,47 +57,52 @@ def __init__(self) -> None:
         self._startup_error: BaseException | None = None
         self._close_requested = False
         self._lock = threading.Lock()
-        self._operation_lock = threading.Lock()
 
     def run(self, awaitable: Awaitable[Any]) -> Any:
-        with self._operation_lock:
+        if threading.current_thread() is self._thread:
+            raise RuntimeError("cannot synchronously run from the runner thread")
+
+        while True:
             loop = self._ensure_loop()
-            future = asyncio.run_coroutine_threadsafe(
-                self._await_result(awaitable), loop
-            )
-            return future.result()
+            with self._lock:
+                if loop is self._loop and not self._close_requested:
+                    future = asyncio.run_coroutine_threadsafe(
+                        self._await_result(awaitable), loop
+                    )
+                    break
+        return future.result()
 
     def close(self) -> None:
         current = threading.current_thread()
         with self._lock:
-            if current is self._thread and self._loop is not None:
-                self._loop.call_soon(self._loop.stop)
+            loop = self._loop
+            thread = self._thread
+            if thread is None:
                 return
-
-        with self._operation_lock:
-            with self._lock:
-                loop = self._loop
-                thread = self._thread
-                if thread is None:
-                    return
-                if loop is None:
-                    self._close_requested = True
-                else:
-                    self._loop = None
-                    self._thread = None
-                    self._closing_threads.add(thread)
-
+            self._close_requested = True
             if loop is None:
+                self._startup_error = RuntimeError("runner closed during startup")
+            else:
+                self._loop = None
+                self._thread = None
+                self._closing_threads.add(thread)
+
+        if loop is None:
+            if thread is not current:
                 thread.join()
-                return
+            return
 
-            if loop.is_closed():
-                with self._lock:
-                    self._closing_threads.discard(thread)
-                return
+        if loop.is_closed():
+            with self._lock:
+                self._closing_threads.discard(thread)
+            return
+
+        if thread is current:
+            loop.call_soon(loop.stop)
+            return
 
-            loop.call_soon_threadsafe(loop.stop)
-            thread.join()
+        loop.call_soon_threadsafe(loop.stop)
+        thread.join()
 
     def owns_thread(self, thread: threading.Thread) -> bool:
         with self._lock:
@@ -152,6 +157,8 @@ def _run_loop(self) -> None:
         with self._lock:
             self._loop = loop
             close_requested = self._close_requested
+            if close_requested and self._startup_error is None:
+                self._startup_error = RuntimeError("runner closed during startup")
             self._started.set()
 
         if close_requested:

posthog/test/test_async_utils.py

@@ -32,9 +32,11 @@ def create_loop():
             return _ContextEventLoop()
 
         def run():
+            awaitable = asyncio.sleep(0)
             try:
-                runner.run(asyncio.sleep(0))
+                runner.run(awaitable)
             except BaseException as error:
+                awaitable.close()
                 run_errors.append(error)
 
         with mock.patch(
@@ -52,6 +54,31 @@ def run():
 
         self.assertFalse(run_thread.is_alive())
         self.assertFalse(close_thread.is_alive())
-        self.assertEqual(run_errors, [])
+        self.assertEqual(len(run_errors), 1)
+        self.assertRegex(str(run_errors[0]), "closed during startup")
         self.assertIsNone(runner._thread)
         self.assertIsNone(runner._loop)
+
+    def test_run_from_runner_thread_fails_instead_of_deadlocking(self):
+        runner = _BackgroundEventLoopRunner()
+
+        async def reenter():
+            awaitable = asyncio.sleep(0)
+            try:
+                with self.assertRaisesRegex(RuntimeError, "runner thread"):
+                    runner.run(awaitable)
+            finally:
+                awaitable.close()
+
+        runner.run(reenter())
+        runner.close()
+
+    def test_close_from_runner_thread_allows_fresh_loop(self):
+        runner = _BackgroundEventLoopRunner()
+
+        async def close_runner():
+            runner.close()
+
+        runner.run(close_runner())
+        runner.run(asyncio.sleep(0))
+        runner.close()