From f0ddcff3ba2465312ff1b312da4823c9c61d909a Mon Sep 17 00:00:00 2001 From: Arthur Tacca Date: Tue, 14 May 2024 17:31:49 +0100 Subject: [PATCH] Allow cancellation to take effect even after a call to reschedule() --- src/trio/_core/_generated_run.py | 15 +++++++++++++-- src/trio/_core/_run.py | 24 ++++++++++++++++++++---- src/trio/_sync.py | 4 ++-- 3 files changed, 35 insertions(+), 8 deletions(-) diff --git a/src/trio/_core/_generated_run.py b/src/trio/_core/_generated_run.py index 47febd3ca1..af776f0d73 100644 --- a/src/trio/_core/_generated_run.py +++ b/src/trio/_core/_generated_run.py @@ -101,7 +101,9 @@ def current_root_task() -> Task | None: raise RuntimeError("must be called from async context") from None -def reschedule(task: Task, next_send: Outcome[Any] = _NO_SEND) -> None: +def reschedule( + task: Task, next_send: Outcome[Any] = _NO_SEND, allow_abort: bool = False +) -> None: """Reschedule the given task with the given :class:`outcome.Outcome`. @@ -112,16 +114,25 @@ def reschedule(task: Task, next_send: Outcome[Any] = _NO_SEND) -> None: returning :data:`Abort.SUCCEEDED` from an abort callback is equivalent to calling :func:`reschedule` once.) + As an exception, if ``allow_abort`` is ``True`` then the abort + callback may still be called up until the rescheduled coroutine + actually resumes. In that case, there are effectively two calls, with + the abort always happening second. Note that any outcome passed to + this function (as the ``next_send`` argument) will be silently + discarded if that happens. + Args: task (trio.lowlevel.Task): the task to be rescheduled. Must be blocked in a call to :func:`wait_task_rescheduled`. next_send (outcome.Outcome): the value (or error) to return (or raise) from :func:`wait_task_rescheduled`. + allow_abort (bool): whether the abort function may still be called + after this function is called. """ locals()[LOCALS_KEY_KI_PROTECTION_ENABLED] = True try: - return GLOBAL_RUN_CONTEXT.runner.reschedule(task, next_send) + return GLOBAL_RUN_CONTEXT.runner.reschedule(task, next_send, allow_abort) except AttributeError: raise RuntimeError("must be called from async context") from None diff --git a/src/trio/_core/_run.py b/src/trio/_core/_run.py index a8b632ce53..f272511b0d 100644 --- a/src/trio/_core/_run.py +++ b/src/trio/_core/_run.py @@ -1441,7 +1441,13 @@ def _attempt_abort(self, raise_cancel: _core.RaiseCancelT) -> None: # whether we succeeded or failed. self._abort_func = None if success is Abort.SUCCEEDED: - self._runner.reschedule(self, capture(raise_cancel)) + if self._next_send is None: + self._runner.reschedule(self, capture(raise_cancel)) + else: + # reschedule(..., allow_abort=True) has already been called, + # so this task is already in the run queue. All we need to do + # is replace the outcome that will be used. + self._next_send = capture(raise_cancel) def _attempt_delivery_of_any_pending_cancel(self) -> None: if self._abort_func is None: @@ -1695,7 +1701,7 @@ def current_root_task(self) -> Task | None: @_public # Type-ignore due to use of Any here. def reschedule( # type: ignore[misc] - self, task: Task, next_send: Outcome[Any] = _NO_SEND + self, task: Task, next_send: Outcome[Any] = _NO_SEND, allow_abort: bool = False ) -> None: """Reschedule the given task with the given :class:`outcome.Outcome`. @@ -1707,11 +1713,20 @@ def reschedule( # type: ignore[misc] returning :data:`Abort.SUCCEEDED` from an abort callback is equivalent to calling :func:`reschedule` once.) + As an exception, if ``allow_abort`` is ``True`` then the abort + callback may still be called up until the rescheduled coroutine + actually resumes. If that happens, and the abort function returns + ``Abort.SUCCEEDED``, then there are effectively two calls, with the + abort always happening second. In that case, any result passed to this + function (as the ``next_send`` argument) will be silently discarded. + Args: task (trio.lowlevel.Task): the task to be rescheduled. Must be blocked in a call to :func:`wait_task_rescheduled`. next_send (outcome.Outcome): the value (or error) to return (or raise) from :func:`wait_task_rescheduled`. + allow_abort (bool): whether the abort function may still be called + after this function is called. """ if next_send is _NO_SEND: @@ -1721,7 +1736,8 @@ def reschedule( # type: ignore[misc] assert task._next_send_fn is None task._next_send_fn = task.coro.send task._next_send = next_send - task._abort_func = None + if not allow_abort: + task._abort_func = None task.custom_sleep_data = None if not self.runq and self.is_guest: self.force_guest_tick_asap() @@ -2602,7 +2618,7 @@ def unrolled_run( next_send_fn = task._next_send_fn next_send = task._next_send - task._next_send_fn = task._next_send = None + task._next_send_fn = task._next_send = task._abort_func = None final_outcome: Outcome[Any] | None = None try: # We used to unwrap the Outcome object here and send/throw diff --git a/src/trio/_sync.py b/src/trio/_sync.py index 6e62eceeff..993164b382 100644 --- a/src/trio/_sync.py +++ b/src/trio/_sync.py @@ -73,7 +73,7 @@ def set(self) -> None: if not self._flag: self._flag = True for task in self._tasks: - _core.reschedule(task) + _core.reschedule(task, allow_abort=True) self._tasks.clear() async def wait(self) -> None: @@ -89,7 +89,7 @@ async def wait(self) -> None: self._tasks.add(task) def abort_fn(_: RaiseCancelT) -> Abort: - self._tasks.remove(task) + self._tasks.discard(task) return _core.Abort.SUCCEEDED await _core.wait_task_rescheduled(abort_fn)