|
|
|
@ -27,6 +27,7 @@
|
|
|
|
|
|
|
|
|
|
from __future__ import absolute_import
|
|
|
|
|
import os
|
|
|
|
|
import threading
|
|
|
|
|
|
|
|
|
|
import mitogen
|
|
|
|
|
import mitogen.master
|
|
|
|
@ -42,13 +43,13 @@ import ansible_mitogen.mixins
|
|
|
|
|
|
|
|
|
|
def wrap_action_loader__get(name, *args, **kwargs):
|
|
|
|
|
"""
|
|
|
|
|
Trap calls to the action plug-in loader, supplementing the type of any
|
|
|
|
|
ActionModule with Mitogen's ActionModuleMixin before constructing it,
|
|
|
|
|
causing the mix-in methods to override any inherited from Ansible's base
|
|
|
|
|
class, replacing most shell use with pure Python equivalents.
|
|
|
|
|
While the mitogen stratey is active, trap action_loader.get() calls,
|
|
|
|
|
augmenting any fetched class with ActionModuleMixin, which replaces various
|
|
|
|
|
helper methods inherited from ActionBase with implementations that avoid
|
|
|
|
|
the use of shell fragments wherever possible.
|
|
|
|
|
|
|
|
|
|
This is preferred to static subclassing as it generalizes to third party
|
|
|
|
|
action modules existing outside the Ansible tree.
|
|
|
|
|
This is used instead of static subclassing as it generalizes to third party
|
|
|
|
|
action modules outside the Ansible tree.
|
|
|
|
|
"""
|
|
|
|
|
klass = action_loader__get(name, class_only=True)
|
|
|
|
|
if klass:
|
|
|
|
@ -59,12 +60,13 @@ def wrap_action_loader__get(name, *args, **kwargs):
|
|
|
|
|
return adorned_klass
|
|
|
|
|
return adorned_klass(*args, **kwargs)
|
|
|
|
|
|
|
|
|
|
action_loader__get = ansible.plugins.action_loader.get
|
|
|
|
|
ansible.plugins.action_loader.get = wrap_action_loader__get
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def wrap_connection_loader__get(name, play_context, new_stdin):
|
|
|
|
|
"""
|
|
|
|
|
While the mitogen strategy is active, rewrite connection_loader.get() calls
|
|
|
|
|
for the 'ssh' and 'local' transports into corresponding requests for the
|
|
|
|
|
'mitogen' connection type, passing the original transport name into it as
|
|
|
|
|
an argument, so that it can emulate the original type.
|
|
|
|
|
"""
|
|
|
|
|
kwargs = {}
|
|
|
|
|
if name in ('ssh', 'local'):
|
|
|
|
@ -72,16 +74,29 @@ def wrap_connection_loader__get(name, play_context, new_stdin):
|
|
|
|
|
name = 'mitogen'
|
|
|
|
|
return connection_loader__get(name, play_context, new_stdin, **kwargs)
|
|
|
|
|
|
|
|
|
|
connection_loader__get = ansible.plugins.connection_loader.get
|
|
|
|
|
ansible.plugins.connection_loader.get = wrap_connection_loader__get
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class ContextProxyService(mitogen.service.Service):
|
|
|
|
|
"""
|
|
|
|
|
Implement a service accessible from worker processes connecting back into
|
|
|
|
|
the top-level process. The service yields an existing context matching a
|
|
|
|
|
connection configuration if it exists, otherwise it constructs a new
|
|
|
|
|
conncetion before returning it.
|
|
|
|
|
Used by worker processes connecting back into the top-level process to
|
|
|
|
|
fetch the single Context instance corresponding to the supplied connection
|
|
|
|
|
configuration, creating a matching connection if it does not exist.
|
|
|
|
|
|
|
|
|
|
For connection methods and their parameters, refer to:
|
|
|
|
|
http://mitogen.readthedocs.io/en/latest/api.html#context-factories
|
|
|
|
|
|
|
|
|
|
This concentrates all SSH connections in the top-level process, which may
|
|
|
|
|
become a bottleneck. There are multiple ways to fix that:
|
|
|
|
|
* creating one .local() child context per CPU and sharding connections
|
|
|
|
|
between them, using the master process to route messages, or
|
|
|
|
|
* as above, but having each child create a unique UNIX listener and
|
|
|
|
|
having workers connect in directly.
|
|
|
|
|
|
|
|
|
|
:param dict dct:
|
|
|
|
|
Parameters passed to mitogen.master.Router.[method](). One key,
|
|
|
|
|
"method", is popped from the dictionary and used to look up the method.
|
|
|
|
|
|
|
|
|
|
:returns mitogen.master.Context:
|
|
|
|
|
Corresponding Context instance.
|
|
|
|
|
"""
|
|
|
|
|
well_known_id = 500
|
|
|
|
|
max_message_size = 1000
|
|
|
|
@ -104,9 +119,9 @@ class ContextProxyService(mitogen.service.Service):
|
|
|
|
|
class StrategyModule(ansible.plugins.strategy.linear.StrategyModule):
|
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
|
|
|
super(StrategyModule, self).__init__(*args, **kwargs)
|
|
|
|
|
self.add_connection_plugin_path()
|
|
|
|
|
self._add_connection_plugin_path()
|
|
|
|
|
|
|
|
|
|
def add_connection_plugin_path(self):
|
|
|
|
|
def _add_connection_plugin_path(self):
|
|
|
|
|
"""
|
|
|
|
|
Automatically add the connection plug-in directory to the ModuleLoader
|
|
|
|
|
path, slightly reduces end-user configuration.
|
|
|
|
@ -116,23 +131,69 @@ class StrategyModule(ansible.plugins.strategy.linear.StrategyModule):
|
|
|
|
|
conn_dir = os.path.join(basedir, 'connection')
|
|
|
|
|
ansible.plugins.connection_loader.add_directory(conn_dir)
|
|
|
|
|
|
|
|
|
|
def run(self, iterator, play_context, result=0):
|
|
|
|
|
def _setup_logging(self):
|
|
|
|
|
"""
|
|
|
|
|
Setup Mitogen's logging. Eventually this should be redirected into
|
|
|
|
|
Ansible's logging.
|
|
|
|
|
"""
|
|
|
|
|
log_level = os.environ.get('MITOGEN_LOG_LEVEL', 'INFO')
|
|
|
|
|
log_io = 'MITOGEN_LOG_IO' in os.environ
|
|
|
|
|
mitogen.utils.log_to_file(level=log_level, io=log_io)
|
|
|
|
|
|
|
|
|
|
def _setup_master(self):
|
|
|
|
|
"""
|
|
|
|
|
Construct a Router, Broker, mitogen.unix listener thread, and thread
|
|
|
|
|
serving connection requests from worker processes.
|
|
|
|
|
"""
|
|
|
|
|
self.router = mitogen.master.Router()
|
|
|
|
|
self.router.responder.whitelist_prefix('ansible')
|
|
|
|
|
self.router.responder.whitelist_prefix('ansible_mitogen')
|
|
|
|
|
self.listener = mitogen.unix.Listener(self.router)
|
|
|
|
|
os.environ['LISTENER_SOCKET_PATH'] = self.listener.path
|
|
|
|
|
|
|
|
|
|
# TODO: gracefully shutdown and join on this at exist.
|
|
|
|
|
self.service = ContextProxyService(self.router)
|
|
|
|
|
#mitogen.utils.log_to_file(level='DEBUG', io=False)
|
|
|
|
|
|
|
|
|
|
import threading
|
|
|
|
|
th = threading.Thread(target=self.service.run)
|
|
|
|
|
th.setDaemon(True)
|
|
|
|
|
th.start()
|
|
|
|
|
self.service_thread = threading.Thread(target=self.service.run)
|
|
|
|
|
self.service_thread.setDaemon(True)
|
|
|
|
|
self.service_thread.start()
|
|
|
|
|
|
|
|
|
|
def _run_with_master(self, iterator, play_context, result):
|
|
|
|
|
"""
|
|
|
|
|
Arrange for a mitogen.master.Router to be available for the duration of
|
|
|
|
|
the strategy's real run() method.
|
|
|
|
|
"""
|
|
|
|
|
self._setup_logging()
|
|
|
|
|
self._setup_master()
|
|
|
|
|
try:
|
|
|
|
|
return super(StrategyModule, self).run(iterator, play_context)
|
|
|
|
|
finally:
|
|
|
|
|
self.router.broker.shutdown()
|
|
|
|
|
os.unlink(self.listener.path)
|
|
|
|
|
|
|
|
|
|
def _install_wrappers(self):
|
|
|
|
|
"""
|
|
|
|
|
Install our PluginLoader monkey patches and update global variables
|
|
|
|
|
with references to the real functions.
|
|
|
|
|
"""
|
|
|
|
|
global action_loader__get
|
|
|
|
|
action_loader__get = ansible.plugins.action_loader.get
|
|
|
|
|
ansible.plugins.action_loader.get = wrap_action_loader__get
|
|
|
|
|
|
|
|
|
|
global connection_loader__get
|
|
|
|
|
connection_loader__get = ansible.plugins.connection_loader.get
|
|
|
|
|
ansible.plugins.connection_loader.get = wrap_connection_loader__get
|
|
|
|
|
|
|
|
|
|
def _remove_wrappers(self):
|
|
|
|
|
"""
|
|
|
|
|
Uninstall the PluginLoader monkey patches.
|
|
|
|
|
"""
|
|
|
|
|
ansible.plugins.action_loader.get = action_loader__get
|
|
|
|
|
ansible.plugins.connection_loader.get = connection_loader__get
|
|
|
|
|
|
|
|
|
|
def run(self, iterator, play_context, result=0):
|
|
|
|
|
self._install_wrappers()
|
|
|
|
|
try:
|
|
|
|
|
return self._run_with_master(iterator, play_context, result)
|
|
|
|
|
finally:
|
|
|
|
|
self._remove_wrappers()
|
|
|
|
|
self._setup_master()
|
|
|
|
|