gh-96471: Add queue shutdown, next step.

Lib/asyncio/queues.py

@@ -1,6 +1,14 @@
-__all__ = ('Queue', 'PriorityQueue', 'LifoQueue', 'QueueFull', 'QueueEmpty')
+__all__ = (
+    'Queue',
+    'PriorityQueue',
+    'LifoQueue',
+    'QueueFull',
+    'QueueEmpty',
+    'QueueShutDown',
+)
 
 import collections
+import enum
 import heapq
 from types import GenericAlias
 
@@ -18,6 +26,17 @@ class QueueFull(Exception):
     pass
 
 
+class QueueShutDown(Exception):
+    """Raised when putting on to or getting from a shut-down Queue."""
+    pass
+
+
+class _QueueState(enum.Enum):
+    ALIVE = "alive"
+    SHUTDOWN = "shutdown"
+    SHUTDOWN_IMMEDIATE = "shutdown-immediate"
+
+
 class Queue(mixins._LoopBoundMixin):
     """A queue, useful for coordinating producer and consumer coroutines.
 
@@ -41,6 +60,7 @@ def __init__(self, maxsize=0):
         self._finished = locks.Event()
         self._finished.set()
         self._init(maxsize)
+        self._shutdown_state = _QueueState.ALIVE
 
     # These three are overridable in subclasses.
 
@@ -81,6 +101,8 @@ def _format(self):
             result += f' _putters[{len(self._putters)}]'
         if self._unfinished_tasks:
             result += f' tasks={self._unfinished_tasks}'
+        if not self._is_alive():
+            result += f' state={self._shutdown_state.value}'
         return result
 
     def qsize(self):
@@ -112,7 +134,11 @@ async def put(self, item):
 
         Put an item into the queue. If the queue is full, wait until a free
         slot is available before adding item.
+
+        Raise a QueueShutDown if _shutdown_state is not set to `ALIVE`.
         """
+        if not self._is_alive():
+            raise QueueShutDown
         while self.full():
             putter = self._get_loop().create_future()
             self._putters.append(putter)
@@ -125,20 +151,28 @@ async def put(self, item):
                     self._putters.remove(putter)
                 except ValueError:
                     # The putter could be removed from self._putters by a
-                    # previous get_nowait call.
+                    # previous get_nowait call,
+                    # or a shutdown call.
                     pass
                 if not self.full() and not putter.cancelled():
                     # We were woken up by get_nowait(), but can't take
                     # the call.  Wake up the next in line.
                     self._wakeup_next(self._putters)
                 raise
+            if not self._is_alive():
+                raise QueueShutDown
         return self.put_nowait(item)
 
     def put_nowait(self, item):
         """Put an item into the queue without blocking.
 
         If no free slot is immediately available, raise QueueFull.
+
+        Raise a QueueShutDown if _shutdown_state is not set to
+        `ALIVE`.
         """
+        if not self._is_alive():
+            raise QueueShutDown
         if self.full():
             raise QueueFull
         self._put(item)
@@ -150,8 +184,18 @@ async def get(self):
         """Remove and return an item from the queue.
 
         If queue is empty, wait until an item is available.
+
+        Raise a QueueShutDown if _shutdown_state is set to
+        `SHUTDOWN_IMMEDIATE`.
+
+        Raise a QueueShutDown if _shutdown_state is set to
+        `SHUTDOWN` and queue is empty.
         """
+        if self._is_shutdown_immediate():
+            raise QueueShutDown
         while self.empty():
+            if self._is_shutdown():
+                raise QueueShutDown
             getter = self._get_loop().create_future()
             self._getters.append(getter)
             try:
@@ -163,21 +207,34 @@ async def get(self):
                     self._getters.remove(getter)
                 except ValueError:
                     # The getter could be removed from self._getters by a
-                    # previous put_nowait call.
+                    # previous put_nowait call,
+                    # or a shutdown call.
                     pass
                 if not self.empty() and not getter.cancelled():
                     # We were woken up by put_nowait(), but can't take
                     # the call.  Wake up the next in line.
                     self._wakeup_next(self._getters)
                 raise
