refactor: Replace AsyncWorkerPool with BoundedTaskSet

The async event delivery concurrency limiter no longer mirrors the sync
FixedThreadPool's shape. BoundedTaskSet drops the unused name parameter and
the thread vocabulary, reserves a slot synchronously at spawn (so a full set
rejects rather than queues), and uses a done-callback plus asyncio.gather for
cleanup and draining.

Author: jsonbailey

ldclient/impl/aio/concurrency.py

@@ -2,7 +2,7 @@
 Async concurrency primitives used by the async data source, event processor, and
 data system. Each wraps a piece of fiddly asyncio plumbing (timeout-aware waits,
 queue exception normalization, an interval-from-start repeating task, a bounded
-task pool) that callers would otherwise inline repeatedly. The sync code uses the
+task set) that callers would otherwise inline repeatedly. The sync code uses the
 equivalent stdlib/SDK primitives (``threading.Event``/``Lock``, ``queue.Queue``,
 ``RepeatingTask``, ``FixedThreadPool``) directly, so these have no sync twin.
 """
@@ -12,7 +12,7 @@
 import time
 from queue import Empty as QueueEmpty  # noqa: F401  (shared timeout exception)
 from queue import Full as QueueFull  # noqa: F401  (shared capacity exception)
-from typing import Any, Callable, Optional, Set
+from typing import Any, Awaitable, Callable, Optional, Set
 
 from ldclient.impl.util import log
 
@@ -250,50 +250,37 @@ async def _run(self):
             pass
 
 
-class AsyncWorkerPool:
-    """A fixed-size pool of concurrent tasks that rejects jobs when its limit
-    is reached. Matches the contract of
-    ``ldclient.impl.fixed_thread_pool.FixedThreadPool``."""
+class BoundedTaskSet:
+    """Runs up to ``limit`` coroutines concurrently as background tasks. When the
+    limit is reached, ``try_spawn`` rejects new work (returning False) rather than
+    queuing it, so callers can apply their own backpressure. ``drain`` awaits all
+    in-flight tasks; ``stop`` prevents any further work from being accepted."""
 
-    def __init__(self, size: int, name: str):
-        self._size = size
-        self._name = name
-        self._busy: Set[asyncio.Task] = set()
-        self._event = AsyncEvent()
-        self._stopped = False
+    def __init__(self, limit: int):
+        self._limit = limit
+        self._tasks: Set[asyncio.Task] = set()
+        self._accepting = True
 
-    def execute(self, jobFn: Callable) -> bool:
-        """Schedules a job for execution if the pool is not already at its
-        limit, and returns True if successful; returns False if all workers
-        are busy."""
-        if self._stopped or len(self._busy) >= self._size:
+    def try_spawn(self, job: Callable[[], Awaitable]) -> bool:
+        """Starts ``job()`` as a background task if fewer than ``limit`` tasks are
+        already running and the set is still accepting work. Returns True if the
+        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.ensure_future(self._run_job(jobFn))
-        self._busy.add(task)
+        task = asyncio.ensure_future(job())
+        self._tasks.add(task)
+        task.add_done_callback(self._on_done)
         return True
 
-    async def _run_job(self, jobFn: Callable) -> None:
-        try:
-            result = jobFn()
-            if inspect.isawaitable(result):
-                await result
-        except Exception:
-            log.warning('Unhandled exception in worker thread', exc_info=True)
-        finally:
-            task = asyncio.current_task()
-            if task is not None:
-                self._busy.discard(task)
-            self._event.set()
-
-    async def wait(self) -> None:
-        """Waits until all currently busy workers have completed their jobs."""
-        while len(self._busy) > 0:
-            self._event.clear()
-            if len(self._busy) == 0:
-                return
-            await self._event.wait()
+    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 drain(self) -> None:
+        """Waits for all currently running tasks to complete."""
+        await asyncio.gather(*self._tasks, return_exceptions=True)
 
     def stop(self) -> None:
-        """Tells the pool to reject any further jobs; active jobs run to
-        completion."""
-        self._stopped = True
+        """Rejects any further work; tasks already running finish normally."""
+        self._accepting = False

ldclient/impl/events/async_event_processor.py

@@ -19,7 +19,7 @@
     AsyncQueue,
     AsyncRepeatingTask,
     AsyncTaskRunner,
-    AsyncWorkerPool
+    BoundedTaskSet
 )
 from ldclient.impl.aio.transport import AsyncHTTPTransport
 from ldclient.impl.events.diagnostics import create_diagnostic_init
@@ -38,7 +38,7 @@
 )
 from ldclient.interfaces import AsyncEventProcessor
 
-__MAX_FLUSH_THREADS__ = 5
+__MAX_FLUSH_CONCURRENCY__ = 5
 __CURRENT_EVENT_SCHEMA__ = 4
 
 
