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

ansible: hook strategy and worker processes into profiler

This commit is contained in:
David Wilson 2019-01-30 00:39:13 +00:00
parent 5dc85c05f9
commit e587396e70
1 changed files with 26 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
ansible_mitogen/strategy.py
View File

@ -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()