[Performance] Optimized UserEvent.decay to resolve N+1 queries. Closes #3460

Author: ayushgupta704

api_app/user_events_manager/queryset.py

@@ -12,9 +12,7 @@
 class UserEventQuerySet(QuerySet):
     def decay(self):
         from collections import defaultdict
-
         from django.db import transaction
-
         objects = (
             self.exclude(decay_progression=DecayProgressionEnum.FIXED.value)
             .exclude(next_decay__isnull=True)
@@ -29,22 +27,17 @@ def decay(self):
             objects = objects.select_related("data_model")
         else:
             objects = objects.prefetch_related("data_model")
-
         # Load into memory so we can mutate fields and bulk-write back.
         events = list(objects)
         if not events:
             return 0
-
         # Group by concrete class since bulk_update works per-table.
         data_models_by_class = defaultdict(list)
-
         for obj in events:
             obj.decay_times += 1
             data_model = obj.data_model
-
             if data_model is not None:
                 data_model.reliability -= 1
-
             if data_model is None or data_model.reliability == 0:
                 obj.next_decay = None
             else:
@@ -54,17 +47,15 @@ def decay(self):
                     obj.next_decay += datetime.timedelta(
                         days=obj.decay_timedelta_days ** (obj.decay_times + 1)
                     )
-
             if data_model is not None:
                 data_models_by_class[data_model.__class__].append(data_model)
-
         # Bulk-write instead of per-object .save() to avoid O(N) queries.
         # Atomic so partial failures don't leave inconsistent state.
         with transaction.atomic():
             for model_class, models_list in data_models_by_class.items():
-                model_class.objects.bulk_update(models_list, ["reliability"])
-            self.model.objects.bulk_update(events, ["decay_times", "next_decay"])
-
+                unique_models = {m.pk: m for m in models_list}.values()
+                model_class.objects.bulk_update(unique_models, ["reliability"], batch_size=1000)
+            self.model.objects.bulk_update(events, ["decay_times", "next_decay"], batch_size=1000)
         return len(events)
 
     def visible_for_user(self, user):

tests/api_app/user_events_manager/test_queryset.py

@@ -151,6 +151,38 @@ def test_decay_reliability_reaches_zero(self):
         ua.delete()
         an.delete()
 
+    def test_decay_performance(self):
+        from django.db import connection
+        from django.test.utils import CaptureQueriesContext
+        an = Analyzable.objects.create(
+            name="test_perf.com",
+            classification=Classification.DOMAIN,
+        )
+        N = 5
+        for i in range(N):
+            an_i = Analyzable.objects.create(
+                name=f"test_perf{i}.com",
+                classification=Classification.DOMAIN,
+            )
+            ue_ser = UserAnalyzableEventSerializer(
+                data={
+                    "analyzable": {"name": an_i.name},
+                    "decay_progression": 0,
+                    "decay_timedelta_days": 1,
+                    "data_model_content": {"evaluation": "malicious", "reliability": 8},
+                },
+                context={"request": MockUpRequest(self.user)},
+            )
+            ue_ser.is_valid()
+            ua = ue_ser.save()
+            ua.next_decay = now() - datetime.timedelta(days=1)
+            ua.save()
+        with CaptureQueriesContext(connection) as queries:
+            number = UserAnalyzableEvent.objects.decay()
+            self.assertEqual(number, N)
+        self.assertLessEqual(len(queries), 10)
+        Analyzable.objects.filter(name__startswith="test_perf").delete()
+
     def test_decay_multiple_events(self):
         analyzables = []
         events = []