fix: Release BoundedTaskSet slots synchronously to prevent a livelock

BoundedTaskSet freed a slot in a done-callback, which runs on a later
loop turn. When flush tasks finished and a flush_and_wait or shutdown
message arrived in the same event-loop batch, the dispatcher's
retry loop (while not self._trigger_flush(): await self._flush_workers.wait())
could spin forever: try_run still counted the finished tasks, and wait()
over already-complete tasks did not yield, so the done-callbacks never ran
to free the slots. This wedged the event loop at ~100% CPU.

Release the slot from within the task, synchronously in a finally, like the
sync FixedThreadPool. Add regression tests at the BoundedTaskSet and event
processor levels, and drop an unused import.

Author: jsonbailey

ldclient/impl/aio/concurrency.py

@@ -267,15 +267,19 @@ def try_run(self, job: Callable[[], Coroutine]) -> bool:
         task was started, or False if the set is full or has been stopped."""
         if not self._accepting or len(self._tasks) >= self._limit:
             return False
-        task = asyncio.create_task(job())
+        task = asyncio.create_task(self._run(job))
         self._tasks.add(task)
-        task.add_done_callback(self._on_done)
         return True
 
-    def _on_done(self, task: asyncio.Task) -> None:
-        self._tasks.discard(task)
-        if not task.cancelled() and task.exception() is not None:
-            log.warning('Unhandled exception in background task', exc_info=task.exception())
+    async def _run(self, job: Callable[[], Coroutine]) -> None:
+        try:
+            await job()
+        except Exception:
+            log.warning('Unhandled exception in background task', exc_info=True)
+        finally:
+            task = asyncio.current_task()
+            if task is not None:
+                self._tasks.discard(task)
 
     async def wait(self) -> None:
         """Waits for the tasks currently running to complete. This does not stop

ldclient/impl/events/async_event_processor.py

@@ -6,7 +6,6 @@
 import gzip
 import json
 import queue
-import time
 import uuid
 from collections import namedtuple
 from random import Random

ldclient/testing/impl/events/test_async_event_processor.py

@@ -10,6 +10,7 @@
 import asyncio
 import gzip
 import json
+import threading
 import uuid
 from contextlib import asynccontextmanager
 from email.utils import formatdate
@@ -491,3 +492,88 @@ async def test_event_payload_is_gzip_compressed_when_enabled():
         output = json.loads(gzip.decompress(mock_http.request_data))
         assert len(output) == 1
         assert output[0]['kind'] == 'identify'
+
+
+async def test_flush_and_wait_completes_when_in_flight_posts_finish_together():
+    # Regression test for a livelock: when all in-flight flush POSTs finish and
+    # the flush_and_wait message arrives in the same event-loop batch, the
+    # dispatcher's retry loop must still make progress and reply, not spin.
+    class GatedMockAioHttp(MockAioHttp):
+        """Parks each request until the gate is set, then completes without yielding."""
+
+        def __init__(self):
+            super().__init__()
+            self.gate = asyncio.Event()
+
+        def request(self, method, uri, headers=None, data=None, timeout=None, proxy=None):
+            self._recorded_requests.append((headers, data))
+            outer = self
+
+            class _Ctx:
+                async def __aenter__(self):
+                    await outer.gate.wait()
+                    return MockAioResponse(200, {})
+
+                async def __aexit__(self, exc_type, exc_value, traceback):
+                    return False
+
+            return _Ctx()
+
+    # DefaultAsyncEventProcessor discards the dispatcher, so capture it here to
+    # let the watchdog recover the loop if it wedges.
+    dispatcher_holder = []
+
+    def capture_dispatcher(inbox, config, http, diagnostic_accumulator):
+        dispatcher = EventDispatcher(inbox, config, http, diagnostic_accumulator)
+        dispatcher_holder.append(dispatcher)
+        return dispatcher
+
+    mock_http = GatedMockAioHttp()
+    config = AsyncConfig(sdk_key='SDK_KEY', diagnostic_opt_out=True)
+    ep = DefaultAsyncEventProcessor(config, mock_http, dispatcher_class=capture_dispatcher)
+    try:
+        # Saturate all five flush workers with parked POSTs.
+        for i in range(5):
+            ep.send_event(EventInputIdentify(timestamp, Context.create('user%d' % i)))
+            ep.flush()
+        deadline = asyncio.get_running_loop().time() + 2
+        while len(mock_http.recorded_requests) < 5:
+            assert asyncio.get_running_loop().time() < deadline, 'workers never saturated'
+            await asyncio.sleep(0.01)
+
+        # Buffer one more event so the awaited flush has work to hand off.
+        ep.send_event(EventInputIdentify(timestamp, context))
+        await asyncio.sleep(0.05)
+
+        # Complete all in-flight POSTs and enqueue flush_and_wait in one loop
+        # callback, so the dispatcher sees the message in the batch the workers
+        # finished in.
+        reply = AsyncEvent()
+        mock_http.gate.set()
+        ep._inbox.put_nowait(EventProcessorMessage('flush_and_wait', reply))
+
+        # A wedged loop cannot fire asyncio timeouts, so a watchdog thread
+        # detects the wedge and frees the task set to fail with an assertion
+        # rather than hang the run.
+        finished = threading.Event()
+        wedged = threading.Event()
+
+        def watchdog():
+            if not finished.wait(3):
+                wedged.set()
+                dispatcher_holder[0]._flush_workers._tasks.clear()
+
+        threading.Thread(target=watchdog, daemon=True).start()
+
+        replied = await reply.wait(10)
+        finished.set()
+
+        assert replied is True
+        assert not wedged.is_set(), (
+            "flush_and_wait wedged the event loop: the dispatcher spun in "
+            "'while not self._trigger_flush(): await self._flush_workers.wait()' "
+            "without yielding, and only the watchdog un-stuck it"
+        )
+    finally:
+        mock_http.gate.set()
+        await ep.stop()

ldclient/testing/test_aio.py

@@ -223,6 +223,38 @@ async def test_wait_returns_when_idle(self):
         await tasks.wait()
         tasks.stop()
 
+    @pytest.mark.asyncio
+    async def test_retry_pattern_makes_progress_when_task_finished_in_same_batch(self):
+        tasks = aio.BoundedTaskSet(1)
+        finished = asyncio.Event()
+
+        async def job():
+            finished.set()
+
+        assert tasks.try_run(job) is True
+
+        # A raw asyncio.Event wakes this coroutine in the same loop batch the
+        # job's task finished in, before try_run sees the slot freed.
+        await finished.wait()
+
+        async def noop():
+            pass
+
+        # Mirrors the retry loop in EventDispatcher._run_main_loop: try_run must
+        # free the finished task's slot, or the real (unbounded) loop spins.
+        accepted = False
+        for _ in range(100):
+            if tasks.try_run(noop):
+                accepted = True
+                break
+            await tasks.wait()
+
+        assert accepted, (
+            "try_run never freed capacity: the finished task was still counted "
+            "against the limit, so the dispatcher's retry loop would spin forever"
+        )
+        await tasks.wait()
+
 
 class TestTaskRunnerParity:
     @pytest.mark.asyncio