New trio.lowlevel.hold_autojump() context manager, and use in to_thread.run_sync()

docs/source/reference-lowlevel.rst

@@ -556,6 +556,8 @@ Low-level blocking
 .. autofunction:: wait_task_rescheduled
 .. autoclass:: Abort
 .. autofunction:: reschedule
+.. autofunction:: hold_autojump
+   :with:
 
 Here's an example lock class implemented using
 :func:`wait_task_rescheduled` directly. This implementation has a

newsfragments/3553.feature.rst

@@ -0,0 +1,9 @@
+Added `trio.lowlevel.hold_autojump`, a context manager that keeps the
+autojump clock (`trio.testing.MockClock` with a finite
+``autojump_threshold``) from jumping while the block runs. It is for code
+that parks a task with `trio.lowlevel.wait_task_rescheduled` until work
+outside the run loop finishes -- in a thread, or on another event loop --
+and wakes it through `trio.lowlevel.TrioToken.run_sync_soon`. Such a task
+otherwise looks idle, so virtual time could jump past its timeouts while the
+work is still in progress. `trio.to_thread.run_sync` gained a matching
+``hold_autojump`` argument.

newsfragments/456.bugfix.rst

@@ -0,0 +1,7 @@
+The autojump clock (`trio.testing.MockClock` with a finite
+``autojump_threshold``) no longer jumps virtual time ahead while a task is
+waiting for a thread started by `trio.to_thread.run_sync`, which made timeouts
+in concurrent tasks fire spuriously. A thread that waits until a task has
+slept in virtual time (for example on a `queue.Queue` that a task fills after
+a `trio.sleep`) must now be started with ``hold_autojump=False``, or it
+deadlocks.

src/trio/_core/__init__.py

@@ -46,6 +46,7 @@
     current_task,
     current_time,
     current_trio_token,
+    hold_autojump,
     in_trio_run,
     in_trio_task,
     notify_closing,

src/trio/_core/_mock_clock.py

@@ -61,6 +61,14 @@ class MockClock(Clock):
           in `wait_all_tasks_blocked`, then the autojump feature treats that
           as active task and does *not* jump the clock.
 
+          Likewise, the clock does *not* jump while a task is waiting for a
+          thread started by `trio.to_thread.run_sync`, unless that thread
+          was started with ``hold_autojump=False``; a thread that waits
+          until a task has slept in virtual time needs that, or the two wait
+          for each other forever. Code that waits for other work outside
+          the run loop can get the same treatment with
+          `trio.lowlevel.hold_autojump`.
+
     """
 
     def __init__(self, rate: float = 0.0, autojump_threshold: float = inf) -> None:

src/trio/_core/_run.py

@@ -1810,6 +1810,10 @@ class Runner:  # type: ignore[explicit-any]
 
     # If everything goes idle for this long, we call clock._autojump()
     clock_autojump_threshold: float = inf
+    # Nonzero while work that will wake a task is still running off the event
+    # loop; the autojump above is not armed until it returns to zero. See
+    # hold_autojump() below, and _AutojumpThreadHold in trio._threads.
+    autojump_holds: int = 0
 
     # Guest mode stuff
     is_guest: bool = False
@@ -2752,7 +2756,16 @@ def unrolled_run(
             # We use 'elif' here because if there are tasks in
             # wait_all_tasks_blocked, then those tasks will wake up without
             # jumping the clock, so we don't need to autojump.
-            elif runner.clock_autojump_threshold < timeout:
+            #
+            # Similarly, while anything holds the autojump (a
+            # to_thread.run_sync thread, or a hold_autojump() block), the
+            # system isn't really idle: real work is happening off the event
+            # loop, and it will eventually wake up a task. So we don't arm
+            # the autojump, and the IO wait simply sleeps until that work
+            # reports back.
+            elif (
+                runner.clock_autojump_threshold < timeout and runner.autojump_holds == 0
+            ):
                 timeout = runner.clock_autojump_threshold
                 idle_primed = IdlePrimedTypes.AUTOJUMP_CLOCK
 
@@ -2987,6 +3000,52 @@ def current_task() -> Task:
         raise RuntimeError("must be called from async context") from None
 
 
+@contextmanager
+def hold_autojump() -> Generator[None, None, None]:
+    """Keep the autojump clock from jumping while the ``with`` block runs.
+
+    Use this around a wait for work that happens outside the run loop, where
+    Trio cannot see it: for example, a task parked in
+    :func:`wait_task_rescheduled` until a thread or another event loop wakes
+    it through :meth:`TrioToken.run_sync_soon`. Such a task looks idle, so
+    without this a :class:`trio.testing.MockClock` with a finite
+    ``autojump_threshold`` may jump ahead, and fire timeouts, while the
+    work is still in progress. :func:`trio.to_thread.run_sync` does the
+    equivalent for its worker threads, unless given ``hold_autojump=False``.
+
+    Only autojumping is affected. This does nothing under a real clock, and
+    tasks blocked in :func:`trio.testing.wait_all_tasks_blocked` still wake
+    up. Holds are counted, so they nest, and several tasks can hold at once.
+    While anything holds, the clock does not autojump: if the work never
+    finishes, a timeout that relies on autojumping never fires.
+
+    The block must be entered and exited from inside a task of the same run.
+    :meth:`TrioToken.run_sync_soon` callbacks qualify, so one callback may
+    enter the block and a later one exit it; but such a hold only starts when
+    the callback runs, and the clock may jump first. A task that enters the
+    block itself and then waits has no such gap.
+
+    Raises:
+      RuntimeError: if entered outside a task.
+
+    """
+    # Only tasks may take a hold. The run loop re-checks the count after every
+    # batch of task steps, and runs none between arming the autojump and
+    # firing it, so the count cannot rise in between.
+    current_task()
+    runner = GLOBAL_RUN_CONTEXT.runner
+    runner.autojump_holds += 1
+    try:
+        yield
+        # Checked only on normal exit, so that an exception already on its
+        # way out is never replaced; the finally releases the hold either way.
+        assert (
+            getattr(GLOBAL_RUN_CONTEXT, "runner", None) is runner
+        ), "hold_autojump() block exited outside the run it was entered in"
+    finally:
+        runner.autojump_holds -= 1
+
+
 def current_effective_deadline() -> float:
     """Returns the current effective deadline for the current task.
 

