Add message with debugging info to Cancelled (#3256)

---------

Co-authored-by: CoolCat467 <[email protected]>
Co-authored-by: A5Rocks <[email protected]>

Author: 3 people

newsfragments/3232.feature.rst

@@ -0,0 +1 @@
+:exc:`Cancelled` strings can now display the source and reason for a cancellation. Trio-internal sources of cancellation will set this string, and :meth:`CancelScope.cancel` now has a ``reason`` string parameter that can be used to attach info to any :exc:`Cancelled` to help in debugging.

src/trio/_channel.py

@@ -547,7 +547,9 @@ async def context_manager(
                     yield wrapped_recv_chan
                 # User has exited context manager, cancel to immediately close the
                 # abandoned generator if it's still alive.
-                nursery.cancel_scope.cancel()
+                nursery.cancel_scope.cancel(
+                    "exited trio.as_safe_channel context manager"
+                )
         except BaseExceptionGroup as eg:
             try:
                 raise_single_exception_from_group(eg)

src/trio/_core/_asyncgens.py

@@ -230,7 +230,9 @@ async def _finalize_one(
             # with an exception, not even a Cancelled. The inside
             # is cancelled so there's no deadlock risk.
             with _core.CancelScope(shield=True) as cancel_scope:
-                cancel_scope.cancel()
+                cancel_scope.cancel(
+                    reason="disallow async work when closing async generators during trio shutdown"
+                )
                 await agen.aclose()
         except BaseException:
             ASYNCGEN_LOGGER.exception(

src/trio/_core/_exceptions.py

@@ -1,12 +1,26 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from functools import partial
+from typing import TYPE_CHECKING, Literal
+
+import attrs
 
 from trio._util import NoPublicConstructor, final
 
 if TYPE_CHECKING:
     from collections.abc import Callable
 
+    from typing_extensions import Self, TypeAlias
+
+CancelReasonLiteral: TypeAlias = Literal[
+    "KeyboardInterrupt",
+    "deadline",
+    "explicit",
+    "nursery",
+    "shutdown",
+    "unknown",
+]
+
 
 class TrioInternalError(Exception):
     """Raised by :func:`run` if we encounter a bug in Trio, or (possibly) a
@@ -34,6 +48,7 @@ class WouldBlock(Exception):
 
 
 @final
+@attrs.define(eq=False, kw_only=True)
 class Cancelled(BaseException, metaclass=NoPublicConstructor):
     """Raised by blocking calls if the surrounding scope has been cancelled.
 
@@ -67,11 +82,42 @@ class Cancelled(BaseException, metaclass=NoPublicConstructor):
 
     """
 
+    source: CancelReasonLiteral
+    # repr(Task), so as to avoid gc troubles from holding a reference
+    source_task: str | None = None
+    reason: str | None = None
+
     def __str__(self) -> str:
-        return "Cancelled"
+        return (
+            f"cancelled due to {self.source}"
+            + ("" if self.reason is None else f" with reason {self.reason!r}")
+            + ("" if self.source_task is None else f" from task {self.source_task}")
+        )
 
     def __reduce__(self) -> tuple[Callable[[], Cancelled], tuple[()]]:
-        return (Cancelled._create, ())
+        # The `__reduce__` tuple does not support directly passing kwargs, and the
+        # kwargs are required so we can't use the third item for adding to __dict__,
+        # so we use partial.
+        return (
+            partial(
+                Cancelled._create,
+                source=self.source,
+                source_task=self.source_task,
+                reason=self.reason,
+            ),
+            (),
+        )
+
+    if TYPE_CHECKING:
+        # for type checking on internal code
+        @classmethod
+        def _create(
+            cls,
+            *,
+            source: CancelReasonLiteral,
+            source_task: str | None = None,
+            reason: str | None = None,
+        ) -> Self: ...
 
 
 class BusyResourceError(Exception):

src/trio/_core/_run.py

@@ -36,7 +36,12 @@
 from ._asyncgens import AsyncGenerators
 from ._concat_tb import concat_tb
 from ._entry_queue import EntryQueue, TrioToken
-from ._exceptions import Cancelled, RunFinishedError, TrioInternalError
+from ._exceptions import (
+    Cancelled,
+    CancelReasonLiteral,
+    RunFinishedError,
+    TrioInternalError,
+)
 from ._instrumentation import Instruments
 from ._ki import KIManager, enable_ki_protection
 from ._parking_lot import GLOBAL_PARKING_LOT_BREAKER
@@ -305,7 +310,7 @@ def expire(self, now: float) -> bool:
                 did_something = True
                 # This implicitly calls self.remove(), so we don't need to
                 # decrement _active here
-                cancel_scope.cancel()
+                cancel_scope._cancel(CancelReason(source="deadline"))
         # If we've accumulated too many stale entries, then prune the heap to
         # keep it under control. (We only do this occasionally in a batch, to
         # keep the amortized cost down)
@@ -314,6 +319,20 @@ def expire(self, now: float) -> bool:
         return did_something
 
 
+@attrs.define
+class CancelReason:
+    """Attached to a :class:`CancelScope` upon cancellation with details of the source of the
+    cancellation, which is then used to construct the string in a :exc:`Cancelled`.
+    Users can pass a ``reason`` str to :meth:`CancelScope.cancel` to set it.
+
+    Not publicly exported or documented.
+    """
+
+    source: CancelReasonLiteral
+    source_task: str | None = None
+    reason: str | None = None
+
+
 @attrs.define(eq=False)
 class CancelStatus:
     """Tracks the cancellation status for a contiguous extent
@@ -468,6 +487,14 @@ def recalculate(self) -> None:
                 or current.parent_cancellation_is_visible_to_us
             )
             if new_state != current.effectively_cancelled:
+                if (
+                    current._scope._cancel_reason is None
+                    and current.parent_cancellation_is_visible_to_us
+                ):
+                    assert current._parent is not None
+                    current._scope._cancel_reason = (
+                        current._parent._scope._cancel_reason
+                    )
                 current.effectively_cancelled = new_state
                 if new_state:
                     for task in current._tasks:
@@ -558,6 +585,8 @@ class CancelScope:
     _cancel_called: bool = attrs.field(default=False, init=False)
     cancelled_caught: bool = attrs.field(default=False, init=False)
 
+    _cancel_reason: CancelReason | None = attrs.field(default=None, init=False)
+
     # Constructor arguments:
     _relative_deadline: float = attrs.field(
         default=inf,
@@ -594,7 +623,7 @@ def __enter__(self) -> Self:
             self._relative_deadline = inf
 
         if current_time() >= self._deadline:
-            self.cancel()
+            self._cancel(CancelReason(source="deadline"))
         with self._might_change_registered_deadline():
             self._cancel_status = CancelStatus(scope=self, parent=task._cancel_status)
             task._activate_cancel_status(self._cancel_status)
@@ -883,19 +912,42 @@ def shield(self, new_value: bool) -> None:
             self._cancel_status.recalculate()
 
     @enable_ki_protection
-    def cancel(self) -> None:
-        """Cancels this scope immediately.
-
-        This method is idempotent, i.e., if the scope was already
-        cancelled then this method silently does nothing.
+    def _cancel(self, cancel_reason: CancelReason | None) -> None:
+        """Internal sources of cancellation should use this instead of :meth:`cancel`
+        in order to set a more detailed :class:`CancelReason`
+        Helper or high-level functions can use `cancel`.
         """
         if self._cancel_called:
             return
+
+        if self._cancel_reason is None:
+            self._cancel_reason = cancel_reason
+
         with self._might_change_registered_deadline():
             self._cancel_called = True
+
         if self._cancel_status is not None:
             self._cancel_status.recalculate()
 
+    @enable_ki_protection
+    def cancel(self, reason: str | None = None) -> None:
+        """Cancels this scope immediately.
+
+        The optional ``reason`` argument accepts a string, which will be attached to
+        any resulting :exc:`Cancelled` exception to help you understand where that
+        cancellation is coming from and why it happened.
+
+        This method is idempotent, i.e., if the scope was already
+        cancelled then this method silently does nothing.
+        """
+        try:
+            current_task = repr(_core.current_task())
+        except RuntimeError:
+            current_task = None
+        self._cancel(
+            CancelReason(reason=reason, source="explicit", source_task=current_task)
+        )
+
     @property
     def cancel_called(self) -> bool:
         """Readonly :class:`bool`. Records whether cancellation has been
@@ -924,7 +976,7 @@ def cancel_called(self) -> bool:
             # but it makes the value returned by cancel_called more
             # closely match expectations.
             if not self._cancel_called and current_time() >= self._deadline:
-                self.cancel()
+                self._cancel(CancelReason(source="deadline"))
         return self._cancel_called
 
 
@@ -1192,9 +1244,9 @@ def parent_task(self) -> Task:
         "(`~trio.lowlevel.Task`):  The Task that opened this nursery."
         return self._parent_task
 
-    def _add_exc(self, exc: BaseException) -> None:
+    def _add_exc(self, exc: BaseException, reason: CancelReason | None) -> None:
         self._pending_excs.append(exc)
-        self.cancel_scope.cancel()
+        self.cancel_scope._cancel(reason)
 
     def _check_nursery_closed(self) -> None:
         if not any([self._nested_child_running, self._children, self._pending_starts]):
@@ -1210,7 +1262,14 @@ def _child_finished(
     ) -> None:
         self._children.remove(task)
         if isinstance(outcome, Error):
-            self._add_exc(outcome.error)
+            self._add_exc(
+                outcome.error,
+                CancelReason(
+                    source="nursery",
+                    source_task=repr(task),
+                    reason=f"child task raised exception {outcome.error!r}",
+                ),
+            )
         self._check_nursery_closed()
 
     async def _nested_child_finished(
@@ -1220,7 +1279,14 @@ async def _nested_child_finished(
         # Returns ExceptionGroup instance (or any exception if the nursery is in loose mode
         # and there is just one contained exception) if there are pending exceptions
         if nested_child_exc is not None:
-            self._add_exc(nested_child_exc)
+            self._add_exc(
+                nested_child_exc,
+                reason=CancelReason(
+                    source="nursery",
+                    source_task=repr(self._parent_task),
+                    reason=f"Code block inside nursery contextmanager raised exception {nested_child_exc!r}",
+                ),
+            )
         self._nested_child_running = False
         self._check_nursery_closed()
 
@@ -1231,7 +1297,13 @@ async def _nested_child_finished(
             def aborted(raise_cancel: _core.RaiseCancelT) -> Abort:
                 exn = capture(raise_cancel).error
                 if not isinstance(exn, Cancelled):
-                    self._add_exc(exn)
+                    self._add_exc(
+                        exn,
+                        CancelReason(
+                            source="KeyboardInterrupt",
+                            source_task=repr(self._parent_task),
+                        ),
+                    )
                 # see test_cancel_scope_exit_doesnt_create_cyclic_garbage
                 del exn  # prevent cyclic garbage creation
                 return Abort.FAILED
@@ -1245,7 +1317,8 @@ def aborted(raise_cancel: _core.RaiseCancelT) -> Abort:
             try:
                 await cancel_shielded_checkpoint()
             except BaseException as exc:
-                self._add_exc(exc)
+                # there's no children to cancel, so don't need to supply cancel reason
+                self._add_exc(exc, reason=None)
 
         popped = self._parent_task._child_nurseries.pop()
         assert popped is self
@@ -1575,8 +1648,17 @@ def _attempt_delivery_of_any_pending_cancel(self) -> None:
         if not self._cancel_status.effectively_cancelled:
             return
 
+        reason = self._cancel_status._scope._cancel_reason
+
         def raise_cancel() -> NoReturn:
-            raise Cancelled._create()
+            if reason is None:
+                raise Cancelled._create(source="unknown", reason="misnesting")
+            else:
+                raise Cancelled._create(
+                    source=reason.source,
+                    reason=reason.reason,
+                    source_task=reason.source_task,
+                )
 
         self._attempt_abort(raise_cancel)
 
@@ -2075,15 +2157,27 @@ async def init(
                     )
 
                 # Main task is done; start shutting down system tasks
-                self.system_nursery.cancel_scope.cancel()
+                self.system_nursery.cancel_scope._cancel(
+                    CancelReason(
+                        source="shutdown",
+                        reason="main task done, shutting down system tasks",
+                        source_task=repr(self.init_task),
+                    )
+                )
 
             # System nursery is closed; finalize remaining async generators
             await self.asyncgens.finalize_remaining(self)
 
             # There are no more asyncgens, which means no more user-provided
             # code except possibly run_sync_soon callbacks. It's finally safe
             # to stop the run_sync_soon task and exit run().
-            run_sync_soon_nursery.cancel_scope.cancel()
+            run_sync_soon_nursery.cancel_scope._cancel(
+                CancelReason(
+                    source="shutdown",
+                    reason="main task done, shutting down run_sync_soon callbacks",
+                    source_task=repr(self.init_task),
+                )
+            )
 
     ################
     # Outside context problems
@@ -2926,7 +3020,18 @@ async def checkpoint() -> None:
     if task._cancel_status.effectively_cancelled or (
         task is task._runner.main_task and task._runner.ki_pending
     ):
-        with CancelScope(deadline=-inf):
+        cs = CancelScope(deadline=-inf)
+        if (
+            task._cancel_status._scope._cancel_reason is None
+            and task is task._runner.main_task
+            and task._runner.ki_pending
+        ):
+            task._cancel_status._scope._cancel_reason = CancelReason(
+                source="KeyboardInterrupt"
+            )
+        assert task._cancel_status._scope._cancel_reason is not None
+        cs._cancel_reason = task._cancel_status._scope._cancel_reason
+        with cs:
             await _core.wait_task_rescheduled(lambda _: _core.Abort.SUCCEEDED)