From 78b00652bdc7c853289d910c05b533b53b576d4b Mon Sep 17 00:00:00 2001
From: Ryan Petrello <rpetrell@redhat.com>
Date: Mon, 27 Jan 2020 10:41:54 -0500
Subject: [PATCH] add the ability to enable profiling for the callback receiver
 workers

---
 awx/main/dispatch/worker/callback.py | 27 +++++++++++++++++++++++++++
 awx/settings/defaults.py             |  3 +++
 2 files changed, 30 insertions(+)

diff --git a/awx/main/dispatch/worker/callback.py b/awx/main/dispatch/worker/callback.py
index b80c2bce86..739f1538ea 100644
--- a/awx/main/dispatch/worker/callback.py
+++ b/awx/main/dispatch/worker/callback.py
@@ -1,8 +1,14 @@
+import cProfile
 import logging
+import os
+import pstats
+import signal
+import tempfile
 import time
 import traceback
 from queue import Empty as QueueEmpty
 
+from django.conf import settings
 from django.utils.timezone import now as tz_now
 from django.db import DatabaseError, OperationalError, connection as django_connection
 from django.db.utils import InterfaceError, InternalError, IntegrityError
@@ -31,6 +37,7 @@ class CallbackBrokerWorker(BaseWorker):
     '''
 
     MAX_RETRIES = 2
+    prof = None
 
     def __init__(self):
         self.buff = {}
@@ -41,6 +48,26 @@ class CallbackBrokerWorker(BaseWorker):
         except QueueEmpty:
             return {'event': 'FLUSH'}
 
+    def toggle_profiling(self, *args):
+        if self.prof:
+            self.prof.disable()
+            filename = f'callback-{os.getpid()}.pstats'
+            filepath = os.path.join(tempfile.gettempdir(), filename)
+            with open(filepath, 'w') as f:
+                pstats.Stats(self.prof, stream=f).sort_stats('cumulative').print_stats()
+            pstats.Stats(self.prof).dump_stats(filepath + '.raw')
+            self.prof = False
+            logger.error(f'profiling is disabled, wrote {filepath}')
+        else:
+            self.prof = cProfile.Profile()
+            self.prof.enable()
+            logger.error('profiling is enabled')
+
+    def work_loop(self, *args, **kw):
+        if settings.AWX_CALLBACK_PROFILE:
+            signal.signal(signal.SIGUSR1, self.toggle_profiling)
+        return super(CallbackBrokerWorker, self).work_loop(*args, **kw)
+
     def flush(self, force=False):
         now = tz_now()
         if (
diff --git a/awx/settings/defaults.py b/awx/settings/defaults.py
index 1835bf57ac..4efc70b0f2 100644
--- a/awx/settings/defaults.py
+++ b/awx/settings/defaults.py
@@ -1222,6 +1222,9 @@ AWX_REQUEST_PROFILE = False
 #
 AWX_REQUEST_PROFILE_WITH_DOT = False
 
+# Allow profiling callback workers via SIGUSR1
+AWX_CALLBACK_PROFILE = False
+
 # Delete temporary directories created to store playbook run-time
 AWX_CLEANUP_PATHS = True