Customize application_name for different connections in dispatcher service (#13074)

* Introduce new method in settings, import in-line w NOQA mark

* Further refine the app_name to use shorter service names like dispatcher

* Clean up listener logic, change some names

awx/main/dispatch/__init__.py

@@ -4,6 +4,8 @@ import select
 
 from contextlib import contextmanager
 
+from awx.settings.application_name import get_application_name
+
 from django.conf import settings
 from django.db import connection as pg_connection
 
@@ -83,10 +85,11 @@ def pg_bus_conn(new_connection=False):
     '''
 
     if new_connection:
-        conf = settings.DATABASES['default']
+        conf = settings.DATABASES['default'].copy()
-        conn = psycopg2.connect(
+        conf['OPTIONS'] = conf.get('OPTIONS', {}).copy()
-            dbname=conf['NAME'], host=conf['HOST'], user=conf['USER'], password=conf['PASSWORD'], port=conf['PORT'], **conf.get("OPTIONS", {})
+        # Modify the application name to distinguish from other connections the process might use
-        )
+        conf['OPTIONS']['application_name'] = get_application_name(settings.CLUSTER_HOST_ID, function='listener')
+        conn = psycopg2.connect(dbname=conf['NAME'], host=conf['HOST'], user=conf['USER'], password=conf['PASSWORD'], port=conf['PORT'], **conf['OPTIONS'])
         # Django connection.cursor().connection doesn't have autocommit=True on by default
         conn.set_session(autocommit=True)
     else:

awx/main/dispatch/periodic.py

@@ -10,6 +10,7 @@ from django_guid import set_guid
 from django_guid.utils import generate_guid
 
 from awx.main.dispatch.worker import TaskWorker
+from awx.main.utils.db import set_connection_name
 
 logger = logging.getLogger('awx.main.dispatch.periodic')
 
@@ -21,6 +22,9 @@ class Scheduler(Scheduler):
         def run():
             ppid = os.getppid()
             logger.warning('periodic beat started')
+
+            set_connection_name('periodic')  # set application_name to distinguish from other dispatcher processes
+
             while True:
                 if os.getppid() != ppid:
                     # if the parent PID changes, this process has been orphaned

awx/main/dispatch/worker/base.py

@@ -18,6 +18,7 @@ from django.conf import settings
 from awx.main.dispatch.pool import WorkerPool
 from awx.main.dispatch import pg_bus_conn
 from awx.main.utils.common import log_excess_runtime
+from awx.main.utils.db import set_connection_name
 
 if 'run_callback_receiver' in sys.argv:
     logger = logging.getLogger('awx.main.commands.run_callback_receiver')
@@ -219,6 +220,7 @@ class BaseWorker(object):
     def work_loop(self, queue, finished, idx, *args):
         ppid = os.getppid()
         signal_handler = WorkerSignalHandler()
+        set_connection_name('worker')  # set application_name to distinguish from other dispatcher processes
         while not signal_handler.kill_now:
             # if the parent PID changes, this process has been orphaned
             # via e.g., segfault or sigkill, we should exit too

awx/main/utils/db.py

@@ -3,6 +3,9 @@
 
 from itertools import chain
 
+from awx.settings.application_name import set_application_name
+from django.conf import settings
+
 
 def get_all_field_names(model):
     # Implements compatibility with _meta.get_all_field_names
@@ -18,3 +21,7 @@ def get_all_field_names(model):
             )
         )
     )
+
+
+def set_connection_name(function):
+    set_application_name(settings.DATABASES, settings.CLUSTER_HOST_ID, function=function)

awx/settings/application_name.py

@@ -0,0 +1,31 @@
+import os
+import sys
+
+
+def get_service_name(argv):
+    '''
+    Return best-effort guess as to the name of this service
+    '''
+    for arg in argv:
+        if arg == '-m':
+            continue
+        if 'python' in arg:
+            continue
+        if 'manage' in arg:
+            continue
+        if arg.startswith('run_'):
+            return arg[len('run_') :]
+        return arg
+
+
+def get_application_name(CLUSTER_HOST_ID, function=''):
+    if function:
+        function = f'_{function}'
+    return f'awx-{os.getpid()}-{get_service_name(sys.argv)}{function}-{CLUSTER_HOST_ID}'[:63]
+
+
+def set_application_name(DATABASES, CLUSTER_HOST_ID, function=''):
+    if 'sqlite3' in DATABASES['default']['ENGINE']:
+        return
+    options_dict = DATABASES['default'].setdefault('OPTIONS', dict())
+    options_dict['application_name'] = get_application_name(CLUSTER_HOST_ID, function)

awx/settings/development.py

@@ -105,8 +105,11 @@ AWX_CALLBACK_PROFILE = True
 AWX_DISABLE_TASK_MANAGERS = False
 # ======================!!!!!!! FOR DEVELOPMENT ONLY !!!!!!!=================================
 
-if 'sqlite3' not in DATABASES['default']['ENGINE']:  # noqa
+from .application_name import set_application_name
-    DATABASES['default'].setdefault('OPTIONS', dict()).setdefault('application_name', f'{CLUSTER_HOST_ID}-{os.getpid()}-{" ".join(sys.argv)}'[:63])  # noqa
+
+set_application_name(DATABASES, CLUSTER_HOST_ID)
+
+del set_application_name
 
 # If any local_*.py files are present in awx/settings/, use them to override
 # default settings for development.  If not present, we can still run using

awx/settings/production.py

@@ -100,6 +100,8 @@ except IOError:
 
 # The below runs AFTER all of the custom settings are imported.
 
-DATABASES.setdefault('default', dict()).setdefault('OPTIONS', dict()).setdefault(
+from .application_name import set_application_name
-    'application_name', f'{CLUSTER_HOST_ID}-{os.getpid()}-{" ".join(sys.argv)}'[:63]  # NOQA
+
-)  # noqa
+set_application_name(DATABASES, CLUSTER_HOST_ID)  # NOQA
+
+del set_application_name