diff --git a/ansible_mitogen/strategy.py b/ansible_mitogen/strategy.py index e82e6d58..1ecd6912 100644 --- a/ansible_mitogen/strategy.py +++ b/ansible_mitogen/strategy.py @@ -30,10 +30,13 @@ from __future__ import absolute_import import os import threading +import mitogen.core import ansible_mitogen.loaders import ansible_mitogen.mixins import ansible_mitogen.process +import ansible.executor.process.worker + def _patch_awx_callback(): """ @@ -93,6 +96,20 @@ def wrap_connection_loader__get(name, *args, **kwargs): return connection_loader__get(name, *args, **kwargs) +def wrap_worker__run(*args, **kwargs): + """ + While the strategy is active, rewrite connection_loader.get() calls for + some transports into requests for a compatible Mitogen transport. + """ + if mitogen.core._profile_hook.__name__ != '_profile_hook': + import signal + signal.signal(signal.SIGTERM, signal.SIG_IGN) + + return mitogen.core._profile_hook('WorkerProcess', + lambda: worker__run(*args, **kwargs) + ) + + class StrategyMixin(object): """ This mix-in enhances any built-in strategy by arranging for various Mitogen @@ -167,12 +184,17 @@ class StrategyMixin(object): connection_loader__get = ansible_mitogen.loaders.connection_loader.get ansible_mitogen.loaders.connection_loader.get = wrap_connection_loader__get + global worker__run + worker__run = ansible.executor.process.worker.WorkerProcess.run + ansible.executor.process.worker.WorkerProcess.run = wrap_worker__run + def _remove_wrappers(self): """ Uninstall the PluginLoader monkey patches. """ ansible_mitogen.loaders.action_loader.get = action_loader__get ansible_mitogen.loaders.connection_loader.get = connection_loader__get + ansible.executor.process.worker.WorkerProcess.run = worker__run def _add_plugin_paths(self): """ @@ -193,9 +215,12 @@ class StrategyMixin(object): the strategy's real run() method. """ ansible_mitogen.process.MuxProcess.start() + run = super(StrategyMixin, self).run self._add_plugin_paths() self._install_wrappers() try: - return super(StrategyMixin, self).run(iterator, play_context) + return mitogen.core._profile_hook('Strategy', + lambda: run(iterator, play_context) + ) finally: self._remove_wrappers()