fixes

Author: lynnagara

src/sentry/conf/types/kafka_definition.py

@@ -86,6 +86,8 @@ class ConsumerDefinition(TypedDict, total=False):
     dlq_max_invalid_ratio: float | None
     dlq_max_consecutive_count: int | None
 
+    stale_topic: Topic
+
 
 def validate_consumer_definition(consumer_definition: ConsumerDefinition) -> None:
     if "dlq_topic" not in consumer_definition and (

src/sentry/consumers/__init__.py

@@ -569,19 +569,20 @@ def build_consumer_config(group_id: str):
     else:
         stale_topic = None
 
+    dlq_policy = None
+
     if enable_dlq or stale_threshold_sec:
         dlq_producer = build_dlq_producer(dlq_topic=dlq_topic, stale_topic=stale_topic)
 
-        dlq_policy = DlqPolicy(
-            dlq_producer,
-            DlqLimit(
-                max_invalid_ratio=consumer_definition.get("dlq_max_invalid_ratio"),
-                max_consecutive_count=consumer_definition.get("dlq_max_consecutive_count"),
-            ),
-            None,
-        )
-    else:
-        dlq_policy = None
+        if dlq_producer:
+            dlq_policy = DlqPolicy(
+                dlq_producer,
+                DlqLimit(
+                    max_invalid_ratio=consumer_definition.get("dlq_max_invalid_ratio"),
+                    max_consecutive_count=consumer_definition.get("dlq_max_consecutive_count"),
+                ),
+                None,
+            )
 
     return StreamProcessor(
         consumer=consumer,

src/sentry/consumers/dlq.py

@@ -1,19 +1,22 @@
+import logging
 import time
-from collections.abc import Callable, Mapping, MutableMapping
+from collections.abc import Mapping, MutableMapping
 from concurrent.futures import Future
 from datetime import datetime, timedelta, timezone
 from enum import Enum
 
 from arroyo.backends.kafka import KafkaPayload, KafkaProducer
 from arroyo.dlq import InvalidMessage, KafkaDlqProducer
 from arroyo.processing.strategies.abstract import ProcessingStrategy, ProcessingStrategyFactory
-from arroyo.types import FILTERED_PAYLOAD, BrokerValue, Message, Partition
+from arroyo.types import FILTERED_PAYLOAD, BrokerValue, Commit, FilteredPayload, Message, Partition
 from arroyo.types import Topic as ArroyoTopic
 from arroyo.types import Value
 
 from sentry.conf.types.kafka_definition import Topic
 from sentry.utils.kafka_config import get_kafka_producer_cluster_options, get_topic_definition
 
+logger = logging.getLogger(__name__)
+
 
 class RejectReason(Enum):
     STALE = "stale"
@@ -27,16 +30,30 @@ class MultipleDestinationDlqProducer(KafkaDlqProducer):
 
     def __init__(
         self,
-        producers: Mapping[RejectReason, KafkaDlqProducer],
-        topic_selector: Callable[[BrokerValue[KafkaPayload], str], RejectReason],
+        producers: Mapping[RejectReason, KafkaDlqProducer | None],
     ) -> None:
         self.producers = producers
-        self.topic_selector = topic_selector
 
     def produce(
-        self, value: BrokerValue[KafkaPayload], reason: str
+        self,
+        value: BrokerValue[KafkaPayload],
+        reason: str | None = None,
     ) -> Future[BrokerValue[KafkaPayload]]:
-        return self.producers[self.topic_selector(value, reason)].produce(value)
+        try:
+            reject_reason = RejectReason(reason)
+            producer = self.producers.get(reject_reason)
+        except ValueError:
+            producer = None
+
+        if producer:
+            return producer.produce(value)
+        else:
+            # No DLQ producer configured for the reason.
+            logger.error("No DLQ producer configured for reason %s", reason)
+            future: Future[BrokerValue[KafkaPayload]] = Future()
+            future.set_running_or_notify_cancel()
+            future.set_result(value)
+            return future
 
 
 def _get_dlq_producer(topic: Topic | None) -> KafkaDlqProducer | None:
@@ -50,7 +67,8 @@ def _get_dlq_producer(topic: Topic | None) -> KafkaDlqProducer | None:
 
 
 def build_dlq_producer(
-    dlq_topic: Topic | None, stale_topic: Topic | None
+    dlq_topic: Topic | None,
+    stale_topic: Topic | None,
 ) -> MultipleDestinationDlqProducer | None:
     if dlq_topic is None and stale_topic is None:
         return None
@@ -63,9 +81,11 @@ def build_dlq_producer(
     return MultipleDestinationDlqProducer(producers)
 
 
-class DlqStaleMessages(ProcessingStrategy):
+class DlqStaleMessages(ProcessingStrategy[KafkaPayload]):
     def __init__(
-        self, stale_threshold_sec: int, next_step: ProcessingStrategy[KafkaPayload]
+        self,
+        stale_threshold_sec: int,
+        next_step: ProcessingStrategy[KafkaPayload | FilteredPayload],
     ) -> None:
         self.stale_threshold_sec = stale_threshold_sec
         self.next_step = next_step
@@ -80,17 +100,18 @@ def submit(self, message: Message[KafkaPayload]) -> None:
         )
 
         if isinstance(message.value, BrokerValue):
-            message_timestamp = message.timestamp.astimezone(timezone.utc)
-            if message_timestamp < min_accepted_timestamp:
-                self.offsets_to_forward[message.value.partition, message.value.next_offset]
+            if message.value.timestamp < min_accepted_timestamp:
+                self.offsets_to_forward[message.value.partition] = message.value.next_offset
                 raise InvalidMessage(
-                    message.value.partition, message.value.offset, RejectReason.STALE.value
+                    message.value.partition, message.value.offset, reason=RejectReason.STALE.value
                 )
 
         if self.offsets_to_forward and time.time() > self.last_forwarded_offsets + 1:
-            message = Message(Value(FILTERED_PAYLOAD), self.offsets_to_forward)
+            filtered_message = Message(Value(FILTERED_PAYLOAD, self.offsets_to_forward))
             self.offsets_to_forward = {}
-            self.next_step.submit(message)
+            self.next_step.submit(filtered_message)
+
+        self.next_step.submit(message)
 
     def poll(self) -> None:
         self.next_step.poll()
@@ -105,17 +126,23 @@ def terminate(self) -> None:
         self.next_step.terminate()
 
 
-class DlqStaleMessagesStrategyFactoryWrapper(ProcessingStrategyFactory):
+class DlqStaleMessagesStrategyFactoryWrapper(ProcessingStrategyFactory[KafkaPayload]):
     """
     Wrapper used to dlq a message with a stale timestamp before it is passed to
     the rest of the pipeline. The InvalidMessage is raised with a
     "stale" reason so it can be routed to a separate stale topic.
     """
 
-    def __init__(self, stale_threshold_sec: int, inner: ProcessingStrategyFactory) -> None:
+    def __init__(
+        self,
+        stale_threshold_sec: int,
+        inner: ProcessingStrategyFactory[KafkaPayload | FilteredPayload],
+    ) -> None:
         self.stale_threshold_sec = stale_threshold_sec
         self.inner = inner
 
-    def create_with_partitions(self, commit, partitions) -> ProcessingStrategy:
+    def create_with_partitions(
+        self, commit: Commit, partitions: Mapping[Partition, int]
+    ) -> ProcessingStrategy[KafkaPayload]:
         rv = self.inner.create_with_partitions(commit, partitions)
         return DlqStaleMessages(self.stale_threshold_sec, rv)