ansible: hook strategy and worker processes into profiler

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