issue #217: ansible: working, if extremely inefficient implementation

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
     )

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,

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

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()
 

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]