ansible: connection plugin v2

Now featuring one roundtrip per module invocation.

mitogen/ansible/action.py

+# Copyright 2017, David Wilson
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its contributors
+# may be used to endorse or promote products derived from this software without
+# specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+import json
+import os
+
+import ansible
+import ansible.plugins
+import ansible.plugins.action.normal
+import mitogen.ansible.helpers
+
+
+ANSIBLE_BASEDIR = os.path.dirname(ansible.__file__)
+
+
+class ActionModule(ansible.plugins.action.normal.ActionModule):
+    def get_py_module_name(self, module_name):
+        path = ansible.plugins.module_loader.find_plugin(module_name, '')
+        relpath = os.path.relpath(path, ANSIBLE_BASEDIR)
+        root, _ = os.path.splitext(relpath)
+        return 'ansible.' + root.replace('/', '.')
+
+    def _execute_module(self, module_name=None, module_args=None, tmp=None,
+                        task_vars=None, persist_files=False,
+                        delete_remote_tmp=True, wrap_async=False):
+
+        module_name = module_name or self._task.action
+        module_args = module_args or self._task.args
+        task_vars = task_vars or {}
+
+        self._update_module_args(module_name, module_args, task_vars)
+
+        #####################################################################
+
+        py_module_name = self.get_py_module_name(module_name)
+        js = self._connection.py_call(mitogen.ansible.helpers.run_module, py_module_name,
+                                      args=json.loads(json.dumps(module_args)))
+
+        #####################################################################
+
+        data = self._parse_returned_data({
+            'rc': 0,
+            'stdout': js,
+            'stdout_lines': [js],
+            'stderr': ''
+        })
+
+        if wrap_async:
+            data['changed'] = True
+
+        # pre-split stdout/stderr into lines if needed
+        if 'stdout' in data and 'stdout_lines' not in data:
+            # if the value is 'False', a default won't catch it.
+            txt = data.get('stdout', None) or u''
+            data['stdout_lines'] = txt.splitlines()
+        if 'stderr' in data and 'stderr_lines' not in data:
+            # if the value is 'False', a default won't catch it.
+            txt = data.get('stderr', None) or u''
+            data['stderr_lines'] = txt.splitlines()
+
+        return data

mitogen/ansible/connection.py

@@ -43,8 +43,7 @@ Enable it by:
 """
 
 import mitogen.master
-import mitogen.ssh
-import mitogen.utils
+import mitogen.unix
 from mitogen.ansible import helpers
 
 import ansible.errors
@@ -52,7 +51,7 @@ import ansible.plugins.connection
 
 
 class Connection(ansible.plugins.connection.ConnectionBase):
-    broker = None
+    router = None
     context = None
 
     become_methods = []
@@ -60,35 +59,35 @@ class Connection(ansible.plugins.connection.ConnectionBase):
 
     @property
     def connected(self):
-        return self.broker is not None
+        return self.router is not None
 
     def _connect(self):
         if self.connected:
             return
-        self.broker = mitogen.master.Broker()
-        self.router = mitogen.master.Router(self.broker)
-        if self._play_context.remote_addr == 'localhost':
-            self.context = self.router.connect(mitogen.master.Stream)
-        else:
-            self.context = self.router.connect(mitogen.ssh.Stream,
-                hostname=self._play_context.remote_addr,
-            )
+
+        self.router, self.parent = mitogen.unix.connect('/tmp/mitosock')
+        self.context = mitogen.service.call(self.parent, 500, {
+            'hostname': self._play_context.remote_addr,
+        })
+
+    def py_call(self, func, *args, **kwargs):
+        self._connect()
+        return self.context.call(func, *args, **kwargs)
 
     def exec_command(self, cmd, in_data=None, sudoable=True):
         super(Connection, self).exec_command(cmd, in_data=in_data, sudoable=sudoable)
         if in_data:
             raise ansible.errors.AnsibleError("does not support module pipelining")
-
-        return self.context.call(helpers.exec_command, cmd, in_data)
+        return self.py_call(helpers.exec_command, cmd, in_data)
 
     def fetch_file(self, in_path, out_path):
-        output = self.context.call(helpers.read_path, in_path)
+        output = self.py_call(helpers.read_path, in_path)
         helpers.write_path(out_path, output)
 
     def put_file(self, in_path, out_path):
-        self.context.call(helpers.write_path, out_path,
+        self.py_call(helpers.write_path, out_path,
                      helpers.read_path(in_path))
 
     def close(self):
-        self.broker.shutdown()
-        self.broker.join()
+        self.router.broker.shutdown()
+        self.router.broker.join()

mitogen/ansible/helpers.py

@@ -121,7 +121,7 @@ def run_module(module, raw_params=None, args=None):
         # invocation, so we need to invoke main explicitly.
         mod.main()
     except Exit, e:
-        return e.dct
+        return json.dumps(e.dct)
 
 
 def exec_command(cmd, in_data=None):

mitogen/ansible/strategy.py

+# Copyright 2017, David Wilson
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its contributors
+# may be used to endorse or promote products derived from this software without
+# specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+import mitogen
+import mitogen.master
+import mitogen.service
+import mitogen.unix
+import mitogen.utils
+import mitogen.ansible.action
+
+import ansible.errors
+import ansible.plugins.strategy.linear
+from ansible.plugins import action_loader
+
+
+class ContextProxyService(mitogen.service.Service):
+    well_known_id = 500
+    max_message_size = 1000
+
+    def __init__(self, router):
+        super(ContextProxyService, self).__init__(router)
+        self._context_by_id = {}
+
+    def validate_args(self, args):
+        return (isinstance(args, dict) and
+                isinstance(args.get('hostname'), basestring))
+
+    def dispatch(self, dct, msg):
+        key = repr(sorted(dct.items()))
+        if key not in self._context_by_id:
+            self._context_by_id[key] = self.router.ssh(**dct)
+        return self._context_by_id[key]
+
+
+class StrategyModule(ansible.plugins.strategy.linear.StrategyModule):
+    def run(self, *args, **kwargs):
+        self.router = mitogen.master.Router()
+        self.listener = mitogen.unix.Listener(self.router, path='/tmp/mitosock')
+        self.service = ContextProxyService(self.router)
+        mitogen.utils.log_to_file()
+
+        import threading
+        th = threading.Thread(target=self.service.run)
+        th.setDaemon(True)
+        th.start()
+
+        real_get = action_loader.get
+        def get(name, *args, **kwargs):
+            if name == 'normal':
+                return mitogen.ansible.action.ActionModule(*args, **kwargs)
+            return real_get(name, *args, **kwargs)
+        action_loader.get = get
+
+        try:
+            return super(StrategyModule, self).run(*args, **kwargs)
+        finally:
+            self.router.broker.shutdown()