issue #217: ansible: beginnings of ModuleDepService.

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):
     """

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')),
         )

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)

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]