Skip to content

Commit 0ed8330

Browse files
committed
gh-155477: Fix multiprocessing.Pool deadlock on close() with a buffersize imap
close() did not release the buffersize semaphores that throttle the task generator, so a partially-consumed imap left the task handler blocked and join() deadlocked. Release them in close() and stop the generator once the pool leaves the RUN state.
1 parent 7c653e2 commit 0ed8330

3 files changed

Lines changed: 33 additions & 0 deletions

File tree

Lib/multiprocessing/pool.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -403,6 +403,11 @@ def _guarded_task_generation(self, result_job, func, iterable, sema=None):
403403
enumerated_iter = iter(enumerate(iterable))
404404
while True:
405405
sema.acquire()
406+
if self._state != RUN:
407+
# The pool is closing or terminating; stop submitting
408+
# the still-throttled tasks so the task handler can
409+
# finish instead of blocking here forever.
410+
break
406411
try:
407412
i, x = next(enumerated_iter)
408413
except StopIteration:
@@ -661,6 +666,10 @@ def close(self):
661666
self._state = CLOSE
662667
self._worker_handler._state = CLOSE
663668
self._change_notifier.put(None)
669+
# Wake any task generator throttled on a buffersize semaphore so
670+
# it observes the CLOSE state and stops submitting.
671+
for sema in tuple(self._taskqueue_buffersize_semaphores):
672+
sema.release()
664673

665674
def terminate(self):
666675
util.debug('terminating pool')

Lib/test/_test_multiprocessing.py

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3227,6 +3227,27 @@ def produce_args():
32273227
p.terminate()
32283228
p.join()
32293229

3230+
@warnings_helper.ignore_fork_in_thread_deprecation_warnings()
3231+
@support.subTests('method_name', ("imap", "imap_unordered"))
3232+
def test_imap_with_buffersize_close_after_partial_consumption(
3233+
self, method_name
3234+
):
3235+
# close()/join() must not deadlock when a buffersize iterator is
3236+
# only partially consumed (the throttled task generator must stop).
3237+
p = self.Pool(2)
3238+
method = getattr(p, method_name)
3239+
it = method(sqr, range(1000), buffersize=2)
3240+
next(it)
3241+
finished = threading.Event()
3242+
def finalize():
3243+
p.close()
3244+
p.join()
3245+
finished.set()
3246+
t = threading.Thread(target=finalize)
3247+
t.start()
3248+
t.join(support.SHORT_TIMEOUT)
3249+
self.assertTrue(finished.is_set(), "close()/join() deadlocked")
3250+
32303251
@support.subTests('method_name', ("imap", "imap_unordered"))
32313252
def test_imap_and_imap_unordered_with_buffersize_on_empty_iterable(
32323253
self, method_name
Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
Fix a deadlock when ``multiprocessing.pool.Pool.close()`` is followed by
2+
``join()`` after only partially consuming an ``imap()`` or ``imap_unordered()``
3+
iterator that was created with a *buffersize*. Patch by tonghuaroot.

0 commit comments

Comments
 (0)