FazBrowse GitHub Viewer | Trending |
URL:
| Home
Tools: [Download Repo ZIP]   [Original HTTPS Page]

gh-159689: Fix ProcessPoolExecutor hang when a worker is killed mid-result by junaidkhan80121 · Pull Request #159697 · python/cpython · GitHub

Repository navigation

11 changes: 10 additions & 1 deletion Lib/concurrent/futures/process.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,7 @@
# This import is required to load the multiprocessing.connection submodule
# so that it can be accessed later as `mp.connection`
import multiprocessing.connection
from multiprocessing.connection import _SentinelReady
from multiprocessing.queues import Queue
import threading
import weakref
Expand Down Expand Up @@ -518,8 +519,16 @@ def wait_result_broken_or_wakeup(self):
result_item = None
if result_reader in ready:
try:
result_item = result_reader.recv()
# Keep watching worker sentinels for the rest of the
# message. recv() otherwise blocks in plain reads, and the
# parent's write end of this pipe prevents a killed worker
# from delivering EOF. gh-159689.
result_item = result_reader.recv(sentinels=worker_sentinels)
is_broken = False
except _SentinelReady:
# A worker exited while its result was only partly written.
# Leave is_broken set so the exit code is reported.
pass
except BaseException as exc:
cause = format_exception(exc)

Expand Down
124 changes: 104 additions & 20 deletions Lib/multiprocessing/connection.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,20 @@ def _init_timeout(timeout=CONNECTION_TIMEOUT):
def _check_timeout(t):
return time.monotonic() > t


class _SentinelReady(Exception):
"""A wait sentinel became readable while a connection read was in progress.

Used by ProcessPoolExecutor so a worker killed mid-result does not leave
the manager thread blocked in a read that can never see EOF. The parent's
copy of the result pipe's write end stays open for the executor's
lifetime, so the death of the writer is not delivered as end-of-file.
"""

def __init__(self, ready):
super().__init__(ready)
self.ready = ready

#
#
#
Expand Down Expand Up @@ -253,11 +267,18 @@ def recv_bytes_into(self, buf, offset=0):
(offset + size) // itemsize])
return size

def recv(self):
"""Receive a (picklable) object"""
def recv(self, *, sentinels=()):
"""Receive a (picklable) object.

*sentinels* is an optional sequence of objects accepted by
:func:`wait`, such as process sentinels. Before each read, those
objects are waited on together with this connection. If one of them
is ready and this connection is not, :class:`_SentinelReady` is
raised instead of blocking until the rest of the message arrives.
"""
self._check_closed()
self._check_readable()
buf = self._recv_bytes()
buf = self._recv_bytes(sentinels=sentinels)
return _ForkingPickler.loads(buf.getbuffer())

def poll(self, timeout=0.0):
Expand Down Expand Up @@ -317,7 +338,31 @@ def _send_bytes(self, buf):
assert err == 0
assert nwritten == len(buf)

def _recv_bytes(self, maxsize=None):
def _wait_overlapped_or_sentinels(self, ov, sentinels):
"""Block until *ov* completes, or a sentinel is signaled."""
if not sentinels:
waitres = _winapi.WaitForMultipleObjects(
[ov.event], False, INFINITE)
assert waitres == WAIT_OBJECT_0
return
sentinels = list(sentinels)
handles = [ov.event, *(int(item) for item in sentinels)]
waitres = _winapi.WaitForMultipleObjects(handles, False, INFINITE)
if waitres == WAIT_OBJECT_0:
return
ov.cancel()
try:
ov.GetOverlappedResult(True)
except OSError:
pass
idx = waitres - WAIT_OBJECT_0
if not 0 < idx <= len(sentinels):
idx = waitres - WAIT_ABANDONED_0
if 0 < idx <= len(sentinels):
raise _SentinelReady([sentinels[idx - 1]])
raise OSError("failed waiting for pipe data")

def _recv_bytes(self, maxsize=None, sentinels=()):
if self._got_empty_message:
self._got_empty_message = False
return io.BytesIO()
Expand All @@ -329,23 +374,29 @@ def _recv_bytes(self, maxsize=None):

