ref(aci): handle sharding project list key for buffer processing (#97638)

Author: cathteng

src/sentry/buffer/base.py

@@ -58,9 +58,14 @@ def get_hash(self, model: type[models.Model], field: dict[str, BufferField]) ->
     def get_hash_length(self, model: type[models.Model], field: dict[str, BufferField]) -> int:
         raise NotImplementedError
 
-    def get_sorted_set(self, key: str, min: float, max: float) -> list[tuple[int, datetime]]:
+    def get_sorted_set(self, key: str, min: float, max: float) -> list[tuple[int, float]]:
         return []
 
+    def bulk_get_sorted_set(
+        self, keys: list[str], min: float, max: float
+    ) -> dict[int, list[float]]:
+        return {}
+
     def push_to_sorted_set(self, key: str, value: list[int] | int) -> None:
         return None
 
@@ -92,6 +97,9 @@ def delete_hash(
     def delete_key(self, key: str, min: float, max: float) -> None:
         return None
 
+    def delete_keys(self, keys: list[str], min: float, max: float) -> None:
+        return None
+
     def incr(
         self,
         model: type[models.Model],

src/sentry/buffer/redis.py

@@ -2,6 +2,7 @@
 
 import logging
 import pickle
+from collections import defaultdict
 from collections.abc import Callable
 from dataclasses import dataclass
 from datetime import date, datetime, timezone
@@ -368,6 +369,26 @@ def _execute_redis_operation(
             pipe.expire(key, self.key_expire)
         return pipe.execute()[0]
 
+    def _execute_sharded_redis_operation(
+        self,
+        keys: list[str],
+        operation: RedisOperation,
+        *args: Any,
+        **kwargs: Any,
+    ) -> Any:
+        """
+        Execute a Redis operation on a list of keys, using the same args and kwargs for each key.
+        """
+
+        metrics_str = f"redis_buffer.{operation.value}"
+        metrics.incr(metrics_str, amount=len(keys))
+        pipe = self.get_redis_connection(self.pending_key)
+        for key in keys:
+            getattr(pipe, operation.value)(key, *args, **kwargs)
+            if args:
+                pipe.expire(key, self.key_expire)
+        return pipe.execute()
+
     def push_to_sorted_set(self, key: str, value: list[int] | int) -> None:
         now = time()
         if isinstance(value, list):
@@ -376,7 +397,7 @@ def push_to_sorted_set(self, key: str, value: list[int] | int) -> None:
             value_dict = {value: now}
         self._execute_redis_operation(key, RedisOperation.SORTED_SET_ADD, value_dict)
 
-    def get_sorted_set(self, key: str, min: float, max: float) -> list[tuple[int, datetime]]:
+    def get_sorted_set(self, key: str, min: float, max: float) -> list[tuple[int, float]]:
         redis_set = self._execute_redis_operation(
             key,
             RedisOperation.SORTED_SET_GET_RANGE,
@@ -393,9 +414,38 @@ def get_sorted_set(self, key: str, min: float, max: float) -> list[tuple[int, da
             decoded_set.append(data_and_timestamp)
         return decoded_set
 
+    def bulk_get_sorted_set(
+        self, keys: list[str], min: float, max: float
+    ) -> dict[int, list[float]]:
+        data_to_timestamps: dict[int, list[float]] = defaultdict(list)
+
+        redis_set = self._execute_sharded_redis_operation(
+            keys,
+            RedisOperation.SORTED_SET_GET_RANGE,
+            min=min,
+            max=max,
+            withscores=True,
+        )
+        for result in redis_set:
+            for items in result:
+                item = items[0]
+                if isinstance(item, bytes):
+                    item = item.decode("utf-8")
+                data_to_timestamps[int(item)].append(items[1])
+
+        return data_to_timestamps
+
     def delete_key(self, key: str, min: float, max: float) -> None:
         self._execute_redis_operation(key, RedisOperation.SORTED_SET_DELETE_RANGE, min=min, max=max)
 
+    def delete_keys(self, keys: list[str], min: float, max: float) -> None:
+        self._execute_sharded_redis_operation(
+            keys,
+            RedisOperation.SORTED_SET_DELETE_RANGE,
+            min=min,
+            max=max,
+        )
+
     def delete_hash(
         self,
         model: type[models.Model],

src/sentry/rules/processing/buffer_processing.py

@@ -32,6 +32,8 @@ class BufferHashKeys:
 
 class DelayedProcessingBase(ABC):
     buffer_key: ClassVar[str]
+    buffer_shards: ClassVar[int] = 1  # 1 shard will use the original buffer key
+    buffer_separator: ClassVar[str] = ":"
     option: ClassVar[str | None]
 
     def __init__(self, project_id: int):
@@ -47,6 +49,13 @@ def hash_args(self) -> BufferHashKeys:
     def processing_task(self) -> Task:
         raise NotImplementedError
 
+    @classmethod
+    def get_buffer_keys(cls) -> list[str]:
+        return [
+            f"{cls.buffer_key}{cls.buffer_separator}{shard}" if shard > 0 else cls.buffer_key
+            for shard in range(cls.buffer_shards)
+        ]
+
 
 delayed_processing_registry = Registry[type[DelayedProcessingBase]]()
 
@@ -156,21 +165,31 @@ def process_buffer() -> None:
             # The staler this timestamp, the more likely it'll miss some recently updated projects,
             # and the more likely we'll have frequently updated projects that are never actually
             # retrieved and processed here.
-            fetch_time = datetime.now(tz=timezone.utc)
-            project_ids = buffer.backend.get_sorted_set(
-                handler.buffer_key, min=0, max=fetch_time.timestamp()
+            fetch_time = datetime.now(tz=timezone.utc).timestamp()
+            buffer_keys = handler.get_buffer_keys()
+            all_project_ids_and_timestamps = buffer.backend.bulk_get_sorted_set(
+                buffer_keys,
+                min=0,
+                max=fetch_time,
             )
+
             if should_emit_logs:
                 log_str = ", ".join(
-                    f"{project_id}: {timestamp}" for project_id, timestamp in project_ids
+                    f"{project_id}: {timestamps}"
+                    for project_id, timestamps in all_project_ids_and_timestamps.items()
                 )
                 log_name = f"{processing_type}.project_id_list"
                 logger.info(log_name, extra={"project_ids": log_str})
 
-            for project_id, _ in project_ids:
+            project_ids = list(all_project_ids_and_timestamps.keys())
+            for project_id in project_ids:
                 process_in_batches(project_id, processing_type)
 
-            buffer.backend.delete_key(handler.buffer_key, min=0, max=fetch_time.timestamp())
+            buffer.backend.delete_keys(
+                buffer_keys,
+                min=0,
+                max=fetch_time,
+            )
 
 
 if not redis_buffer_registry.has(BufferHookEvent.FLUSH):

tests/sentry/buffer/test_redis.py

@@ -1,6 +1,7 @@
 import copy
 import datetime
 import pickle
+import random
 from collections import defaultdict
 from collections.abc import Mapping
 from unittest import mock
@@ -281,6 +282,54 @@ def test_enqueue(self) -> None:
         result = json.loads(project_ids_to_rule_data[project_id2][0].get(f"{rule2_id}:{group3_id}"))
         assert result.get("event_id") == event4_id
 
+    def test_get_bulk_sorted_set(self) -> None:
+        shards = 3
+        project_ids = [1, 2, 2, 3, 3, 3, 4, 4, 4, 4]
+        for id in project_ids:
+            shard = random.randrange(shards)
+            if shard == 0:
+                key = PROJECT_ID_BUFFER_LIST_KEY
+            else:
+                key = f"{PROJECT_ID_BUFFER_LIST_KEY}:{shard}"
+            self.buf.push_to_sorted_set(key=key, value=id)
+
+        buffer_keys = [
+            f"{PROJECT_ID_BUFFER_LIST_KEY}:{shard}" if shard > 0 else PROJECT_ID_BUFFER_LIST_KEY
+            for shard in range(shards)
+        ]
+
+        project_ids_and_timestamps = self.buf.bulk_get_sorted_set(
+            buffer_keys,
+            min=0,
+            max=datetime.datetime.now().timestamp(),
+        )
+        assert len(project_ids_and_timestamps) == 4
+        assert set(project_ids_and_timestamps.keys()) == set(project_ids)
+
+        self.buf.delete_keys(
+            buffer_keys,
+            min=0,
+            max=datetime.datetime.now().timestamp(),
+        )
+        project_ids_and_timestamps = self.buf.bulk_get_sorted_set(
+            buffer_keys,
+            min=0,
+            max=datetime.datetime.now().timestamp(),
+        )
+        assert len(project_ids_and_timestamps) == 0
+
+    def test_bulk_sorted_set_single_key(self) -> None:
+        project_ids = [1, 2, 2, 3, 3, 3, 4, 4, 4, 4]
+        for id in project_ids:
+            self.buf.push_to_sorted_set(key=PROJECT_ID_BUFFER_LIST_KEY, value=id)
+        project_ids_and_timestamps = self.buf.bulk_get_sorted_set(
+            [PROJECT_ID_BUFFER_LIST_KEY],
+            min=0,
+            max=datetime.datetime.now().timestamp(),
+        )
+        assert len(project_ids_and_timestamps) == 4
+        assert set(project_ids_and_timestamps.keys()) == set(project_ids)
+
     def test_buffer_hook_registry(self) -> None:
         """Test that we can add an event to the registry and that the callback is invoked"""
         mock = Mock()

tests/sentry/rules/processing/test_buffer_processing.py

@@ -245,6 +245,20 @@ def test_skips_processing_with_option(self, mock_process_in_batches) -> None:
             self.project_two.id,
         }
 
+    @patch("sentry.rules.processing.buffer_processing.process_in_batches")
+    @patch("sentry.rules.processing.delayed_processing.DelayedRule.buffer_shards", 3)
+    def test_fetches_with_shards(self, mock_process_in_batches: MagicMock) -> None:
+        project = self.create_project()
+        buffer.backend.push_to_sorted_set(key=f"{PROJECT_ID_BUFFER_LIST_KEY}", value=project.id)
+        buffer.backend.push_to_sorted_set(
+            key=f"{PROJECT_ID_BUFFER_LIST_KEY}:1", value=self.project_two.id
+        )
+        buffer.backend.push_to_sorted_set(key=f"{PROJECT_ID_BUFFER_LIST_KEY}:2", value=project.id)
+
+        process_buffer()
+
+        assert mock_process_in_batches.call_count == 3
+
 
 class ProcessInBatchesTest(CreateEventTestCase):
     def setUp(self) -> None: