Repository: PrefectHQ/prefect
Pull request: Adds a dead-letter queue to memory subscriptions to avoid infinite loops
<!-- 
Thanks for opening a pull request to Prefect! 
If this is your first contribution, please make sure to review our contribution guidelines: https://docs.prefect.io/latest/contributing/overview/
-->

<!-- Include an overview of the proposed changes here -->
In https://github.com/PrefectHQ/prefect/issues/15607, users are seeing the task run recorder get stuck in an infinite loop when trying to insert a task run state update with the same task run ID and timestamp as another task run state update, but with a different ID. This PR doesn't solve the underlying problem (which appears to be a race condition in the client-side code), but it does add a dead-letter queue to the `Subscription` class to avoid messages being retried ad-infinitum. By default, messages will attempt to be processed 3 times before being removed from the subscription's queue.

Closes https://github.com/PrefectHQ/prefect/issues/15607

Problem statement:
Infinite Loop Symptoms on Prefect Server
### Bug summary

I ran into an issue with Prefect while working locally that looks like an endless loop happening on the backend.

I've observed the following behavior locally several times:
1. After running several dozen manual script executions, the workflows stop displaying correctly in the web UI. Specifically, the green squares on new flows are no longer visible.
2. At the same time, I see a large stream of logs in the PostgreSQL terminal.
3. Restarting Prefect resolves the issue.

Logs:
```
prefect-postgres-1  | 2024-10-07 23:57:37.075 UTC [2893] STATEMENT:  INSERT INTO task_run_state (task_run_id, type, timestamp, name, message, state_details, data, id, created, updated) VALUES ($1::UUID, $2::state_type, $3::TIMESTAMP WITH TIME ZONE, $4::VARCHAR, $5::VARCHAR, $6, $7::JSON, $8::UUID, $9::TIMESTAMP WITH TIME ZONE, $10::TIMESTAMP WITH TIME ZONE) ON CONFLICT (id) DO NOTHING
prefect-postgres-1  | 2024-10-07 23:57:37.082 UTC [2893] ERROR:  duplicate key value violates unique constraint "uq_task_run_state__task_run_id_timestamp_desc"
prefect-postgres-1  | 2024-10-07 23:57:37.082 UTC [2893] DETAIL:  Key (task_run_id, "timestamp")=(3649e07c-9883-42ac-8819-76c821f5f8eb, 2024-10-07 23:45:17.096643+00) already exists.
prefect-postgres-1  | 2024-10-07 23:57:37.082 UTC [2893] STATEMENT:  INSERT INTO task_run_state (task_run_id, type, timestamp, name, message, state_details, data, id, created, updated) VALUES ($1::UUID, $2::state_type, $3::TIMESTAMP WITH TIME ZONE, $4::VARCHAR, $5::VARCHAR, $6, $7::JSON, $8::UUID, $9::TIMESTAMP WITH TIME ZONE, $10::TIMESTAMP WITH TIME ZONE) ON CONFLICT (id) DO NOTHING
....
```

### Version info (`prefect version` output)

```Text
Version:             3.0.4
API version:         0.8.4
Python version:      3.12.6
Git commit:          c068d7e2
Built:               Tue, Oct 1, 2024 11:54 AM
OS/Arch:             linux/x86_64
Profile:             ephemeral
Server type:         server
Pydantic version:    2.9.2
```


### Additional context

PostgreSQL version: 16

Hints from issue discussion:
Here is an example of the error in the UI. You can see that tasks are not displayed, but there are saved artifacts. The Task Runs list is also empty.

