From 99815f8962bfe88f8a485b16736599445219fe7d Mon Sep 17 00:00:00 2001
From: Elijah DeLee <kdelee@redhat.com>
Date: Fri, 26 Aug 2022 11:40:36 -0400
Subject: [PATCH] calcuate consumed capacity in same way in metrics

We should be consistent about this. Also this takes us from doing a as
many queries to the UnifiedJob table as we have instances to doing 1
query to the UnifiedJob table (and both do 1 query to Instances table)
---
 awx/main/analytics/collectors.py          | 31 ++++++++++++-----------
 awx/main/scheduler/task_manager_models.py |  8 +++---
 2 files changed, 20 insertions(+), 19 deletions(-)

diff --git a/awx/main/analytics/collectors.py b/awx/main/analytics/collectors.py
index 27a41e9dff..e6ee155a83 100644
--- a/awx/main/analytics/collectors.py
+++ b/awx/main/analytics/collectors.py
@@ -16,6 +16,7 @@ from awx.conf.license import get_license
 from awx.main.utils import get_awx_version, camelcase_to_underscore, datetime_hook
 from awx.main import models
 from awx.main.analytics import register
+from awx.main.scheduler.task_manager_models import TaskManagerInstances
 
 """
 This module is used to define metrics collected by awx.main.analytics.gather()
@@ -235,25 +236,25 @@ def projects_by_scm_type(since, **kwargs):
 @register('instance_info', '1.2', description=_('Cluster topology and capacity'))
 def instance_info(since, include_hostnames=False, **kwargs):
     info = {}
-    instances = models.Instance.objects.values_list('hostname').values(
-        'uuid', 'version', 'capacity', 'cpu', 'memory', 'managed_by_policy', 'hostname', 'enabled'
-    )
-    for instance in instances:
-        consumed_capacity = sum(x.task_impact for x in models.UnifiedJob.objects.filter(execution_node=instance['hostname'], status__in=('running', 'waiting')))
+    # Use same method that the TaskManager does to compute consumed capacity without querying all running jobs for each Instance
+    active_tasks = models.UnifiedJob.objects.filter(status__in=['running', 'waiting']).only('task_impact', 'controller_node', 'execution_node')
+    tm_instances = TaskManagerInstances(active_tasks, instance_fields=['uuid', 'version', 'capacity', 'cpu', 'memory', 'managed_by_policy', 'enabled'])
+    for instance in tm_instances.instance_objects:
+        consumed_capacity = tm_instances[instance.hostname].consumed_capacity
         instance_info = {
-            'uuid': instance['uuid'],
-            'version': instance['version'],
-            'capacity': instance['capacity'],
-            'cpu': instance['cpu'],
-            'memory': instance['memory'],
-            'managed_by_policy': instance['managed_by_policy'],
-            'enabled': instance['enabled'],
+            'uuid': instance.uuid,
+            'version': instance.version,
+            'capacity': instance.capacity,
+            'cpu': instance.cpu,
+            'memory': instance.memory,
+            'managed_by_policy': instance.managed_by_policy,
+            'enabled': instance.enabled,
             'consumed_capacity': consumed_capacity,
-            'remaining_capacity': instance['capacity'] - consumed_capacity,
+            'remaining_capacity': instance.capacity - consumed_capacity,
         }
         if include_hostnames is True:
-            instance_info['hostname'] = instance['hostname']
-        info[instance['uuid']] = instance_info
+            instance_info['hostname'] = instance.hostname
+        info[instance.uuid] = instance_info
     return info
 
 
diff --git a/awx/main/scheduler/task_manager_models.py b/awx/main/scheduler/task_manager_models.py
index cade939343..b5071266ec 100644
--- a/awx/main/scheduler/task_manager_models.py
+++ b/awx/main/scheduler/task_manager_models.py
@@ -34,14 +34,14 @@ class TaskManagerInstance:
 
 
 class TaskManagerInstances:
-    def __init__(self, active_tasks, instances=None):
+    def __init__(self, active_tasks, instances=None, instance_fields=['node_type', 'capacity', 'hostname', 'enabled']):
         self.instances_by_hostname = dict()
+        self.instance_objects = []
         if instances is None:
-            instances = (
-                Instance.objects.filter(hostname__isnull=False, enabled=True).exclude(node_type='hop').only('node_type', 'capacity', 'hostname', 'enabled')
-            )
+            instances = Instance.objects.filter(hostname__isnull=False, enabled=True).exclude(node_type='hop').only(*instance_fields)
         for instance in instances:
             self.instances_by_hostname[instance.hostname] = TaskManagerInstance(instance)
+            self.instance_objects.append(instance)
 
         # initialize remaining capacity based on currently waiting and running tasks
         for task in active_tasks: