tests/unit/test_shutdown_queue.py

#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import asyncio
import logging
import time
import unittest
from queue import Queue
from threading import Event, Thread

from skywalking.agent import (
    _abandon_async_queue,
    _abandon_sync_queue,
    _await_shutdown_or_background_failure,
    _cancel_pending_tasks,
    _join_sync_queue,
    _shutdown_async_queue,
    _shutdown_sync_queue,
)


class TestShutdownQueueHelpers(unittest.TestCase):

    def test_abandon_sync_queue_unblocks_join(self):
        q = Queue()
        q.put('a')
        q.put('b')
        self.assertEqual(q.unfinished_tasks, 2)

        hung = Event()

        def _join():
            q.join()
            hung.set()

        Thread(target=_join, daemon=True).start()
        time.sleep(0.05)
        self.assertFalse(hung.is_set())

        abandoned = _abandon_sync_queue(q)
        self.assertEqual(abandoned, 2)
        self.assertTrue(hung.wait(1.0))
        self.assertTrue(_join_sync_queue(q, 0.5))

    def test_shutdown_sync_skips_flush_when_not_ready(self):
        q = Queue()
        q.put('x')
        called = []

        def report():
            called.append(1)
            raise AssertionError('must not flush when may_send is False')

        _shutdown_sync_queue(report, q, 'test', may_send=False)
        self.assertEqual(called, [])
        self.assertTrue(q.empty())
        self.assertEqual(q.unfinished_tasks, 0)

    def test_shutdown_sync_flush_timeout_then_abandon(self):
        q = Queue()
        q.put('x')
        started = Event()

        def report():
            started.set()
            time.sleep(10)  # longer than flush budget

        # Temporarily shrink budget via monkeypatch on module constant
        import skywalking.agent as agent_mod

        previous = agent_mod._SHUTDOWN_FLUSH_TIMEOUT_SEC
        try:
            agent_mod._SHUTDOWN_FLUSH_TIMEOUT_SEC = 0.2
            t0 = time.monotonic()
            _shutdown_sync_queue(report, q, 'test', may_send=True)
            elapsed = time.monotonic() - t0
        finally:
            agent_mod._SHUTDOWN_FLUSH_TIMEOUT_SEC = previous

        self.assertTrue(started.wait(1.0))
        self.assertLess(elapsed, 2.0)
        self.assertTrue(q.empty())
        self.assertEqual(q.unfinished_tasks, 0)

    def test_abandon_async_queue_unblocks_join(self):
        async def _run():
            q = asyncio.Queue()
            await q.put('a')
            await q.put('b')
            join_task = asyncio.create_task(q.join())
            await asyncio.sleep(0.05)
            self.assertFalse(join_task.done())
            abandoned = await _abandon_async_queue(q)
            self.assertEqual(abandoned, 2)
            await asyncio.wait_for(join_task, timeout=1.0)

        asyncio.run(_run())

    def test_cancel_pending_tasks_excludes_caller(self):
        """Regression: gathering the caller's own task made shutdown hang until the outer budget."""
        loop = asyncio.new_event_loop()
        thread = Thread(target=loop.run_forever, daemon=True)
        thread.start()

        async def _forever():
            while True:
                await asyncio.sleep(0.05)

        try:
            async def _spawn():
                return [asyncio.create_task(_forever()) for _ in range(2)]

            reporter_tasks = asyncio.run_coroutine_threadsafe(_spawn(), loop).result(timeout=2.0)
            time.sleep(0.1)

            t0 = time.monotonic()
            future = asyncio.run_coroutine_threadsafe(_cancel_pending_tasks(reporter_tasks), loop)
            future.result(timeout=2.0)
            self.assertLess(time.monotonic() - t0, 2.0)
            for task in reporter_tasks:
                self.assertTrue(task.cancelled() or task.done())
        finally:
            loop.call_soon_threadsafe(loop.stop)
            thread.join(timeout=2.0)
            loop.close()

    def test_async_shutdown_cleanup_runs_inside_asyncio_run_root(self):
        """
        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()

        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()

            async def reporter():
                while not finished.is_set():
                    await asyncio.sleep(0.02)

            tasks = {asyncio.create_task(reporter()) for _ in range(2)}
            loop_ready.set()
            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))
        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))

    def test_background_task_failure_is_logged_not_silenced(self):
        """Regression: failing background tasks must be observed and logged exactly once."""
        holder = {'errors': []}
        error_logged = Event()
        aclose_done = Event()

        class _Handler(logging.Handler):
            def emit(self, record):
                msg = record.getMessage()
                if 'Error in Python agent asyncio event loop' in msg:
                    holder['errors'].append(msg)
                    error_logged.set()

        agent_logger = logging.getLogger('skywalking')
        handler = _Handler()
        agent_logger.addHandler(handler)
        previous_level = agent_logger.level
        agent_logger.setLevel(logging.ERROR)

        async def failing_background():
            await asyncio.sleep(0.02)
            raise ValueError('command dispatch failed')

        async def yielding_aclose():
            await asyncio.sleep(0.02)
            aclose_done.set()

        async def root():
            finished = asyncio.Event()
            holder['finished'] = finished
            failing_task = asyncio.create_task(failing_background())

            async def reporter():
                while not finished.is_set():
                    await asyncio.sleep(0.05)

            reporter_task = asyncio.create_task(reporter())
            tasks = {failing_task, reporter_task}
            await _await_shutdown_or_background_failure(finished, tasks)
            await _cancel_pending_tasks(tasks)
            await yielding_aclose()
            holder['after_wait'] = True
            holder['failing_task'] = failing_task

        try:
            asyncio.run(root())
            self.assertTrue(holder.get('after_wait'))
            self.assertTrue(error_logged.wait(2.0))
            self.assertTrue(aclose_done.wait(2.0))
            self.assertEqual(len(holder['errors']), 1)
            self.assertIn('command dispatch failed', holder['errors'][0])
            self.assertTrue(holder['finished'].is_set())
            self.assertTrue(holder['failing_task'].done())
            self.assertIsInstance(holder['failing_task'].exception(), ValueError)
        finally:
            agent_logger.removeHandler(handler)
            agent_logger.setLevel(previous_level)

    def test_clean_shutdown_does_not_log_normal_background_completion(self):
        """Clean shutdown: reporters finish after _finished; cleanup must stay silent."""
        holder = {'errors': []}

        class _Handler(logging.Handler):
            def emit(self, record):
                if 'Error in Python agent asyncio event loop' in record.getMessage():
                    holder['errors'].append(record.getMessage())

        agent_logger = logging.getLogger('skywalking')
        handler = _Handler()
        agent_logger.addHandler(handler)
        previous_level = agent_logger.level
        agent_logger.setLevel(logging.ERROR)

        async def root():
            finished = asyncio.Event()

            async def reporter():
                while not finished.is_set():
                    await asyncio.sleep(0.01)

            tasks = {asyncio.create_task(reporter()) for _ in range(2)}
            finished.set()
            await _await_shutdown_or_background_failure(finished, tasks)
            # Give reporters a turn to observe _finished and return normally.
            await asyncio.sleep(0.05)
            await _cancel_pending_tasks(tasks)
            holder['all_done'] = all(task.done() for task in tasks)

        try:
            asyncio.run(root())
            self.assertTrue(holder.get('all_done'))
            self.assertEqual(holder['errors'], [])
        finally:
            agent_logger.removeHandler(handler)
            agent_logger.setLevel(previous_level)

    def test_unexpected_pre_shutdown_cancellation_is_logged(self):
        """Cancellation before shutdown must be logged and trigger orderly cleanup."""
        holder = {'errors': []}
        error_logged = Event()
        aclose_done = Event()

        class _Handler(logging.Handler):
            def emit(self, record):
                msg = record.getMessage()
                if 'Error in Python agent asyncio event loop' in msg:
                    holder['errors'].append(msg)
                    error_logged.set()

        agent_logger = logging.getLogger('skywalking')
        handler = _Handler()
        agent_logger.addHandler(handler)
        previous_level = agent_logger.level
        agent_logger.setLevel(logging.ERROR)

        async def cancelled_background():
            asyncio.current_task().cancel()
            await asyncio.sleep(0)

        async def yielding_aclose():
            await asyncio.sleep(0.02)
            aclose_done.set()

        async def root():
            finished = asyncio.Event()
            holder['finished'] = finished

            async def reporter():
                while not finished.is_set():
                    await asyncio.sleep(0.05)

            cancelled_task = asyncio.create_task(cancelled_background())
            reporter_task = asyncio.create_task(reporter())
            tasks = {cancelled_task, reporter_task}
            await _await_shutdown_or_background_failure(finished, tasks)
            await _cancel_pending_tasks(tasks)
            await yielding_aclose()
            holder['after_wait'] = True

        try:
            asyncio.run(root())
            self.assertTrue(holder.get('after_wait'))
            self.assertTrue(error_logged.wait(2.0))
            self.assertTrue(aclose_done.wait(2.0))
            self.assertEqual(len(holder['errors']), 1)
            self.assertIn('cancelled unexpectedly', holder['errors'][0])
            self.assertTrue(holder['finished'].is_set())
        finally:
            agent_logger.removeHandler(handler)
            agent_logger.setLevel(previous_level)

    def test_intentional_cleanup_cancellation_is_silent(self):
        """Tasks cancelled by _cancel_pending_tasks during cleanup must not ERROR."""
        holder = {'errors': []}

        class _Handler(logging.Handler):
            def emit(self, record):
                if 'Error in Python agent asyncio event loop' in record.getMessage():
                    holder['errors'].append(record.getMessage())

        agent_logger = logging.getLogger('skywalking')
        handler = _Handler()
        agent_logger.addHandler(handler)
        previous_level = agent_logger.level
        agent_logger.setLevel(logging.ERROR)

        async def forever():
            while True:
                await asyncio.sleep(0.05)

        async def root():
            tasks = {asyncio.create_task(forever()) for _ in range(2)}
            await _cancel_pending_tasks(tasks)
            holder['all_done'] = all(task.done() for task in tasks)

        try:
            asyncio.run(root())
            self.assertTrue(holder.get('all_done'))
            self.assertEqual(holder['errors'], [])
        finally:
            agent_logger.removeHandler(handler)
            agent_logger.setLevel(previous_level)

    def test_shutdown_async_queue_bounded(self):
        async def _run():
            q = asyncio.Queue()
            await q.put('a')
            t0 = time.monotonic()
            await _shutdown_async_queue(q, 'test')
            self.assertLess(time.monotonic() - t0, 2.0)
            self.assertTrue(q.empty())

        asyncio.run(_run())


if __name__ == '__main__':
    unittest.main()