Shut-down immediate consumes queue

Author: EpicWink

Lib/queue.py

@@ -29,10 +29,6 @@ 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.
 
@@ -63,7 +59,7 @@ def __init__(self, maxsize=0):
         self.unfinished_tasks = 0
 
         # Queue shutdown state
-        self.shutdown_state = _queue_alive
+        self.is_shutdown = False
 
     def task_done(self):
         '''Indicate that a formerly enqueued task is complete.
@@ -82,8 +78,6 @@ def task_done(self):
         Raises ShutDown if the queue has been shut down immediately.
         '''
         with self.all_tasks_done:
-            if self._is_shutdown_immediate():
-                raise ShutDown
             unfinished = self.unfinished_tasks - 1
             if unfinished <= 0:
                 if unfinished < 0:
@@ -103,12 +97,8 @@ def join(self):
         Raises ShutDown if the queue has been shut down immediately.
         '''
         with self.all_tasks_done:
-            if self._is_shutdown_immediate():
-                raise ShutDown
             while self.unfinished_tasks:
                 self.all_tasks_done.wait()
-                if self._is_shutdown_immediate():
-                    raise ShutDown
 
     def qsize(self):
         '''Return the approximate size of the queue (not reliable!).'''
@@ -154,7 +144,7 @@ def put(self, item, block=True, timeout=None):
         Raises ShutDown if the queue has been shut down.
         '''
         with self.not_full:
-            if not self._is_alive():
+            if self.is_shutdown:
                 raise ShutDown
             if self.maxsize > 0:
                 if not block:
@@ -163,7 +153,7 @@ def put(self, item, block=True, timeout=None):
                 elif timeout is None:
                     while self._qsize() >= self.maxsize:
                         self.not_full.wait()
-                        if not self._is_alive():
+                        if self.is_shutdown:
                             raise ShutDown
                 elif timeout < 0:
                     raise ValueError("'timeout' must be a non-negative number")
@@ -174,7 +164,7 @@ def put(self, item, block=True, timeout=None):
                         if remaining <= 0.0:
                             raise Full
                         self.not_full.wait(remaining)
-                        if not self._is_alive():
+                        if self.is_shutdown:
                             raise ShutDown
             self._put(item)
             self.unfinished_tasks += 1
@@ -195,16 +185,15 @@ def get(self, block=True, timeout=None):
         or if the queue has been shut down immediately.
         '''
         with self.not_empty:
-            if self._is_shutdown_immediate() or\
-                (self._is_shutdown() and not self._qsize()):
+            if self.is_shutdown and not self._qsize():
                 raise ShutDown
             if not block:
                 if not self._qsize():
                     raise Empty
             elif timeout is None:
                 while not self._qsize():
                     self.not_empty.wait()
-                    if self._is_shutdown_immediate():
+                    if self.is_shutdown and not self._qsize():
                         raise ShutDown
             elif timeout < 0:
                 raise ValueError("'timeout' must be a non-negative number")
@@ -215,7 +204,7 @@ def get(self, block=True, timeout=None):
                     if remaining <= 0.0:
                         raise Empty
                     self.not_empty.wait(remaining)
-                    if self._is_shutdown_immediate():
+                    if self.is_shutdown and not self._qsize():
                         raise ShutDown
             item = self._get()
             self.not_full.notify()
@@ -247,32 +236,16 @@ def shutdown(self, immediate=False):
         and join() if 'immediate'. The ShutDown exception is raised.
         '''
         with self.mutex:
-            if self._is_shutdown_immediate():
-                return
+            self.is_shutdown = True
             if immediate:
-                self._set_shutdown_immediate()
+                while self._qsize():
+                    self._get()
+                self.unfinished_tasks = 0
                 self.not_empty.notify_all()
                 # release all blocked threads in `join()`
                 self.all_tasks_done.notify_all()
-            else:
-                self._set_shutdown()
             self.not_full.notify_all()
 
-    def _is_alive(self):
-        return self.shutdown_state == _queue_alive
-
-    def _is_shutdown(self):
-        return self.shutdown_state == _queue_shutdown
-
-    def _is_shutdown_immediate(self):
-        return self.shutdown_state == _queue_shutdown_immediate
-
-    def _set_shutdown(self):
-        self.shutdown_state = _queue_shutdown
-
-    def _set_shutdown_immediate(self):
-        self.shutdown_state = _queue_shutdown_immediate
-
 
     # Override these methods to implement other queue organizations
     # (e.g. stack or priority queue).

