Rooba
/
mitogen
mirror of https://github.com/mitogen-hq/mitogen.git
1
1
Fork 0
Code Issues Packages Projects Releases Wiki Activity

issue #106: support custom new-style modules + 'reexec by default' 

Rather than assume any structure about the Python code:

* Delete the exit_json/fail_json monkeypatches.
* Patch SystemExit rather than a magic monkeypatch-thrown exception
* Setup fake cStringIO stdin, stdout, stderr and return those along with
  SystemExit exit status
* Setup _ANSIBLE_ARGS as we used to, since we still want to override
  that with '{}' to prevent accidental import hangs, but also provide
  the same string via sys.stdin.
* Compile the module bytecode once and re-execute it for every
  invocation. May change this back again later, once some benchmarks are
  done.
* Remove the fixups stuff for now, it's handled by ^ above.

Should support any "somewhat new style" Python module, including those
that just give up and dump stuff to stdout directly.
This commit is contained in:
David Wilson 2018-04-01 21:10:04 +01:00
parent a954d54644
commit 8cc20856a8
3 changed files with 74 additions and 131 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
ansible_mitogen/helpers.py
View File

@ -51,6 +51,9 @@ LOG = logging.getLogger(__name__)
#: Caching of fetched file data.
_file_cache = {}
#: Caching of compiled new-style module bytecode.
_bytecode_cache = {}
#: Mapping of job_id<->result dict
_result_by_job_id = {}
@ -84,6 +87,13 @@ def get_file(context, path):
return _file_cache[path]
def get_bytecode(context, path):
if path not in _bytecode_cache:
source = get_file(context, path)
_bytecode_cache[path] = compile(source, path, 'exec')
return _bytecode_cache[path]
def run_module(kwargs):
"""
Set up the process environment in preparation for running an Ansible

26
ansible_mitogen/planner.py
View File

@ -241,7 +241,7 @@ class WantJsonPlanner(ScriptPlanner):
return 'WANT_JSON' in invocation.module_source
class NewStylePlanner(Planner):
class NewStylePlanner(ScriptPlanner):
"""
The Ansiballz framework differs from module replacer in that it uses real
Python imports of things in ansible/module_utils instead of merely
@ -252,30 +252,6 @@ class NewStylePlanner(Planner):
def detect(self, invocation):
return 'from ansible.module_utils.' in invocation.module_source
def get_command_module_name(self, module_name):
"""
Given the name of an Ansible command module, return its canonical
module path within the ansible.
:param module_name:
"shell"
:return:
"ansible.modules.commands.shell"
"""
path = module_loader.find_plugin(module_name, '')
relpath = os.path.relpath(path, os.path.dirname(ansible.__file__))
root, _ = os.path.splitext(relpath)
return 'ansible.' + root.replace('/', '.')
def plan(self, invocation):
return {
'runner_name': self.runner_name,
'module': invocation.module_name,
'mod_name': self.get_command_module_name(invocation.module_name),
'args': invocation.module_args,
'env': invocation.env,
}
class ReplacerPlanner(NewStylePlanner):
runner_name = 'ReplacerRunner'

169
ansible_mitogen/runner.py
View File

