From 81b62d9a1aa19fbdf406f346002e1b6db22f88e1 Mon Sep 17 00:00:00 2001 From: David Wilson Date: Sat, 12 May 2018 14:14:35 +0100 Subject: [PATCH] issue #217: ansible: beginnings of ModuleDepService. --- ansible_mitogen/planner.py | 49 ++++++++++++++++++++++++++++++++++++- ansible_mitogen/process.py | 1 + ansible_mitogen/runner.py | 4 +++ ansible_mitogen/services.py | 42 +++++++++++++++++++++++++++++++ 4 files changed, 95 insertions(+), 1 deletion(-) diff --git a/ansible_mitogen/planner.py b/ansible_mitogen/planner.py index f19d77d4..aa96bd77 100644 --- a/ansible_mitogen/planner.py +++ b/ansible_mitogen/planner.py @@ -35,17 +35,19 @@ files/modules known missing. """ from __future__ import absolute_import -import json import logging import os from ansible.executor import module_common import ansible.errors +import ansible.module_utils try: from ansible.plugins.loader import module_loader + from ansible.plugins.loader import module_utils_loader except ImportError: # Ansible <2.4 from ansible.plugins import module_loader + from ansible.plugins import module_utils_loader import mitogen import mitogen.service @@ -281,6 +283,51 @@ class NewStylePlanner(ScriptPlanner): def detect(self, invocation): return 'from ansible.module_utils.' in invocation.module_source + def get_module_utils_path(self, invocation): + paths = [ + path + for path in module_utils_loader._get_paths(subdirs=False) + if os.path.isdir(path) + ] + paths.append(module_common._MODULE_UTILS_PATH) + return tuple(paths) + + def get_module_utils(self, invocation): + module_utils = mitogen.service.call( + context=invocation.connection.parent, + handle=ansible_mitogen.services.ModuleDepService.handle, + method='scan', + kwargs={ + 'module_name': 'ansible_module_%s' % (invocation.module_name,), + 'module_path': invocation.module_path, + 'search_path': self.get_module_utils_path(invocation), + } + ) + modutils_dir = os.path.dirname(ansible.module_utils.__file__) + has_custom = not all(path.startswith(modutils_dir) + for fullname, path, is_pkg in module_utils) + 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, + should_fork=(self.get_should_fork(invocation) or has_custom), + ) + class ReplacerPlanner(NewStylePlanner): """ diff --git a/ansible_mitogen/process.py b/ansible_mitogen/process.py index f5dc7be5..82b6a59c 100644 --- a/ansible_mitogen/process.py +++ b/ansible_mitogen/process.py @@ -156,6 +156,7 @@ class MuxProcess(object): services=[ ansible_mitogen.services.ContextService(self.router), ansible_mitogen.services.FileService(self.router), + ansible_mitogen.services.ModuleDepService(self.router), ], size=int(os.environ.get('MITOGEN_POOL_SIZE', '16')), ) diff --git a/ansible_mitogen/runner.py b/ansible_mitogen/runner.py index bd2b0cc7..e37c9326 100644 --- a/ansible_mitogen/runner.py +++ b/ansible_mitogen/runner.py @@ -402,6 +402,10 @@ class NewStyleRunner(ScriptRunner): #: path => new-style module bytecode. _code_by_path = {} + def __init__(self, module_utils, **kwargs): + super(NewStyleRunner, self).__init__(**kwargs) + self.module_utils = module_utils + def setup(self): super(NewStyleRunner, self).setup() self._stdio = NewStyleStdio(self.args) diff --git a/ansible_mitogen/services.py b/ansible_mitogen/services.py index 38bca7da..55c0365b 100644 --- a/ansible_mitogen/services.py +++ b/ansible_mitogen/services.py @@ -49,7 +49,9 @@ import threading import zlib import mitogen +import mitogen.master import mitogen.service +import ansible_mitogen.module_finder import ansible_mitogen.target @@ -440,6 +442,17 @@ class FileService(mitogen.service.Service): except KeyError: return None + @mitogen.service.expose(policy=mitogen.service.AllowParents()) + @mitogen.service.arg_spec({ + 'paths': list + }) + def register_many(self, paths): + """ + Batch version of register(). + """ + for path in paths: + self.register(path) + @mitogen.service.expose(policy=mitogen.service.AllowParents()) @mitogen.service.arg_spec({ 'path': basestring @@ -589,3 +602,32 @@ class FileService(mitogen.service.Service): self._schedule_pending_unlocked(state) finally: state.lock.release() + + +class ModuleDepService(mitogen.service.Service): + """ + Scan a new-style module and produce a cached mapping of module_utils names + to their resolved filesystem paths. + """ + max_message_size = 1000 + handle = 502 + + def __init__(self, *args, **kwargs): + super(ModuleDepService, self).__init__(*args, **kwargs) + self._cache = {} + + @mitogen.service.expose(policy=mitogen.service.AllowParents()) + @mitogen.service.arg_spec({ + 'module_name': basestring, + 'module_path': basestring, + 'search_path': tuple, + }) + def scan(self, module_name, module_path, search_path): + if (module_name, search_path) not in self._cache: + resolved = ansible_mitogen.module_finder.scan( + module_name=module_name, + module_path=module_path, + search_path=search_path, + ) + self._cache[module_name, search_path] = resolved + return self._cache[module_name, search_path]