Repository: PrefectHQ/prefect

Issue:
Memory leak
### First check

- [X] I added a descriptive title to this issue.
- [X] I used the GitHub search to find a similar issue and didn't find it.
- [X] I searched the Prefect documentation for this issue.
- [X] I checked that this issue is related to Prefect and not one of its dependencies.

### Bug summary

When running prefect app in an infinite loop (as a job that should listen to e.g. some messaging service and process the events) we are running into infinite memory growth. We tried to inspect our code but even if we run the most minimal example without persisting any data and also garbage collecting the memory still grows.

### Reproduction

```python3
import gc
import asyncio
import time

from prefect import flow, task



@task(
    #retries=3,
    #persist_result=False,
    #cache_result_in_memory=False,
)
async def build_dataframe_async(id, index):
    print(f"Hello World = id = {id}, index = {index}")


@flow(
    # log_prints=True,
    # persist_result=False,
    cache_result_in_memory=False,
)
async def flow_async(id):
    for index in range(30):
        await build_dataframe_async(id, index)


if __name__ == '__main__':
    for index in range(5):
        print(f"____________________ Index = {index} ____________________")
        asyncio.run(flow_async(index))

        gc.collect(0)
        gc.collect(1)
        gc.collect(2)
        time.sleep(0.01)
        gc.collect()
        time.sleep(0.01)
```


### Error

_No response_

### Versions

```Text
python 3.10
prefect 2.10.18
```


### Additional context

_No response_


Discussion:
Also experiencing this with a very basic example flow set to run ~1min:


