ref(rules): Add functionality to RedisBuffer (#68049)

Add methods to `RedisBuffer` to add data about rules that we can later
use to process "slow" rules.

Closes
[#238](getsentry/team-core-product-foundations#238)

Author: ceorourke

src/sentry/buffer/redis.py

@@ -4,6 +4,7 @@
 import pickle
 import threading
 from datetime import date, datetime, timezone
+from enum import Enum
 from time import time
 from typing import Any
 
@@ -32,6 +33,8 @@
 # Debounce our JSON validation a bit in order to not cause too much additional
 # load everywhere
 _last_validation_log: float | None = None
+Pipeline = Any
+# TODO type Pipeline instead of using Any here
 
 
 def _validate_json_roundtrip(value: dict[str, Any], model: type[models.Model]) -> None:
@@ -49,6 +52,13 @@ def _validate_json_roundtrip(value: dict[str, Any], model: type[models.Model]) -
             logger.exception("buffer.invalid_value", extra={"value": value, "model": model})
 
 
+class RedisOperation(Enum):
+    SET_ADD = "sadd"
+    SET_GET = "smembers"
+    HASH_ADD = "hset"
+    HASH_GET_ALL = "hgetall"
+
+
 class PendingBuffer:
     def __init__(self, size: int):
         assert size > 0
@@ -208,6 +218,48 @@ def get(
             col: (int(results[i]) if results[i] is not None else 0) for i, col in enumerate(columns)
         }
 
+    def get_redis_connection(self, key: str) -> Pipeline | None:
+        if is_instance_redis_cluster(self.cluster, self.is_redis_cluster):
+            conn = self.cluster
+        elif is_instance_rb_cluster(self.cluster, self.is_redis_cluster):
+            conn = self.cluster.get_local_client_for_key(key)
+        else:
+            raise AssertionError("unreachable")
+
+        pipe = conn.pipeline()
+        return pipe
+
+    def _execute_redis_operation(self, key: str, operation: RedisOperation, *args: Any) -> Any:
+        pending_key = self._make_pending_key_from_key(key)
+        pipe = self.get_redis_connection(pending_key)
+        if pipe:
+            getattr(pipe, operation.value)(key, *args)
+            if args:
+                pipe.expire(key, self.key_expire)
+            return pipe.execute()
+
+    def push_to_set(self, key: str, value: list[int] | int) -> None:
+        self._execute_redis_operation(key, RedisOperation.SET_ADD, value)
+
+    def get_set(self, key: str) -> list[set[int]]:
+        return self._execute_redis_operation(key, RedisOperation.SET_GET)
+
+    def push_to_hash(
+        self,
+        model: type[models.Model],
+        filters: dict[str, models.Model | str | int],
+        field: str,
+        value: int,
+    ) -> None:
+        key = self._make_key(model, filters)
+        self._execute_redis_operation(key, RedisOperation.HASH_ADD, field, value)
+
+    def get_hash(
+        self, model: type[models.Model], field: dict[str, models.Model | str | int]
+    ) -> dict[str, str]:
+        key = self._make_key(model, field)
+        return self._execute_redis_operation(key, RedisOperation.HASH_GET_ALL)
+
     def incr(
         self,
         model: type[models.Model],
@@ -226,19 +278,13 @@ def incr(
             - Perform a set on signal_only (only if True)
         - Add hashmap key to pending flushes
         """
-
         key = self._make_key(model, filters)
         pending_key = self._make_pending_key_from_key(key)
         # We can't use conn.map() due to wanting to support multiple pending
         # keys (one per Redis partition)
-        if is_instance_redis_cluster(self.cluster, self.is_redis_cluster):
-            conn = self.cluster
-        elif is_instance_rb_cluster(self.cluster, self.is_redis_cluster):
-            conn = self.cluster.get_local_client_for_key(key)
-        else:
-            raise AssertionError("unreachable")
-
-        pipe = conn.pipeline()
+        pipe = self.get_redis_connection(key)
+        if not pipe:
+            return
         pipe.hsetnx(key, "m", f"{model.__module__}.{model.__name__}")
         _validate_json_roundtrip(filters, model)
 

tests/sentry/buffer/test_redis.py

@@ -1,5 +1,6 @@
 import datetime
 import pickle
+from collections import defaultdict
 from unittest import mock
 
 import pytest
@@ -252,6 +253,85 @@ def test_process_pending_partitions_none(self, process_pending, process_incr):
         # Make sure we didn't queue up more
         assert len(process_pending.apply_async.mock_calls) == 2
 
+    def group_rule_data_by_project_id(self, buffer, project_ids):
+        project_ids_to_rule_data = defaultdict(list)
+        for proj_id in project_ids[0]:
+            rule_group_pairs = buffer.get_hash(Project, {"project_id": proj_id})
+            for pair in rule_group_pairs:
+                for k, v in pair.items():
+                    if isinstance(k, bytes):
+                        k = k.decode("utf-8")
+                    if isinstance(v, bytes):
+                        v = v.decode("utf-8")
+                    project_ids_to_rule_data[int(proj_id)].append({k: v})
+        return project_ids_to_rule_data
+
+    def test_enqueue(self):
+        PROJECT_ID_BUFFER_LIST_KEY = "project_id_buffer_list"
+        project_id = 1
+        rule_id = 2
+        group_id = 3
+        event_id = 4
+        group2_id = 5
+        event2_id = 6
+
+        project_id2 = 7
+        rule2_id = 8
+        group3_id = 9
+        event3_id = 10
+        event4_id = 11
+
+        # store the project ids
+        self.buf.push_to_set(key=PROJECT_ID_BUFFER_LIST_KEY, value=project_id)
+        self.buf.push_to_set(key=PROJECT_ID_BUFFER_LIST_KEY, value=project_id2)
+
+        # store the rules and group per project
+        self.buf.push_to_hash(
+            model=Project,
+            filters={"project_id": project_id},
+            field=f"{rule_id}:{group_id}",
+            value=event_id,
+        )
+        self.buf.push_to_hash(
+            model=Project,
+            filters={"project_id": project_id},
+            field=f"{rule_id}:{group2_id}",
+            value=event2_id,
+        )
+        self.buf.push_to_hash(
+            model=Project,
+            filters={"project_id": project_id2},
+            field=f"{rule2_id}:{group3_id}",
+            value=event3_id,
+        )
+
+        project_ids = self.buf.get_set(PROJECT_ID_BUFFER_LIST_KEY)
+        assert project_ids
+
+        project_ids_to_rule_data = self.group_rule_data_by_project_id(self.buf, project_ids)
+        assert project_ids_to_rule_data[project_id][0].get(f"{rule_id}:{group_id}") == str(event_id)
+        assert project_ids_to_rule_data[project_id][1].get(f"{rule_id}:{group2_id}") == str(
+            event2_id
+        )
+        assert project_ids_to_rule_data[project_id2][0].get(f"{rule2_id}:{group3_id}") == str(
+            event3_id
+        )
+
+        # overwrite the value to event4_id
+        self.buf.push_to_hash(
+            model=Project,
+            filters={"project_id": project_id2},
+            field=f"{rule2_id}:{group3_id}",
+            value=event4_id,
+        )
+
+        project_ids_to_rule_data = project_ids_to_rule_data = self.group_rule_data_by_project_id(
+            self.buf, project_ids
+        )
+        assert project_ids_to_rule_data[project_id2][0].get(f"{rule2_id}:{group3_id}") == str(
+            event4_id
+        )
+
     @mock.patch("sentry.buffer.redis.RedisBuffer._make_key", mock.Mock(return_value="foo"))
     @mock.patch("sentry.buffer.base.Buffer.process")
     def test_process_uses_signal_only(self, process):