sentinel = object()
return_value = sentinel
sentinel_ready = False
try:
try:
if err == _winapi.ERROR_IO_PENDING:
waitres = _winapi.WaitForMultipleObjects(
[ov.event], False, INFINITE)
assert waitres == WAIT_OBJECT_0
self._wait_overlapped_or_sentinels(
ov, sentinels)
except _SentinelReady:
# Already cancelled and reaped above.
sentinel_ready = True
raise
except:
ov.cancel()
raise
finally:
nread, err = ov.GetOverlappedResult(True)
if err == 0:
f = io.BytesIO()
f.write(ov.getbuffer())
return_value = f
elif err == _winapi.ERROR_MORE_DATA:
return_value = self._get_more_data(ov, maxsize)
if not sentinel_ready:
nread, err = ov.GetOverlappedResult(True)
if err == 0:
f = io.BytesIO()
f.write(ov.getbuffer())
return_value = f
elif err == _winapi.ERROR_MORE_DATA:
return_value = self._get_more_data(
ov, maxsize, sentinels)
except:
if return_value is sentinel:
raise
Expand All @@ -365,11 +416,28 @@ def _poll(self, timeout):
return True
return bool(wait([self], timeout))

def _get_more_data(self, ov, maxsize):
def _get_more_data(self, ov, maxsize, sentinels=()):
buf = ov.getbuffer()
f = io.BytesIO()
f.write(buf)
left = _winapi.PeekNamedPipe(self._handle)[1]
while left == 0 and sentinels:
# The rest of this message has not been written yet. The
# parent keeps its own write end open, so a killed writer
# does not surface as EOF. Poll the sentinels directly:
# waiting on this pipe here would start another overlapped
# read and consume the message.
handles = [int(item) for item in sentinels]
waitres = _winapi.WaitForMultipleObjects(handles, False, 50)
if waitres == WAIT_TIMEOUT:
left = _winapi.PeekNamedPipe(self._handle)[1]
continue
idx = waitres - WAIT_OBJECT_0
if not 0 <= idx < len(handles):
idx = waitres - WAIT_ABANDONED_0
if 0 <= idx < len(handles):
raise _SentinelReady([sentinels[idx]])
raise OSError("failed waiting for pipe data")
assert left > 0
if maxsize is not None and len(buf) + left > maxsize:
self._bad_message_length()
Expand Down Expand Up @@ -407,11 +475,27 @@ def _send(self, buf, write=_write):
break
buf = buf[n:]

def _recv(self, size, read=_read):
def _abandon_read_if_sentinel_ready(self, sentinels):
"""Raise if a sentinel is ready and this connection has no data.

A non-blocking poll covers the common case where the next chunk is
already buffered, so a large result does not block once per chunk.
"""
ready = wait([self, *sentinels], timeout=0)
if self in ready:
return
if not ready:
ready = wait([self, *sentinels])
if self not in ready:
raise _SentinelReady(ready)

def _recv(self, size, sentinels=(), read=_read):
buf = io.BytesIO()
handle = self._handle
remaining = size
while remaining > 0:
if sentinels:
self._abandon_read_if_sentinel_ready(sentinels)
to_read = min(BUFSIZE, remaining)
chunk = read(handle, to_read)
n = len(chunk)
Expand Down Expand Up @@ -447,15 +531,15 @@ def _send_bytes(self, buf):
# to avoid "broken pipe" errors if the other end closed the pipe.
self._send(header + buf)

def _recv_bytes(self, maxsize=None):
buf = self._recv(4)
def _recv_bytes(self, maxsize=None, sentinels=()):
buf = self._recv(4, sentinels)
size, = struct.unpack("!i", buf.getvalue())
if size == -1:
buf = self._recv(8)
buf = self._recv(8, sentinels)
size, = struct.unpack("!Q", buf.getvalue())
if maxsize is not None and size > maxsize:
return None
return self._recv(size)
return self._recv(size, sentinels)

def _poll(self, timeout):
r = wait([self], timeout)
Expand Down
131 changes: 128 additions & 3 deletions Lib/test/test_concurrent_futures/test_process_pool.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
Original file line number Diff line number Diff line change
@@ -1,5 +1,8 @@
import multiprocessing as mp
import os
import queue
import signal
import struct
import sys
import threading
import time
Expand All @@ -8,16 +11,16 @@
import unittest.mock
import weakref
from concurrent import futures
from concurrent.futures.process import BrokenProcessPool
from concurrent.futures.process import BrokenProcessPool, _check_system_limits

from test import support
from test.support import hashlib_helper, threading_helper, warnings_helper
from test.test_importlib.metadata.fixtures import parameterize

from .executor import ExecutorTest, mul
from .util import (
ProcessPoolForkMixin, ProcessPoolForkserverMixin, ProcessPoolSpawnMixin,
create_executor_tests, setup_module)
BaseTestCase, ProcessPoolForkMixin, ProcessPoolForkserverMixin,
ProcessPoolSpawnMixin, create_executor_tests, setup_module)


class EventfulGCObj():
Expand Down Expand Up @@ -50,6 +53,40 @@ def _report_wait_return(queue, event, value):
return value


def _send_result_header_only(writer):
"""Write a length header and exit before the message body."""
writer._send(struct.pack("!i", 100_000))


