fix(logging_worker): carry queued tasks across event-loop change instead of dropping them

LoggingWorker._ensure_queue nulled self._queue on a loop change, discarding every
pending LoggingTask (each an un-awaited spend-logging coroutine) with no counter and
only a debug log. SDK callers using asyncio.run() per request and mixed sync/async
processes rebind the queue's loop and silently lose spend rows and observability events.

Drain the stale queue and move the pending tasks onto a fresh queue bound to the new
loop, warn with the carried-over count, and keep flush()/join() honest since the queue
is no longer thrown away. Adds a regression test that fills the queue before the loop
change and asserts every task survives and still executes.
This commit is contained in:
mateo-berri 2026-08-24 13:46:09 -07:00
parent 85d5ac2b5c
commit 12a34a10d8
2 changed files with 70 additions and 5 deletions

View file

@ -5,7 +5,7 @@ import asyncio
import atexit import atexit
import contextvars import contextvars
import logging import logging
from collections.abc import Coroutine from collections.abc import Coroutine, Iterator
from typing import Final from typing import Final
from typing_extensions import TypedDict from typing_extensions import TypedDict
@ -61,6 +61,19 @@ class LoggingWorker:
# Register cleanup handler to flush remaining events on exit # Register cleanup handler to flush remaining events on exit
atexit.register(self._flush_on_exit) atexit.register(self._flush_on_exit)
@staticmethod
def _drain_pending(queue: "asyncio.Queue[LoggingTask]") -> tuple[LoggingTask, ...]:
"""Pop every task still queued, without awaiting them, so they can be moved to another queue."""
def _pop_until_empty() -> Iterator[LoggingTask]:
while True:
try:
yield queue.get_nowait()
except asyncio.QueueEmpty:
return
return tuple(_pop_until_empty())
def _ensure_queue(self) -> None: def _ensure_queue(self) -> None:
"""Initialize the queue if it doesn't exist or if event loop has changed.""" """Initialize the queue if it doesn't exist or if event loop has changed."""
try: try:
@ -69,14 +82,27 @@ class LoggingWorker:
# No running loop, can't initialize # No running loop, can't initialize
return return
# Check if we need to reinitialize due to event loop change # The queue, semaphore and worker task are all bound to the loop that created them. On a
# loop change we hand the still-pending tasks to a fresh queue instead of dropping them,
# so queued spend-logging coroutines are not silently discarded (and never left un-awaited).
if self._queue is not None and self._bound_loop is not current_loop: if self._queue is not None and self._bound_loop is not current_loop:
carried_over: Final = self._drain_pending(self._queue)
new_queue: Final[asyncio.Queue[LoggingTask]] = asyncio.Queue(maxsize=self.max_queue_size)
for carried_task in carried_over:
new_queue.put_nowait(carried_task)
if carried_over:
verbose_logger.warning(
"LoggingWorker: event loop changed; carried %d pending logging task(s) onto the new loop",
len(carried_over),
)
else:
verbose_logger.debug("LoggingWorker: Event loop changed, reinitializing queue and worker") verbose_logger.debug("LoggingWorker: Event loop changed, reinitializing queue and worker")
# Clear old state - these are bound to the old loop
self._queue = None
self._sem = None self._sem = None
self._worker_task = None self._worker_task = None
self._running_tasks.clear() self._running_tasks.clear()
self._queue = new_queue
self._bound_loop = current_loop
return
if self._queue is None: if self._queue is None:
self._queue = asyncio.Queue(maxsize=self.max_queue_size) self._queue = asyncio.Queue(maxsize=self.max_queue_size)

View file

@ -413,3 +413,42 @@ class TestLoggingWorker:
assert worker2._bound_loop is not None assert worker2._bound_loop is not None
await worker2.stop() await worker2.stop()
def test_event_loop_change_carries_pending_tasks_over(self):
"""Regression (LIT-6028): a loop change must not silently drop queued coroutines.
Before the fix ``_ensure_queue`` nulled ``self._queue`` on a loop change, discarding
every pending ``LoggingTask`` (each an un-awaited spend-logging coroutine). The tasks
must instead be moved onto the queue bound to the new loop and still execute there.
"""
worker = LoggingWorker(timeout=1.0, max_queue_size=10)
executed: list[int] = []
async def spend_log(index: int) -> None:
executed.append(index)
async def enqueue_on_first_loop() -> None:
worker._ensure_queue()
for i in range(5):
worker.enqueue(spend_log(i))
assert worker._queue is not None
assert worker._queue.qsize() == 5
asyncio.run(enqueue_on_first_loop())
stale_queue = worker._queue
assert stale_queue is not None
async def rebind_on_second_loop() -> None:
worker._ensure_queue()
assert worker._queue is not None
# A fresh queue bound to the new loop, holding every carried-over task (not dropped).
assert worker._queue is not stale_queue
assert worker._queue.qsize() == 5
while not worker._queue.empty():
task = worker._queue.get_nowait()
await task["context"].run(asyncio.create_task, task["coroutine"])
asyncio.run(rebind_on_second_loop())
assert sorted(executed) == [0, 1, 2, 3, 4]