ref(aci): fix N+1 query for DataConditionGroups in delayed workflow (#95058)

Author: cathteng

src/sentry/workflow_engine/processors/delayed_workflow.py

@@ -50,16 +50,17 @@
     SLOW_CONDITIONS,
     Condition,
 )
+from sentry.workflow_engine.models.workflow_data_condition_group import WorkflowDataConditionGroup
 from sentry.workflow_engine.processors.action import filter_recently_fired_workflow_actions
 from sentry.workflow_engine.processors.data_condition_group import (
     evaluate_data_conditions,
     get_slow_conditions_for_groups,
 )
-from sentry.workflow_engine.processors.detector import get_detector_by_event
+from sentry.workflow_engine.processors.detector import get_detectors_by_groupevents_bulk
 from sentry.workflow_engine.processors.log_util import log_if_slow, track_batch_performance
 from sentry.workflow_engine.processors.workflow import (
     WORKFLOW_ENGINE_BUFFER_LIST_KEY,
-    evaluate_workflows_action_filters,
+    evaluate_action_filters,
 )
 from sentry.workflow_engine.processors.workflow_fire_history import create_workflow_fire_histories
 from sentry.workflow_engine.tasks.actions import build_trigger_action_task_params, trigger_action
@@ -567,6 +568,27 @@ def get_group_to_groupevent(
     return group_to_groupevent
 
 
+def get_dcgs_by_group(
+    groups_to_fire: dict[GroupId, set[DataConditionGroup]],
+    event_data: EventRedisData,
+    dcg_type: DataConditionHandler.Group,
+) -> dict[GroupId, set[DataConditionGroup]]:
+    """
+    Extract DataConditionGroups from groups_to_fire, grouped by group ID, for a particular DataConditionGroup type (e.g. workflow trigger)
+    trigger_group_to_dcg_model is the mapping from DataConditionGroup type to DataConditionGroup id to Workflow id
+    Returns a dict mapping GroupId to set of DCGs.
+    """
+    workflow_dcg_ids = set(event_data.trigger_group_to_dcg_model[dcg_type].keys())
+
+    workflow_dcgs_by_group = {}
+    for group_id, dcgs in groups_to_fire.items():
+        workflow_dcgs = {dcg for dcg in dcgs if dcg.id in workflow_dcg_ids}
+        if workflow_dcgs:
+            workflow_dcgs_by_group[group_id] = workflow_dcgs
+
+    return workflow_dcgs_by_group
+
+
 @sentry_sdk.trace
 def fire_actions_for_groups(
     organization: Organization,
@@ -585,6 +607,27 @@ def fire_actions_for_groups(
         },
     )
 
+    workflow_triggers = get_dcgs_by_group(
+        groups_to_fire, event_data, DataConditionHandler.Group.WORKFLOW_TRIGGER
+    )
+    action_filters = get_dcgs_by_group(
+        groups_to_fire, event_data, DataConditionHandler.Group.ACTION_FILTER
+    )
+    all_workflow_triggers = set().union(*list(workflow_triggers.values()))
+
+    # Bulk fetch detectors
+    event_id_to_detector = get_detectors_by_groupevents_bulk(list(group_to_groupevent.values()))
+
+    # Bulk fetch action filters for workflow triggers
+    workflows = Workflow.objects.filter(when_condition_group_id__in=all_workflow_triggers)
+
+    dcg_to_workflow = {
+        wdcg.condition_group: wdcg.workflow
+        for wdcg in WorkflowDataConditionGroup.objects.select_related(
+            "workflow", "condition_group"
+        ).filter(workflow__in=workflows)
+    }
+
     total_actions = 0
     with track_batch_performance(
         "workflow_engine.delayed_workflow.fire_actions_for_groups.loop",
@@ -594,42 +637,45 @@ def fire_actions_for_groups(
         for group, group_event in group_to_groupevent.items():
             with tracker.track(str(group.id)), log_context.new_context(group_id=group.id):
                 workflow_event_data = WorkflowEventData(event=group_event, group=group)
-                detector = get_detector_by_event(workflow_event_data)
-
-                workflow_triggers: set[DataConditionGroup] = set()
-                action_filters: set[DataConditionGroup] = set()
-                for dcg in groups_to_fire[group.id]:
-                    if (
-                        dcg.id
-                        in event_data.trigger_group_to_dcg_model[
-                            DataConditionHandler.Group.WORKFLOW_TRIGGER
-                        ]
-                    ):
-                        workflow_triggers.add(dcg)
-                    elif (
-                        dcg.id
-                        in event_data.trigger_group_to_dcg_model[
-                            DataConditionHandler.Group.ACTION_FILTER
-                        ]
-                    ):
-                        action_filters.add(dcg)
-
-                # process workflow_triggers
-                workflows = set(
-                    Workflow.objects.filter(when_condition_group_id__in=workflow_triggers)
-                )
+                detector = event_id_to_detector.get(group_event.event_id)
+
+                if detector is None:
+                    logger.warning(
+                        "No detector found for event, skipping",
+                        extra={
+                            "event_id": group_event.event_id,
+                            "group_id": group.id,
+                        },
+                    )
+                    continue
+
+                workflow_triggers_for_group = workflow_triggers.get(group.id, set())
+                action_filters_for_group = action_filters.get(group.id, set())
 
                 with log_if_slow(
                     logger,
                     "workflow_engine.delayed_workflow.slow_evaluate_workflows_action_filters",
                     extra={"event_data": workflow_event_data},
                     threshold_seconds=1,
                 ):
-                    workflows_actions = evaluate_workflows_action_filters(
-                        workflows, workflow_event_data
+                    # Process workflow filters for passing trigger groups
+                    triggered_workflow_ids = {
+                        event_data.dcg_to_workflow[dcg.id] for dcg in workflow_triggers_for_group
+                    }
+
+                    filter_dcg_to_workflow: dict[DataConditionGroup, Workflow] = {
+                        dcg: workflow
+                        for dcg, workflow in dcg_to_workflow.items()
+                        if workflow.id in triggered_workflow_ids
+                    }
+
+                    workflows_actions = evaluate_action_filters(
+                        workflow_event_data,
+                        filter_dcg_to_workflow,
                     )
+
                 filtered_actions = filter_recently_fired_workflow_actions(
-                    action_filters | workflows_actions, workflow_event_data
+                    action_filters_for_group | workflows_actions, workflow_event_data
                 )
                 create_workflow_fire_histories(detector, filtered_actions, workflow_event_data)
 
@@ -647,7 +693,7 @@ def fire_actions_for_groups(
                 logger.debug(
                     "workflow_engine.delayed_workflow.triggered_actions",
                     extra={
-                        "workflow_ids": [workflow.id for workflow in workflows],
+                        "workflow_ids": triggered_workflow_ids,
                         "actions": [action.id for action in filtered_actions],
                         "event_data": workflow_event_data,
                         "event_id": event_id,

src/sentry/workflow_engine/processors/detector.py

@@ -1,9 +1,13 @@
 from __future__ import annotations
 
 import logging
+from collections import defaultdict
+from collections.abc import Callable, Mapping
+from typing import NamedTuple
 
 import sentry_sdk
 
+from sentry.db.models.manager.base_query_set import BaseQuerySet
 from sentry.eventstore.models import GroupEvent
 from sentry.issues.issue_occurrence import IssueOccurrence
 from sentry.issues.producer import PayloadType, produce_occurrence_to_kafka
@@ -57,6 +61,142 @@ def get_detector_by_event(event_data: WorkflowEventData) -> Detector:
     return detector
 
 
+class _SplitEvents(NamedTuple):
+    events_with_occurrences: list[tuple[GroupEvent, int]]
+    error_events: list[GroupEvent]
+    events_missing_detectors: list[GroupEvent]
+
+
+def _split_events_by_occurrence(
+    event_list: list[GroupEvent],
+) -> _SplitEvents:
+    from sentry.grouping.grouptype import ErrorGroupType
+
+    events_with_occurrences: list[tuple[GroupEvent, int]] = []
+    error_events: list[GroupEvent] = []  # only error events don't have occurrences
+    events_missing_detectors: list[GroupEvent] = []
+
+    for event in event_list:
+        issue_occurrence = event.occurrence
+        if issue_occurrence is None:
+            assert event.group.issue_type.slug == ErrorGroupType.slug
+            error_events.append(event)
+        elif detector_id := issue_occurrence.evidence_data.get("detector_id"):
+            events_with_occurrences.append((event, detector_id))
+        else:
+            events_missing_detectors.append(event)
+
+    return _SplitEvents(
+        events_with_occurrences,
+        error_events,
+        events_missing_detectors,
+    )
+
+
+def _create_event_detector_map(
+    detectors: BaseQuerySet[Detector],
+    key_event_map: dict[int, list[GroupEvent]],
+    detector_key_extractor: Callable[[Detector], int],
+) -> tuple[dict[str, Detector], set[int]]:
+    result: dict[str, Detector] = {}
+
+    # used to track existing keys (detector_id or project_id) to log missing keys
+    keys = set()
+
+    for detector in detectors:
+        key = detector_key_extractor(detector)
+        keys.add(key)
+        detector_events = key_event_map[key]
+        result.update({event.event_id: detector for event in detector_events})
+
+    return result, keys
+
+
+def get_detectors_by_groupevents_bulk(
+    event_list: list[GroupEvent],
+) -> Mapping[str, Detector]:
+    """
+    Given a list of GroupEvents, return a mapping of event_id to Detector.
+    """
+    from sentry.grouping.grouptype import ErrorGroupType
+
+    if not event_list:
+        return {}
+
+    result: dict[str, Detector] = {}
+
+    # Separate events by whether they have occurrences or not
+    events_with_occurrences, error_events, events_missing_detectors = _split_events_by_occurrence(
+        event_list
+    )
+
+    # Fetch detectors for events with occurrences (by detector_id)
+    missing_detector_ids = set()
+    if events_with_occurrences:
+        detector_id_to_events: dict[int, list[GroupEvent]] = defaultdict(list)
+
+        for event, detector_id in events_with_occurrences:
+            detector_id_to_events[detector_id].append(event)
+
+        def _extract_events_lookup_key(detector: Detector) -> int:
+            return detector.id
+
+        if detector_id_to_events:
+            detectors = Detector.objects.filter(id__in=list(detector_id_to_events.keys()))
+            mapping, found_detector_ids = _create_event_detector_map(
+                detectors,
+                key_event_map=detector_id_to_events,
+                detector_key_extractor=_extract_events_lookup_key,
+            )
+            result.update(mapping)
+
+            missing_detector_ids = set(detector_id_to_events.keys()) - found_detector_ids
+
+    # Fetch detectors for events without occurrences (by project_id)
+    projects_missing_detectors = set()
+    if error_events:
+        # Group events by project_id
+        project_to_events: dict[int, list[GroupEvent]] = defaultdict(list)
+
+        for event in error_events:
+            project_to_events[event.project_id].append(event)
+
+        def _extract_events_lookup_key(detector: Detector) -> int:
+            return detector.project_id
+
+        detectors = Detector.objects.filter(
+            project_id__in=project_to_events.keys(),
+            type=ErrorGroupType.slug,
+        )
+        mapping, projects_with_error_detectors = _create_event_detector_map(
+            detectors,
+            key_event_map=project_to_events,
+            detector_key_extractor=_extract_events_lookup_key,
+        )
+        result.update(mapping)
+
+        projects_missing_detectors = set(project_to_events.keys()) - projects_with_error_detectors
+
+    # Log all missing detectors
+    if missing_detector_ids or projects_missing_detectors or events_missing_detectors:
+        metrics.incr(
+            "workflow_engine.detectors.error",
+            amount=len(projects_missing_detectors) + len(missing_detector_ids),
+        )
+        logger.error(
+            "Detectors not found for events",
+            extra={
+                "projects_missing_error_detectors": projects_missing_detectors,
+                "missing_detectors": missing_detector_ids,
+                "events_missing_detectors": [
+                    (event.event_id, event.group_id) for event in events_missing_detectors
+                ],
+            },
+        )
+
+    return result
+
+
 def create_issue_platform_payload(result: DetectorEvaluationResult) -> None:
     occurrence, status_change = None, None