feat(dlq): Added produce policy

arroyo/processing/strategies/dead_letter_queue/__init__.py

@@ -2,6 +2,7 @@
 from .policies.abstract import DeadLetterQueuePolicy, InvalidMessages
 from .policies.count import CountInvalidMessagePolicy
 from .policies.ignore import IgnoreInvalidMessagePolicy
+from .policies.produce import ProduceInvalidMessagePolicy
 from .policies.raise_e import RaiseInvalidMessagePolicy
 
 __all__ = [
@@ -11,4 +12,5 @@
     "CountInvalidMessagePolicy",
     "IgnoreInvalidMessagePolicy",
     "RaiseInvalidMessagePolicy",
+    "ProduceInvalidMessagePolicy",
 ]

arroyo/processing/strategies/dead_letter_queue/dead_letter_queue.py

@@ -54,5 +54,6 @@ def terminate(self) -> None:
         self.__next_step.terminate()
 
     def join(self, timeout: Optional[float] = None) -> None:
+        self.__policy.join(timeout)
         self.__next_step.close()
         self.__next_step.join(timeout)

arroyo/processing/strategies/dead_letter_queue/policies/abstract.py

@@ -1,5 +1,53 @@
+import base64
 from abc import ABC, abstractmethod
-from typing import Any, Sequence
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Mapping, Optional, Sequence, Union
+
+Serializable = Union[str, bytes]
+
+
+@dataclass(frozen=True)
+class InvalidMessage:
+    """
+    A dataclass to generally represent a bad Kafka message.
+    The Kafka specific fields (offset etc) are all optional
+    since this can be used to pass almost anything to the
+    DLQ.
+
+    A `reason` can be provided for debugging purposes. If a
+    produce policy is configured on the relevant DLQ, a message
+    in the form returned by `to_dict()` will be produced via
+    the policy.
+
+    An `InvalidMessages` exception should be raised containing
+    one or more objects of this class in order to actually
+    pass data to the DLQ.
+    """
+
+    payload: Serializable
+    timestamp: datetime
+    reason: Optional[str] = None
+    consumer_group: Optional[str] = None
+    partition: Optional[int] = None
+    offset: Optional[int] = None
+
+    def to_dict(self) -> Mapping[str, str]:
+        if isinstance(self.payload, bytes):
+            try:
+                decoded = self.payload.decode("utf-8")
+            except UnicodeDecodeError:
+                decoded = base64.b64encode(self.payload).decode("utf-8")
+        else:
+            decoded = self.payload
+        return {
+            "payload": decoded,
+            "timestamp": str(self.timestamp),
+            "reason": str(self.reason),
+            "consumer_group": str(self.consumer_group),
+            "partition": str(self.partition),
+            "offset": str(self.offset),
+        }
 
 
 class InvalidMessages(Exception):
@@ -8,12 +56,12 @@ class InvalidMessages(Exception):
     so they are handled correctly.
     """
 
-    def __init__(self, messages: Sequence[Any]):
+    def __init__(
+        self,
+        messages: Sequence[InvalidMessage],
+    ):
         self.messages = messages
 
-    def __str__(self) -> str:
-        return f"Invalid Message(s): {self.messages}"
-
 
 class DeadLetterQueuePolicy(ABC):
     """
@@ -25,4 +73,11 @@ def handle_invalid_messages(self, e: InvalidMessages) -> None:
         """
         Decide what to do with invalid messages.
         """
-        pass
+        raise NotImplementedError()
+
+    @abstractmethod
+    def join(self, timeout: Optional[float]) -> None:
+        """
+        Cleanup any asynchronous tasks that may be running.
+        """
+        raise NotImplementedError()

arroyo/processing/strategies/dead_letter_queue/policies/count.py

@@ -6,7 +6,6 @@
     DeadLetterQueuePolicy,
     InvalidMessages,
 )
-from arroyo.utils.metrics import get_metrics
 
 
 class _Bucket(NamedTuple):
