From b2d51f6dbc020a3b30b6de1cf99dc6f2724162bc Mon Sep 17 00:00:00 2001 From: Sanjay Santhanam <51058514+Sanjays2402@users.noreply.github.com> Date: Mon, 13 Jul 2026 11:32:05 -0700 Subject: [PATCH] queues: Preserve items when waiting consumers are canceled Queue.put_nowait resolves a waiting getter before its consumer task resumes. If that task is canceled in the gap, cancellation replaces the delivered result and the item disappears. Use a cancellation-aware getter future to return a delivered item to the queue when cancellation wins that race. Add regression coverage for the put-then-cancel ordering. Fixes #2826. --- tornado/queues.py | 28 ++++++++++++++++++++++++++-- tornado/test/queues_test.py | 16 ++++++++++++++++ 2 files changed, 42 insertions(+), 2 deletions(-) diff --git a/tornado/queues.py b/tornado/queues.py index 3eeb13e505..d5e05f2cbb 100644 --- a/tornado/queues.py +++ b/tornado/queues.py @@ -27,10 +27,11 @@ from __future__ import annotations +import asyncio import collections import datetime import heapq -from collections.abc import Awaitable +from collections.abc import Awaitable, Generator from typing import Any, Generic, TypeVar from tornado import gen, ioloop @@ -66,6 +67,20 @@ def on_timeout() -> None: future.add_done_callback(lambda _: io_loop.remove_timeout(timeout_handle)) +class _QueueGetter(Future[_T]): + def __init__(self, queue: Queue[_T]) -> None: + super().__init__() + self.queue = queue + + def __await__(self) -> Generator[Any, None, _T]: + try: + return (yield from super().__await__()) + except asyncio.CancelledError: + if self.done() and not self.cancelled(): + self.queue._put_cancelled(self.result()) + raise + + class _QueueIterator(Generic[_T]): def __init__(self, q: Queue[_T]) -> None: self.q = q @@ -248,7 +263,7 @@ def get(self, timeout: float | datetime.timedelta | None = None) -> Awaitable[_T A ``timeout`` argument of zero will either return or raise immediately. Previously, zero was treated equivalent to ``None`` (wait forever). """ - future: Future[_T] = Future() + future: Future[_T] = _QueueGetter(self) try: future.set_result(self.get_nowait()) except QueueEmpty: @@ -326,6 +341,15 @@ def __put_internal(self, item: _T) -> None: self._finished.clear() self._put(item) + def _put_cancelled(self, item: _T) -> None: + self._consume_expired() + if self._getters: + assert self.empty(), "queue non-empty, why are getters waiting?" + getter = self._getters.popleft() + future_set_result_unless_cancelled(getter, item) + else: + self._put(item) + def _consume_expired(self) -> None: # Remove timed-out waiters. while self._putters and self._putters[0][1].done(): diff --git a/tornado/test/queues_test.py b/tornado/test/queues_test.py index c29eea9c06..35b94dc364 100644 --- a/tornado/test/queues_test.py +++ b/tornado/test/queues_test.py @@ -133,6 +133,22 @@ def test_get_timeout_preempted(self): yield gen.sleep(0.02) self.assertEqual(0, (yield get)) + @gen_test + async def test_cancelled_get_does_not_lose_item(self): + q: queues.Queue[int] = queues.Queue() + + async def consume(): + return await q.get() + + consumer = asyncio.ensure_future(consume()) + await asyncio.sleep(0) + q.put_nowait(0) + consumer.cancel() + + with self.assertRaises(asyncio.CancelledError): + await consumer + self.assertEqual(0, q.get_nowait()) + @gen_test def test_get_clears_timed_out_putters(self): q: queues.Queue[int] = queues.Queue(1)