[stream-refactor] parent: fix crash on graceful shutdown

Now it's possible for stream.protocol to not refer to MitogenProtocol,
move the signal handler to a MitogenProtocol subclass instead.

Fixes a crash where CTRL+C during child bootstrap would print
AttributeError.
pull/607/head
David Wilson 5 years ago
parent 2ccdeeeb87
commit 6b8a7cbcc4

@ -638,10 +638,10 @@ def _upgrade_broker(broker):
# This function is deadly! The act of calling start_receive() generates log # This function is deadly! The act of calling start_receive() generates log
# messages which must be silenced as the upgrade progresses, otherwise the # messages which must be silenced as the upgrade progresses, otherwise the
# poller state will change as it is copied, resulting in write fds that are # poller state will change as it is copied, resulting in write fds that are
# lost. (Due to LogHandler->Router->Stream->Broker->Poller, where Stream # lost. (Due to LogHandler->Router->Stream->Protocol->Broker->Poller, where
# only calls start_transmit() when transitioning from empty to non-empty # Stream only calls start_transmit() when transitioning from empty to
# buffer. If the start_transmit() is lost, writes from the child hang # non-empty buffer. If the start_transmit() is lost, writes from the child
# permanently). # hang permanently).
root = logging.getLogger() root = logging.getLogger()
old_level = root.level old_level = root.level
root.setLevel(logging.CRITICAL) root.setLevel(logging.CRITICAL)
@ -810,7 +810,8 @@ class CallSpec(object):
class PollPoller(mitogen.core.Poller): class PollPoller(mitogen.core.Poller):
""" """
Poller based on the POSIX poll(2) interface. Not available on some versions Poller based on the POSIX poll(2) interface. Not available on some versions
of OS X, otherwise it is the preferred poller for small FD counts. of OS X, otherwise it is the preferred poller for small FD counts, as there
is no setup/teardown/configuration system call overhead.
""" """
SUPPORTED = hasattr(select, 'poll') SUPPORTED = hasattr(select, 'poll')
_repr = 'PollPoller()' _repr = 'PollPoller()'
@ -1106,8 +1107,8 @@ class BootstrapProtocol(RegexProtocol):
""" """
Respond to stdout of a child during bootstrap. Wait for EC0_MARKER to be Respond to stdout of a child during bootstrap. Wait for EC0_MARKER to be
written by the first stage to indicate it can receive the bootstrap, then written by the first stage to indicate it can receive the bootstrap, then
await EC1_MARKER to indicate success, and await EC1_MARKER to indicate success, and :class:`MitogenProtocol` can be
:class:`mitogen.core.MitogenProtocol` can be enabled. enabled.
""" """
#: Sentinel value emitted by the first stage to indicate it is ready to #: Sentinel value emitted by the first stage to indicate it is ready to
#: receive the compressed bootstrap. For :mod:`mitogen.ssh` this must have #: receive the compressed bootstrap. For :mod:`mitogen.ssh` this must have
@ -1161,6 +1162,26 @@ class LogProtocol(LineLoggingProtocolMixin, mitogen.core.DelimitedProtocol):
LOG.info(u'%s: %s', self.stream.name, line.decode('utf-8', 'replace')) LOG.info(u'%s: %s', self.stream.name, line.decode('utf-8', 'replace'))
class MitogenProtocol(mitogen.core.MitogenProtocol):
"""
Extend core.MitogenProtocol to cause SHUTDOWN to be sent to the child
during graceful shutdown.
"""
def on_shutdown(self, broker):
"""
Respond to the broker's request for the stream to shut down by sending
SHUTDOWN to the child.
"""
LOG.debug('%r: requesting child shutdown', self)
self._send(
mitogen.core.Message(
src_id=mitogen.context_id,
dst_id=self.remote_id,
handle=mitogen.core.SHUTDOWN,
)
)
class Options(object): class Options(object):
name = None name = None
@ -1407,7 +1428,7 @@ class Connection(object):
if not self.exception: if not self.exception:
self._router.register(self.context, self.stdio_stream) self._router.register(self.context, self.stdio_stream)
self.stdio_stream.set_protocol( self.stdio_stream.set_protocol(
mitogen.core.MitogenProtocol( MitogenProtocol(
router=self._router, router=self._router,
remote_id=self.context.context_id, remote_id=self.context.context_id,
) )
@ -1428,19 +1449,6 @@ class Connection(object):
stream.on_disconnect(self._router.broker) stream.on_disconnect(self._router.broker)
self._complete_connection() self._complete_connection()
def on_stream_shutdown(self):
"""
Request the slave gracefully shut itself down.
"""
LOG.debug('%r: requesting child shutdown', self)
self.stdio_stream.protocol._send(
mitogen.core.Message(
src_id=mitogen.context_id,
dst_id=self.stdio_stream.protocol.remote_id,
handle=mitogen.core.SHUTDOWN,
)
)
eof_error_msg = 'EOF on stream; last 100 lines received:\n' eof_error_msg = 'EOF on stream; last 100 lines received:\n'
def on_stdio_disconnect(self): def on_stdio_disconnect(self):
@ -1511,7 +1519,6 @@ class Connection(object):
stream.name = self.options.name or self._get_name() stream.name = self.options.name or self._get_name()
stream.accept(self.proc.stdout, self.proc.stdin) stream.accept(self.proc.stdout, self.proc.stdin)
mitogen.core.listen(stream, 'shutdown', self.on_stream_shutdown)
mitogen.core.listen(stream, 'disconnect', self.on_stdio_disconnect) mitogen.core.listen(stream, 'disconnect', self.on_stdio_disconnect)
self._router.broker.start_receive(stream) self._router.broker.start_receive(stream)
return stream return stream

Loading…
Cancel
Save