+            if self._is_shutdown_immediate():
+                raise QueueShutDown
         return self.get_nowait()
 
     def get_nowait(self):
         """Remove and return an item from the queue.
 
         Return an item if one is immediately available, else raise QueueEmpty.
+
+        Raise a QueueShutDown if _shutdown_state is set to
+        `SHUTDOWN_IMMEDIATE`.
+
+        Raise a QueueShutDown if _shutdown_state is set to
+        `SHUTDOWN` and queue is empty.
         """
+        if self._is_shutdown_immediate():
+            raise QueueShutDown
         if self.empty():
+            if self._is_shutdown():
+                raise QueueShutDown
             raise QueueEmpty
         item = self._get()
         self._wakeup_next(self._putters)
@@ -196,7 +253,12 @@ def task_done(self):
 
         Raises ValueError if called more times than there were items placed in
         the queue.
+
+        Raise a QueueShutDown if _shutdown_state is set to
+        `SHUTDOWN_IMMEDIATE`.
         """
+        if self._is_shutdown_immediate():
+            raise QueueShutDown
         if self._unfinished_tasks <= 0:
             raise ValueError('task_done() called too many times')
         self._unfinished_tasks -= 1
@@ -210,9 +272,58 @@ async def join(self):
         queue. The count goes down whenever a consumer calls task_done() to
         indicate that the item was retrieved and all work on it is complete.
         When the count of unfinished tasks drops to zero, join() unblocks.
+
+        Raise a QueueShutDown if _shutdown_state is set to
+        `SHUTDOWN_IMMEDIATE`.
         """
+        if self._is_shutdown_immediate():
+            raise QueueShutDown
         if self._unfinished_tasks > 0:
             await self._finished.wait()
+            if self._is_shutdown_immediate():
+                raise QueueShutDown
+
+    def shutdown(self, immediate=False):
+        """Shut-down the queue, making queue gets and puts raise.
+
+        By default, gets will only raise once the queue is empty. Set
+        'immediate' to True to make gets raise immediately instead.
+
+        All blocked callers of put() will be unblocked, and also get()
+        and join() if 'immediate'. The QueueShutDown exception is raised.
+        """
+        if self._is_shutdown_immediate():
+            return
+        # here _shutdown_state is ALIVE or SHUTDOWN
+        if immediate:
+            self._set_shutdown_immediate()
+            while self._getters:
+                getter = self._getters.popleft()
+                if not getter.done():
+                    getter.set_result(None)
+            # Release all 'blocked' tasks/coros in `join()`
+            self._finished.set()
+        else:
+            self._set_shutdown()
+        while self._putters:
+            putter = self._putters.popleft()
+            if not putter.done():
+                putter.set_result(None)
+
+    def _is_alive(self):
+        return self._shutdown_state is _QueueState.ALIVE
+
+    def _is_shutdown(self):
+        return self._shutdown_state is _QueueState.SHUTDOWN
+
+    def _is_shutdown_immediate(self):
+        return self._shutdown_state is _QueueState.SHUTDOWN_IMMEDIATE
+
+    def _set_shutdown(self):
+        self._shutdown_state = _QueueState.SHUTDOWN
+
+    def _set_shutdown_immediate(self):
+        self._shutdown_state = _QueueState.SHUTDOWN_IMMEDIATE
 
 
 class PriorityQueue(Queue):

Lib/multiprocessing/queues.py

@@ -18,7 +18,7 @@
 import weakref
 import errno
 
-from queue import Empty, Full
+from queue import Empty, Full, ShutDown
 
 import _multiprocessing
 
@@ -28,6 +28,10 @@
 
 from .util import debug, info, Finalize, register_after_fork, is_exiting
 
+_queue_alive = 0
+_queue_shutdown = 1
+_queue_shutdown_immediate = 2
+
 #
 # Queue type using a pipe, buffer and thread
 #
@@ -50,18 +54,21 @@ def __init__(self, maxsize=0, *, ctx):
         # For use by concurrent.futures
         self._ignore_epipe = False
         self._reset()