@@ -108,13 +108,13 @@ def __init__(self, inbox: AsyncQueue, config: AsyncConfig, http_client, diagnost
         self._sampler = Sampler(Random())
         self._omit_anonymous_contexts = config.omit_anonymous_contexts
 
-        self._flush_workers = AsyncWorkerPool(__MAX_FLUSH_THREADS__, "ldclient.flush")
-        self._diagnostic_flush_workers: Optional[AsyncWorkerPool] = None
+        self._flush_workers = BoundedTaskSet(__MAX_FLUSH_CONCURRENCY__)
+        self._diagnostic_flush_workers: Optional[BoundedTaskSet] = None
         if self._diagnostic_accumulator is not None:
-            self._diagnostic_flush_workers = AsyncWorkerPool(1, "ldclient.events.diag_flush")
+            self._diagnostic_flush_workers = BoundedTaskSet(1)
             init_event = create_diagnostic_init(self._diagnostic_accumulator.data_since_date, self._diagnostic_accumulator.diagnostic_id, config)
             task = DiagnosticEventSendTask(self._http, self._config, init_event)
-            self._diagnostic_flush_workers.execute(task.run)
+            self._diagnostic_flush_workers.try_spawn(task.run)
 
         self._runner = AsyncTaskRunner()
         self._runner.spawn("ldclient.events.processor", self._run_main_loop)
@@ -134,12 +134,12 @@ async def _run_main_loop(self):
                     self._send_and_reset_diagnostics()
                 elif message.type == 'flush_and_wait':
                     self._trigger_flush()
-                    await self._flush_workers.wait()
+                    await self._flush_workers.drain()
                     message.param.set()
                 elif message.type == 'test_sync':
-                    await self._flush_workers.wait()
+                    await self._flush_workers.drain()
                     if self._diagnostic_flush_workers is not None:
-                        await self._diagnostic_flush_workers.wait()
+                        await self._diagnostic_flush_workers.drain()
                     message.param.set()
                 elif message.type == 'stop':
                     await self._do_shutdown()
@@ -156,7 +156,7 @@ def _trigger_flush(self):
             self._diagnostic_accumulator.record_events_in_batch(len(payload.events))
         if len(payload.events) > 0 or not payload.summary.is_empty():
             task = EventPayloadSendTask(self._http, self._config, self._formatter, payload, self._handle_response)
-            if self._flush_workers.execute(task.run):
+            if self._flush_workers.try_spawn(task.run):
                 # The events have been handed off to a flush worker; clear them from our buffer.
                 self._outbox.clear()
             else:
@@ -169,15 +169,15 @@ def _send_and_reset_diagnostics(self):
             stats_event = self._diagnostic_accumulator.create_event_and_reset(dropped_event_count, self._deduplicated_contexts)
             self._deduplicated_contexts = 0
             task = DiagnosticEventSendTask(self._http, self._config, stats_event)
-            self._diagnostic_flush_workers.execute(task.run)
+            self._diagnostic_flush_workers.try_spawn(task.run)
 
     async def _do_shutdown(self):
         self._flush_workers.stop()
-        await self._flush_workers.wait()
+        await self._flush_workers.drain()
 
         if self._diagnostic_flush_workers is not None:
             self._diagnostic_flush_workers.stop()
-            await self._diagnostic_flush_workers.wait()
+            await self._diagnostic_flush_workers.drain()
 
         await self._http.close()
 

ldclient/testing/test_aio.py

@@ -203,38 +203,43 @@ async def action():
 
 
 # ---------------------------------------------------------------------------
-# WorkerPool
+# BoundedTaskSet
 # ---------------------------------------------------------------------------
 
-class TestWorkerPoolParity:
+class TestBoundedTaskSet:
     @pytest.mark.asyncio
-    async def test_async_saturation_returns_false(self):
-        pool = aio.AsyncWorkerPool(1, "test.pool")
+    async def test_saturation_returns_false(self):
+        tasks = aio.BoundedTaskSet(1)
         release = aio.AsyncEvent()
         started = aio.AsyncEvent()
 
         async def job():
             started.set()
             await release.wait(2)
 
-        assert pool.execute(job) is True
+        async def noop():
+            pass
+
+        assert tasks.try_spawn(job) is True
         await started.wait(2)
-        assert pool.execute(lambda: None) is False
+        # Set is full (limit 1), so the next spawn is rejected rather than queued.
+        assert tasks.try_spawn(noop) is False
         release.set()
-        await pool.wait()
+        await tasks.drain()
 
-        async def noop():
-            pass
+        # A slot is free again once the first task drained.
+        assert tasks.try_spawn(noop) is True
+        await tasks.drain()
 
-        assert pool.execute(noop) is True
-        await pool.wait()
-        pool.stop()
+        # After stop(), further work is rejected.
+        tasks.stop()
+        assert tasks.try_spawn(noop) is False
 
     @pytest.mark.asyncio
-    async def test_async_wait_returns_when_idle(self):
-        pool = aio.AsyncWorkerPool(2, "test.pool")
-        await pool.wait()
-        pool.stop()
+    async def test_drain_returns_when_idle(self):
+        tasks = aio.BoundedTaskSet(2)
+        await tasks.drain()
+        tasks.stop()
 
 
 # ---------------------------------------------------------------------------