From 00edf0d66d04f637f26f843d5cb54a12934b252e Mon Sep 17 00:00:00 2001 From: David Wilson Date: Wed, 23 May 2018 16:23:22 +0100 Subject: [PATCH] core: have ExternalContext accept a config dict rather than kwargs. The parameter lists had gotten out of control. --- mitogen/core.py | 78 ++++++++++++++++++++++++-------------------- mitogen/fakessh.py | 39 ++++++++++++---------- mitogen/fork.py | 10 +++--- mitogen/parent.py | 6 ++-- tests/parent_test.py | 1 + 5 files changed, 74 insertions(+), 60 deletions(-) diff --git a/mitogen/core.py b/mitogen/core.py index c5c31360..ad948853 100644 --- a/mitogen/core.py +++ b/mitogen/core.py @@ -1596,11 +1596,14 @@ class Broker(object): class ExternalContext(object): detached = False + def __init__(self, config): + self.config = config + def _on_broker_shutdown(self): self.channel.close() def _on_broker_exit(self): - if not self.profiling: + if not self.config['profiling']: os.kill(os.getpid(), signal.SIGTERM) def _on_shutdown_msg(self, msg): @@ -1638,29 +1641,31 @@ class ExternalContext(object): LOG.error('Stream had %d bytes after 2000ms', pending) self.broker.defer(stream.on_disconnect, self.broker) - def _setup_master(self, max_message_size, profiling, unidirectional, - parent_id, context_id, in_fd, out_fd): - Router.max_message_size = max_message_size - self.profiling = profiling - if profiling: + def _setup_master(self): + Router.max_message_size = self.config['max_message_size'] + if self.config['profiling']: enable_profiling() self.broker = Broker() self.router = Router(self.broker) - self.router.undirectional = unidirectional + self.router.undirectional = self.config['unidirectional'] self.router.add_handler( fn=self._on_shutdown_msg, handle=SHUTDOWN, policy=has_parent_authority, ) self.master = Context(self.router, 0, 'master') + parent_id = self.config['parent_ids'][0] if parent_id == 0: self.parent = self.master else: self.parent = Context(self.router, parent_id, 'parent') - self.channel = Receiver(router=self.router, - handle=CALL_FUNCTION, - policy=has_parent_authority) + in_fd = self.config.get('in_fd', 100) + out_fd = self.config.get('out_fd', 1) + + self.recv = Receiver(router=self.router, + handle=CALL_FUNCTION, + policy=has_parent_authority) self.stream = Stream(self.router, parent_id) self.stream.name = 'parent' self.stream.accept(in_fd, out_fd) @@ -1678,20 +1683,22 @@ class ExternalContext(object): except OSError: pass # No first stage exists (e.g. fakessh) - def _setup_logging(self, debug, log_level): + def _setup_logging(self): root = logging.getLogger() - root.setLevel(log_level) + root.setLevel(self.config['log_level']) root.handlers = [LogHandler(self.master)] - if debug: + if self.config['debug']: enable_debug_logging() - def _setup_importer(self, importer, core_src_fd, whitelist, blacklist): + def _setup_importer(self): + importer = self.config.get('importer') if importer: importer._install_handler(self.router) importer._context = self.parent else: + core_src_fd = self.config.get('core_src_fd', 101) if core_src_fd: - fp = os.fdopen(101, 'r', 1) + fp = os.fdopen(core_src_fd, 'r', 1) try: core_size = int(fp.readline()) core_src = fp.read(core_size) @@ -1702,8 +1709,13 @@ class ExternalContext(object): else: core_src = None - importer = Importer(self.router, self.parent, - core_src, whitelist, blacklist) + importer = Importer( + self.router, + self.parent, + core_src, + self.config.get('whitelist', ()), + self.config.get('blacklist', ()), + ) self.importer = importer self.router.importer = importer @@ -1723,12 +1735,12 @@ class ExternalContext(object): sys.modules['mitogen.core'] = mitogen.core del sys.modules['__main__'] - def _setup_globals(self, version, context_id, parent_ids): - mitogen.__version__ = version + def _setup_globals(self): mitogen.is_master = False - mitogen.context_id = context_id - mitogen.parent_ids = parent_ids - mitogen.parent_id = parent_ids[0] + mitogen.__version__ = self.config['version'] + mitogen.context_id = self.config['context_id'] + mitogen.parent_ids = self.config['parent_ids'][:] + mitogen.parent_id = mitogen.parent_ids[0] def _setup_stdio(self): # We must open this prior to closing stdout, otherwise it will recycle @@ -1774,7 +1786,7 @@ class ExternalContext(object): return fn(*args, **kwargs) def _dispatch_calls(self): - for msg in self.channel: + for msg in self.recv: try: msg.reply(self._dispatch_one(msg)) except Exception: @@ -1783,28 +1795,24 @@ class ExternalContext(object): msg.reply(CallError(e)) self.dispatch_stopped = True - def main(self, parent_ids, context_id, debug, profiling, log_level, - unidirectional, max_message_size, version, in_fd=100, out_fd=1, - core_src_fd=101, setup_stdio=True, setup_package=True, - importer=None, whitelist=(), blacklist=()): - self._setup_master(max_message_size, profiling, unidirectional, - parent_ids[0], context_id, in_fd, out_fd) + def main(self): + self._setup_master() try: try: - self._setup_logging(debug, log_level) - self._setup_importer(importer, core_src_fd, whitelist, blacklist) + self._setup_logging() + self._setup_importer() self._reap_first_stage() - if setup_package: + if self.config.get('setup_package', True): self._setup_package() - self._setup_globals(version, context_id, parent_ids) - if setup_stdio: + self._setup_globals() + if self.config.get('setup_stdio', True): self._setup_stdio() self.router.register(self.parent, self.stream) sys.executable = os.environ.pop('ARGV0', sys.executable) _v and LOG.debug('Connected to %s; my ID is %r, PID is %r', - self.parent, context_id, os.getpid()) + self.parent, mitogen.context_id, os.getpid()) _v and LOG.debug('Recovered sys.executable: %r', sys.executable) _profile_hook('main', self._dispatch_calls) diff --git a/mitogen/fakessh.py b/mitogen/fakessh.py index 3ee91015..27cb4acd 100644 --- a/mitogen/fakessh.py +++ b/mitogen/fakessh.py @@ -304,6 +304,25 @@ def _fakessh_main(dest_context_id, econtext): process.control.put(('exit', None)) +def _get_econtext_config(context, sock2): + parent_ids = mitogen.parent_ids[:] + parent_ids.insert(0, mitogen.context_id) + return { + 'context_id': context.context_id, + 'core_src_fd': None, + 'debug': getattr(context.router, 'debug', False), + 'in_fd': sock2.fileno(), + 'log_level': mitogen.parent.get_log_level(), + 'max_message_size': context.router.max_message_size, + 'out_fd': sock2.fileno(), + 'parent_ids': parent_ids, + 'profiling': getattr(context.router, 'profiling', False), + 'unidirectional': getattr(context.router, 'unidirectional', False), + 'setup_stdio': False, + 'version': mitogen.__version__, + } + + # # Public API. # @@ -328,9 +347,6 @@ def run(dest, router, args, deadline=None, econtext=None): # Held in socket buffer until process is booted. fakessh.call_async(_fakessh_main, dest.context_id) - parent_ids = mitogen.parent_ids[:] - parent_ids.insert(0, mitogen.context_id) - tmp_path = tempfile.mkdtemp(prefix='mitogen_fakessh') try: ssh_path = os.path.join(tmp_path, 'ssh') @@ -339,20 +355,9 @@ def run(dest, router, args, deadline=None, econtext=None): fp.write('#!%s\n' % (sys.executable,)) fp.write(inspect.getsource(mitogen.core)) fp.write('\n') - fp.write('ExternalContext().main(**%r)\n' % ({ - 'context_id': context_id, - 'core_src_fd': None, - 'debug': getattr(router, 'debug', False), - 'in_fd': sock2.fileno(), - 'log_level': mitogen.parent.get_log_level(), - 'max_message_size': router.max_message_size, - 'out_fd': sock2.fileno(), - 'parent_ids': parent_ids, - 'profiling': getattr(router, 'profiling', False), - 'unidirectional': getattr(router, 'unidirectional', False), - 'setup_stdio': False, - 'version': mitogen.__version__, - },)) + fp.write('ExternalContext(%r).main()\n' % ( + _get_econtext_config(context, sock2), + )) finally: fp.close() diff --git a/mitogen/fork.py b/mitogen/fork.py index 4a5627dc..dc013fa5 100644 --- a/mitogen/fork.py +++ b/mitogen/fork.py @@ -148,12 +148,12 @@ class Stream(mitogen.parent.Stream): os.close(devnull) childfp.close() - kwargs = self.get_main_kwargs() - kwargs['core_src_fd'] = None - kwargs['importer'] = self.importer - kwargs['setup_package'] = False + config = self.get_econtext_config() + config['core_src_fd'] = None + config['importer'] = self.importer + config['setup_package'] = False try: - mitogen.core.ExternalContext().main(**kwargs) + mitogen.core.ExternalContext(config).main() finally: # Don't trigger atexit handlers, they were copied from the parent. os._exit(0) diff --git a/mitogen/parent.py b/mitogen/parent.py index 692adf24..9d66c736 100644 --- a/mitogen/parent.py +++ b/mitogen/parent.py @@ -896,7 +896,7 @@ class Stream(mitogen.core.Stream): 'exec(_(_("%s".encode(),"base64"),"zip"))' % (encoded,) ] - def get_main_kwargs(self): + def get_econtext_config(self): assert self.max_message_size is not None parent_ids = mitogen.parent_ids[:] parent_ids.insert(0, mitogen.context_id) @@ -915,8 +915,8 @@ class Stream(mitogen.core.Stream): def get_preamble(self): source = inspect.getsource(mitogen.core) - source += '\nExternalContext().main(**%r)\n' % ( - self.get_main_kwargs(), + source += '\nExternalContext(%r).main()\n' % ( + self.get_econtext_config(), ) return zlib.compress(minimize_source(source), 9) diff --git a/tests/parent_test.py b/tests/parent_test.py index 4cbdad38..c4dafb62 100644 --- a/tests/parent_test.py +++ b/tests/parent_test.py @@ -1,6 +1,7 @@ import errno import os import subprocess +import sys import tempfile import time