ENH independent com channel for wakeup to avoid deadlocks in shutdown

Author: tomMoral

Lib/concurrent/futures/process.py

@@ -76,12 +76,36 @@
 _threads_queues = weakref.WeakKeyDictionary()
 _global_shutdown = False
 
+
+# This constants control the maximal wakeup. If a job is submitted to the
+# Executor, it might take up to _POLL_TIMEOUT for the executor to notice and
+# start launching the job. This _POLL_TIMEOUT is to be cumulated with the
+# communication overhead.
+_POLL_TIMEOUT = .001
+
+
+class _Sentinel:
+    __slot__ = ["_state"]
+
+    def __init__(self):
+        self._state = False
+
+    def set(self):
+        self._state = True
+
+    def get_and_unset(self):
+        s = self._state
+        if s:
+            self._state = False
+        return s
+
+
 def _python_exit():
     global _global_shutdown
     _global_shutdown = True
     items = list(_threads_queues.items())
-    for t, q in items:
-        q.put(None)
+    for t, wakeup in items:
+        wakeup.set()
     for t, q in items:
         t.join()
 
@@ -266,7 +290,8 @@ def _queue_management_worker(executor_reference,
                              pending_work_items,
                              work_ids_queue,
                              call_queue,
-                             result_queue):
+                             result_queue,
+                             wakeup):
     """Manages the communication between this process and the worker processes.
 
     This function is run in a local thread.
@@ -284,6 +309,8 @@ def _queue_management_worker(executor_reference,
             derived from _WorkItems for processing by the process workers.
         result_queue: A ctx.SimpleQueue of _ResultItems generated by the
             process workers.
+        wakeup: A _Sentinel to allow waking up the queue_manager_thread from
+            the main Thread and avoid deadlocks caused by broken queues.
     """
     executor = None
 
@@ -292,38 +319,56 @@ def shutting_down():
                 or executor._shutdown_thread)
 
     def shutdown_worker():
-        # This is an upper bound
-        nb_children_alive = sum(p.is_alive() for p in processes.values())
-        for i in range(0, nb_children_alive):
-            try:
-                call_queue.put_nowait(None)
-            except Full:
-                pass
+        # This is an upper bound on the number of children alive.
+        n_children_alive = sum(p.is_alive() for p in processes.values())
+        n_children_to_stop = n_children_alive
+        n_sentinels_sent = 0
+        # Sent the right number of sentinels, to make sure all children are
+        # properly terminated.
+        while n_sentinels_sent < n_children_to_stop and n_children_alive > 0:
+            for i in range(n_children_to_stop - n_sentinels_sent):
+                try:
+                    call_queue.put_nowait(None)
+                    n_sentinels_sent += 1
+                except Full:
+                    break
+            n_children_alive = sum(p.is_alive() for p in processes.values())
+
         # Release the queue's resources as soon as possible.
         call_queue.close()
         # If .join() is not called on the created processes then
         # some ctx.Queue methods may deadlock on Mac OS X.
         for p in processes.values():
             p.join()
 
-    reader = result_queue._reader
+    result_reader = result_queue._reader
 
     while True:
         _add_call_item_to_queue(pending_work_items,
                                 work_ids_queue,
                                 call_queue)
 
-        sentinels = [p.sentinel for p in processes.values()]
-        assert sentinels
-        ready = wait([reader] + sentinels)
-
+        # Wait for a result to be ready in the result_queue while checking
+        # that worker process are still running.
+        worker_sentinels = [p.sentinel for p in processes.values()]
         received_item = False
-        if reader in ready:
+        while not wakeup.get_and_unset():
+            ready = wait([result_reader] + worker_sentinels,
+                         timeout=_POLL_TIMEOUT)
+            if len(ready) > 0:
+                break
+        else:
+            # The thread has been woken up by the main thread or the gc.
+            ready = []
+            result_item = None
+            received_item = True
+
+        if result_reader in ready:
             try:
-                result_item = reader.recv()
+                result_item = result_reader.recv()
                 received_item = True
-            except:
-                pass
+            except BaseException as e:
+                traceback.print_exc()
         if not received_item:
             # Mark the process pool broken so that submits fail right now.
             executor = executor_reference()
@@ -499,12 +544,20 @@ def __init__(self, max_workers=None, mp_context=None,
         self._result_queue = mp_context.SimpleQueue()
         self._work_ids = queue.Queue()
 
+        # Permits to wake_up the queue_manager_thread independently of
+        # result_queue state. This avoid deadlocks caused by the non
+        # transmission of wakeup signal when a worker died with the
+        # _result_queue write lock.
+        self._wakeup = _Sentinel()
+
     def _start_queue_management_thread(self):
-        # When the executor gets lost, the weakref callback will wake up
-        # the queue management thread.
-        def weakref_cb(_, q=self._result_queue):
-            q.put(None)
         if self._queue_management_thread is None:
+            # When the executor gets lost, the weakref callback will wake up
+            # the queue management thread.
+            def weakref_cb(_, wakeup=self._wakeup):
+                mp.util.debug('Executor collected: triggering callback for'
+                              ' QueueManager wakeup')
+                wakeup.set()
             # Start the processes so that their sentinels are known.
             self._adjust_process_count()
             self._queue_management_thread = threading.Thread(
@@ -514,11 +567,12 @@ def weakref_cb(_, q=self._result_queue):
                       self._pending_work_items,
                       self._work_ids,
                       self._call_queue,
-                      self._result_queue),
+                      self._result_queue,
+                      self._wakeup),
                 name="QueueManagerThread")
             self._queue_management_thread.daemon = True
             self._queue_management_thread.start()
-            _threads_queues[self._queue_management_thread] = self._result_queue
+            _threads_queues[self._queue_management_thread] = self._wakeup
 
     def _adjust_process_count(self):
         for _ in range(len(self._processes), self._max_workers):
@@ -545,7 +599,7 @@ def submit(self, fn, *args, **kwargs):
             self._work_ids.put(self._queue_count)
             self._queue_count += 1
             # Wake up queue management thread
-            self._result_queue.put(None)
+            self._wakeup.set()
 
             self._start_queue_management_thread()
             return f
@@ -585,7 +639,7 @@ def shutdown(self, wait=True):
             self._shutdown_thread = True
         if self._queue_management_thread:
             # Wake up queue management thread
-            self._result_queue.put(None)
+            self._wakeup.set()
             if wait:
                 self._queue_management_thread.join()
         # To reduce the risk of opening too many files, remove references to

Lib/test/test_concurrent_futures.py

@@ -988,22 +988,10 @@ def test_shutdown_deadlock(self):
             executor.shutdown(wait=True)
 
 
-class ProcessPoolForkExecutorDeadlockTest(ProcessPoolForkMixin,
-                                          ExecutorDeadlockTest,
-                                          unittest.TestCase):
-    pass
-
-
-class ProcessPoolForkserverExecutorDeadlockTest(ProcessPoolForkserverMixin,
-                                                ExecutorDeadlockTest,
-                                                unittest.TestCase):
-    pass
-
-
-class ProcessPoolSpawnExecutorDeadlockTest(ProcessPoolSpawnMixin,
-                                           ExecutorDeadlockTest,
-                                           unittest.TestCase):
-    pass
+create_executor_tests(ExecutorDeadlockTest,
+                      executor_mixins=(ProcessPoolForkMixin,
+                                       ProcessPoolForkserverMixin,
+                                       ProcessPoolSpawnMixin))
 
 
 class FutureTests(BaseTestCase):