perf(consumers): Route messages based on kafka headers

snuba/datasets/message_filters.py

@@ -21,3 +21,19 @@ def should_drop(self, message: Message[TPayload]) -> bool:
 class PassthroughKafkaFilter(StreamMessageFilter[KafkaPayload]):
     def should_drop(self, message: Message[KafkaPayload]) -> bool:
         return False
+
+
+class KafkaHeaderFilter(StreamMessageFilter[KafkaPayload]):
+    def __init__(self, header_key: str, header_value: str):
+        self.__header_key = header_key
+        self.__header_value = header_value
+
+    def should_drop(self, message: Message[KafkaPayload]) -> bool:
+        for key, value in message.payload.headers:
+            if key != self.__header_key:
+                continue
+
+            str_value = value.decode("utf-8")
+            return True if str_value == self.__header_value else False
+
+        return False

snuba/datasets/storages/errors.py

@@ -2,6 +2,7 @@
 from snuba.clusters.storage_sets import StorageSetKey
 from snuba.datasets.errors_processor import ErrorsProcessor
 from snuba.datasets.errors_replacer import ErrorsReplacer, ReplacerState
+from snuba.datasets.message_filters import KafkaHeaderFilter
 from snuba.datasets.schemas.tables import WritableTableSchema
 from snuba.datasets.storage import WritableTableStorage
 from snuba.datasets.storages import StorageKey
@@ -36,6 +37,7 @@
     mandatory_condition_checkers=[ProjectIdEnforcer()],
     stream_loader=build_kafka_stream_loader_from_settings(
         processor=ErrorsProcessor(promoted_tag_columns),
+        pre_filter=KafkaHeaderFilter("transaction_forwarder", "1"),
         default_topic=Topic.EVENTS,
         replacement_topic=Topic.EVENT_REPLACEMENTS,
         commit_log_topic=Topic.COMMIT_LOG,

snuba/datasets/storages/transactions.py

@@ -12,6 +12,7 @@
 from snuba.clickhouse.columns import SchemaModifiers as Modifiers
 from snuba.clickhouse.columns import String, UInt
 from snuba.clusters.storage_sets import StorageSetKey
+from snuba.datasets.message_filters import KafkaHeaderFilter
 from snuba.datasets.schemas.tables import WritableTableSchema
 from snuba.datasets.storage import WritableTableStorage
 from snuba.datasets.storages import StorageKey
@@ -151,6 +152,7 @@
     ],
     stream_loader=build_kafka_stream_loader_from_settings(
         processor=TransactionsMessageProcessor(),
+        pre_filter=KafkaHeaderFilter("transaction_forwarder", "0"),
         default_topic=Topic.EVENTS,
         commit_log_topic=Topic.COMMIT_LOG,
         subscription_scheduler_mode=SchedulingWatermarkMode.PARTITION,

tests/datasets/test_message_filters.py

@@ -0,0 +1,52 @@
+from datetime import datetime
+
+import pytest
+from arroyo import Message, Partition, Topic
+from arroyo.backends.kafka import KafkaPayload
+
+from snuba.datasets.message_filters import KafkaHeaderFilter
+
+test_data = [
+    pytest.param(
+        KafkaHeaderFilter("should_drop", "1"),
+        Message(
+            Partition(Topic("random"), 1),
+            1,
+            KafkaPayload(b"key", b"value", [("should_drop", b"1")]),
+            datetime.now(),
+        ),
+        True,
+        id="matching-headers",
+    ),
+    pytest.param(
+        KafkaHeaderFilter("should_drop", "0"),
+        Message(
+            Partition(Topic("random"), 1),
+            1,
+            KafkaPayload(b"key", b"value", [("should_drop", b"1")]),
+            datetime.now(),
+        ),
+        False,
+        id="mismatched-headers",
+    ),
+    pytest.param(
+        KafkaHeaderFilter("should_drop", "1"),
+        Message(
+            Partition(Topic("random"), 1),
+            1,
+            KafkaPayload(b"key", b"value", [("missing", b"0")]),
+            datetime.now(),
+        ),
+        False,
+        id="missing-headers",
+    ),
+]
+
+
+@pytest.mark.parametrize("header_filter, message, expected_result", test_data)
+def test_kafka_filter_header_should_drop(
+    header_filter: KafkaHeaderFilter,
+    message: Message[KafkaPayload],
+    expected_result: bool,
+) -> None:
+    assert header_filter.should_drop(message) == expected_result