diff --git a/ansible_mitogen/loaders.py b/ansible_mitogen/loaders.py index 21a2c145..08c59278 100644 --- a/ansible_mitogen/loaders.py +++ b/ansible_mitogen/loaders.py @@ -34,7 +34,6 @@ from __future__ import absolute_import try: from ansible.plugins.loader import action_loader - from ansible.plugins.loader import callback_loader from ansible.plugins.loader import connection_loader from ansible.plugins.loader import module_loader from ansible.plugins.loader import module_utils_loader @@ -42,7 +41,6 @@ try: from ansible.plugins.loader import strategy_loader except ImportError: # Ansible <2.4 from ansible.plugins import action_loader - from ansible.plugins import callback_loader from ansible.plugins import connection_loader from ansible.plugins import module_loader from ansible.plugins import module_utils_loader diff --git a/ansible_mitogen/strategy.py b/ansible_mitogen/strategy.py index 217a6def..3cdf85f9 100644 --- a/ansible_mitogen/strategy.py +++ b/ansible_mitogen/strategy.py @@ -28,12 +28,41 @@ from __future__ import absolute_import import os +import threading import ansible_mitogen.loaders import ansible_mitogen.mixins import ansible_mitogen.process +def _patch_awx_callback(): + """ + issue #400: AWX loads a display callback that suffers from thread-safety + issues. Detect the presence of older AWX versions and patch the bug. + """ + # AWX uses sitecustomize.py to force-load this package. If it exists, we're + # running under AWX. + try: + from awx_display_callback.events import EventContext + from awx_display_callback.events import event_context + except ImportError: + return + + if hasattr(EventContext(), '_local'): + # Patched version. + return + + def patch_add_local(self, **kwargs): + tls = vars(self._local) + ctx = tls.setdefault('_ctx', {}) + ctx.update(kwargs) + + EventContext._local = threading.local() + EventContext.add_local = patch_add_local + +_patch_awx_callback() + + def wrap_action_loader__get(name, *args, **kwargs): """ While the mitogen strategy is active, trap action_loader.get() calls, @@ -65,29 +94,6 @@ def wrap_connection_loader__get(name, *args, **kwargs): return connection_loader__get(name, *args, **kwargs) -def patch_awx_callback(): - """ - issue #400: AWX loads a display callback that suffers from thread-safety - issues. Detect the presence of older AWX versions and patch the bug. - """ - cls = ansible_mitogen.loaders.callback_loader.get( - 'awx_display', - class_only=True, - ) - if cls is None: # callback does not exist. - return - - # Callback load will have updated sys.path. Now import its guts. - try: - from awx_display_callback.events import event_context - except ImportError: - return # Newer or ancient AWX. - - # Executing this before starting additional threads avoids race. - with event_context.set_global(): - pass - - class StrategyMixin(object): """ This mix-in enhances any built-in strategy by arranging for various Mitogen @@ -187,7 +193,6 @@ class StrategyMixin(object): ansible_mitogen.process.MuxProcess.start() self._add_connection_plugin_path() self._install_wrappers() - patch_awx_callback() try: return super(StrategyMixin, self).run(iterator, play_context) finally: