You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
mitogen/mitogen/parent.py

703 lines
23 KiB
Python

# Copyright 2017, David Wilson
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its contributors
# may be used to endorse or promote products derived from this software without
# specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
import getpass
import inspect
import logging
import os
import re
import select
import signal
import socket
import sys
import termios
import textwrap
import threading
import time
import zlib
import mitogen.core
from mitogen.core import LOG
from mitogen.core import IOLOG
DOCSTRING_RE = re.compile(r'""".+?"""', re.M | re.S)
COMMENT_RE = re.compile(r'^[ ]*#[^\n]*$', re.M)
class Argv(object):
def __init__(self, argv):
self.argv = argv
def escape(self, x):
s = '"'
for c in x:
if c in '\\$"`':
s += '\\'
s += c
s += '"'
return s
def __str__(self):
return ' '.join(map(self.escape, self.argv))
def get_log_level():
return (LOG.level or logging.getLogger().level or logging.INFO)
def minimize_source(source):
subber = lambda match: '""' + ('\n' * match.group(0).count('\n'))
source = DOCSTRING_RE.sub(subber, source)
source = COMMENT_RE.sub('', source)
return source.replace(' ', '\t')
def flags(names):
"""Return the result of ORing a set of (space separated) :py:mod:`termios`
module constants together."""
return sum(getattr(termios, name) for name in names.split())
def cfmakeraw(tflags):
"""Given a list returned by :py:func:`termios.tcgetattr`, return a list
that has been modified in the same manner as the `cfmakeraw()` C library
function."""
iflag, oflag, cflag, lflag, ispeed, ospeed, cc = tflags
iflag &= ~flags('IGNBRK BRKINT PARMRK ISTRIP INLCR IGNCR ICRNL IXON')
oflag &= ~flags('OPOST IXOFF')
lflag &= ~flags('ECHO ECHOE ECHONL ICANON ISIG IEXTEN')
cflag &= ~flags('CSIZE PARENB')
cflag |= flags('CS8')
# TODO: one or more of the above bit twiddles sets or omits a necessary
# flag. Forcing these fields to zero, as shown below, gets us what we want
# on Linux/OS X, but it is possibly broken on some other OS.
iflag = 0
oflag = 0
lflag = 0
return [iflag, oflag, cflag, lflag, ispeed, ospeed, cc]
def disable_echo(fd):
old = termios.tcgetattr(fd)
new = cfmakeraw(old)
flags = (
termios.TCSAFLUSH |
getattr(termios, 'TCSASOFT', 0)
)
termios.tcsetattr(fd, flags, new)
def close_nonstandard_fds():
for fd in xrange(3, 1024):
try:
os.close(fd)
except OSError:
pass
def create_child(*args):
parentfp, childfp = socket.socketpair()
parentfp.setsockopt(socket.SOL_SOCKET,
socket.SO_SNDBUF,
mitogen.core.CHUNK_SIZE)
childfp.setsockopt(socket.SOL_SOCKET,
socket.SO_RCVBUF,
mitogen.core.CHUNK_SIZE)
pid = os.fork()
if not pid:
mitogen.core.set_block(childfp.fileno())
os.dup2(childfp.fileno(), 0)
os.dup2(childfp.fileno(), 1)
childfp.close()
parentfp.close()
os.execvp(args[0], args)
childfp.close()
# Decouple the socket from the lifetime of the Python socket object.
fd = os.dup(parentfp.fileno())
parentfp.close()
LOG.debug('create_child() child %d fd %d, parent %d, cmd: %s',
pid, fd, os.getpid(), Argv(args))
return pid, fd
def tty_create_child(*args):
master_fd, slave_fd = os.openpty()
disable_echo(master_fd)
disable_echo(slave_fd)
pid = os.fork()
if not pid:
mitogen.core.set_block(slave_fd)
os.dup2(slave_fd, 0)
os.dup2(slave_fd, 1)
os.dup2(slave_fd, 2)
close_nonstandard_fds()
os.setsid()
os.close(os.open(os.ttyname(1), os.O_RDWR))
os.execvp(args[0], args)
os.close(slave_fd)
LOG.debug('tty_create_child() child %d fd %d, parent %d, cmd: %s',
pid, master_fd, os.getpid(), Argv(args))
return pid, master_fd
def write_all(fd, s, deadline=None):
timeout = None
written = 0
while written < len(s):
if deadline is not None:
timeout = max(0, deadline - time.time())
if timeout == 0:
raise mitogen.core.TimeoutError('write timed out')
_, wfds, _ = select.select([], [fd], [], timeout)
if not wfds:
continue
n, disconnected = mitogen.core.io_op(os.write, fd, buffer(s, written))
if disconnected:
raise mitogen.core.StreamError('EOF on stream during write')
written += n
def iter_read(fd, deadline=None):
bits = []
timeout = None
while True:
if deadline is not None:
timeout = max(0, deadline - time.time())
if timeout == 0:
break
rfds, _, _ = select.select([fd], [], [], timeout)
if not rfds:
continue
s, disconnected = mitogen.core.io_op(os.read, fd, 4096)
IOLOG.debug('iter_read(%r) -> %r', fd, s)
if disconnected or not s:
raise mitogen.core.StreamError(
'EOF on stream; last 300 bytes received: %r' %
(''.join(bits)[-300:],)
)
bits.append(s)
yield s
raise mitogen.core.TimeoutError('read timed out')
def discard_until(fd, s, deadline):
for buf in iter_read(fd, deadline):
if IOLOG.level == logging.DEBUG:
for line in buf.splitlines():
IOLOG.debug('discard_until: discarding %r', line)
if buf.endswith(s):
return
def upgrade_router(econtext):
if not isinstance(econtext.router, Router): # TODO
econtext.router.__class__ = Router # TODO
econtext.router.upgrade(
importer=econtext.importer,
parent=econtext.parent,
)
def _docker_method():
import mitogen.docker
return mitogen.docker.Stream
def _local_method():
return mitogen.parent.Stream
def _ssh_method():
import mitogen.ssh
return mitogen.ssh.Stream
def _sudo_method():
import mitogen.sudo
return mitogen.sudo.Stream
METHOD_NAMES = {
'docker': _docker_method,
'local': _local_method,
'ssh': _ssh_method,
'sudo': _sudo_method,
}
@mitogen.core.takes_econtext
def _proxy_connect(name, method_name, kwargs, econtext):
mitogen.parent.upgrade_router(econtext)
context = econtext.router._connect(
klass=METHOD_NAMES[method_name](),
name=name,
**kwargs
)
return context.context_id, context.name
class Stream(mitogen.core.Stream):
"""
Base for streams capable of starting new slaves.
"""
#: The path to the remote Python interpreter.
python_path = 'python2.7'
#: Maximum time to wait for a connection attempt.
connect_timeout = 30.0
#: Derived from :py:attr:`connect_timeout`; absolute floating point
#: UNIX timestamp after which the connection attempt should be abandoned.
connect_deadline = None
#: True to cause context to write verbose /tmp/mitogen.<pid>.log.
debug = False
#: True to cause context to write /tmp/mitogen.stats.<pid>.<thread>.log.
profiling = False
def __init__(self, *args, **kwargs):
super(Stream, self).__init__(*args, **kwargs)
self.sent_modules = set(['mitogen', 'mitogen.core'])
#: List of contexts reachable via this stream; used to cleanup routes
#: during disconnection.
self.routes = set([self.remote_id])
def construct(self, remote_name=None, python_path=None, debug=False,
connect_timeout=None, profiling=False, **kwargs):
"""Get the named context running on the local machine, creating it if
it does not exist."""
super(Stream, self).construct(**kwargs)
if python_path:
self.python_path = python_path
if sys.platform == 'darwin' and self.python_path == '/usr/bin/python':
# OS X installs a craptacular argv0-introspecting Python version
# switcher as /usr/bin/python. Override attempts to call it with an
# explicit call to python2.7
self.python_path = '/usr/bin/python2.7'
if connect_timeout:
self.connect_timeout = connect_timeout
if remote_name is None:
remote_name = '%s@%s:%d'
remote_name %= (getpass.getuser(), socket.gethostname(), os.getpid())
if '/' in remote_name or '\\' in remote_name:
raise ValueError('remote_name= cannot contain slashes')
self.remote_name = remote_name
self.debug = debug
self.profiling = profiling
self.connect_deadline = time.time() + self.connect_timeout
def on_shutdown(self, broker):
"""Request the slave gracefully shut itself down."""
LOG.debug('%r closing CALL_FUNCTION channel', self)
self.send(
mitogen.core.Message(
src_id=mitogen.context_id,
dst_id=self.remote_id,
handle=mitogen.core.SHUTDOWN,
)
)
# Minimised, gzipped, base64'd and passed to 'python -c'. It forks, dups
# file descriptor 0 as 100, creates a pipe, then execs a new interpreter
# with a custom argv.
# * Optimized for minimum byte count after minification & compression.
# * 'CONTEXT_NAME', 'PREAMBLE_COMPRESSED_LEN', and 'PREAMBLE_LEN' are
# substituted with their respective values.
# * CONTEXT_NAME must be prefixed with the name of the Python binary in
# order to allow virtualenvs to detect their install prefix.
@staticmethod
def _first_stage():
R,W=os.pipe()
r,w=os.pipe()
if os.fork():
os.dup2(0,100)
os.dup2(R,0)
os.dup2(r,101)
os.close(R)
os.close(r)
os.close(W)
os.close(w)
os.environ['ARGV0']=sys.executable
os.execl(sys.executable,sys.executable+'(mitogen:CONTEXT_NAME)')
os.write(1,'EC0\n')
C=_(os.fdopen(0,'rb').read(PREAMBLE_COMPRESSED_LEN),'zip')
os.fdopen(W,'w',0).write(C)
os.fdopen(w,'w',0).write('PREAMBLE_LEN\n'+C)
os.write(1,'EC1\n')
def get_boot_command(self):
source = inspect.getsource(self._first_stage)
source = textwrap.dedent('\n'.join(source.strip().split('\n')[2:]))
source = source.replace(' ', '\t')
source = source.replace('CONTEXT_NAME', self.remote_name)
preamble_compressed = self.get_preamble()
source = source.replace('PREAMBLE_COMPRESSED_LEN', str(len(preamble_compressed)))
source = source.replace('PREAMBLE_LEN', str(len(zlib.decompress(preamble_compressed))))
encoded = zlib.compress(source, 9).encode('base64').replace('\n', '')
# We can't use bytes.decode() in 3.x since it was restricted to always
# return unicode, so codecs.decode() is used instead. In 3.x
# codecs.decode() requires a bytes object. Since we must be compatible
# with 2.4 (no bytes literal), an extra .encode() either returns the
# same str (2.x) or an equivalent bytes (3.x).
return [
self.python_path, '-c',
'import codecs,os,sys;_=codecs.decode;'
'exec(_(_("%s".encode(),"base64"),"zip"))' % (encoded,)
]
def get_preamble(self):
parent_ids = mitogen.parent_ids[:]
parent_ids.insert(0, mitogen.context_id)
source = inspect.getsource(mitogen.core)
source += '\nExternalContext().main(**%r)\n' % ({
'parent_ids': parent_ids,
'context_id': self.remote_id,
'debug': self.debug,
'profiling': self.profiling,
'log_level': get_log_level(),
'whitelist': self._router.get_module_whitelist(),
'blacklist': self._router.get_module_blacklist(),
},)
return zlib.compress(minimize_source(source), 9)
create_child = staticmethod(create_child)
def connect(self):
LOG.debug('%r.connect()', self)
pid, fd = self.create_child(*self.get_boot_command())
self.name = 'local.%s' % (pid,)
self.receive_side = mitogen.core.Side(self, fd)
self.transmit_side = mitogen.core.Side(self, os.dup(fd))
LOG.debug('%r.connect(): child process stdin/stdout=%r',
self, self.receive_side.fd)
self._connect_bootstrap()
def _ec0_received(self):
LOG.debug('%r._ec0_received()', self)
write_all(self.transmit_side.fd, self.get_preamble())
discard_until(self.receive_side.fd, 'EC1\n', time.time() + 10.0)
def _connect_bootstrap(self):
deadline = time.time() + self.connect_timeout
discard_until(self.receive_side.fd, 'EC0\n', deadline)
self._ec0_received()
class ChildIdAllocator(object):
def __init__(self, router):
self.router = router
self.lock = threading.Lock()
self.it = iter(xrange(0))
def allocate(self):
self.lock.acquire()
try:
for id_ in self.it:
return id_
master = mitogen.core.Context(self.router, 0)
start, end = master.send_await(
mitogen.core.Message(dst_id=0, handle=mitogen.core.ALLOCATE_ID)
)
self.it = iter(xrange(start, end))
finally:
self.lock.release()
return self.allocate()
class RouteMonitor(object):
def __init__(self, router, parent=None):
self.router = router
self.parent = parent
self.router.add_handler(
fn=self._on_add_route,
handle=mitogen.core.ADD_ROUTE,
persist=True,
)
self.router.add_handler(
fn=self._on_del_route,
handle=mitogen.core.DEL_ROUTE,
persist=True,
)
def propagate(self, handle, target_id):
# self.parent is None in the master.
if self.parent:
self.parent.send(
mitogen.core.Message(
handle=handle,
data=str(target_id),
)
)
def notice_stream(self, stream):
"""
When this parent is responsible for a new directly connected child
stream, we're also responsible for broadcasting DEL_ROUTE upstream
if/when that child disconnects.
"""
self.propagate(mitogen.core.ADD_ROUTE, stream.remote_id)
mitogen.core.listen(
obj=stream,
name='disconnect',
func=lambda: self._on_stream_disconnect(stream),
)
def _on_stream_disconnect(self, stream):
"""
Respond to disconnection of a local stream by
"""
for target_id in stream.routes:
LOG.debug('%r is gone; propagating DEL_ROUTE for ID %d',
stream, target_id)
self.router.del_route(target_id)
self.propagate(mitogen.core.DEL_ROUTE, target_id)
def _on_add_route(self, msg):
if msg == mitogen.core._DEAD:
return
target_id = int(msg.data)
stream = self.router.stream_by_id(msg.auth_id)
current = self.router.stream_by_id(target_id)
if current and current.remote_id != mitogen.parent_id:
LOG.error('Cannot add duplicate route to %r via %r, '
'already have existing route via %r',
target_id, stream, current)
return
LOG.debug('Adding route to %d via %r', target_id, stream)
stream.routes.add(target_id)
self.router.add_route(target_id, stream)
self.propagate(mitogen.core.ADD_ROUTE, target_id)
def _on_del_route(self, msg):
if msg == mitogen.core._DEAD:
return
target_id = int(msg.data)
registered_stream = self.router.stream_by_id(target_id)
stream = self.router.stream_by_id(msg.auth_id)
if registered_stream != stream:
LOG.error('Received DEL_ROUTE for %d from %r, expected %r',
target_id, stream, registered_stream)
return
LOG.debug('Deleting route to %d via %r', target_id, stream)
stream.routes.discard(target_id)
self.router.del_route(target_id)
self.propagate(mitogen.core.DEL_ROUTE, target_id)
class Router(mitogen.core.Router):
context_class = mitogen.core.Context
id_allocator = None
responder = None
log_forwarder = None
route_monitor = None
def upgrade(self, importer, parent):
LOG.debug('%r.upgrade()', self)
self.id_allocator = ChildIdAllocator(router=self)
self.responder = ModuleForwarder(
router=self,
parent_context=parent,
importer=importer,
)
self.route_monitor = RouteMonitor(self, parent)
def stream_by_id(self, dst_id):
return self._stream_by_id.get(dst_id,
self._stream_by_id.get(mitogen.parent_id))
def add_route(self, target_id, stream):
LOG.debug('%r.add_route(%r, %r)', self, target_id, stream)
assert isinstance(target_id, int)
assert isinstance(stream, Stream)
try:
self._stream_by_id[target_id] = stream
except KeyError:
LOG.error('%r: cant add route to %r via %r: no such stream',
self, target_id, stream)
def del_route(self, target_id):
LOG.debug('%r.del_route(%r)', self, target_id)
try:
del self._stream_by_id[target_id]
except KeyError:
LOG.error('%r: cant delete route to %r: no such stream',
self, target_id)
def get_module_blacklist(self):
if mitogen.context_id == 0:
return self.responder.blacklist
return self.importer.blacklist
def get_module_whitelist(self):
if mitogen.context_id == 0:
return self.responder.whitelist
return self.importer.whitelist
def allocate_id(self):
return self.id_allocator.allocate()
def context_by_id(self, context_id, via_id=None):
context = self._context_by_id.get(context_id)
if context is None:
context = self.context_class(self, context_id)
if via_id is not None:
context.via = self.context_by_id(via_id)
self._context_by_id[context_id] = context
return context
def _connect(self, klass, name=None, **kwargs):
context_id = self.allocate_id()
context = self.context_class(self, context_id)
stream = klass(self, context_id, **kwargs)
if name is not None:
stream.name = name
stream.connect()
context.name = stream.name
self.route_monitor.notice_stream(stream)
self.route_monitor.propagate(mitogen.core.ADD_ROUTE, context_id)
self.register(context, stream)
return context
def connect(self, method_name, name=None, **kwargs):
klass = METHOD_NAMES[method_name]()
kwargs.setdefault('debug', self.debug)
kwargs.setdefault('profiling', self.profiling)
via = kwargs.pop('via', None)
if via is not None:
return self.proxy_connect(via, method_name, name=name, **kwargs)
return self._connect(klass, name=name, **kwargs)
def proxy_connect(self, via_context, method_name, name=None, **kwargs):
context_id, name = via_context.call(_proxy_connect,
name=name,
method_name=method_name,
kwargs=kwargs
)
name = '%s.%s' % (via_context.name, name)
context = self.context_class(self, context_id, name=name)
context.via = via_context
self._context_by_id[context.context_id] = context
return context
class ProcessMonitor(object):
def __init__(self):
# pid -> callback()
self.callback_by_pid = {}
signal.signal(signal.SIGCHLD, self._on_sigchld)
def _on_sigchld(self, _signum, _frame):
for pid, callback in self.callback_by_pid.items():
pid, status = os.waitpid(pid, os.WNOHANG)
if pid:
callback(status)
del self.callback_by_pid[pid]
def add(self, pid, callback):
self.callback_by_pid[pid] = callback
_instance = None
@classmethod
def instance(cls):
if cls._instance is None:
cls._instance = cls()
return cls._instance
class ModuleForwarder(object):
"""
Respond to GET_MODULE requests in a slave by forwarding the request to our
parent context, or satisfying the request from our local Importer cache.
"""
def __init__(self, router, parent_context, importer):
self.router = router
self.parent_context = parent_context
self.importer = importer
router.add_handler(self._on_get_module, mitogen.core.GET_MODULE)
def __repr__(self):
return 'ModuleForwarder(%r)' % (self.router,)
def _on_get_module(self, msg):
LOG.debug('%r._on_get_module(%r)', self, msg)
if msg == mitogen.core._DEAD:
return
fullname = msg.data
callback = lambda: self._on_cache_callback(msg, fullname)
self.importer._request_module(fullname, callback)
def _send_one_module(self, msg, tup):
self.router._async_route(
mitogen.core.Message.pickled(
tup,
dst_id=msg.src_id,
handle=mitogen.core.LOAD_MODULE,
)
)
def _on_cache_callback(self, msg, fullname):
LOG.debug('%r._on_get_module(): sending %r', self, fullname)
tup = self.importer._cache[fullname]
if tup is not None:
for related in tup[4]:
rtup = self.importer._cache.get(related)
if not rtup:
LOG.debug('%r._on_get_module(): skipping absent %r',
self, related)
continue
self._send_one_module(msg, rtup)
self._send_one_module(msg, tup)