fix: async shutdown root lifecycle and queue_timeout=0 dequeue

Run async cleanup on the asyncio.run root after _finished is set so
protocol aclose() completes before Runner teardown. Preserve the first
Queue.get attempt when SW_AGENT_QUEUE_TIMEOUT=0. Add regression tests.

Author: songzhendong

skywalking/agent/__init__.py

@@ -151,9 +151,8 @@ async def _cancel_pending_tasks(tasks) -> None:
     """
     Cancel agent-owned reporter / connectivity-watch tasks only.
 
-    Must not cancel the asyncio.run root (__start_event_loop_async): that tears down
-    the Runner and cancels this fini coroutine before protocol aclose() runs.
-    The current task (fini) is also excluded so we never await ourselves.
+    Cancel only the given reporter / watch tasks; never the asyncio.run root.
+    The current task is excluded so we never await ourselves.
     """
     current = asyncio.current_task()
     pending = [
@@ -820,11 +819,12 @@ async def __start_event_loop_async(self) -> None:
 
         await self.__bootstrap()  # gather all coroutines
 
-        # Track Tasks explicitly so shutdown cancels only reporters/watchers, not the
-        # asyncio.run root task that is awaiting this gather.
         self.background_tasks = {asyncio.create_task(coro) for coro in self.background_coroutines}
         logger.debug('All background coroutines started')
-        await asyncio.gather(*self.background_tasks)
+        # Wait for shutdown inside the asyncio.run root, then clean up here so the
+        # Runner stays alive through protocol aclose() before asyncio.run returns.
+        await self._finished.wait()
+        await self.__async_shutdown_cleanup()
 
     def __start_event_loop(self) -> None:
         try:
@@ -872,17 +872,13 @@ def start(self) -> None:
         self.event_loop_thread = Thread(name='event_loop_thread', target=self.__start_event_loop, daemon=True)
         self.event_loop_thread.start()
 
-    async def __fini_async(self):
+    async def __async_shutdown_cleanup(self) -> None:
         """
-        This method is called when the agent is shutting down.
-        Clean up all the queues and stop all the asyncio tasks.
+        Async shutdown body that must run on the asyncio.run root task.
 
         Do not await report_* here: aio generators use unbounded queue.get() and can
         hang forever. Abandon + timed join, then cancel tasks.
         """
-        if self._finished is not None:
-            self._finished.set()
-
         await _shutdown_async_queue(self.__segment_queue, 'segment')
 
         if config.agent_log_reporter_active:
@@ -908,14 +904,13 @@ async def __fini_async(self):
                 close()
 
     def __fini(self):
-        if not self.loop.is_closed():
-            future = asyncio.run_coroutine_threadsafe(self.__fini_async(), self.loop)
-            try:
-                future.result(timeout=_SHUTDOWN_LOOP_JOIN_TIMEOUT_SEC + _SHUTDOWN_JOIN_TIMEOUT_SEC * 4)
-            except Exception:  # noqa: BLE001 - never block process exit on shutdown
-                logger.exception('async agent shutdown exceeded budget or failed')
-                future.cancel()
-        self.event_loop_thread.join(timeout=_SHUTDOWN_LOOP_JOIN_TIMEOUT_SEC)
+        loop = getattr(self, 'loop', None)
+        if loop is not None and not loop.is_closed() and self._finished is not None:
+            loop.call_soon_threadsafe(self._finished.set)
+        if self.event_loop_thread is not None:
+            self.event_loop_thread.join(
+                timeout=_SHUTDOWN_LOOP_JOIN_TIMEOUT_SEC + _SHUTDOWN_JOIN_TIMEOUT_SEC * 4,
+            )
         if self.event_loop_thread.is_alive():
             logger.warning(
                 'Python agent event_loop thread still alive after %.1fs shutdown budget',

skywalking/agent/protocol/grpc.py

@@ -45,20 +45,24 @@
 from skywalking.trace.segment import Segment
 
 
-def _queue_get_within_batch(queue: Queue, block: bool, batch_deadline: float):
+def _queue_get_within_batch(queue: Queue, block: bool, batch_deadline: float, *, allow_immediate: bool = False):
     """
     Get one item within an absolute batch window (monotonic deadline).
 
     Avoids int(elapsed) truncation that could let queue waits approach
     agent_queue_timeout + 1s and collide with a tight RPC deadline.
+    When allow_immediate is True (first generator iteration), still attempt
+    Queue.get once so SW_AGENT_QUEUE_TIMEOUT=0 can drain an immediately
+    available item via get(timeout=0).
     Returns None when the window is exhausted or the queue is empty.
     """
     remaining = batch_deadline - monotonic()
-    if remaining <= 0:
+    if remaining <= 0 and not allow_immediate:
         return None
     try:
         if block:
-            return queue.get(block=True, timeout=remaining)
+            timeout = remaining if remaining > 0 else 0
+            return queue.get(block=True, timeout=timeout)
         return queue.get(block=False)
     except Empty:
         return None
@@ -172,8 +176,10 @@ def generator():
             nonlocal sent
 
             batch_deadline = monotonic() + float(config.agent_queue_timeout)
+            first_get = True
             while True:
-                segment = _queue_get_within_batch(queue, block, batch_deadline)  # type: Segment
+                segment = _queue_get_within_batch(queue, block, batch_deadline, allow_immediate=first_get)  # type: Segment
+                first_get = False
                 if segment is None:
                     return
 
@@ -240,8 +246,10 @@ def generator():
             nonlocal sent
 
             batch_deadline = monotonic() + float(config.agent_queue_timeout)
+            first_get = True
             while True:
-                log_data = _queue_get_within_batch(queue, block, batch_deadline)  # type: LogData
+                log_data = _queue_get_within_batch(queue, block, batch_deadline, allow_immediate=first_get)  # type: LogData
+                first_get = False
                 if log_data is None:
                     return
 
@@ -270,8 +278,10 @@ def generator():
             nonlocal sent
 
             batch_deadline = monotonic() + float(config.agent_queue_timeout)
+            first_get = True
             while True:
-                meter_data = _queue_get_within_batch(queue, block, batch_deadline)  # type: MeterData
+                meter_data = _queue_get_within_batch(queue, block, batch_deadline, allow_immediate=first_get)  # type: MeterData
+                first_get = False
                 if meter_data is None:
                     return
 
@@ -299,8 +309,10 @@ def generator():
             nonlocal sent
 
             batch_deadline = monotonic() + float(config.agent_queue_timeout)
+            first_get = True
             while True:
-                snapshot = _queue_get_within_batch(queue, block, batch_deadline)  # type: TracingThreadSnapshot
+                snapshot = _queue_get_within_batch(queue, block, batch_deadline, allow_immediate=first_get)  # type: TracingThreadSnapshot
+                first_get = False
                 if snapshot is None:
                     return
 

tests/unit/test_grpc_channel.py

@@ -19,6 +19,8 @@
 import json
 import unittest
 import asyncio
+from queue import Queue
+from time import monotonic
 from unittest.mock import MagicMock, patch
 
 import grpc
@@ -646,6 +648,28 @@ def test_sync_report_uses_timeout_with_margin(self):
             config.agent_queue_timeout = prev
 
 
+class TestQueueGetWithinBatch(unittest.TestCase):
+
+    def test_queue_timeout_zero_drains_immediately_available_item(self):
+        from skywalking.agent.protocol.grpc import _queue_get_within_batch
+
+        q = Queue()
+        q.put('segment')
+        batch_deadline = monotonic()
+        item = _queue_get_within_batch(q, True, batch_deadline, allow_immediate=True)
+        self.assertEqual(item, 'segment')
+        self.assertTrue(q.empty())
+
+    def test_queue_timeout_zero_skips_when_empty(self):
+        from skywalking.agent.protocol.grpc import _queue_get_within_batch
+
+        q = Queue()
+        batch_deadline = monotonic()
+        self.assertIsNone(
+            _queue_get_within_batch(q, True, batch_deadline, allow_immediate=True),
+        )
+
+
 class TestCollectorChannelNotInstrumented(unittest.TestCase):
 
     def test_multi_address_collector_channel_skips_sw_interceptor(self):

tests/unit/test_shutdown_queue.py

@@ -136,43 +136,44 @@ async def _spawn():
             thread.join(timeout=2.0)
             loop.close()
 
-    def test_cancel_preserves_asyncio_run_root_so_fini_completes(self):
+    def test_async_shutdown_cleanup_runs_inside_asyncio_run_root(self):
         """
-        Production path uses asyncio.run(root). Cancelling the root tears down the Runner
-        and cancels fini before aclose(). Only background tasks may be cancelled.
+        Production topology: root waits on _finished, then cancels reporters and
+        awaits a yielding protocol aclose() before asyncio.run returns.
         """
+        holder = {}
+        aclose_entered = Event()
         aclose_done = Event()
         loop_ready = Event()
-        holder = {}
+
+        async def yielding_aclose():
+            aclose_entered.set()
+            await asyncio.sleep(0.05)
+            aclose_done.set()
 
         async def root():
+            finished = asyncio.Event()
+            holder['finished'] = finished
             holder['loop'] = asyncio.get_running_loop()
-            holder['root'] = asyncio.current_task()
 
             async def reporter():
-                while True:
-                    await asyncio.sleep(0.05)
+                while not finished.is_set():
+                    await asyncio.sleep(0.02)
 
-            holder['tasks'] = {asyncio.create_task(reporter()) for _ in range(2)}
+            tasks = {asyncio.create_task(reporter()) for _ in range(2)}
             loop_ready.set()
-            # return_exceptions: cancelled reporters must not surface as root CancelledError
-            await asyncio.gather(*holder['tasks'], return_exceptions=True)
+            await finished.wait()
+            await _cancel_pending_tasks(tasks)
+            await yielding_aclose()
             holder['root_finished'] = True
 
         thread = Thread(target=lambda: asyncio.run(root()), daemon=True)
         thread.start()
         self.assertTrue(loop_ready.wait(3.0))
-
-        async def fini():
-            # Must not cancel holder["root"]; doing so would cancel this fini via Runner teardown.
-            await _cancel_pending_tasks(holder['tasks'])
-            aclose_done.set()
-            return 'ok'
-
-        result = asyncio.run_coroutine_threadsafe(fini(), holder['loop']).result(timeout=3.0)
-        self.assertEqual(result, 'ok')
-        self.assertTrue(aclose_done.wait(1.0))
-        thread.join(timeout=3.0)
+        holder['loop'].call_soon_threadsafe(holder['finished'].set)
+        thread.join(timeout=5.0)
+        self.assertTrue(aclose_entered.wait(2.0))
+        self.assertTrue(aclose_done.wait(2.0))
         self.assertTrue(holder.get('root_finished', False))