@ -36,10 +36,13 @@ how to build arguments for it, preseed related data, etc.
"""
from __future__ import absolute_import
import cStringIO
import json
import logging
import os
import sys
import tempfile
import types
import ansible_mitogen.helpers # TODO: circular import
@ -131,124 +134,37 @@ class TemporaryEnvironment(object):
os.environ.update(self.original)
class NewStyleModuleExit(Exception):
"""
Capture the result of a call to `.exit_json()` or `.fail_json()` by a
native Ansible module.
"""
def __init__(self, ansible_module, **kwargs):
ansible_module.add_path_info(kwargs)
kwargs.setdefault('invocation', {
'module_args': ansible_module.params
})
ansible_module.do_cleanup_files()
self.dct = ansible.module_utils.basic.remove_values(
kwargs,
ansible_module.no_log_values,
)
class NewStyleMethodOverrides(object):
@staticmethod
def exit_json(self, **kwargs):
"""
Raise exit_json() output as the `.dct` attribute of a
:class:`NewStyleModuleExit` exception`.
"""
kwargs.setdefault('changed', False)
raise NewStyleModuleExit(self, **kwargs)
@staticmethod
def fail_json(self, **kwargs):
"""
Raise fail_json() output as the `.dct` attribute of a
:class:`NewStyleModuleExit` exception`.
"""
kwargs.setdefault('failed', True)
raise NewStyleModuleExit(self, **kwargs)
klass = ansible.module_utils.basic.AnsibleModule
def __init__(self):
self._original_exit_json = self.klass.exit_json
self._original_fail_json = self.klass.fail_json
self.klass.exit_json = self.exit_json
self.klass.fail_json = self.fail_json
class TemporaryArgv(object):
def __init__(self, argv):
self.original = sys.argv[:]
sys.argv[:] = argv
def revert(self):
"""
Restore prior state.
"""
self.klass.exit_json = self._original_exit_json
self.klass.fail_json = self._original_fail_json
sys.argv[:] = self.original
class NewStyleModuleArguments(object):
class NewStyleStdio(object):
"""
Patch ansible.module_utils.basic argument globals.
"""
def __init__(self, args):
self.original = ansible.module_utils.basic._ANSIBLE_ARGS
self.original_stdout = sys.stdout
self.original_stderr = sys.stderr
self.original_stdin = sys.stdin
sys.stdout = cStringIO.StringIO()
sys.stderr = cStringIO.StringIO()
ansible.module_utils.basic._ANSIBLE_ARGS = json.dumps({
'ANSIBLE_MODULE_ARGS': args
})
sys.stdin = cStringIO.StringIO(
ansible.module_utils.basic._ANSIBLE_ARGS
)
def revert(self):
"""
Restore prior state.
"""
ansible.module_utils.basic._ANSIBLE_ARGS = self.original
class NewStyleRunner(Runner):
"""
Execute a new-style Ansible module, where Module Replacer-related tricks
aren't required.
"""
def __init__(self, mod_name, **kwargs):
super(NewStyleRunner, self).__init__(**kwargs)
self.mod_name = mod_name
def setup(self):
super(NewStyleRunner, self).setup()
self._overrides = NewStyleMethodOverrides()
self._args = NewStyleModuleArguments(self.args)
def revert(self):
super(NewStyleRunner, self).revert()
self._args.revert()
self._overrides.revert()
def _fixup__default(self, mod):
pass
def _fixup__yum_repository(self, mod):
# https://github.com/dw/mitogen/issues/154
mod.YumRepo.repofile = mod.configparser.RawConfigParser()
def _run(self):
fixup = getattr(self, '_fixup__' + self.module, self._fixup__default)
try:
mod = __import__(self.mod_name, {}, {}, [''])
fixup(mod)
# Ansible modules begin execution on import. Thus the above
# __import__ will cause either Exit or ModuleError to be raised. If
# we reach the line below, the module did not execute and must
# already have been imported for a previous invocation, so we need
# to invoke main explicitly.
mod.main()
except NewStyleModuleExit, e:
return {
'rc': 0,
'stdout': json.dumps(e.dct),
'stderr': '',
}
return {
'rc': 1,
'stdout': '',
'stderr': 'ansible_mitogen: module did not exit normally.',
}
sys.stdout = self.original_stdout
sys.stderr = self.original_stderr
sys.stdin = self.original_stdin
ansible.module_utils.basic._ANSIBLE_ARGS = '{}'
class ProgramRunner(Runner):
@ -385,6 +301,47 @@ class ScriptRunner(ProgramRunner):
return '\n'.join(new)
class NewStyleRunner(ScriptRunner):
"""
Execute a new-style Ansible module, where Module Replacer-related tricks
aren't required.
"""
def setup(self):
super(NewStyleRunner, self).setup()
self._stdio = NewStyleStdio(self.args)
self._argv = TemporaryArgv([self.path])
def revert(self):
super(NewStyleRunner, self).revert()
self._stdio.revert()
def _get_bytecode(self):
"""
Fetch the module binary from the master if necessary.
"""
return ansible_mitogen.helpers.get_bytecode(
context=self.service_context,
path=self.path,
)
def _run(self):
bytecode = self._get_bytecode()
mod = types.ModuleType('__main__')
d = vars(mod)
e = None
try:
exec bytecode in d, d
except SystemExit, e:
pass
return {
'rc': e[0] if e else 2,
'stdout': sys.stdout.getvalue(),
'stderr': sys.stderr.getvalue(),
}
class JsonArgsFileRunner(ArgsFileRunner, ScriptRunner):
JSON_ARGS = '<<INCLUDE_ANSIBLE_MODULE_JSON_ARGS>>'
@ -411,4 +368,4 @@ class OldStyleRunner(ArgsFileRunner, ScriptRunner):
return ' '.join(
'%s=%s' % (key, shlex_quote(str(self.args[key])))
for key in self.args
)
) + ' ' # Bug-for-bug :(