src/trio/_core/_tests/test_mock_clock.py

@@ -1,11 +1,13 @@
+import contextlib
+import threading
 import time
 from math import inf
 
 import pytest
 
-from trio import sleep
+from trio import Event, move_on_after, sleep, sleep_forever
 
-from ... import _core
+from ... import _abc, _core
 from .. import wait_all_tasks_blocked
 from .._mock_clock import MockClock
 from .._run import GLOBAL_RUN_CONTEXT
@@ -191,3 +193,173 @@ async def test_initialization_doesnt_mutate_runner() -> None:
         GLOBAL_RUN_CONTEXT.runner.clock_autojump_threshold,
     )
     assert before == after
+
+
+@pytest.mark.parametrize("hold", [False, True])
+async def test_hold_autojump_while_waiting_for_thread(
+    autojump_clock: MockClock,
+    hold: bool,
+) -> None:
+    # A task parked in wait_task_rescheduled until a non-Trio thread wakes it
+    # through run_sync_soon looks idle to the run loop. Without a hold the
+    # clock jumps while the thread is still working; with one, virtual time
+    # waits for the thread. The thread works until ``proceed`` is set: by the
+    # sleeper (so only once the clock has jumped) or, when the clock is held
+    # and cannot jump, as soon as every task is blocked.
+    record = []
+    proceed = threading.Event()
+
+    async def wait_for_thread() -> None:
+        task = _core.current_task()
+        token = _core.current_trio_token()
+
+        def thread_fn() -> None:
+            proceed.wait()
+            # Not needed by a correct clock; it gives a broken one time to jump.
+            time.sleep(0.05)
+            token.run_sync_soon(_core.reschedule, task)
+
+        thread = threading.Thread(target=thread_fn)
+        thread.start()
+        try:
+            with _core.hold_autojump() if hold else contextlib.nullcontext():
+                await _core.wait_task_rescheduled(lambda _: _core.Abort.FAILED)
+        finally:
+            proceed.set()
+            thread.join()
+        record.append(("thread", _core.current_time()))
+
+    async def sleeper() -> None:
+        await sleep(10)
+        record.append(("sleeper", _core.current_time()))
+        proceed.set()
+
+    async def release_thread_if_held() -> None:
+        await wait_all_tasks_blocked()
+        if hold:
+            proceed.set()
+
+    async with _core.open_nursery() as nursery:
+        nursery.start_soon(wait_for_thread)
+        nursery.start_soon(sleeper)
+        nursery.start_soon(release_thread_if_held)
+
+    if hold:
+        assert record == [("thread", 0), ("sleeper", 10)]
+    else:
+        assert record == [("sleeper", 10), ("thread", 10)]
+
+
+async def test_hold_autojump_across_run_sync_soon_callbacks(
+    autojump_clock: MockClock,
+) -> None:
+    # run_sync_soon callbacks run inside a task, so one callback can enter the
+    # block and a later one, here sent from another thread, can exit it.
+    token = _core.current_trio_token()
+    hold = _core.hold_autojump()
+    done = Event()
+
+    def finished() -> None:
+        hold.__exit__(None, None, None)
+        done.set()
+
+    def thread_fn() -> None:
+        time.sleep(0.2)
+        token.run_sync_soon(finished)
+
+    # Queued before this task blocks, so it runs before the loop can go idle.
+    token.run_sync_soon(hold.__enter__)
+    thread = threading.Thread(target=thread_fn)
+    thread.start()
+    try:
+        with move_on_after(60) as scope:
+            await done.wait()
+    finally:
+        thread.join()
+    assert not scope.cancelled_caught
+    assert _core.current_time() == 0
+    assert GLOBAL_RUN_CONTEXT.runner.autojump_holds == 0
+
+
+async def test_hold_autojump_nests_and_always_releases(
+    autojump_clock: MockClock,
+) -> None:
+    runner = GLOBAL_RUN_CONTEXT.runner
+    with _core.hold_autojump():
+        with _core.hold_autojump():
+            assert runner.autojump_holds == 2
+        assert runner.autojump_holds == 1
+    assert runner.autojump_holds == 0
+
+    with pytest.raises(KeyError):
+        with _core.hold_autojump():
+            raise KeyError
+    assert runner.autojump_holds == 0
+
+    with _core.CancelScope() as scope:
+        with _core.hold_autojump():
+            scope.cancel()
+            await sleep_forever()
+    assert scope.cancelled_caught
+    assert runner.autojump_holds == 0
+
+    # ...and the clock autojumps again afterwards
+    start = _core.current_time()
+    await sleep(3)
+    assert _core.current_time() - start == 3
+
+
+def test_hold_autojump_with_real_clock() -> None:
+    # Under a real clock there is no autojump to hold; the block is harmless.
+    async def main() -> None:
+        with _core.hold_autojump():
+            await sleep(0.01)
+
+    _core.run(main)
+
+
+def test_hold_autojump_outside_task() -> None:
+    with pytest.raises(RuntimeError, match=r"^must be called from async context$"):
+        with _core.hold_autojump():
+            pass  # pragma: no cover
+
+    # Inside a run but between task steps is refused too: that is what keeps
+    # the count from changing between the run loop arming the autojump and
+    # firing it. (Exceptions in instrument hooks are swallowed, hence record.)
+    record = []
+
+    class TryToHold(_abc.Instrument):
+        def before_io_wait(self, timeout: float) -> None:
+            try:
+                with _core.hold_autojump():
+                    record.append("allowed")  # pragma: no cover
+            except RuntimeError:
+                record.append("refused")
+
+    _core.run(_core.checkpoint, instruments=[TryToHold()])
+    assert record
+    assert set(record) == {"refused"}
+
+
+async def test_hold_autojump_exit_inside_another_run(
+    autojump_clock: MockClock,
+) -> None:
+    # A hold belongs to the run it was entered in. Exiting the block from
+    # inside a different run is a misuse, which fails an assertion; but it
+    # still releases the original run's unit and leaves the other run's count
+    # alone.
+    runner = GLOBAL_RUN_CONTEXT.runner
+    hold = _core.hold_autojump()
+    hold.__enter__()
+    other_run_count = []
+
+    async def exit_here() -> None:
+        with pytest.raises(AssertionError, match="outside the run"):
+            hold.__exit__(None, None, None)
+        other_run_count.append(GLOBAL_RUN_CONTEXT.runner.autojump_holds)
+
+    thread = threading.Thread(target=_core.run, args=(exit_here,))
+    thread.start()
+    thread.join()
+    assert other_run_count == [0]
+    assert runner.autojump_holds == 0

