Adds tracking of which process in the pool takes which job from the queue; adds test for issue22393/issue38084.

Author: applio

Lib/multiprocessing/pool.py

@@ -121,6 +121,7 @@ def worker(inqueue, outqueue, initializer=None, initargs=(), maxtasks=None,
             break
 
         job, i, func, args, kwds = task
+        put((job, i, (None, os.getpid())))  # Provide info on who took job
         try:
             result = (True, func(*args, **kwds))
         except Exception as e:
@@ -220,12 +221,14 @@ def __init__(self, processes=None, initializer=None, initargs=(),
 
         sentinels = self._get_sentinels()
 
+        self._job_assignments = {}
         self._worker_handler = threading.Thread(
             target=Pool._handle_workers,
             args=(self._cache, self._taskqueue, self._ctx, self.Process,
                   self._processes, self._pool, self._inqueue, self._outqueue,
                   self._initializer, self._initargs, self._maxtasksperchild,
-                  self._wrap_exception, sentinels, self._change_notifier)
+                  self._wrap_exception, sentinels, self._change_notifier,
+                  self._job_assignments)
             )
         self._worker_handler.daemon = True
         self._worker_handler._state = RUN
@@ -243,7 +246,8 @@ def __init__(self, processes=None, initializer=None, initargs=(),
 
         self._result_handler = threading.Thread(
             target=Pool._handle_results,
-            args=(self._outqueue, self._quick_get, self._cache)
+            args=(self._outqueue, self._quick_get, self._cache,
+                  self._job_assignments)
             )
         self._result_handler.daemon = True
         self._result_handler._state = RUN
@@ -284,7 +288,7 @@ def _get_worker_sentinels(workers):
                 workers if hasattr(worker, "sentinel")]
 
     @staticmethod
-    def _join_exited_workers(pool):
+    def _join_exited_workers(pool, outqueue, job_assignments):
         """Cleanup after any worker processes which have exited due to reaching
         their specified lifetime.  Returns True if any workers were cleaned up.
         """
@@ -294,8 +298,15 @@ def _join_exited_workers(pool):
             if worker.exitcode is not None:
                 # worker exited
                 util.debug('cleaning up worker %d' % i)
+                pid =  worker.ident
                 worker.join()
                 cleaned = True
+                if pid in job_assignments:
+                    # If the worker process died without communicating back
+                    # while running a job, add a default result for it.
+                    job = job_assignments[pid]
+                    outqueue.put((job, i, (False, RuntimeError("Worker died"))))
+                    del job_assignments[pid]
                 del pool[i]
         return cleaned
 
@@ -330,10 +341,10 @@ def _repopulate_pool_static(ctx, Process, processes, pool, inqueue,
     @staticmethod
     def _maintain_pool(ctx, Process, processes, pool, inqueue, outqueue,
                        initializer, initargs, maxtasksperchild,
-                       wrap_exception):
+                       wrap_exception, job_assignments):
         """Clean up any exited workers and start replacements for them.
         """
-        if Pool._join_exited_workers(pool):
+        if Pool._join_exited_workers(pool, outqueue, job_assignments):
             Pool._repopulate_pool_static(ctx, Process, processes, pool,
                                          inqueue, outqueue, initializer,
                                          initargs, maxtasksperchild,
@@ -504,15 +515,16 @@ def _wait_for_updates(sentinels, change_notifier, timeout=None):
     def _handle_workers(cls, cache, taskqueue, ctx, Process, processes,
                         pool, inqueue, outqueue, initializer, initargs,
                         maxtasksperchild, wrap_exception, sentinels,
-                        change_notifier):
+                        change_notifier, job_assignments):
         thread = threading.current_thread()
 
         # Keep maintaining workers until the cache gets drained, unless the pool
         # is terminated.
         while thread._state == RUN or (cache and thread._state != TERMINATE):
             cls._maintain_pool(ctx, Process, processes, pool, inqueue,
                                outqueue, initializer, initargs,
-                               maxtasksperchild, wrap_exception)
+                               maxtasksperchild, wrap_exception,
+                               job_assignments)
 
             current_sentinels = [*cls._get_worker_sentinels(pool), *sentinels]
 
@@ -568,7 +580,7 @@ def _handle_tasks(taskqueue, put, outqueue, pool, cache):
         util.debug('task handler exiting')
 
     @staticmethod
-    def _handle_results(outqueue, get, cache):
+    def _handle_results(outqueue, get, cache, job_assignments):
         thread = threading.current_thread()
 
         while 1:
@@ -587,12 +599,18 @@ def _handle_results(outqueue, get, cache):
                 util.debug('result handler got sentinel')
                 break
 
-            job, i, obj = task
-            try:
-                cache[job]._set(i, obj)
-            except KeyError:
-                pass
-            task = job = obj = None
+            job, i, (task_info, value) = task
+            if task_info is None:
+                # task_info is True or False when a task has completed but
+                # None indicates information about which process has
+                # accepted a job from the queue.
+                job_assignments[value] = job
+            else:
+                try:
+                    cache[job]._set(i, (task_info, value))
+                except KeyError:
+                    pass
+            task = job = task_info = value = None
 
         while cache and thread._state != TERMINATE:
             try:
@@ -604,12 +622,16 @@ def _handle_results(outqueue, get, cache):
             if task is None:
                 util.debug('result handler ignoring extra sentinel')
                 continue
-            job, i, obj = task
-            try:
-                cache[job]._set(i, obj)
-            except KeyError:
-                pass
-            task = job = obj = None
+
+            job, i, (task_info, value) = task
+            if task_info is None:
+                job_assignments[value] = job
+            else:
+                try:
+                    cache[job]._set(i, (task_info, value))
+                except KeyError:
+                    pass
+            task = job = task_info = value = None
 
         if hasattr(outqueue, '_reader'):
             util.debug('ensuring that outqueue is not full')

Lib/test/_test_multiprocessing.py

@@ -2712,6 +2712,27 @@ def errback(exc):
         p.close()
         p.join()
 
+    def test_pool_worker_died_without_communicating(self):
+        # Issue22393: test fix of indefinite hang caused by worker processes
+        # exiting abruptly (such as via os._exit()) without communicating
+        # back to the pool at all.
+        prog = (
+            "import os, multiprocessing as mp; "
+            "is_main = (__name__ == '__main__'); "
+            "p = mp.Pool(1) if is_main else print('worker'); "
+            "p.map(os._exit, [1]) if is_main else None; "
+            "(p.close() or p.join()) if is_main else None"
+        )
+        # Only if there is a regression will this ever trigger a
+        # subprocess.TimeoutExpired.
+        completed_process = subprocess.run(
+            [sys.executable, '-E', '-S', '-O', '-c', prog],
+            check=False,
+            timeout=100,
+            capture_output=True
+        )
+        self.assertNotEqual(0, completed_process.returncode)
+
 class _TestPoolWorkerLifetime(BaseTestCase):
     ALLOWED_TYPES = ('processes', )