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 meant 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`. Without it
+such a task looks idle, and virtual time can jump past timeouts while the
+work is still in progress. `trio.to_thread.run_sync` already holds the
+autojump for its own worker threads.

newsfragments/456.bugfix.rst

@@ -0,0 +1,21 @@
+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`. Such a task is now
+treated as active, like a task blocked in
+`trio.testing.wait_all_tasks_blocked`. Previously the run loop looked idle
+while work was happening in the thread, so virtual time could race ahead of it
+and fire timeouts in concurrent tasks spuriously.
+
+Not waited for: threads that have been abandoned (``abandon_on_cancel=True``,
+after the waiting task was cancelled), threads that are currently inside
+`trio.from_thread.run`, and the threads Trio itself uses only to wait for a
+child process to exit or for a Windows handle, so virtual time can still pass
+while a child process runs.
+
+This can change how existing tests behave. A virtual-time timeout no longer
+fires while a `trio.to_thread.run_sync` call is still blocking; the call is
+waited for in real time. A thread can also deadlock now: if it blocks, other
+than through `trio.from_thread.run`, waiting for a task that first has to sleep
+in virtual time (for example ``queue.get()`` on a queue that a task fills after
+a `trio.sleep`), the clock waits for the thread and the thread for the clock.
+Have the thread wait through `trio.from_thread.run` instead.

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,21 @@ class MockClock(Clock):
           in `wait_all_tasks_blocked`, then the autojump feature treats that
           as active task and does *not* jump the clock.
 
+          Tasks waiting on a thread started by `trio.to_thread.run_sync` are
+          also treated as active: while the thread is running, real work is
+          still happening in real time, so the clock does *not* jump.
+          Virtual time advances again once the thread finishes. The threads
+          Trio itself uses only to wait for a child process to exit, or for
+          a Windows handle, do not count: virtual time can pass while a
+          child process runs. Beware of a thread that blocks, other than
+          through `trio.from_thread.run`, waiting for a task that first has
+          to sleep in virtual time (``queue.get()`` on a queue that a task
+          fills after a sleep): the clock waits for the thread and the
+          thread waits for the clock, a deadlock.
+
+          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,47 @@ 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` already does
+    the equivalent for its worker threads.
+
+    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
+    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,171 @@ 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, but it must release the original
+    # run's unit and leave 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:
+        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

@@ -6,7 +6,7 @@
 
 from .. import _core, _subprocess
 from .._sync import CapacityLimiter, Event
-from .._threads import to_thread_run_sync
+from .._threads import _to_thread_run_sync
 
 assert (sys.platform != "win32" and sys.platform != "darwin") or not TYPE_CHECKING
 
@@ -78,11 +78,13 @@ async def _waitid_system_task(pid: int, event: Event) -> None:
     # call to trio.run is shutting down.
 
     try:
-        await to_thread_run_sync(
+        await _to_thread_run_sync(
             sync_wait_reapable,
-            pid,
+            (pid,),
+            thread_name=None,
             abandon_on_cancel=True,
             limiter=waitid_limiter,
+            hold_autojump=False,
         )
     except OSError:
         # If waitid fails, waitpid will fail too, so it still makes