src/trio/_subprocess_platform/waitid.py

@@ -83,6 +83,10 @@ async def _waitid_system_task(pid: int, event: Event) -> None:
             pid,
             abandon_on_cancel=True,
             limiter=waitid_limiter,
+            # This thread only waits for the child to exit. Where the platform
+            # allows, that is an I/O wait, which does not hold the autojump
+            # clock either.
+            hold_autojump=False,
         )
     except OSError:
         # If waitid fails, waitpid will fail too, so it still makes

src/trio/_tests/test_subprocess.py

@@ -34,6 +34,7 @@
     sleep_forever,
 )
 from .._core._tests.tutil import skip_if_fbsd_pipes_broken, slow
+from .._subprocess_platform import wait_child_exiting
 from ..lowlevel import open_process
 from ..testing import MockClock, assert_no_checkpoints, wait_all_tasks_blocked
 
@@ -623,6 +624,44 @@ def on_alarm(sig: int, frame: FrameType | None) -> None:
         signal.signal(signal.SIGALRM, old_sigalrm)
 
 
+@pytest.mark.skipif(
+    not wait_child_exiting.__module__.endswith("waitid"),
+    reason="needs the waitid backend",
+)
+async def test_autojump_clock_does_not_wait_for_waitid_thread(
+    autojump_clock: MockClock,
+) -> None:
+    # The thread that the waitid backend uses only waits for the child to
+    # exit. The autojump clock must not wait for it: a virtual-time timeout
+    # has to fire while the child is still running, as it does on backends
+    # where waiting for a child is an I/O wait.
+    proc = await open_process(SLEEP(3600))
+    try:
+        with move_on_after(5) as scope:
+            await wait_child_exiting(proc)
+        assert scope.cancelled_caught
+        assert proc.returncode is None
+    finally:
+        proc.kill()
+        await proc.wait()
+
+
+async def test_autojump_clock_does_not_wait_for_child(
+    autojump_clock: MockClock,
+) -> None:
+    # On every backend, a virtual-time timeout around a child process fires
+    # while the child is still running.
+    proc = await open_process(SLEEP(3600))
+    try:
+        with move_on_after(5) as scope:
+            await proc.wait()
+        assert scope.cancelled_caught
+        assert proc.returncode is None
+    finally:
+        proc.kill()
+        await proc.wait()
+
+
 async def test_custom_deliver_cancel() -> None:
     custom_deliver_cancel_called = False