ref(aci): enqueue project ids to sharded list in process_workflows (#97641)

Author: cathteng

src/sentry/workflow_engine/processors/workflow.py

@@ -1,3 +1,4 @@
+import random
 from dataclasses import asdict, dataclass, replace
 from datetime import datetime
 from enum import StrEnum
@@ -9,9 +10,9 @@
 from django.utils import timezone
 
 from sentry import buffer, features
+from sentry.eventstore.models import GroupEvent
 from sentry.models.activity import Activity
 from sentry.models.environment import Environment
-from sentry.services.eventstore.models import GroupEvent
 from sentry.utils import json
 from sentry.workflow_engine.models import Action, DataConditionGroup, Detector, Workflow
 from sentry.workflow_engine.models.workflow_data_condition_group import WorkflowDataConditionGroup
@@ -30,7 +31,6 @@
 
 logger = log_context.get_logger(__name__)
 
-WORKFLOW_ENGINE_BUFFER_LIST_KEY = "workflow_engine_delayed_processing_buffer"
 DetectorId = int | None
 DataConditionGroupId = int
 
@@ -99,6 +99,8 @@ def buffer_value(self) -> str:
 def enqueue_workflows(
     items_by_workflow: dict[Workflow, DelayedWorkflowItem],
 ) -> None:
+    from sentry.workflow_engine.tasks.delayed_workflows import DelayedWorkflow
+
     items_by_project_id = DefaultDict[int, list[DelayedWorkflowItem]](list)
     for queue_item in items_by_workflow.values():
         if not queue_item.delayed_if_group_ids and not queue_item.passing_if_group_ids:
@@ -125,9 +127,8 @@ def enqueue_workflows(
 
     sentry_sdk.set_tag("delayed_workflow_items", items)
 
-    buffer.backend.push_to_sorted_set(
-        key=WORKFLOW_ENGINE_BUFFER_LIST_KEY, value=list(items_by_project_id.keys())
-    )
+    sharded_key = random.choice(DelayedWorkflow.get_buffer_keys())
+    buffer.backend.push_to_sorted_set(key=sharded_key, value=list(items_by_project_id.keys()))
 
     logger.debug(
         "workflow_engine.workflows.enqueued",

src/sentry/workflow_engine/tasks/delayed_workflows.py

@@ -17,7 +17,6 @@
 from sentry.taskworker.namespaces import workflow_engine_tasks
 from sentry.taskworker.retry import Retry
 from sentry.workflow_engine.models import Workflow
-from sentry.workflow_engine.processors.workflow import WORKFLOW_ENGINE_BUFFER_LIST_KEY
 from sentry.workflow_engine.utils import log_context
 
 logger = log_context.get_logger("sentry.workflow_engine.tasks.delayed_workflows")
@@ -52,7 +51,8 @@ def process_delayed_workflows_shim(
 
 @delayed_processing_registry.register("delayed_workflow")
 class DelayedWorkflow(DelayedProcessingBase):
-    buffer_key = WORKFLOW_ENGINE_BUFFER_LIST_KEY
+    buffer_key = "workflow_engine_delayed_processing_buffer"
+    buffer_shards = 8
     option = "delayed_workflow.rollout"
 
     @property

tests/sentry/workflow_engine/processors/test_delayed_workflow.py

@@ -60,7 +60,7 @@
     get_group_to_groupevent,
     get_groups_to_fire,
 )
-from sentry.workflow_engine.processors.workflow import WORKFLOW_ENGINE_BUFFER_LIST_KEY
+from sentry.workflow_engine.tasks.delayed_workflows import DelayedWorkflow
 from tests.sentry.workflow_engine.test_base import BaseWorkflowTest
 from tests.snuba.rules.conditions.test_event_frequency import BaseEventFrequencyPercentTest
 
@@ -119,12 +119,8 @@ def setUp(self) -> None:
         self.mock_redis_buffer = mock_redis_buffer()
         self.mock_redis_buffer.__enter__()
 
-        buffer.backend.push_to_sorted_set(
-            key=WORKFLOW_ENGINE_BUFFER_LIST_KEY, value=self.project.id
-        )
-        buffer.backend.push_to_sorted_set(
-            key=WORKFLOW_ENGINE_BUFFER_LIST_KEY, value=self.project2.id
-        )
+        buffer.backend.push_to_sorted_set(key=DelayedWorkflow.buffer_key, value=self.project.id)
+        buffer.backend.push_to_sorted_set(key=DelayedWorkflow.buffer_key, value=self.project2.id)
 
     def tearDown(self):
         self.mock_redis_buffer.__exit__(None, None, None)

tests/sentry/workflow_engine/processors/test_workflow.py

@@ -24,14 +24,14 @@
 )
 from sentry.workflow_engine.models.data_condition import Condition
 from sentry.workflow_engine.processors.workflow import (
-    WORKFLOW_ENGINE_BUFFER_LIST_KEY,
     DelayedWorkflowItem,
     delete_workflow,
     enqueue_workflows,
     evaluate_workflow_triggers,
     evaluate_workflows_action_filters,
     process_workflows,
 )
+from sentry.workflow_engine.tasks.delayed_workflows import DelayedWorkflow
 from sentry.workflow_engine.types import WorkflowEventData
 from tests.sentry.workflow_engine.test_base import BaseWorkflowTest
 
@@ -414,6 +414,8 @@ def setUp(self) -> None:
         )
         self.event_data = WorkflowEventData(event=self.group_event, group=self.group)
         self.action_group, _ = self.create_workflow_action(self.workflow)
+
+        self.buffer_keys = DelayedWorkflow.get_buffer_keys()
         self.mock_redis_buffer = mock_redis_buffer()
         self.mock_redis_buffer.__enter__()
 
@@ -438,11 +440,12 @@ def test_enqueues_workflow_all_logic_type(self) -> None:
 
         process_workflows(self.event_data)
 
-        project_ids = buffer.backend.get_sorted_set(
-            WORKFLOW_ENGINE_BUFFER_LIST_KEY, 0, self.buffer_timestamp
+        project_ids = buffer.backend.bulk_get_sorted_set(
+            self.buffer_keys,
+            min=0,
+            max=self.buffer_timestamp,
         )
-        assert project_ids
-        assert project_ids[0][0] == self.project.id
+        assert list(project_ids.keys()) == [self.project.id]
 
     def test_enqueues_workflow_any_logic_type(self) -> None:
         assert self.workflow.when_condition_group
@@ -469,10 +472,12 @@ def test_enqueues_workflow_any_logic_type(self) -> None:
 
         process_workflows(self.event_data)
 
-        project_ids = buffer.backend.get_sorted_set(
-            WORKFLOW_ENGINE_BUFFER_LIST_KEY, 0, self.buffer_timestamp
+        project_ids = buffer.backend.bulk_get_sorted_set(
+            self.buffer_keys,
+            min=0,
+            max=self.buffer_timestamp,
         )
-        assert project_ids[0][0] == self.project.id
+        assert list(project_ids.keys()) == [self.project.id]
 
     def test_skips_enqueuing_any(self) -> None:
         # skips slow conditions if the condition group evaluates to True without evaluating them
@@ -549,10 +554,12 @@ def test_enqueues_with_when_and_if_slow_conditions(self) -> None:
 
         process_workflows(self.event_data)
 
-        project_ids = buffer.backend.get_sorted_set(
-            WORKFLOW_ENGINE_BUFFER_LIST_KEY, 0, self.buffer_timestamp
+        project_ids = buffer.backend.bulk_get_sorted_set(
+            self.buffer_keys,
+            min=0,
+            max=self.buffer_timestamp,
         )
-        assert project_ids[0][0] == self.project.id
+        assert list(project_ids.keys()) == [self.project.id]
 
     def test_enqueues_event_if_meets_fast_conditions(self) -> None:
         assert self.workflow.when_condition_group
@@ -576,8 +583,10 @@ def test_enqueues_event_if_meets_fast_conditions(self) -> None:
 
         process_workflows(self.event_data)
 
-        project_ids = buffer.backend.get_sorted_set(
-            WORKFLOW_ENGINE_BUFFER_LIST_KEY, 0, self.buffer_timestamp
+        project_ids = buffer.backend.bulk_get_sorted_set(
+            self.buffer_keys,
+            min=0,
+            max=self.buffer_timestamp,
         )
         assert not project_ids
 
@@ -586,10 +595,12 @@ def test_enqueues_event_if_meets_fast_conditions(self) -> None:
 
         process_workflows(self.event_data)
 
-        project_ids = buffer.backend.get_sorted_set(
-            WORKFLOW_ENGINE_BUFFER_LIST_KEY, 0, self.buffer_timestamp
+        project_ids = buffer.backend.bulk_get_sorted_set(
+            self.buffer_keys,
+            min=0,
+            max=self.buffer_timestamp,
         )
-        assert project_ids[0][0] == self.project.id
+        assert list(project_ids.keys()) == [self.project.id]
 
 
 @freeze_time(FROZEN_TIME)
@@ -609,6 +620,7 @@ def setUp(self) -> None:
             occurrence=self.build_occurrence(evidence_data={"detector_id": self.detector.id})
         )
         self.event_data = WorkflowEventData(event=self.group_event, group=self.group)
+        self.buffer_keys = DelayedWorkflow.get_buffer_keys()
 
     @patch("sentry.utils.metrics.incr")
     def test_metrics_issue_dual_processing_metrics(self, mock_incr: MagicMock) -> None:
@@ -736,11 +748,12 @@ def test_enqueues_when_slow_conditions(self):
 
         enqueue_workflows(queue_items)
 
-        project_ids = buffer.backend.get_sorted_set(
-            WORKFLOW_ENGINE_BUFFER_LIST_KEY, 0, timezone.now().timestamp()
+        project_ids = buffer.backend.bulk_get_sorted_set(
+            self.buffer_keys,
+            min=0,
+            max=timezone.now().timestamp(),
         )
-        assert project_ids
-        assert project_ids[0][0] == self.project.id
+        assert list(project_ids.keys()) == [self.project.id]
 
 
 class TestEnqueueWorkflows(BaseWorkflowTest):
@@ -774,9 +787,11 @@ def setUp(self) -> None:
 
     @patch("sentry.buffer.backend.push_to_sorted_set")
     @patch("sentry.buffer.backend.push_to_hash_bulk")
+    @patch("random.choice")
     def test_enqueue_workflows__adds_to_workflow_engine_buffer(
-        self, mock_push_to_hash_bulk, mock_push_to_sorted_set
+        self, mock_randchoice, mock_push_to_hash_bulk, mock_push_to_sorted_set
     ):
+        mock_randchoice.return_value = f"{DelayedWorkflow.buffer_key}:{5}"
         enqueue_workflows(
             {
                 self.workflow: DelayedWorkflowItem(
@@ -791,7 +806,7 @@ def test_enqueue_workflows__adds_to_workflow_engine_buffer(
         )
 
         mock_push_to_sorted_set.assert_called_once_with(
-            key=WORKFLOW_ENGINE_BUFFER_LIST_KEY,
+            key=f"{DelayedWorkflow.buffer_key}:{5}",
             value=[self.group_event.project_id],
         )
 

tests/sentry/workflow_engine/test_integration.py

@@ -24,7 +24,7 @@
 from sentry.workflow_engine.models.data_condition import Condition
 from sentry.workflow_engine.processors import process_data_source, process_detectors
 from sentry.workflow_engine.processors.delayed_workflow import process_delayed_workflows
-from sentry.workflow_engine.processors.workflow import WORKFLOW_ENGINE_BUFFER_LIST_KEY
+from sentry.workflow_engine.tasks.delayed_workflows import DelayedWorkflow
 from sentry.workflow_engine.types import DetectorPriorityLevel
 from tests.sentry.workflow_engine.test_base import BaseWorkflowTest
 
@@ -213,6 +213,7 @@ def setUp(self) -> None:
         )
         self.workflow_triggers.conditions.all().delete()
         self.action_group, self.action = self.create_workflow_action(workflow=self.workflow)
+        self.buffer_keys = DelayedWorkflow.get_buffer_keys()
 
     @pytest.fixture(autouse=True)
     def with_feature_flags(self):
@@ -409,8 +410,8 @@ def test_slow_condition_workflow_with_conditions(self, mock_trigger: MagicMock)
         self.post_process_error(event_1, is_new=True)
         assert not mock_trigger.called
 
-        project_ids = buffer.backend.get_sorted_set(
-            WORKFLOW_ENGINE_BUFFER_LIST_KEY, 0, timezone.now().timestamp()
+        project_ids = buffer.backend.bulk_get_sorted_set(
+            self.buffer_keys, 0, timezone.now().timestamp()
         )
         assert not project_ids
 
@@ -419,8 +420,8 @@ def test_slow_condition_workflow_with_conditions(self, mock_trigger: MagicMock)
         self.post_process_error(event_2, is_new=True)
         assert not mock_trigger.called
 
-        project_ids = buffer.backend.get_sorted_set(
-            WORKFLOW_ENGINE_BUFFER_LIST_KEY, 0, timezone.now().timestamp()
+        project_ids = buffer.backend.bulk_get_sorted_set(
+            self.buffer_keys, 0, timezone.now().timestamp()
         )
         assert not project_ids
 
@@ -433,11 +434,13 @@ def test_slow_condition_workflow_with_conditions(self, mock_trigger: MagicMock)
         self.post_process_error(event_5)
         assert not mock_trigger.called
 
-        project_ids = buffer.backend.get_sorted_set(
-            WORKFLOW_ENGINE_BUFFER_LIST_KEY, 0, timezone.now().timestamp()
+        project_ids = buffer.backend.bulk_get_sorted_set(
+            self.buffer_keys,
+            min=0,
+            max=timezone.now().timestamp(),
         )
 
-        process_delayed_workflows(project_ids[0][0])
+        process_delayed_workflows(list(project_ids.keys())[0])
         mock_trigger.assert_called_once()
 
     def test_slow_condition_subqueries(self, mock_trigger: MagicMock) -> None:
@@ -474,11 +477,13 @@ def test_slow_condition_subqueries(self, mock_trigger: MagicMock) -> None:
             self.post_process_error(event_3)
             assert not mock_trigger.called
 
-            project_ids = buffer.backend.get_sorted_set(
-                WORKFLOW_ENGINE_BUFFER_LIST_KEY, 0, timezone.now().timestamp()
+            project_ids = buffer.backend.bulk_get_sorted_set(
+                self.buffer_keys,
+                min=0,
+                max=timezone.now().timestamp(),
             )
 
-            process_delayed_workflows(project_ids[0][0])
+            process_delayed_workflows(list(project_ids.keys())[0])
             assert not mock_trigger.called
 
         with freeze_time(now + timedelta(minutes=1)):
@@ -487,9 +492,11 @@ def test_slow_condition_subqueries(self, mock_trigger: MagicMock) -> None:
             self.post_process_error(event_4)
             assert not mock_trigger.called
 
-            project_ids = buffer.backend.get_sorted_set(
-                WORKFLOW_ENGINE_BUFFER_LIST_KEY, 0, timezone.now().timestamp()
+            project_ids = buffer.backend.bulk_get_sorted_set(
+                self.buffer_keys,
+                min=0,
+                max=timezone.now().timestamp(),
             )
 
-            process_delayed_workflows(project_ids[0][0])
+            process_delayed_workflows(list(project_ids.keys())[0])
             mock_trigger.assert_called_once()