queues: Preserve items when waiting consumers are canceled

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

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)