From 05aab5da4cbbd02c8830402fb00db4235eedb5e1 Mon Sep 17 00:00:00 2001
From: AlanCoding <arominge@redhat.com>
Date: Mon, 10 Jun 2019 14:03:17 -0400
Subject: [PATCH] cache smart inventory memberships to avoid slow query

---
 awx/main/models/events.py | 16 ++++++++++--
 awx/main/tasks.py         | 51 ++++++++++++++++++++++++++-------------
 2 files changed, 48 insertions(+), 19 deletions(-)

diff --git a/awx/main/models/events.py b/awx/main/models/events.py
index aec07f4f4340..5b424353eb13 100644
--- a/awx/main/models/events.py
+++ b/awx/main/models/events.py
@@ -346,8 +346,20 @@ def save(self, *args, **kwargs):
 
             # Update host related field from host_name.
             if hasattr(self, 'job') and not self.host_id and self.host_name:
-                host_qs = self.job.inventory.hosts.filter(name=self.host_name)
-                host_id = host_qs.only('id').values_list('id', flat=True).first()
+                if self.job.inventory.kind == 'smart':
+                    # optimization to avoid calling inventory.hosts, which
+                    # can take a long time to run under some circumstances
+                    from awx.main.models.inventory import SmartInventoryMembership
+                    membership = SmartInventoryMembership.objects.filter(
+                        inventory=self.job.inventory, host__name=self.host_name
+                    ).first()
+                    if membership:
+                        host_id = membership.host_id
+                    else:
+                        host_id = None
+                else:
+                    host_qs = self.job.inventory.hosts.filter(name=self.host_name)
+                    host_id = host_qs.only('id').values_list('id', flat=True).first()
                 if host_id != self.host_id:
                     self.host_id = host_id
                     if 'host_id' not in update_fields:
diff --git a/awx/main/tasks.py b/awx/main/tasks.py
index 5ff4f9d4c31d..f899fadb0513 100644
--- a/awx/main/tasks.py
+++ b/awx/main/tasks.py
@@ -601,26 +601,39 @@ def update_inventory_computed_fields(inventory_id, should_update_hosts=True):
         raise
 
 
-@task()
-def update_host_smart_inventory_memberships():
-    try:
+def update_smart_memberships_for_inventory(smart_inventory):
+    current = set(SmartInventoryMembership.objects.filter(inventory=smart_inventory).values_list('host_id', flat=True))
+    new = set(smart_inventory.hosts.values_list('id', flat=True))
+    additions = new - current
+    removals = current - new
+    if additions or removals:
         with transaction.atomic():
-            smart_inventories = Inventory.objects.filter(kind='smart', host_filter__isnull=False, pending_deletion=False)
-            SmartInventoryMembership.objects.all().delete()
-            memberships = []
-            changed_inventories = set([])
-            for smart_inventory in smart_inventories:
+            if removals:
+                SmartInventoryMembership.objects.filter(inventory=smart_inventory, host_id__in=removals).delete()
+            if additions:
                 add_for_inventory = [
-                    SmartInventoryMembership(inventory_id=smart_inventory.id, host_id=host_id[0])
-                    for host_id in smart_inventory.hosts.values_list('id')
+                    SmartInventoryMembership(inventory_id=smart_inventory.id, host_id=host_id)
+                    for host_id in additions
                 ]
-                memberships.extend(add_for_inventory)
-                if add_for_inventory:
-                    changed_inventories.add(smart_inventory)
-            SmartInventoryMembership.objects.bulk_create(memberships)
-    except IntegrityError as e:
-        logger.error("Update Host Smart Inventory Memberships failed due to an exception: {}".format(e))
-        return
+                SmartInventoryMembership.objects.bulk_create(add_for_inventory)
+        logger.debug('Smart host membership cached for {}, {} additions, {} removals, {} total count.'.format(
+            smart_inventory.pk, len(additions), len(removals), len(new)
+        ))
+        return True  # changed
+    return False
+
+
+@task()
+def update_host_smart_inventory_memberships():
+    smart_inventories = Inventory.objects.filter(kind='smart', host_filter__isnull=False, pending_deletion=False)
+    changed_inventories = set([])
+    for smart_inventory in smart_inventories:
+        try:
+            changed = update_smart_memberships_for_inventory(smart_inventory)
+            if changed:
+                changed_inventories.add(smart_inventory)
+        except IntegrityError:
+            logger.exception('Failed to update smart inventory memberships for {}'.format(smart_inventory.pk))
     # Update computed fields for changed inventories outside atomic action
     for smart_inventory in changed_inventories:
         smart_inventory.update_computed_fields(update_groups=False, update_hosts=False)
@@ -1588,6 +1601,10 @@ def pre_run_hook(self, job):
                                             job_explanation=('Previous Task Failed: {"job_type": "%s", "job_name": "%s", "job_id": "%s"}' %
                                                              ('project_update', local_project_sync.name, local_project_sync.id)))
                     raise
+        if job.inventory.kind == 'smart':
+            # cache smart inventory memberships so that the host_filter query is not
+            # ran inside of the event saving code
+            update_smart_memberships_for_inventory(job.inventory)
 
     def final_run_hook(self, job, status, private_data_dir, fact_modification_times, isolated_manager_instance=None):
         super(RunJob, self).final_run_hook(job, status, private_data_dir, fact_modification_times)