Add threading queue shutdown

* Include docs

Author: EpicWink

Doc/library/queue.rst

@@ -93,6 +93,14 @@ The :mod:`queue` module defines the following classes and exceptions:
    on a :class:`Queue` object which is full.
 
 
+.. exception:: ShutDown
+
+   Exception raised when :meth:`~Queue.put` or :meth:`~Queue.get` is called on
+   a :class:`Queue` object which has been shut down.
+
+   .. versionadded:: 3.12
+
+
 .. _queueobjects:
 
 Queue Objects
@@ -135,6 +143,8 @@ provide the public methods described below.
    immediately available, else raise the :exc:`Full` exception (*timeout* is
    ignored in that case).
 
+   Raises :exc:`ShutDown` if the queue has been shut down.
+
 
 .. method:: Queue.put_nowait(item)
 
@@ -155,6 +165,9 @@ provide the public methods described below.
    an uninterruptible wait on an underlying lock.  This means that no exceptions
    can occur, and in particular a SIGINT will not trigger a :exc:`KeyboardInterrupt`.
 
+   Raises :exc:`ShutDown` if the queue has been shut down and is empty, or if
+   the queue has been shut down immediately.
+
 
 .. method:: Queue.get_nowait()
 
@@ -177,6 +190,8 @@ fully processed by daemon consumer threads.
    Raises a :exc:`ValueError` if called more times than there were items placed in
    the queue.
 
+   Raises :exc:`ShutDown` if the queue has been shut down immediately.
+
 
 .. method:: Queue.join()
 
@@ -187,6 +202,8 @@ fully processed by daemon consumer threads.
    indicate that the item was retrieved and all work on it is complete.  When the
    count of unfinished tasks drops to zero, :meth:`join` unblocks.
 
+   Raises :exc:`ShutDown` if the queue has been shut down immediately.
+
 
 Example of how to wait for enqueued tasks to be completed::
 
@@ -214,6 +231,25 @@ Example of how to wait for enqueued tasks to be completed::
     print('All work completed')
 
 
+Terminating queues
+^^^^^^^^^^^^^^^^^^
+
+:class:`Queue` objects can be made to prevent further interaction by shutting
+them down.
+
+.. method:: Queue.shutdown(immediate=False)
+
+   Shut-down the queue, making queue gets and puts raise :exc:`ShutDown`.
+
+   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* is true.
+
+   .. versionadded:: 3.12
+
+
 SimpleQueue Objects
 -------------------
 

Lib/queue.py

@@ -25,6 +25,15 @@ class Full(Exception):
     pass
 
 
+class ShutDown(Exception):
+    '''Raised when put/get with shut-down queue.'''
+
+
+_queue_alive = "alive"
+_queue_shutdown = "shutdown"
+_queue_shutdown_immediate = "shutdown-immediate"
+
+
 class Queue:
     '''Create a queue object with a given maximum size.
 
@@ -54,6 +63,9 @@ def __init__(self, maxsize=0):
         self.all_tasks_done = threading.Condition(self.mutex)
         self.unfinished_tasks = 0
 
+        # Queue shut-down state
+        self.shutdown_state = _queue_alive
+
     def task_done(self):
         '''Indicate that a formerly enqueued task is complete.
 
@@ -87,6 +99,8 @@ def join(self):
         '''
         with self.all_tasks_done:
             while self.unfinished_tasks:
+                if self.shutdown_state == _queue_shutdown_immediate:
+                    return
                 self.all_tasks_done.wait()
 
     def qsize(self):
@@ -130,6 +144,8 @@ def put(self, item, block=True, timeout=None):
         is immediately available, else raise the Full exception ('timeout'
         is ignored in that case).
         '''
+        if self.shutdown_state != _queue_alive:
+            raise ShutDown
         with self.not_full:
             if self.maxsize > 0:
                 if not block:
@@ -138,6 +154,8 @@ def put(self, item, block=True, timeout=None):
                 elif timeout is None:
                     while self._qsize() >= self.maxsize:
                         self.not_full.wait()
+                        if self.shutdown_state != _queue_alive:
+                            raise ShutDown
                 elif timeout < 0:
                     raise ValueError("'timeout' must be a non-negative number")
                 else:
@@ -147,6 +165,8 @@ def put(self, item, block=True, timeout=None):
                         if remaining <= 0.0:
                             raise Full
                         self.not_full.wait(remaining)
+                        if self.shutdown_state != _queue_alive:
+                            raise ShutDown
             self._put(item)
             self.unfinished_tasks += 1
             self.not_empty.notify()
@@ -162,22 +182,36 @@ def get(self, block=True, timeout=None):
         available, else raise the Empty exception ('timeout' is ignored
         in that case).
         '''
+        if self.shutdown_state == _queue_shutdown_immediate:
+            raise ShutDown
         with self.not_empty:
             if not block:
                 if not self._qsize():
+                    if self.shutdown_state != _queue_alive:
+                        raise ShutDown
                     raise Empty
             elif timeout is None:
                 while not self._qsize():
+                    if self.shutdown_state != _queue_alive:
+                        raise ShutDown
                     self.not_empty.wait()
+                    if self.shutdown_state != _queue_alive:
+                        raise ShutDown
             elif timeout < 0:
                 raise ValueError("'timeout' must be a non-negative number")
             else:
                 endtime = time() + timeout
                 while not self._qsize():
+                    if self.shutdown_state != _queue_alive:
+                        raise ShutDown
                     remaining = endtime - time()
                     if remaining <= 0.0:
                         raise Empty
                     self.not_empty.wait(remaining)
+                    if self.shutdown_state != _queue_alive:
+                        raise ShutDown
+            if self.shutdown_state == _queue_shutdown_immediate:
+                raise ShutDown
             item = self._get()
             self.not_full.notify()
             return item
@@ -198,6 +232,24 @@ def get_nowait(self):
         '''
         return self.get(block=False)
 
+    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 ShutDown exception is raised.
+        '''
+        with self.mutex:
+            if immediate:
+                self.shutdown_state = _queue_shutdown_immediate
+                self.not_empty.notify_all()
+                self.all_tasks_done.notify_all()
+            else:
+                self.shutdown_state = _queue_shutdown
+            self.not_full.notify_all()
+
     # Override these methods to implement other queue organizations
     # (e.g. stack or priority queue).
     # These will only be called with appropriate locks held

Lib/test/test_queue.py

@@ -241,6 +241,41 @@ def test_shrinking_queue(self):
         with self.assertRaises(self.queue.Full):
             q.put_nowait(4)
 
+    def test_shutdown_empty(self):
+        q = self.type2test()
+        q.shutdown()
+        try:
+            q.put("data")
+            self.fail("Didn't appear to shut-down queue")
+        except self.queue.ShutDown:
+            pass
+        try:
+            q.get()
+            self.fail("Didn't appear to shut-down queue")
+        except self.queue.ShutDown:
+            pass
+
+    def test_shutdown_nonempty(self):
+        q = self.type2test()
+        q.put("data")
+        q.shutdown()
+        q.get()
+        try:
+            q.get()
+            self.fail("Didn't appear to shut-down queue")
+        except self.queue.ShutDown:
+            pass
+
+    def test_shutdown_immediate(self):
+        q = self.type2test()
+        q.put("data")
+        q.shutdown(immediate=True)
+        try:
+            q.get()
+            self.fail("Didn't appear to shut-down queue")
+        except self.queue.ShutDown:
+            pass
+
 class QueueTest(BaseQueueTestMixin):
 
     def setUp(self):