Lib/test/test_queue.py

@@ -267,16 +267,15 @@ def test_shutdown_immediate(self):
     def test_shutdown_allowed_transitions(self):
         # allowed transitions would be from alive via shutdown to immediate
         q = self.type2test()
-        self.assertEqual("alive", q.shutdown_state)
+        self.assertFalse(q.is_shutdown)
 
         q.shutdown()
-        self.assertEqual("shutdown", q.shutdown_state)
+        self.assertTrue(q.is_shutdown)
 
         q.shutdown(immediate=True)
-        self.assertEqual("shutdown-immediate", q.shutdown_state)
+        self.assertTrue(q.is_shutdown)
 
         q.shutdown(immediate=False)
-        self.assertNotEqual("shutdown", q.shutdown_state)
 
     def _shutdown_all_methods_in_one_thread(self, immediate):
         q = self.type2test(2)
@@ -293,10 +292,9 @@ def _shutdown_all_methods_in_one_thread(self, immediate):
                 q.get()
             with self.assertRaises(self.queue.ShutDown):
                 q.get_nowait()
-            with self.assertRaises(self.queue.ShutDown):
+            with self.assertRaises(ValueError):
                 q.task_done()
-            with self.assertRaises(self.queue.ShutDown):
-                q.join()
+            q.join()
         else:
             self.assertIn(q.get(), "LO")
             q.task_done()
@@ -333,10 +331,7 @@ def _write_msg_thread(self, q, n, results, delay,
                 event_end.set()
             time.sleep(delay)
         # end of all puts
-        try:
-            q.join()
-        except self.queue.ShutDown:
-            pass
+        q.join()
 
     def _read_msg_thread(self, q, nb, results, delay, event_start):
         event_start.wait()
@@ -355,26 +350,17 @@ def _read_msg_thread(self, q, nb, results, delay, event_start):
                 nb -= 1
             except self.queue.Empty:
                 pass
-        try:
-            q.join()
-        except self.queue.ShutDown:
-            pass
+        q.join()
 
     def _shutdown_thread(self, q, event_end, immediate):
         event_end.wait()
         q.shutdown(immediate)
-        try:
-            q.join()
-        except self.queue.ShutDown:
-            pass
+        q.join()
 
     def _join_thread(self, q, delay, event_start):
         event_start.wait()
         time.sleep(delay)
-        try:
-            q.join()
-        except self.queue.ShutDown:
-            pass
+        q.join()
 
     def _shutdown_all_methods_in_many_threads(self, immediate):
         q = self.type2test()
@@ -413,6 +399,9 @@ def _shutdown_all_methods_in_many_threads(self, immediate):
             assert(len(res_gets) <= len(res_puts))
             assert(res_gets.count(True) <= res_puts.count(True))
 
+        for thread in ps[1:]:
+            thread.join()
+
     def test_shutdown_all_methods_in_many_threads(self):
         return self._shutdown_all_methods_in_many_threads(False)
 
@@ -544,15 +533,9 @@ def _shutdown_join(self, immediate):
         go = threading.Event()
         nb = q.qsize()
 
-        if immediate:
-            thrds = (
-                (self._join_shutdown, (q, results)),
-                (self._join_shutdown, (q, results)),
-            )
-        else:
-            thrds = (
-                (self._join, (q, results)),
-                (self._join, (q, results)),
+        thrds = (
+            (self._join, (q, results)),
+            (self._join, (q, results)),
         )
         threads = []
         for func, params in thrds:
@@ -584,21 +567,15 @@ def _shutdown_put_join(self, immediate):
         nb = q.qsize()
         # queue not fulled
 
-        if immediate:
-            thrds = (
-                (self._put_shutdown, (q, "E", go, results)),
-                (self._join_shutdown, (q, results)),
-            )
-        else:
-            thrds = (
-                (self._put_shutdown, (q, "E", go, results)),
-                (self._join, (q, results)),
-            )
+        thrds = (
+            (self._put_shutdown, (q, "E", go, results)),
+            (self._join, (q, results)),
+        )
         threads = []
         for func, params in thrds:
             threads.append(threading.Thread(target=func, args=params))
             threads[-1].start()
-        if not immediate:
+        if not immediate or immediate:  # TODO: dedent (minimising Git diff)
             self.assertEqual(q.unfinished_tasks, nb)
             for i in range(nb):
                 t = threading.Thread(target=q.task_done)