@@ -20,9 +19,14 @@ class _Bucket(NamedTuple):
 
 class CountInvalidMessagePolicy(DeadLetterQueuePolicy):
     """
-    Ignore invalid messages up to a certain limit per time unit window in seconds.
-    This window is 1 minute by default. The exception associated with the invalid
-    messages is raised for all incoming invalid messages which go past this limit.
+    Does not raise invalid messages up to a certain limit per time unit window
+    in seconds. This window is 1 minute by default. The exception associated with
+    the invalid messages is raised for all incoming invalid messages which go past
+    this limit.
+
+    A `next_policy` (a DLQ Policy) must be passed to handle any exception which
+    remains within the per second limit. This gives full control over what happens
+    to exceptions within the limit.
 
     A saved state in the form `[(<timestamp: int>, <hits: int>), ...]` can be passed
     on init to load a previously saved state. This state should be aggregated to
@@ -31,13 +35,14 @@ class CountInvalidMessagePolicy(DeadLetterQueuePolicy):
 
     def __init__(
         self,
+        next_policy: DeadLetterQueuePolicy,
         limit: int,
         seconds: int = 60,
         load_state: Optional[Sequence[Tuple[int, int]]] = None,
     ) -> None:
         self.__limit = limit
         self.__seconds = seconds
-        self.__metrics = get_metrics()
+        self.__next_policy = next_policy
         if load_state is None:
             load_state = []
         self.__hits = deque(
@@ -49,7 +54,7 @@ def handle_invalid_messages(self, e: InvalidMessages) -> None:
         self._add(e)
         if self._count() > self.__limit:
             raise e
-        self.__metrics.increment("dlq.dropped_messages", len(e.messages))
+        self.__next_policy.handle_invalid_messages(e)
 
     def _add(self, e: InvalidMessages) -> None:
         now = int(time())
@@ -63,3 +68,6 @@ def _add(self, e: InvalidMessages) -> None:
     def _count(self) -> int:
         start = int(time()) - self.__seconds
         return sum(bucket.hits for bucket in self.__hits if bucket.timestamp >= start)
+
+    def join(self, timeout: Optional[float]) -> None:
+        self.__next_policy.join(timeout)

arroyo/processing/strategies/dead_letter_queue/policies/ignore.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 from arroyo.processing.strategies.dead_letter_queue.policies.abstract import (
     DeadLetterQueuePolicy,
     InvalidMessages,
@@ -11,3 +13,6 @@ def __init__(self) -> None:
 
     def handle_invalid_messages(self, e: InvalidMessages) -> None:
         self.__metrics.increment("dlq.dropped_messages", len(e.messages))
+
+    def join(self, timeout: Optional[float]) -> None:
+        return

arroyo/processing/strategies/dead_letter_queue/policies/produce.py

@@ -0,0 +1,71 @@
+import json
+import time
+from collections import deque
+from concurrent.futures import Future
+from typing import Deque, Optional
+
+from arroyo.backends.abstract import Producer
+from arroyo.backends.kafka.consumer import KafkaPayload
+from arroyo.processing.strategies.dead_letter_queue.policies.abstract import (
+    DeadLetterQueuePolicy,
+    InvalidMessage,
+    InvalidMessages,
+)
+from arroyo.types import Message, Topic
+from arroyo.utils.metrics import get_metrics
+
+MAX_QUEUE_SIZE = 5000
+
+
+class ProduceInvalidMessagePolicy(DeadLetterQueuePolicy):
+    """
+    Produces given InvalidMessages to a dead letter topic.
+    """
+
+    def __init__(
+        self, producer: Producer[KafkaPayload], dead_letter_topic: Topic
+    ) -> None:
+        self.__metrics = get_metrics()
+        self.__dead_letter_topic = dead_letter_topic
+        self.__producer = producer
+        self.__futures: Deque[Future[Message[KafkaPayload]]] = deque()
+
+    def handle_invalid_messages(self, e: InvalidMessages) -> None:
+        """
+        Produces a message to the given dead letter topic for each
+        invalid message. Produced message is in the form provided
+        by `InvalidMessage.to_dict()`
+        """
+        for message in e.messages:
+            payload = self._build_payload(message)
+            self._produce(payload)
+        self.__metrics.increment("dlq.produced_messages", len(e.messages))
+
+    def _build_payload(self, message: InvalidMessage) -> KafkaPayload:
+        data = json.dumps(message.to_dict()).encode("utf-8")
+        return KafkaPayload(key=None, value=data, headers=[])
+
+    def _produce(self, payload: KafkaPayload) -> None:
+        """
+        Prune done futures from queue if filled, then asynchronously
+        produce the message, adding the process (future) to the queue.
+        """
+        self._prune_done_future()
+        self.__futures.append(
+            self.__producer.produce(
+                destination=self.__dead_letter_topic, payload=payload
+            )
+        )
+
+    def _prune_done_future(self) -> None:
+        if len(self.__futures) >= MAX_QUEUE_SIZE:
+            self.__futures[0].result()
+            self.__futures.popleft()
+
+    def join(self, timeout: Optional[float] = None) -> None:
+        start = time.perf_counter()
+        while self.__futures:
+            if self.__futures[0].done():
+                self.__futures.popleft()
+            if timeout is not None and time.perf_counter() - start > timeout:
+                break

arroyo/processing/strategies/dead_letter_queue/policies/raise_e.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 from arroyo.processing.strategies.dead_letter_queue.policies.abstract import (
     DeadLetterQueuePolicy,
     InvalidMessages,
@@ -7,3 +9,6 @@
 class RaiseInvalidMessagePolicy(DeadLetterQueuePolicy):
     def handle_invalid_messages(self, e: InvalidMessages) -> None:
         raise e
+
+    def join(self, timeout: Optional[float]) -> None:
+        return