[WIP] Allow cancellation between reschedule() called and task resumed

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
 

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

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)