# Artifacts
![image](https://github.com/user-attachments/assets/87a0a0c3-6d2b-4b54-8337-fa1467fdaa64)
# Task Runs 
![image](https://github.com/user-attachments/assets/954f6828-0182-4328-85f4-8b4128774e4a)

Workaround:
Execute a command on the PostgreSQL server of Prefect
```sql
DROP INDEX IF EXISTS uq_task_run_state__task_run_id_timestamp_desc;
CREATE INDEX uq_task_run_state__task_run_id_timestamp_desc
    ON task_run_state (task_run_id, timestamp DESC);
```
Thanks for the bug report @a14e! Do you have a flow that you can reliably reproduce this with? I suspect that there's something strange happening with the task run recorder, but it will be easier to confirm with a clear way to reproduce the issue.
@desertaxle Thank you for your reply!
Yes, I can reproduce the issue, but there is a significant amount of randomness involved. Below, I've provided the flow code, and on my machine, it took around 420 flow runs to reproduce the issue twice (I was running it from the UI).

My flow looks simplified like this:
```python
from prefect import flow, task, serve
from typing import List


class DirectusLoaderTasks:

    @task(log_prints=True, name="Directus. Load Topic")
    async def get_topic_by_id(self, topic_id: str) -> str:
        return "123"

class GrammarDBTasks:

    @task(log_prints=True, name="DB. Load items for topic")
    async def load_all_items_from_db(self, topic_id: str) -> list[str]:
        return []

    @task(log_prints=True, name="DB. Load item groups for topic")
    async def load_all_groups_from_db(self, topic_id: str) -> list[str]:
        return []


@task(log_prints=True, name="Simple filter")
async def simple_filter(from_db: List[str]) -> List[str]:
    return []



class OpenAiTasks:
    @task(name="Open AI. filter3", retries=3)
    async def filter3(self,
                      groups: List[str],
                      topic: str) -> List[str]:
        return []

    @task(name="Open AI. filter1", retries=3)
    async def filter1(self,
                      groups: List[str],
                      topic: str) -> List[str]:

        return []

    @task(name="Open AI. filter2", retries=3)
    async def filter2(self,
                      groups: List[str],
                      groups_in_db: List[int],
                      topic: str) -> List[str]:

        return []

    @task(name="Open AI. Normalize", retries=3)
    async def normalize(self,
                        groups: List[str],
                        topic: str) -> List[str]:
        return []




@flow(log_prints=True,name="Item Generation flow")
async def generate_items_flow():

    directus_tasks = DirectusLoaderTasks()
    grammar_tasks = GrammarDBTasks()
    open_ai_tasks = OpenAiTasks()

    topic_id = await directus_tasks.get_topic_by_id("123")

    item_groups_from_db = await grammar_tasks.load_all_groups_from_db(topic_id)
    items_from_db = await grammar_tasks.load_all_items_from_db(topic_id)
    new_item_groups = await simple_filter(items_from_db)

    deduplicated_response: List[str] = await open_ai_tasks.filter1(
        new_item_groups,
        topic_id
    )

    filtered_response: List[str] = await open_ai_tasks.filter2(
        deduplicated_response,
        item_groups_from_db,
        topic_id
    )

    filtered_response = await open_ai_tasks.filter3(filtered_response,
                                                    topic_id)

    normalize_forms = await open_ai_tasks.normalize(filtered_response,
                                                    topic_id)

    return


if __name__ == "__main__":
    example_deploy = generate_items_flow.to_deployment(
        "Generate Grammar Item Groups",
        tags=["tag1", "tag2", "tag3"]
    )
    serve(example_deploy)


```

Here’s a video showing what the PostgreSQL logs look like in the terminal:

https://github.com/user-attachments/assets/85c3dc76-eddd-49bf-9501-61d3427c2a02


I have the exact problem that occurs in the local setup with the PostgreSQL database and prefect server running in a simple docker-compose project. The occurrence rate is quite high, it happens even with such a simple pipeline:

```python 
from prefect import flow, task
import pandas as pd
import time
import os
import random


@task(name="Transform Data 1")
def transform_data_1(path) -> pd.DataFrame:
    print(path)
    data = {"id": [1, 2, 3], "value": [10, 20, 30]}
    time.sleep(random.randint(1, 10))
    return pd.DataFrame(data)


@task(name="Transform Data 2")
def transform_data_2(path) -> pd.DataFrame:
    print(path)
    data = {"id": [4, 5, 6], "value": [40, 50, 60]}
    time.sleep(random.randint(1, 10))
    return pd.DataFrame(data)


@flow
def process_data(file_path: str):
    df1 = transform_data_1(file_path)
    print("Done 1")
    df2 = transform_data_2(file_path)
    print("Done 2")


@flow(name="Example Pipeline")
def example_pipeline():
    for i in range(5):
        process_data(
            file_path=f"file_{i}.csv",
        )


if __name__ == "__main__":
    os.environ["PREFECT_API_URL"] = "http://localhost:4200"
    example_pipeline()

```

I have tried rewriting this into async calls, but had no effect. Tried both postgres:16-alpine and postgres:17-alpine docker images for the postgres, different versions of python images for the prefect server as well. My version of prefect:

```
Version:             3.1.2
API version:         0.8.4
Python version:      3.10.15
Git commit:          02b99f0a
Built:               Tue, Nov 12, 2024 1:38 PM
OS/Arch:             linux/x86_64
Profile:             ephemeral
Server type:         ephemeral
Pydantic version:    2.9.2
Server:
  Database:          postgresql
```

The proposed workaround with the constraint modification on the database side works, but the container with the server is consuming an enormous amount of CPU resources and we need to restart the container to break the loop.
Thanks for the additional info @a14e and @rejner! I haven't been able to reproduce this on my end, but it looks like there's a race condition where the a state change for a task run can be emitted twice with the exact same timestamp. I haven't yet tracked down where that's happening, but the issue is exacerbated because we keep retrying event processing when there's a failure, which causes this infinite loop.

I think I can address the infinite loop issue by adding a dead-letter queue to the task run recorder. The race condition will still exist, but it won't overwhelm the server CPU when it happens, and maybe we'll get some more info about the states that are causing this underlying issue.

Commit message: Merge branch 'main' into add-dlq-to-subscriptions

Patch to review:
diff --git a/src/prefect/server/utilities/messaging/memory.py b/src/prefect/server/utilities/messaging/memory.py
index 1d151cb6a660..e1e04c7ae2e6 100644
--- a/src/prefect/server/utilities/messaging/memory.py
+++ b/src/prefect/server/utilities/messaging/memory.py
@@ -1,4 +1,5 @@
 import asyncio
+import copy
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
 from datetime import timedelta
@@ -29,6 +30,7 @@
 class MemoryMessage:
     data: Union[bytes, str]
     attributes: Dict[str, Any]
+    retry_count: int = 0
 
 
 class Subscription:
@@ -36,8 +38,10 @@ class Subscription:
     _queue: asyncio.Queue
     _retry: asyncio.Queue
 
-    def __init__(self, topic: "Topic") -> None:
+    def __init__(self, topic: "Topic", max_retries: int = 3) -> None:
         self.topic = topic
+        self.max_retries = max_retries
+        self.dead_letter_queue = list()
         self._queue = asyncio.Queue()
         self._retry = asyncio.Queue()
 
@@ -45,7 +49,16 @@ async def deliver(self, message: MemoryMessage) -> None:
         await self._queue.put(message)
 
     async def retry(self, message: MemoryMessage) -> None:
-        await self._retry.put(message)
+        message.retry_count += 1
+        if message.retry_count > self.max_retries:
+            logger.warning(
+                "Message failed after %d retries and will be moved to the dead letter queue",
+                message.retry_count,
+                extra={"event_message": message},
+            )
+            self.dead_letter_queue.append(message)
+        else:
+            await self._retry.put(message)
 
     async def get(self) -> MemoryMessage:
         if self._retry.qsize() > 0:
@@ -93,7 +106,8 @@ def clear(self):
 
     async def publish(self, message: MemoryMessage) -> None:
         for subscription in self._subscriptions:
-            await subscription.deliver(message)
+            # Ensure that each subscription gets its own copy of the message
+            await subscription.deliver(copy.deepcopy(message))
 
 
 @asynccontextmanager
diff --git a/tests/server/services/test_task_run_recorder.py b/tests/server/services/test_task_run_recorder.py
index e373a76c972e..6e6ad72887e3 100644
--- a/tests/server/services/test_task_run_recorder.py
+++ b/tests/server/services/test_task_run_recorder.py
@@ -18,7 +18,7 @@
 from prefect.server.schemas.core import FlowRun, TaskRunPolicy
 from prefect.server.schemas.states import StateDetails, StateType
 from prefect.server.services import task_run_recorder
-from prefect.server.utilities.messaging import MessageHandler
+from prefect.server.utilities.messaging import MessageHandler, create_publisher
 from prefect.server.utilities.messaging.memory import MemoryMessage
 
 
@@ -29,6 +29,7 @@ async def test_start_and_stop_service():
 
     await service.started_event.wait()
     assert service.consumer_task is not None
+    assert service.consumer is not None
 
     await service.stop()
     assert service.consumer_task is None
@@ -753,3 +754,45 @@ async def test_task_run_recorder_handles_all_out_of_order_permutations(
 
     state_types = set(state.type for state in states)
     assert state_types == {StateType.PENDING, StateType.RUNNING, StateType.COMPLETED}
+
+
+async def test_task_run_recorder_sends_repeated_failed_messages_to_dead_letter(
+    session: AsyncSession,
+    pending_event: ReceivedEvent,
+):
+    """
+    Test to ensure situations like the one described in https://github.com/PrefectHQ/prefect/issues/15607
+    don't overwhelm the task run recorder.
+    """
+    pending_transition_time = pendulum.datetime(2024, 1, 1, 0, 0, 0, 0, "UTC")
+    assert pending_event.occurred == pending_transition_time
+
+    service = task_run_recorder.TaskRunRecorder()
+
+    service_task = asyncio.create_task(service.start())
+    await service.started_event.wait()
+
+    async with create_publisher("events") as publisher:
+        await publisher.publish_data(
+            message(pending_event).data, message(pending_event).attributes
+        )
+        # Sending a task run event with the same task run id and timestamp but
+        # a different id will raise an error when trying to insert it into the
+        # database
+        duplicate_pending_event = pending_event.model_copy()
+        duplicate_pending_event.id = UUID("bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb")
+        await publisher.publish_data(
+            message(duplicate_pending_event).data,
+            message(duplicate_pending_event).attributes,
+        )
+
+    while not service.consumer.subscription.dead_letter_queue:
+        await asyncio.sleep(0.1)
+
+    assert len(service.consumer.subscription.dead_letter_queue) == 1
+
+    service_task.cancel()
+    try:
+        await service_task
+    except asyncio.CancelledError:
+        pass
diff --git a/tests/server/utilities/test_messaging.py b/tests/server/utilities/test_messaging.py
index 2ee6bf8e5325..7df580a5d1ea 100644
--- a/tests/server/utilities/test_messaging.py
+++ b/tests/server/utilities/test_messaging.py
@@ -24,6 +24,7 @@
     create_publisher,
     ephemeral_subscription,
 )
+from prefect.server.utilities.messaging.memory import Consumer as MemoryConsumer
 from prefect.settings import (
     PREFECT_MESSAGING_BROKER,
     PREFECT_MESSAGING_CACHE,
@@ -441,3 +442,43 @@ async def handler(message: Message):
     # TODO: is there a way we can test that ephemeral subscriptions really have cleaned
     # up after themselves after they have exited?  This will differ significantly by
     # each broker implementation, so it's hard to write a generic test.
+
+
+async def test_repeatedly_failed_message_is_moved_to_dead_letter_queue(
+    deduplicating_publisher: Publisher,
+    consumer: MemoryConsumer,
+):
+    captured_messages: List[Message] = []
+
+    async def handler(message: Message):
+        captured_messages.append(message)
+        raise ValueError("Simulated failure")
+
+    consumer_task = asyncio.create_task(consumer.run(handler))
+
+    async with deduplicating_publisher as p:
+        await p.publish_data(b"hello, world", {"howdy": "partner"})
+
+    while not consumer.subscription.dead_letter_queue:
+        await asyncio.sleep(0.1)
+    try:
+        consumer_task.cancel()
+        await consumer_task
+    except asyncio.CancelledError:
+        pass
+
+    # Message should have been moved to DLQ after multiple retries
+    assert len(captured_messages) == 4  # Original attempt + 3 retries
+    for message in captured_messages:
+        assert message.data == b"hello, world"
+        assert message.attributes == {"howdy": "partner"}
+
+    # Verify message is in DLQ
+    assert len(consumer.subscription.dead_letter_queue) == 1
+    dlq_message = next(iter(consumer.subscription.dead_letter_queue))
+    assert dlq_message.data == b"hello, world"
+    assert dlq_message.attributes == {"howdy": "partner"}
+    assert dlq_message.retry_count > 3
+
+    remaining_message = await drain_one(consumer)
+    assert not remaining_message


Code review comments:
diff_hunk: @@ -29,23 +30,35 @@
 class MemoryMessage:
     data: Union[bytes, str]
     attributes: Dict[str, Any]
+    retry_count: int = 0
 
 
 class Subscription:
     topic: "Topic"
     _queue: asyncio.Queue
     _retry: asyncio.Queue
 
-    def __init__(self, topic: "Topic") -> None:
+    def __init__(self, topic: "Topic", max_retries: int = 3) -> None:
         self.topic = topic
+        self.max_retries = max_retries
+        self.dead_letter_queue = list()
         self._queue = asyncio.Queue()
         self._retry = asyncio.Queue()
 
     async def deliver(self, message: MemoryMessage) -> None:
         await self._queue.put(message)
 
     async def retry(self, message: MemoryMessage) -> None:
-        await self._retry.put(message)
+        message.retry_count += 1
+        if message.retry_count > self.max_retries:
+            logger.warning(
+                "Message failed after %d retries and will be moved to the dead letter queue",
+                message.retry_count,
+                extra={"event_message": message},
+            )
+            self.dead_letter_queue.append(message)
path: src/prefect/server/utilities/messaging/memory.py
text: @user1:
The DLQ here could end up growing without bounds.  What do you think about writing these to a file in `$PREFECT_HOME` instead?  That gives the user a chance to review them and decide whether there's something that needs to happen, or ignore them and drop them.  This also gives us an opening to implement some kind of redrive process for them later.

Either way, we probably don't want to store them in-memory indefinitely, so we need a way to purge these if we don't write them to disk.

@author:
That's a good idea! For other message brokers, I think we might want to have the DLQ in the same infra (e.g. Redis), but that doesn't work as well for the in-memory implementation. I'll update the current implementation to write to disk with an eye toward writing other places in the future.

@user1:
yep exactly ^