diff --git a/ansible_mitogen/mixins.py b/ansible_mitogen/mixins.py index fdd104b2..18d29da3 100644 --- a/ansible_mitogen/mixins.py +++ b/ansible_mitogen/mixins.py @@ -313,6 +313,7 @@ class ActionModuleMixin(ansible.plugins.action.ActionBase): env = {} self._compute_environment_string(env) + self._connection._connect() return ansible_mitogen.planner.invoke( ansible_mitogen.planner.Invocation( action=self, diff --git a/ansible_mitogen/planner.py b/ansible_mitogen/planner.py index 6605686c..29a1670a 100644 --- a/ansible_mitogen/planner.py +++ b/ansible_mitogen/planner.py @@ -181,7 +181,6 @@ class BinaryPlanner(Planner): return module_common._is_binary(invocation.module_source) def _grant_file_service_access(self, invocation): - invocation.connection._connect() invocation.connection.parent.call_service( service_name='ansible_mitogen.services.FileService', method_name='register', @@ -297,7 +296,6 @@ class NewStylePlanner(ScriptPlanner): ) def get_module_utils(self, invocation): - invocation.connection._connect() return invocation.connection.parent.call_service( service_name='ansible_mitogen.services.ModuleDepService', method_name='scan', @@ -306,6 +304,7 @@ class NewStylePlanner(ScriptPlanner): module_path=invocation.module_path, search_path=self.get_search_path(invocation), builtin_path=module_common._MODULE_UTILS_PATH, + context=invocation.connection.context, ) def plan(self, invocation): diff --git a/ansible_mitogen/services.py b/ansible_mitogen/services.py index ee4ca010..916a4c1a 100644 --- a/ansible_mitogen/services.py +++ b/ansible_mitogen/services.py @@ -624,32 +624,50 @@ class ModuleDepService(mitogen.service.Service): self._file_service = file_service self._cache = {} + def _get_builtin_names(self, builtin_path, resolved): + return [ + fullname + for fullname, path, is_pkg in resolved + if os.path.abspath(path).startswith(builtin_path) + ] + + def _get_custom_tups(self, builtin_path, resolved): + return [ + (fullname, path, is_pkg) + for fullname, path, is_pkg in resolved + if not os.path.abspath(path).startswith(builtin_path) + ] + @mitogen.service.expose(policy=mitogen.service.AllowParents()) @mitogen.service.arg_spec({ 'module_name': basestring, 'module_path': basestring, 'search_path': tuple, 'builtin_path': basestring, + 'context': mitogen.core.Context, }) - def scan(self, module_name, module_path, search_path, builtin_path): - if (module_name, search_path) not in self._cache: + def scan(self, module_name, module_path, search_path, builtin_path, context): + key = (module_name, search_path) + if key not in self._cache: resolved = ansible_mitogen.module_finder.scan( module_name=module_name, module_path=module_path, search_path=tuple(search_path) + (builtin_path,), ) builtin_path = os.path.abspath(builtin_path) - filtered = [ - (fullname, path, is_pkg) - for fullname, path, is_pkg in resolved - if not os.path.abspath(path).startswith(builtin_path) - ] - self._cache[module_name, search_path] = filtered + builtin = self._get_builtin_names(builtin_path, resolved) + custom = self._get_custom_tups(builtin_path, resolved) + self._cache[key] = { + 'builtin': builtin, + 'custom': custom, + } # 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 filtered: + for fullname, path, is_pkg in custom: self._file_service.register(path=path) - return self._cache[module_name, search_path] + for name in self._cache[key]['builtin']: + self.router.responder.forward_module(context, name) + return self._cache[key]['custom']