fix: Await in-flight poll before closing the transport on stop

AsyncRepeatingTask gains wait_stopped() to await the cancelled task; the
polling processor's stop() now waits for the in-flight poll to unwind before
closing the requester's transport, so awaiting stop() guarantees background
work has stopped and the transport isn't closed under a live request.

Author: jsonbailey

ldclient/impl/aio/concurrency.py

@@ -223,6 +223,16 @@ def stop(self):
         if task is not None and task is not asyncio.current_task():
             task.cancel()
 
+    async def wait_stopped(self):
+        """Waits for the task to finish unwinding after ``stop()``. A no-op if
+        the task never started or is the current task."""
+        task = self.__task
+        if task is not None and task is not asyncio.current_task():
+            try:
+                await task
+            except asyncio.CancelledError:
+                pass
+
     async def _run(self):
         try:
             if self.__initial_delay > 0:

ldclient/impl/datasource/async_polling.py

@@ -46,6 +46,10 @@ def initialized(self):
 
     async def stop(self):
         self.__stop_with_error_info(None)
+        # Wait for the in-flight poll to finish unwinding before closing the
+        # transport, so awaiting stop() guarantees background work has stopped
+        # and we don't close the HTTP transport out from under a live request.
+        await self._task.wait_stopped()
         await self._requester.close()
 
     def __stop_with_error_info(self, error: Optional[DataSourceErrorInfo]):

ldclient/testing/impl/datasource/test_async_polling.py

@@ -298,6 +298,36 @@ async def test_stop_closes_requester(self):
 
         processor._requester.close.assert_awaited_once()
 
+    @pytest.mark.asyncio
+    async def test_stop_awaits_poll_before_closing_transport(self):
+        # The in-flight poll must finish unwinding before the transport is
+        # closed, so we never close it out from under a live request.
+        order = []
+        started = asyncio.Event()
+
+        async def slow_poll():
+            started.set()
+            try:
+                await asyncio.sleep(60)
+            except asyncio.CancelledError:
+                order.append('poll_done')
+                raise
+
+        async def close():
+            order.append('transport_closed')
+
+        requester = MagicMock()
+        requester.get_all_data = slow_poll
+        requester.close = close
+
+        processor = make_processor(requester=requester)
+        processor.start()
+        await asyncio.wait_for(started.wait(), timeout=1.0)
+
+        await processor.stop()
+
+        assert order == ['poll_done', 'transport_closed']
+
     @pytest.mark.asyncio
     async def test_stop_cancels_polling_task_cleanly(self):
         store = MockAsyncFeatureStore()