![image](https://github.com/PrefectHQ/prefect/assets/13764397/a93445de-7eae-4451-89e1-d4fa5b22bd68)

Python 3.10
prefect 2.13.0
I think this is related to https://github.com/pydantic/pydantic/pull/6681 which was fixed in pydantic v2. I'm linking this to our ticket for updating to pydantic v2 since I expect this issue will go away once we support the upgrade.

Related to https://github.com/PrefectHQ/prefect/issues/10145
on 2.14.16 memory still leaks
To reproduce:

`pip install matplotlib prefect pylab-sdk memory-profiler`

`test_flow.py`
```python
import asyncio
import collections
import contextlib
import gc
import os
import pprint
import time
import tracemalloc

import anyio
import prefect._internal.concurrency.threads
import prefect.client.base
import prefect.engine
import prefect.events.related
import prefect.utilities.asyncutils
from prefect import context, flow, task, task_runners


@task(
    retries=3,
    persist_result=False,
    cache_result_in_memory=False,
)
async def build_dataframe_async(id, index):
    print(f"Hello World = id = {id}, index = {index}")


@flow(
    cache_result_in_memory=False,
)
async def flow_async(id):
    for index in range(80):
        await build_dataframe_async(id, index)


if __name__ == "__main__":
    for index in range(100):
        print(f"____________________ Index = {index} ____________________")
        asyncio.run(flow_async(index))

        gc.collect(0)
        gc.collect(1)
        gc.collect(2)
        time.sleep(0.01)
        gc.collect()
        time.sleep(0.01)
```

`execute.sh`
```shell
#!/bin/bash

export PYTHONTRACEMALLOC=1

rm mprofile_*
# Current supported backends: 'psutil', 'psutil_pss', 'psutil_uss', 'posix', 'tracemalloc'. Defaults to 'psutil'
mprof run --backend psutil -C --python python ./test_flow.py
mprof plot
```

`./execute.sh`

produces

![obrazek](https://github.com/PrefectHQ/prefect/assets/31843161/f9d72a61-9554-4c38-97f8-4df27e3164f3)

I investigated this issue to see if there is any relation to the work in #12019, and I believe the leak here is different.

I've located code that reveals the leak, but I haven't understood why yet: https://github.com/PrefectHQ/prefect/blob/e1936a860061a6bfb0ec82e1bd8b43a9f7ea4d53/src/prefect/_internal/concurrency/calls.py#L377

If we use this `set_current_call` helper during async function call execution, we get a leak; remove that line and the leak vanishes. However, all that code does is set a contextvar: https://github.com/PrefectHQ/prefect/blob/e1936a860061a6bfb0ec82e1bd8b43a9f7ea4d53/src/prefect/_internal/concurrency/calls.py#L48-L54

I've invalidated several working theories about why this happens but will have to regroup and try to identify some different unknowns / possible conditions. Meanwhile, #12019 will help some egregious leaks with task arguments.
Several theories I have disproved:
- `Call` calls `_run_sync()`, entering its current call context manager, then calls `_run_async()` afterward, which enters *its* context manager (the same one). Is that the issue? 
  - No, if I turn the sync one off and leave the async on, it fails, sync on and async off, it works
  - So the problem is specifically the async use of a contextvar
- The problem is not the lack of an async context manager, same problem happens with async context manager, that doesn't matter anyway with a  contextvar
- The problem is not that async and sync calls use the same contextvar, I tried separate contextvars and context managers, same issue, again only with async calls (i.e. `_call_async()`).
- Doesn’t matter if the call happens within or outside of the cancel scope
- Running _run_async() in a copy of the call’s context doesn’t solve it either
In #12019 I switched to storing a weakref for our `current_call` contextvar because setting the contextvar within an async call appears to never free that memory. We don't need a strong reference at this point, so this works. Doing so improves the memory growth of the example script but doesn't halt it, so I'll continue to investigate other leaks.

Resolving patch:
diff --git a/src/prefect/_internal/concurrency/services.py b/src/prefect/_internal/concurrency/services.py
--- a/src/prefect/_internal/concurrency/services.py
+++ b/src/prefect/_internal/concurrency/services.py
@@ -91,6 +91,11 @@ def _stop(self, at_exit: bool = False):
             self._remove_instance()
             self._stopped = True
 
+            # Allow asyncio task to be garbage-collected. Its context may contain
+            # references to all Prefect Task calls made during a flow run, through
+            # EngineContext. Issue #10338.
+            self._task = None
+
             # Signal completion to the loop
             self._queue.put_nowait(None)
 
diff --git a/src/prefect/_internal/concurrency/threads.py b/src/prefect/_internal/concurrency/threads.py
--- a/src/prefect/_internal/concurrency/threads.py
+++ b/src/prefect/_internal/concurrency/threads.py
@@ -214,6 +214,9 @@ async def _run_until_shutdown(self):
         for call in self._on_shutdown:
             await self._run_call(call)
 
+        # Empty the list to allow calls to be garbage collected. Issue #10338.
+        self._on_shutdown = []
+
     async def _run_call(self, call: Call) -> None:
         task = call.run()
         if task is not None:


Test patch:
diff --git a/tests/_internal/concurrency/test_services.py b/tests/_internal/concurrency/test_services.py
--- a/tests/_internal/concurrency/test_services.py
+++ b/tests/_internal/concurrency/test_services.py
@@ -188,6 +188,15 @@ def test_drain_safe_to_call_multiple_times():
     )
 
 
+def test_drain_clears_asyncio_task():
+    instance = MockService.instance()
+    instance.send(1)
+
+    assert instance._task is not None
+    MockService.drain_all()
+    assert instance._task is None
+
+
 def test_send_many_threads():
     def on_thread(i):
         MockService.instance().send(i)
diff --git a/tests/_internal/concurrency/test_threads.py b/tests/_internal/concurrency/test_threads.py
--- a/tests/_internal/concurrency/test_threads.py
+++ b/tests/_internal/concurrency/test_threads.py
@@ -66,6 +66,18 @@ def test_event_loop_thread_with_on_shutdown_hooks():
     mock.assert_has_awaits(call(i) for i in range(5))
 
 
+def test_event_loop_thread_clears_shutdown_hooks_after_shutdown():
+    event_loop_thread = EventLoopThread()
+
+    event_loop_thread.start()
+    for i in range(5):
+        event_loop_thread.add_shutdown_call(Call.new(AsyncMock(), i))
+    event_loop_thread.shutdown()
+    event_loop_thread.thread.join()
+
+    assert not event_loop_thread._on_shutdown
+
+
 @pytest.mark.parametrize("thread_cls", [WorkerThread, EventLoopThread])
 @pytest.mark.parametrize("daemon", [True, False])
 def test_thread_daemon(daemon, thread_cls):