fakessh: Make ExternalContext.main() more flexible

- Allow specifying non-default in/out FDs
- Allow not passing econtext.core source code FD
- Allow starting directly without a first stage
- Allow disabling stdio redirection
This commit is contained in:
David Wilson 2017-08-30 17:50:49 +05:30
parent 69737701f6
commit 1f49d24cd6
1 changed files with 27 additions and 18 deletions

View File

@ -255,13 +255,13 @@ class Importer(object):
'econtext.utils', 'econtext.utils',
]} ]}
self.tls = threading.local() self.tls = threading.local()
self._cache = { self._cache = {}
'econtext.core': ( if core_src:
self._cache['econtext.core'] = (
None, None,
'econtext/core.py', 'econtext/core.py',
zlib.compress(core_src), zlib.compress(core_src),
) )
}
def __repr__(self): def __repr__(self):
return 'Importer()' return 'Importer()'
@ -963,7 +963,7 @@ class ExternalContext(object):
The :py:class:`IoLogger` connected to ``stderr``. The :py:class:`IoLogger` connected to ``stderr``.
""" """
def _setup_master(self, parent_id, context_id, key): def _setup_master(self, parent_id, context_id, key, in_fd, out_fd):
self.broker = Broker() self.broker = Broker()
self.router = Router(self.broker) self.router = Router(self.broker)
self.master = Context(self.router, 0, 'master') self.master = Context(self.router, 0, 'master')
@ -975,10 +975,13 @@ class ExternalContext(object):
self.channel = Channel(self.master, CALL_FUNCTION) self.channel = Channel(self.master, CALL_FUNCTION)
self.stream = Stream(self.router, parent_id, key) self.stream = Stream(self.router, parent_id, key)
self.stream.name = 'parent' self.stream.name = 'parent'
self.stream.accept(100, 1) self.stream.accept(in_fd, out_fd)
os.wait() # Reap first stage. os.close(in_fd)
os.close(100) try:
os.wait() # Reap first stage.
except OSError:
pass # No first stage exists (e.g. fakessh)
def _setup_logging(self, debug, log_level): def _setup_logging(self, debug, log_level):
root = logging.getLogger() root = logging.getLogger()
@ -987,13 +990,16 @@ class ExternalContext(object):
if debug: if debug:
enable_debug_logging() enable_debug_logging()
def _setup_importer(self): def _setup_importer(self, core_src_fd):
with os.fdopen(101, 'r', 1) as fp: if core_src_fd:
core_size = int(fp.readline()) with os.fdopen(101, 'r', 1) as fp:
core_src = fp.read(core_size) core_size = int(fp.readline())
# Strip "ExternalContext.main()" call from last line. core_src = fp.read(core_size)
core_src = '\n'.join(core_src.splitlines()[:-1]) # Strip "ExternalContext.main()" call from last line.
fp.close() core_src = '\n'.join(core_src.splitlines()[:-1])
fp.close()
else:
core_src = None
self.importer = Importer(self.parent, core_src) self.importer = Importer(self.parent, core_src)
sys.meta_path.append(self.importer) sys.meta_path.append(self.importer)
@ -1040,17 +1046,20 @@ class ExternalContext(object):
ret = fn(*args, **kwargs) ret = fn(*args, **kwargs)
self.master.send(Message.pickled(ret, handle=msg.reply_to)) self.master.send(Message.pickled(ret, handle=msg.reply_to))
except Exception, e: except Exception, e:
LOG.debug('_dispatch_calls: %s', e)
e = CallError(str(e)) e = CallError(str(e))
self.master.send(Message.pickled(e, handle=msg.reply_to)) self.master.send(Message.pickled(e, handle=msg.reply_to))
def main(self, parent_id, context_id, key, debug, log_level): def main(self, parent_id, context_id, key, debug, log_level,
self._setup_master(parent_id, context_id, key) in_fd=100, out_fd=1, core_src_fd=101, setup_stdio=True):
self._setup_master(parent_id, context_id, key, in_fd, out_fd)
try: try:
try: try:
self._setup_logging(debug, log_level) self._setup_logging(debug, log_level)
self._setup_importer() self._setup_importer(core_src_fd)
self._setup_package(context_id) self._setup_package(context_id)
self._setup_stdio() if setup_stdio:
self._setup_stdio()
self.router.register(self.parent, self.stream) self.router.register(self.parent, self.stream)
self.router.register(self.master, self.stream) self.router.register(self.master, self.stream)