def _send_object(writer):
writer.send({"ok": True})


def _return_large_result():
# Larger than Connection._send_bytes' 16384-byte inline threshold, so the
# length header is written separately from the body.
return b"x" * (32 * 1024)


def _pause_after_result_header(writer):
"""Stop this worker after the result length is written, before the body."""
import multiprocessing.connection as connection

original = connection.Connection._send_bytes

def _send_bytes(self, buf):
n = len(buf)
if n > 16384:
self._send(struct.pack("!i", n))
writer.send_bytes(b"1")
os.kill(os.getpid(), signal.SIGSTOP)
self._send(buf)
else:
original(self, buf)

connection.Connection._send_bytes = _send_bytes


class ProcessPoolExecutorTest(ExecutorTest):

@unittest.skipUnless(sys.platform=='win32', 'Windows-only process limit')
Expand All @@ -58,6 +95,33 @@ def test_max_workers_too_large(self):
"max_workers must be <= 61"):
futures.ProcessPoolExecutor(max_workers=62)

@unittest.skipUnless(hasattr(signal, "SIGSTOP"),
"requires SIGSTOP to pause a worker mid-send")
@warnings_helper.ignore_fork_in_thread_deprecation_warnings()
def test_killed_child_during_result_send(self):
# gh-159689: a worker killed after writing the result length but
# before the body must fail the future. The parent holds the
# result pipe's write end, so this is not delivered as EOF.
ctx = self.get_context()
reader, writer = ctx.Pipe(duplex=False)
executor = self.executor_type(
max_workers=1,
mp_context=ctx,
initializer=_pause_after_result_header,
initargs=(writer,))
try:
future = executor.submit(_return_large_result)
self.assertTrue(reader.poll(support.SHORT_TIMEOUT))
self.assertEqual(reader.recv_bytes(), b"1")
worker = next(iter(executor._processes.values()))
os.kill(worker.pid, signal.SIGKILL)
with self.assertRaises(BrokenProcessPool):
future.result(timeout=support.SHORT_TIMEOUT)
finally:
executor.shutdown(wait=True, cancel_futures=True)
writer.close()
reader.close()

@warnings_helper.ignore_fork_in_thread_deprecation_warnings()
def test_killed_child(self):
# When a child process is abruptly terminated, the whole pool gets
Expand Down Expand Up @@ -543,6 +607,67 @@ def test_force_shutdown_workers_stops_pool(self, function_name):
break


@unittest.skipUnless(sys.platform != "win32",
"partial pipe reads are specific to Unix connections")
class ResultPipeSentinelTest(BaseTestCase):
"""Connection.recv(sentinels=...) when a writer dies mid-message.

The parent keeps its copy of the write end, which is how
ProcessPoolExecutor's result queue is set up. gh-159689.
"""

def setUp(self):
super().setUp()
try:
_check_system_limits()
except NotImplementedError:
self.skipTest("ProcessPoolExecutor unavailable on this system")

def _context(self):
return mp.get_context("spawn")

def test_partial_message_raises_when_writer_exits(self):
ctx = self._context()
reader, writer = ctx.Pipe(duplex=False)
proc = ctx.Process(target=_send_result_header_only, args=(writer,))
started = False
try:
proc.start()
started = True
proc.join(support.SHORT_TIMEOUT)
self.assertEqual(proc.exitcode, 0)
t0 = time.monotonic()
with self.assertRaises(mp.connection._SentinelReady):
reader.recv(sentinels=[proc.sentinel])
self.assertLess(time.monotonic() - t0, support.SHORT_TIMEOUT)
finally:
writer.close()
reader.close()
if started:
proc.join(support.SHORT_TIMEOUT)
proc.close()

def test_complete_message_still_received(self):
ctx = self._context()
reader, writer = ctx.Pipe(duplex=False)
proc = ctx.Process(target=_send_object, args=(writer,))
started = False
try:
proc.start()
started = True
proc.join(support.SHORT_TIMEOUT)
self.assertEqual(proc.exitcode, 0)
writer.close()
self.assertEqual(reader.recv(sentinels=[proc.sentinel]),
{"ok": True})
finally:
writer.close()
reader.close()
if started:
proc.join(support.SHORT_TIMEOUT)
proc.close()


create_executor_tests(globals(), ProcessPoolExecutorTest,
executor_mixins=(ProcessPoolForkMixin,
ProcessPoolForkserverMixin,
Expand Down
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
Fix :class:`concurrent.futures.ProcessPoolExecutor` hanging when a worker
is killed while writing a large result. The pending future now raises
:exc:`~concurrent.futures.process.BrokenProcessPool`, and the interpreter can exit.
Loading

Back | FazBrowse Home | New Git URL