fix: clean up async consumer waiters on cancellation

Author: nickita-khylkouski

.sampo/changesets/async-consumer-cancellation.md

@@ -0,0 +1,5 @@
+---
+pypi/posthog: patch
+---
+
+Clean up queue and flush waiters when the async capture consumer is cancelled.

posthog/_async_consumer.py

@@ -117,15 +117,17 @@ def request_flush(self) -> None:
     async def _get_or_flush(self, timeout: float) -> tuple[Any, bool]:
         get_task = asyncio.create_task(self.queue.get())
         flush_task = asyncio.create_task(self._flush_event.wait())
-        done, pending = await asyncio.wait(
-            {get_task, flush_task},
-            timeout=timeout,
-            return_when=asyncio.FIRST_COMPLETED,
-        )
-        for task in pending:
-            task.cancel()
-        if pending:
-            await asyncio.gather(*pending, return_exceptions=True)
+        try:
+            done, _ = await asyncio.wait(
+                {get_task, flush_task},
+                timeout=timeout,
+                return_when=asyncio.FIRST_COMPLETED,
+            )
+        finally:
+            for task in (get_task, flush_task):
+                if not task.done():
+                    task.cancel()
+            await asyncio.gather(get_task, flush_task, return_exceptions=True)
 
         if get_task in done:
             return get_task.result(), False

posthog/test/test_async_consumer.py

@@ -152,3 +152,58 @@ async def test_request_stops_after_configured_retry_limit():
 
     assert batch_post.await_count == 3
     assert [call.args[0] for call in sleep.await_args_list] == [1, 2]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("run_worker", [False, True], ids=["wait", "worker"])
+async def test_get_or_flush_cancels_waiters_on_cancellation(run_worker):
+    consumer = make_consumer(retries=0)
+    consumer.flush_interval = 60
+    wait_started = asyncio.Event()
+    waiters = []
+    real_wait = asyncio.wait
+
+    async def observe_wait(tasks, **kwargs):
+        waiters.extend(tasks)
+        wait_started.set()
+        return await real_wait(tasks, **kwargs)
+
+    with mock.patch("posthog._async_consumer.asyncio.wait", side_effect=observe_wait):
+        task = asyncio.create_task(
+            consumer.run() if run_worker else consumer._get_or_flush(60)
+        )
+        try:
+            await wait_started.wait()
+            task.cancel()
+            with pytest.raises(asyncio.CancelledError):
+                await task
+
+            assert len(waiters) == 2
+            assert all(waiter.cancelled() for waiter in waiters)
+        finally:
+            task.cancel()
+            for waiter in waiters:
+                waiter.cancel()
+            await asyncio.gather(task, *waiters, return_exceptions=True)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("queued", "flush"), [(True, False), (False, True), (False, False), (True, True)]
+)
+async def test_get_or_flush_preserves_results_and_cleans_up_waiters(queued, flush):
+    consumer = make_consumer(retries=0)
+    event = {"event": "test"}
+    if queued:
+        consumer.queue.put_nowait(event)
+    if flush:
+        consumer.request_flush()
+    tasks_before = asyncio.all_tasks()
+
+    result = await consumer._get_or_flush(60 if queued or flush else 0)
+
+    assert result == (event if queued else None, flush and not queued)
+    assert consumer._flush_event.is_set() == (queued and flush)
+    assert not (asyncio.all_tasks() - tasks_before)
+    if queued:
+        consumer.queue.task_done()