fix: serialize async runner run and close

Prevent close from publishing or stopping a loop during an active startup/run and close partially initialized loops on startup failure.

Author: marandaneto

posthog/_async_utils.py

@@ -57,41 +57,47 @@ 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:
-        loop = self._ensure_loop()
-        future = asyncio.run_coroutine_threadsafe(self._await_result(awaitable), loop)
-        return future.result()
+        with self._operation_lock:
+            loop = self._ensure_loop()
+            future = asyncio.run_coroutine_threadsafe(
+                self._await_result(awaitable), loop
+            )
+            return future.result()
 
     def close(self) -> None:
+        current = threading.current_thread()
         with self._lock:
-            loop = self._loop
-            thread = self._thread
-            if thread is None:
+            if current is self._thread and self._loop is not None:
+                self._loop.call_soon(self._loop.stop)
                 return
-            if loop is None:
-                self._close_requested = True
-            else:
-                self._loop = None
-                self._thread = None
-                self._closing_threads.add(thread)
-
-        if loop is None:
-            if thread is not threading.current_thread():
-                thread.join()
-            return
 
-        if loop.is_closed():
+        with self._operation_lock:
             with self._lock:
-                self._closing_threads.discard(thread)
-            return
+                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)
 
-        if thread is threading.current_thread():
-            loop.call_soon(loop.stop)
-            return
+            if loop is None:
+                thread.join()
+                return
+
+            if loop.is_closed():
+                with self._lock:
+                    self._closing_threads.discard(thread)
+                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:
@@ -130,10 +136,13 @@ def _ensure_loop(self) -> asyncio.AbstractEventLoop:
             return self._loop
 
     def _run_loop(self) -> None:
+        loop = None
         try:
             loop = _ContextEventLoop()
             asyncio.set_event_loop(loop)
         except BaseException as error:
+            if loop is not None and not loop.is_closed():
+                loop.close()
             with self._lock:
                 self._startup_error = error
                 self._thread = None

posthog/test/test_async_utils.py

@@ -20,30 +20,38 @@ def test_startup_error_is_reported(self):
 
         awaitable.close()
 
-    def test_close_during_startup_does_not_orphan_thread(self):
+    def test_close_waits_for_run_during_startup(self):
         runner = _BackgroundEventLoopRunner()
         construction_started = threading.Event()
         release_construction = threading.Event()
+        run_errors = []
 
         def create_loop():
             construction_started.set()
             self.assertTrue(release_construction.wait(2))
             return _ContextEventLoop()
 
+        def run():
+            try:
+                runner.run(asyncio.sleep(0))
+            except BaseException as error:
+                run_errors.append(error)
+
         with mock.patch(
             "posthog._async_utils._ContextEventLoop", side_effect=create_loop
         ):
-            ensure_thread = threading.Thread(target=runner._ensure_loop)
-            ensure_thread.start()
+            run_thread = threading.Thread(target=run)
+            run_thread.start()
             self.assertTrue(construction_started.wait(1))
 
             close_thread = threading.Thread(target=runner.close)
             close_thread.start()
             release_construction.set()
-            ensure_thread.join(2)
+            run_thread.join(2)
             close_thread.join(2)
 
-        self.assertFalse(ensure_thread.is_alive())
+        self.assertFalse(run_thread.is_alive())
         self.assertFalse(close_thread.is_alive())
+        self.assertEqual(run_errors, [])
         self.assertIsNone(runner._thread)
         self.assertIsNone(runner._loop)