+        self._shutdown_state = ctx.Value('i', _queue_alive)
 
         if sys.platform != 'win32':
             register_after_fork(self, Queue._after_fork)
 
     def __getstate__(self):
         context.assert_spawning(self)
         return (self._ignore_epipe, self._maxsize, self._reader, self._writer,
-                self._rlock, self._wlock, self._sem, self._opid)
+                self._rlock, self._wlock, self._sem, self._opid,
+                self._shutdown_state)
 
     def __setstate__(self, state):
         (self._ignore_epipe, self._maxsize, self._reader, self._writer,
-         self._rlock, self._wlock, self._sem, self._opid) = state
+         self._rlock, self._wlock, self._sem, self._opid,
+         self._shutdown_state) = state
         self._reset()
 
     def _after_fork(self):
@@ -83,9 +90,26 @@ def _reset(self, after_fork=False):
         self._recv_bytes = self._reader.recv_bytes
         self._poll = self._reader.poll
 
+    def _is_alive(self):
+        return self._shutdown_state.value == _queue_alive
+
+    def _is_shutdown(self):
+        return self._shutdown_state.value == _queue_shutdown
+
+    def _is_shutdown_immediate(self):
+        return self._shutdown_state.value == _queue_shutdown_immediate
+
+    def _set_shutdown(self):
+        self._shutdown_state.value = _queue_shutdown
+
+    def _set_shutdown_immediate(self):
+        self._shutdown_state.value = _queue_shutdown_immediate
+
     def put(self, obj, block=True, timeout=None):
         if self._closed:
             raise ValueError(f"Queue {self!r} is closed")
+        if not self._is_alive():
+            raise ShutDown
         if not self._sem.acquire(block, timeout):
             raise Full
 
@@ -98,6 +122,9 @@ def put(self, obj, block=True, timeout=None):
     def get(self, block=True, timeout=None):
         if self._closed:
             raise ValueError(f"Queue {self!r} is closed")
+        if self._is_shutdown_immediate() or\
+            (self._is_shutdown() and self.empty()):
+            raise ShutDown
         if block and timeout is None:
             with self._rlock:
                 res = self._recv_bytes()
@@ -118,6 +145,8 @@ def get(self, block=True, timeout=None):
                 self._sem.release()
             finally:
                 self._rlock.release()
+        if self._is_shutdown_immediate():
+            raise ShutDown
         # unserialize the data after having released the lock
         return _ForkingPickler.loads(res)
 
@@ -137,6 +166,19 @@ def get_nowait(self):
     def put_nowait(self, obj):
         return self.put(obj, False)
 
+    def shutdown(self, immediate=False):
+        if self._closed:
+            raise ValueError(f"Queue {self!r} is closed")
+        with self._shutdown_state.get_lock():
+            if self._is_shutdown_immediate():
+                return
+            if immediate:
+                self._set_shutdown_immediate()
+                with self._notempty:
+                    self._notempty.notify_all()
+            else:
+                self._set_shutdown()
+
     def close(self):
         self._closed = True
         close = self._close
@@ -310,6 +352,8 @@ def __setstate__(self, state):
     def put(self, obj, block=True, timeout=None):
         if self._closed:
             raise ValueError(f"Queue {self!r} is closed")
+        if not self._is_alive():
+            raise ShutDown
         if not self._sem.acquire(block, timeout):
             raise Full
 
@@ -322,15 +366,28 @@ def put(self, obj, block=True, timeout=None):
 
     def task_done(self):
         with self._cond:
+            if self._is_shutdown_immediate():
+                raise ShutDown
             if not self._unfinished_tasks.acquire(False):
                 raise ValueError('task_done() called too many times')
             if self._unfinished_tasks._semlock._is_zero():
                 self._cond.notify_all()
 
     def join(self):
         with self._cond:
+            if self._is_shutdown_immediate():
+                raise ShutDown
             if not self._unfinished_tasks._semlock._is_zero():
                 self._cond.wait()
+                if self._is_shutdown_immediate():
+                    raise ShutDown
+
+    def shutdown(self, immediate=False):
+        with self._cond:
+            is_alive = self._is_alive()
+            super().shutdown(immediate)
+            if is_alive:
+                self._cond.notify_all() # here to check YD
 
 #
 # Simplified Queue type -- really just a locked pipe