Repository: PrefectHQ/prefect

Issue:
Task blocking during async concurrent subflows
### 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 async concurrent subflows with tasks, blocking appears to occur and one tasks seems to need to finish before the next one starts:

Example (1s sleep):
![Screenshot 2024-01-30 at 3 19 40 PM](https://github.com/PrefectHQ/prefect/assets/40362401/6abb717b-9221-48d7-94a7-594f36c92a42)

(5s sleep, more pronounced)

![Screenshot 2024-01-30 at 3 24 10 PM](https://github.com/PrefectHQ/prefect/assets/40362401/cb28898d-5119-488d-970b-d266a25596c0)

removing the `@flow` decorator from the subflow seems to give the correct result, indicating it _seems_ to be an issue with async subflows only (although not positive):
![Screenshot 2024-01-30 at 3 43 15 PM](https://github.com/PrefectHQ/prefect/assets/40362401/acf7bcad-f2df-4ac5-ae6d-2bf668bad731)


### Reproduction

```python3
import asyncio
from prefect import task, flow


@task(task_run_name="task-{name}-1")
async def task_1(name: str):
    await asyncio.sleep(1)


@task(task_run_name="task-{name}-2")
async def task_2(name: str):
    await asyncio.sleep(1)


@flow(flow_run_name="my_flow-{name}")
async def sub_flow(name: str):
    await task_1(name)
    await task_2(name)


@flow
async def main_flow():
    futures = []

    for name in ['A', 'B', 'C', 'D', 'E']:
        futures.append(sub_flow(name))

    await asyncio.gather(*futures)

if __name__ == "__main__":
    asyncio.run(main_flow())
```


### Error

_No response_

### Versions

```Text
Version:             2.15.0
API version:         0.8.4
Python version:      3.10.9
Git commit:          e90804b8
Built:               Thu, Feb 15, 2024 2:54 PM
OS/Arch:             darwin/x86_64
Profile:             prod-test
Server type:         cloud
```


### Additional context

_No response_


Discussion:
This is also true of flows run by async subflows:
![image](https://github.com/PrefectHQ/prefect/assets/146098880/6ec0a566-f217-4581-9dc0-f35f524f06c9)

```python
import asyncio
import time
from prefect import flow


@flow
async def subflow_1():
    print("starting subflow 1")
    await asyncio.sleep(30)
    print("subflow 1 finished")


@flow
async def subflow_2():
    print("starting subflow 2")
    await asyncio.sleep(30)
    print("subflow 2 finished")


@flow
async def middle_flow():
    print("starting middle flow")
    await subflow_1()
    await subflow_2()
    print("middle flow finished")


@flow
async def parent_flow():
    print("starting parent flow")
    await asyncio.gather(middle_flow(), middle_flow())
    print("parent flow finished")


if __name__ == "__main__":
    start = time.perf_counter()
    asyncio.run(parent_flow())
    end = time.perf_counter() - start
    print(f"Program finished in {end:0.2f} seconds.")
```

Resolving patch:
diff --git a/src/prefect/_internal/concurrency/api.py b/src/prefect/_internal/concurrency/api.py
--- a/src/prefect/_internal/concurrency/api.py
+++ b/src/prefect/_internal/concurrency/api.py
@@ -1,6 +1,7 @@
 """
 Primary developer-facing API for concurrency management.
 """
+
 import abc
 import asyncio
 import concurrent.futures
@@ -18,6 +19,7 @@
 
 from typing_extensions import ParamSpec
 
+from prefect._internal.concurrency.calls import get_current_call
 from prefect._internal.concurrency.threads import (
     WorkerThread,
     get_global_loop,
@@ -104,40 +106,6 @@ def call_soon_in_loop_thread(
         runner.submit(call)
         return call
 
-    @staticmethod
-    def call_soon_in_waiting_thread(
-        __call: Union[Callable[[], T], Call[T]],
-        thread: threading.Thread,
-        timeout: Optional[float] = None,
-    ) -> Call[T]:
-        """
-        Schedule a call for execution in the thread that is waiting for the current
-        call.
-
-        Returns the submitted call.
-        """
-        call = _cast_to_call(__call)
-        waiter = get_waiter_for_thread(thread)
-        if waiter is None:
-            raise RuntimeError(f"No waiter found for thread {thread}.")
-
-        call.set_timeout(timeout)
-        waiter.submit(call)
-        return call
-
-    @staticmethod
-    def call_in_waiting_thread(
-        __call: Union[Callable[[], T], Call[T]],
-        thread: threading.Thread,
-        timeout: Optional[float] = None,
-    ) -> T:
-        """
-        Run a call in the thread that is waiting for the current call.
-
-        Returns the result of the call.
-        """
-        raise NotImplementedError()
-
     @staticmethod
     def call_in_new_thread(
         __call: Union[Callable[[], T], Call[T]], timeout: Optional[float] = None
@@ -196,13 +164,20 @@ async def wait_for_call_in_new_thread(
         return call.result()
 
     @staticmethod
-    def call_in_waiting_thread(
+    def call_soon_in_waiting_thread(
         __call: Union[Callable[[], T], Call[T]],
         thread: threading.Thread,
         timeout: Optional[float] = None,
-    ) -> Awaitable[T]:
-        call = _base.call_soon_in_waiting_thread(__call, thread, timeout=timeout)
-        return call.aresult()
+    ) -> Call[T]:
+        call = _cast_to_call(__call)
+        parent_call = get_current_call()
+        waiter = get_waiter_for_thread(thread, parent_call)
+        if waiter is None:
+            raise RuntimeError(f"No waiter found for thread {thread}.")
+
+        call.set_timeout(timeout)
+        waiter.submit(call)
+        return call
 
     @staticmethod
     def call_in_new_thread(
@@ -257,13 +232,19 @@ def wait_for_call_in_new_thread(
         return call.result()
 
     @staticmethod
-    def call_in_waiting_thread(
+    def call_soon_in_waiting_thread(
         __call: Union[Callable[[], T], Call[T]],
         thread: threading.Thread,
         timeout: Optional[float] = None,
-    ) -> T:
-        call = _base.call_soon_in_waiting_thread(__call, thread, timeout=timeout)
-        return call.result()
+    ) -> Call[T]:
+        call = _cast_to_call(__call)
+        waiter = get_waiter_for_thread(thread)
+        if waiter is None:
+            raise RuntimeError(f"No waiter found for thread {thread}.")
+
+        call.set_timeout(timeout)
+        waiter.submit(call)
+        return call
 
     @staticmethod
     def call_in_new_thread(
diff --git a/src/prefect/_internal/concurrency/waiters.py b/src/prefect/_internal/concurrency/waiters.py
--- a/src/prefect/_internal/concurrency/waiters.py
+++ b/src/prefect/_internal/concurrency/waiters.py
@@ -9,9 +9,9 @@
 import inspect
 import queue
 import threading
-import weakref
 from collections import deque
 from typing import Awaitable, Generic, List, Optional, TypeVar, Union
+from weakref import WeakKeyDictionary
 
 import anyio
 
@@ -24,34 +24,37 @@
 
 
 # Waiters are stored in a stack for each thread
-_WAITERS_BY_THREAD: "weakref.WeakKeyDictionary[threading.Thread, deque[Waiter]]" = (
-    weakref.WeakKeyDictionary()
+_WAITERS_BY_THREAD: "WeakKeyDictionary[threading.Thread, deque[Waiter]]" = (
+    WeakKeyDictionary()
 )
 
 
-def get_waiter_for_thread(thread: threading.Thread) -> Optional["Waiter"]:
+def get_waiter_for_thread(
+    thread: threading.Thread, parent_call: Optional[Call] = None
+) -> Optional["Waiter"]:
     """
-    Get the current waiter for a thread.
+    Get the current waiter for a thread and an optional parent call.
 
-    Returns `None` if one does not exist.
+    To avoid assigning outer callbacks to inner waiters in the case of nested calls,
+    the parent call is used to determine which waiter to return. If a parent call is
+    not provided, we return the most recently created waiter (last in the stack).
+
+    see https://github.com/PrefectHQ/prefect/issues/12036
+
+    Returns `None` if no active waiter is found for the thread.
     """
-    waiters = _WAITERS_BY_THREAD.get(thread)
-
-    if waiters:
-        idx = -1
-        while abs(idx) <= len(waiters):
-            try:
-                waiter = waiters[idx]
-                if not waiter.call_is_done():
-                    return waiter
-                idx = idx - 1
-            # It is possible that items are being added or removed
-            # from the deque, so the index we're using may not always
-            # be valid.
-            except IndexError:
-                break
 
-    return None
+    waiters: "Optional[deque[Waiter]]" = _WAITERS_BY_THREAD.get(thread)
+
+    if waiters and (active_waiters := [w for w in waiters if not w.call_is_done()]):
+        if parent_call and (
+            matching_waiter := next(
+                (w for w in active_waiters if w._call == parent_call), None
+            )
+        ):  # if exists an active waiter responsible for the parent call, return it
+            return matching_waiter
+        else:  # otherwise, return the most recently created waiter
+            return active_waiters[-1]
 
 
 def add_waiter_for_thread(waiter: "Waiter", thread: threading.Thread):
diff --git a/src/prefect/engine.py b/src/prefect/engine.py
--- a/src/prefect/engine.py
+++ b/src/prefect/engine.py
@@ -857,7 +857,9 @@ async def orchestrate_flow_run(
                     )
                 ):
                     from_async.call_soon_in_waiting_thread(
-                        flow_call, thread=user_thread, timeout=flow.timeout_seconds
+                        flow_call,
+                        thread=user_thread,
+                        timeout=flow.timeout_seconds,
                     )
                 else:
                     from_async.call_soon_in_new_thread(


Test patch:
diff --git a/tests/_internal/concurrency/test_api.py b/tests/_internal/concurrency/test_api.py
--- a/tests/_internal/concurrency/test_api.py
+++ b/tests/_internal/concurrency/test_api.py
@@ -304,34 +304,6 @@ def test_from_sync_call_in_loop_thread(work):
     wait_for_global_loop_exit()
 
 
-@pytest.mark.parametrize("work", [identity, aidentity])
-async def test_from_async_call_in_waiting_thread_from_worker_thread(work):
-    async def worker(parent_thread):
-        result = await from_async.call_in_waiting_thread(
-            create_call(work, 1), parent_thread
-        )
-        assert result == 1
-        return 2
-
-    result = await from_async.wait_for_call_in_new_thread(
-        create_call(worker, threading.current_thread())
-    )
-    assert result == 2
-
-
-@pytest.mark.parametrize("work", [identity, aidentity])
-def test_from_sync_call_in_waiting_thread_from_worker_thread(work):
-    def worker(parent_thread):
-        result = from_sync.call_in_waiting_thread(create_call(work, 1), parent_thread)
-        assert result == 1
-        return 2
-
-    result = from_sync.wait_for_call_in_new_thread(
-        create_call(worker, threading.current_thread())
-    )
-    assert result
-
-
 def test_from_sync_call_in_loop_thread_from_loop_thread():
     def worker():
         # Here, a call is submitted to the loop thread from the loop thread which would
diff --git a/tests/_internal/concurrency/test_waiters.py b/tests/_internal/concurrency/test_waiters.py
--- a/tests/_internal/concurrency/test_waiters.py
+++ b/tests/_internal/concurrency/test_waiters.py
@@ -52,9 +52,29 @@ def call_is_done():
     waiter.call_is_done = call_is_done
 
     waiter_for_thread = get_waiter_for_thread(threading.current_thread())
+
     assert waiter_for_thread is None
 
 
+@pytest.mark.parametrize("waiter_cls", [AsyncWaiter, SyncWaiter])
+def test_get_waiter_for_thread_with_parent_call(waiter_cls):
+    parent_call = Call.new(identity, 1)
+    child_call = Call.new(identity, 2)
+
+    parent_waiter = waiter_cls(parent_call)
+    child_waiter = waiter_cls(child_call)
+
+    parent_waiter.submit(child_call)
+
+    # Assert that get_waiter_for_thread returns the parent waiter when the parent call is provided
+    assert (
+        get_waiter_for_thread(threading.current_thread(), parent_call) == parent_waiter
+    )
+
+    # Assert that get_waiter_for_thread returns the child waiter (most recently created) when no parent call is provided
+    assert get_waiter_for_thread(threading.current_thread()) == child_waiter
+
+
 @pytest.mark.parametrize("cls", [AsyncWaiter, SyncWaiter])
 async def test_waiter_repr(cls):
     current_thread_name = threading.current_thread().name
diff --git a/tests/conftest.py b/tests/conftest.py
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -24,7 +24,7 @@
 import sys
 import tempfile
 from pathlib import Path
-from typing import Generator, Optional
+from typing import AsyncGenerator, Optional
 from urllib.parse import urlsplit, urlunsplit
 
 import asyncpg
@@ -374,7 +374,7 @@ def safety_check_settings():
 @pytest.fixture(scope="session", autouse=True)
 async def generate_test_database_connection_url(
     worker_id: str,
-) -> Generator[Optional[str], None, None]:
+) -> AsyncGenerator[Optional[str], None]:
     """Prepares an alternative test database URL, if necessary, for the current
     connection URL.
 
diff --git a/tests/public/flows/test_flow_crashes.py b/tests/public/flows/test_flow_crashes.py
--- a/tests/public/flows/test_flow_crashes.py
+++ b/tests/public/flows/test_flow_crashes.py
@@ -6,6 +6,7 @@
 be reraised to prevent Prefect from swallowing important signals like keyboard
 interrupts.
 """
+
 import asyncio
 import os
 import signal