Repository: agronholm/anyio

Issue:
Memory object streams deadlock on asyncio
### Things to check first

- [X] I have searched the existing issues and didn't find my bug already reported there

- [X] I have checked that my bug is still present in the latest release


### AnyIO version

master (439951d9babf7615f0720fd43c93c2a695549bd4)

### Python version

3.12.5 (CPython)

### What happened?

Some applications started deadlocking recently when updating AnyIO to 4.4.0. It bisects to #735.

(Aside: Apologies I did not follow-up in the discussions about #728 in May. I got unexpectedly busy at the time :/ )

### How can we reproduce the bug?

(This reproducer could be made smaller but I think it would sacrifice readability)

```python
from __future__ import annotations

from collections.abc import AsyncGenerator
from contextlib import asynccontextmanager
from types import TracebackType
from typing import Final
from typing import TypeVar
from typing import cast

from typing_extensions import Self

import anyio
import anyio.lowlevel
from anyio import CancelScope
from anyio.abc import TaskStatus
from anyio.streams.memory import MemoryObjectSendStream

E = TypeVar("E", bound=BaseException)


class A:
    # `A`'s API is pretty similar to `ObjectSendStream[str]`, except it requires
    # `__aexit__` (it doesn't support `aclose`).

    def __init__(self) -> None:
        super().__init__()
        self._acm: Final = self._acm_impl()

    async def __aenter__(self) -> Self:
        # TODO: Mypy bug?
        return await self._acm.__aenter__()  # type: ignore[return-value]

    async def __aexit__(
        self,
        exc_type: type[E] | None,
        exc_value: E | None,
        traceback: TracebackType | None,
    ) -> bool | None:
        return await self._acm.__aexit__(exc_type, exc_value, traceback)

    @asynccontextmanager
    async def _acm_impl(self) -> AsyncGenerator[Self]:
        async with anyio.create_task_group() as task_group:
            outgoing_service_cancel_scope, self._outgoing_send_stream = cast(
                tuple[CancelScope, MemoryObjectSendStream[str]],
                await task_group.start(self._serve_outgoing),
            )
            with self._outgoing_send_stream:
                try:
                    yield self
                finally:
                    # The `async with` body has exited, so if there is a pending
                    # cancellation request it can be allowed to reach the service task
                    # now.
                    outgoing_service_cancel_scope.shield = False
                    # If the `async with A()` isn't already in a cancelled scope, we also
                    # need to request the service task to shut down.
                    task_group.cancel_scope.cancel()

    async def _serve_outgoing(
        self,
        *,
        task_status: TaskStatus[tuple[CancelScope, MemoryObjectSendStream[str]]],
    ) -> None:
        (
            outgoing_send_stream,
            outgoing_receive_stream,
        ) = anyio.create_memory_object_stream[str](0)
        # The cancel scope here is a service task group pattern
        # (https://github.com/python-trio/trio/issues/1521). This service task must not
        # receive cancellation until the body of the `async with A()` block has finished
        # exiting first, because this service task must remain available in any
        # `finally`/etc. blocks of the `async with A()` body.
        with outgoing_receive_stream, CancelScope(shield=True) as cancel_scope:
            task_status.started((cancel_scope, outgoing_send_stream))
            async for msg in outgoing_receive_stream:
                print(f"outgoing stream driver handling outgoing message: {msg!r}")
                await anyio.lowlevel.checkpoint()

    async def send(self, msg: str, /) -> None:
        # Note: If `_serve_outgoing` crashes due to external causes, this will raise
        # `BrokenResourceError`.
        await self._outgoing_send_stream.send(msg)


async def main() -> None:
    with CancelScope() as cancel_scope:
        async with A() as a:
            try:
                await a.send("message during normal operation!")
                ...
                # Suppose that at some point there's a cancellation request from above
                # (e.g. a signal handler requesting process shutdown):
                cancel_scope.cancel()
            finally:
                # A bit of time is needed here for the scheduling order that triggers
                # the deadlock to be hit. (In the real code, there are more awaits here
                # and the deadlock is hit (empirically) every time.)
                n = 0
                while not a._outgoing_send_stream._state.waiting_receivers:
                    n += 1
                    await anyio.lowlevel.cancel_shielded_checkpoint()
                print(f"{n} scheduling rounds")

                with CancelScope(shield=True):
                    await a.send("message during cleanup!")


# Works on Trio; deadlocks on asyncio.
anyio.run(main, backend="asyncio")
```

Discussion:

Patch:
diff --git a/docs/versionhistory.rst b/docs/versionhistory.rst
index 5a42d0b..f6cbf14 100644
--- a/docs/versionhistory.rst
+++ b/docs/versionhistory.rst
@@ -16,6 +16,8 @@ This library adheres to `Semantic Versioning 2.0 <http://semver.org/>`_.
 - Fixed ``to_process.run_sync()`` failing to initialize if ``__main__.__file__`` pointed
   to a file in a nonexistent directory
   (`#696 <https://github.com/agronholm/anyio/issues/696>`_)
+- Fixed ``TaskInfo.has_pending_cancellation()`` on asyncio not respecting shielded
+  scopes (`#771 <https://github.com/agronholm/anyio/issues/771>`_; PR by @gschaffner)
 
 **4.4.0**
 
diff --git a/src/anyio/_backends/_asyncio.py b/src/anyio/_backends/_asyncio.py
index e191f02..493f8ea 100644
--- a/src/anyio/_backends/_asyncio.py
+++ b/src/anyio/_backends/_asyncio.py
@@ -1861,7 +1861,9 @@ class AsyncIOTaskInfo(TaskInfo):
 
         if task_state := _task_states.get(task):
             if cancel_scope := task_state.cancel_scope:
-                return cancel_scope.cancel_called or cancel_scope._parent_cancelled()
+                return cancel_scope.cancel_called or (
+                    not cancel_scope.shield and cancel_scope._parent_cancelled()
+                )
 
         return False
 


Test patch:
diff --git a/tests/test_taskgroups.py b/tests/test_taskgroups.py
index 2487ff9..a460361 100644
--- a/tests/test_taskgroups.py
+++ b/tests/test_taskgroups.py
@@ -628,9 +628,11 @@ async def test_cancel_from_shielded_scope() -> None:
             assert inner_scope.shield
             tg.cancel_scope.cancel()
             assert current_effective_deadline() == math.inf
+            assert not get_current_task().has_pending_cancellation()
             await checkpoint()
 
         assert current_effective_deadline() == -math.inf
+        assert get_current_task().has_pending_cancellation()
 
         with pytest.raises(get_cancelled_exc_class()):
             await sleep(0.01)
@@ -644,6 +646,7 @@ async def test_cancel_shielded_scope() -> None:
         assert cancel_scope.shield
         cancel_scope.cancel()
         assert current_effective_deadline() == -math.inf
+        assert get_current_task().has_pending_cancellation()
 
         with pytest.raises(get_cancelled_exc_class()):
             await checkpoint()