Add always_abort= option to wait_task_rescheduled

This can be helpful to reduce code duplication between the ordinary
and extraordinary wakeup paths – especially in cases where the code
calling reschedule() doesn't necessarily have all the information
needed to properly perform the wakeup.

newsfragments/XX.feature.rst

@@ -0,0 +1 @@
+New argument ``always_abort`` to :func:`trio.hazmat.wait_task_scheduled`.

trio/_core/_run.py

@@ -489,6 +489,7 @@ class Task:
     _next_send = attr.ib(default=None)
     # ParkingLot modifies this directly
     _abort_func = attr.ib(default=None)
+    _always_abort = attr.ib(default=False)
 
     # For introspection and nursery.start()
     _child_nurseries = attr.ib(default=attr.Factory(list))
@@ -545,6 +546,7 @@ def _attempt_abort(self, raise_cancel):
         # We only attempt to abort once per blocking call, regardless of
         # whether we succeeded or failed.
         self._abort_func = None
+        self._always_abort = False
         if success is Abort.SUCCEEDED:
             self._runner.reschedule(self, outcome.capture(raise_cancel))
 
@@ -718,8 +720,12 @@ def reschedule(self, task, next_send=_NO_SEND):
 
         assert task._runner is self
         assert task._next_send is None
+        if task._always_abort:
+            got = task._abort_func(None)
+            assert got is Abort.SUCCEEDED
         task._next_send = next_send
         task._abort_func = None
+        task._always_abort = False
         self.runq.append(task)
         self.instrument("task_scheduled", task)
 
@@ -1392,6 +1398,7 @@ def run_impl(runner, async_fn, args):
                 elif type(msg) is WaitTaskRescheduled:
                     task._cancel_points += 1
                     task._abort_func = msg.abort_func
+                    task._always_abort = msg.always_abort
                     # KI is "outside" all cancel scopes, so check for it
                     # before checking for regular cancellation:
                     if runner.ki_pending and task is runner.main_task:

trio/_core/_traps.py

@@ -61,9 +61,10 @@ class Abort(enum.Enum):
 @attr.s(frozen=True)
 class WaitTaskRescheduled:
     abort_func = attr.ib()
+    always_abort = attr.ib()
 
 
-async def wait_task_rescheduled(abort_func):
+async def wait_task_rescheduled(abort_func, *, always_abort=False):
     """Put the current task to sleep, with cancellation support.
 
     This is the lowest-level API for blocking in trio. Every time a
@@ -90,6 +91,11 @@ async def wait_task_rescheduled(abort_func):
        :func:`wait_task_rescheduled` returns or raises whatever value or error
        was passed to :func:`reschedule`.
 
+       Normally, this does not call ``abort_func``, but if you pass
+       ``always_abort=True``, then it calls ``abort_func(None)``. Also, in
+       this mode, the ``abort_func`` must *always* return
+       :data:`Abort.SUCCEEDED`.
+
     2. The call's context transitions to a cancelled state (e.g. due to a
        timeout expiring). When this happens, the ``abort_func`` is called. It's
        interface looks like::
@@ -156,4 +162,5 @@ def abort(inner_raise_cancel):
        above about how you should use a higher-level API if at all possible?
 
     """
-    return (await _async_yield(WaitTaskRescheduled(abort_func))).unwrap()
+    msg = WaitTaskRescheduled(abort_func, always_abort)
+    return (await _async_yield(msg)).unwrap()

trio/_core/tests/test_run.py

@@ -1894,3 +1894,35 @@ def test_Cancelled_init():
 
     # private constructor should not raise
     _core.Cancelled._init()
+
+
+async def test_always_abort():
+    args = []
+
+    def abort_fn(arg):
+        args.append(arg)
+        return _core.Abort.SUCCEEDED
+
+    task = _core.current_task()
+
+    async def reschedule_me():
+        _core.reschedule(task)
+
+    args.clear()
+    async with _core.open_nursery() as nursery:
+        nursery.start_soon(reschedule_me)
+        await _core.wait_task_rescheduled(abort_fn)
+    assert args == []
+
+    args.clear()
+    async with _core.open_nursery() as nursery:
+        nursery.start_soon(reschedule_me)
+        await _core.wait_task_rescheduled(abort_fn, always_abort=True)
+    assert args == [None]
+
+    args.clear()
+    async with _core.open_nursery() as nursery:
+        nursery.start_soon(reschedule_me)
+        await _core.wait_task_rescheduled(abort_fn)
+
+    assert args == []