diff --git a/ansible_mitogen/module_finder.py b/ansible_mitogen/module_finder.py index dd20237d..44c06a3c 100644 --- a/ansible_mitogen/module_finder.py +++ b/ansible_mitogen/module_finder.py @@ -1,5 +1,6 @@ from __future__ import absolute_import +import collections import imp import os import sys @@ -11,33 +12,27 @@ import ansible.module_utils PREFIX = 'ansible.module_utils.' -class Module(object): - def __init__(self, name, path, kind=imp.PY_SOURCE, parent=None): - self.name = name - self.path = path - self.kind = kind - self.is_pkg = kind == imp.PKG_DIRECTORY - self.parent = parent +Module = collections.namedtuple('Module', 'name path kind parent') - def __hash__(self): - return hash(self.path) - def __eq__(self, other): - return self.path == other.path +def get_fullname(module): + bits = [str(module.name)] + while module.parent: + bits.append(str(module.parent.name)) + module = module.parent + return '.'.join(reversed(bits)) - def fullname(self): - bits = [str(self.name)] - while self.parent: - bits.append(str(self.parent.name)) - self = self.parent - return '.'.join(reversed(bits)) - def code(self): - fp = open(self.path) - try: - return compile(fp.read(), str(self.name), 'exec') - finally: - fp.close() +def get_code(module): + fp = open(module.path) + try: + return compile(fp.read(), str(module.name), 'exec') + finally: + fp.close() + + +def is_pkg(module): + return module.kind == imp.PKG_DIRECTORY def find(name, path=(), parent=None): @@ -76,13 +71,18 @@ def scan_fromlist(code): def scan(module_name, module_path, search_path): - module = Module(module_name, module_path) + module = Module( + name=module_name, + path=module_path, + kind=imp.PY_SOURCE, + parent=None, + ) stack = [module] seen = set() while stack: module = stack.pop(0) - for level, fromname in scan_fromlist(module.code()): + for level, fromname in scan_fromlist(get_code(module)): if not fromname.startswith(PREFIX): continue @@ -90,17 +90,25 @@ def scan(module_name, module_path, search_path): if imported is None or imported in seen: continue + if imported in seen: + continue + seen.add(imported) + stack.append(imported) parent = imported.parent while parent: - module = Module(name=parent.fullname(), path=parent.path, - kind=parent.kind) + module = Module( + name=get_fullname(parent), + path=parent.path, + kind=parent.kind, + parent=None, + ) if module not in seen: seen.add(module) stack.append(module) parent = parent.parent return sorted( - (PREFIX + module.fullname(), module.path, module.is_pkg) + (PREFIX + get_fullname(module), module.path, is_pkg(module)) for module in seen ) diff --git a/ansible_mitogen/planner.py b/ansible_mitogen/planner.py index aa96bd77..0e4499e9 100644 --- a/ansible_mitogen/planner.py +++ b/ansible_mitogen/planner.py @@ -180,7 +180,7 @@ class BinaryPlanner(Planner): def detect(self, invocation): return module_common._is_binary(invocation.module_source) - def plan(self, invocation, **kwargs): + def _grant_file_service_access(self, invocation): invocation.connection._connect() mitogen.service.call( context=invocation.connection.parent, @@ -190,6 +190,9 @@ class BinaryPlanner(Planner): 'path': invocation.module_path } ) + + def plan(self, invocation, **kwargs): + self._grant_file_service_access(invocation) return super(BinaryPlanner, self).plan( invocation=invocation, runner_name=self.runner_name, @@ -270,6 +273,12 @@ class NewStylePlanner(ScriptPlanner): def _get_interpreter(self, invocation): return None, None + def _grant_file_service_access(self, invocation): + """ + Stub out BinaryPlanner's method since ModuleDepService makes internal + calls to grant file access, avoiding 2 IPCs per task invocation. + """ + def get_should_fork(self, invocation): """ In addition to asynchronous tasks, new-style modules should be forked @@ -293,6 +302,7 @@ class NewStylePlanner(ScriptPlanner): return tuple(paths) def get_module_utils(self, invocation): + invocation.connection._connect() module_utils = mitogen.service.call( context=invocation.connection.parent, handle=ansible_mitogen.services.ModuleDepService.handle, @@ -309,19 +319,7 @@ class NewStylePlanner(ScriptPlanner): return module_utils, has_custom def plan(self, invocation): - invocation.connection._connect() module_utils, has_custom = self.get_module_utils(invocation) - mitogen.service.call( - context=invocation.connection.parent, - handle=ansible_mitogen.services.FileService.handle, - method='register_many', - kwargs={ - 'paths': [ - path - for fullname, path, is_pkg in module_utils - ] - } - ) return super(NewStylePlanner, self).plan( invocation, module_utils=module_utils, diff --git a/ansible_mitogen/process.py b/ansible_mitogen/process.py index 82b6a59c..4946aa29 100644 --- a/ansible_mitogen/process.py +++ b/ansible_mitogen/process.py @@ -151,12 +151,16 @@ class MuxProcess(object): Construct a ContextService and a thread to service requests for it arriving from worker processes. """ + file_service = ansible_mitogen.services.FileService(router=self.router) self.pool = mitogen.service.Pool( router=self.router, services=[ + file_service, ansible_mitogen.services.ContextService(self.router), - ansible_mitogen.services.FileService(self.router), - ansible_mitogen.services.ModuleDepService(self.router), + ansible_mitogen.services.ModuleDepService( + router=self.router, + file_service=file_service, + ), ], size=int(os.environ.get('MITOGEN_POOL_SIZE', '16')), ) diff --git a/ansible_mitogen/runner.py b/ansible_mitogen/runner.py index e37c9326..807623f7 100644 --- a/ansible_mitogen/runner.py +++ b/ansible_mitogen/runner.py @@ -38,6 +38,7 @@ how to build arguments for it, preseed related data, etc. from __future__ import absolute_import import cStringIO import ctypes +import imp import json import logging import os @@ -182,6 +183,46 @@ class Runner(object): self.revert() +class ModuleUtilsImporter(object): + """ + :param list module_utils: + List of `(fullname, path, is_pkg)` tuples. + """ + def __init__(self, context, module_utils): + self._context = context + self._by_fullname = { + fullname: (path, is_pkg) + for fullname, path, is_pkg in module_utils + } + self._loaded = set() + sys.meta_path.insert(0, self) + + def revert(self): + sys.meta_path.remove(self) + for fullname in self._loaded: + sys.modules.pop(fullname, None) + + def find_module(self, fullname, path=None): + if fullname in self._by_fullname: + return self + + def load_module(self, fullname): + path, is_pkg = self._by_fullname[fullname] + source = ansible_mitogen.target.get_file(self._context, path) + code = compile(source, path, 'exec') + mod = sys.modules.setdefault(fullname, imp.new_module(fullname)) + mod.__file__ = "master:%s" % (path,) + mod.__loader__ = self + if is_pkg: + mod.__path__ = [] + mod.__package__ = fullname + else: + mod.__package__ = fullname.rpartition('.')[0] + exec(code, mod.__dict__) + self._loaded.add(fullname) + return mod + + class TemporaryEnvironment(object): def __init__(self, env=None): self.original = os.environ.copy() @@ -413,6 +454,10 @@ class NewStyleRunner(ScriptRunner): # module, but this has never been a bug report. Instead act like an # interpreter that had its script piped on stdin. self._argv = TemporaryArgv(['']) + self._importer = ModuleUtilsImporter( + context=self.service_context, + module_utils=self.module_utils, + ) if libc__res_init: libc__res_init() diff --git a/ansible_mitogen/services.py b/ansible_mitogen/services.py index 55c0365b..e7d6b2d7 100644 --- a/ansible_mitogen/services.py +++ b/ansible_mitogen/services.py @@ -612,8 +612,9 @@ class ModuleDepService(mitogen.service.Service): max_message_size = 1000 handle = 502 - def __init__(self, *args, **kwargs): - super(ModuleDepService, self).__init__(*args, **kwargs) + def __init__(self, file_service, **kwargs): + super(ModuleDepService, self).__init__(**kwargs) + self._file_service = file_service self._cache = {} @mitogen.service.expose(policy=mitogen.service.AllowParents()) @@ -630,4 +631,11 @@ class ModuleDepService(mitogen.service.Service): search_path=search_path, ) self._cache[module_name, search_path] = resolved + + # Grant FileService access to paths in here to avoid another 2 IPCs + # from WorkerProcess. + self._file_service.register(path=module_path) + for fullname, path, is_pkg in resolved: + self._file_service.register(path=path) + return self._cache[module_name, search_path]