Expand cancellation usability from native trio threads

docs/source/reference-core.rst

@@ -1823,6 +1823,11 @@ to spawn a child thread, and then use a :ref:`memory channel
 
 .. literalinclude:: reference-core/from-thread-example.py
 
+You can also perform a non-blocking check for cancellation from threads spawned
+by `trio.to_thread.run_sync`.
+
+.. autofunction:: trio.from_thread.check_cancelled
+
 Threads and task-local storage
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

newsfragments/2392.feature.rst

@@ -0,0 +1 @@
+If called from a thread spawned by `trio.to_thread.run_sync`, `trio.from_thread.run` and `trio.from_thread.run_sync` now reuse the task and cancellation status of the host task and have `trio.from_thread.check_cancelled` to efficiently poll for cancellation.

trio/_tests/test_threads.py

@@ -13,13 +13,12 @@
 import pytest
 import sniffio
 
-from trio._core import TrioToken, current_trio_token
-
-from .. import CapacityLimiter, Event, _core, sleep
+from .. import CapacityLimiter, Event, _core, fail_after, sleep, sleep_forever
 from .._core._tests.test_ki import ki_self
 from .._core._tests.tutil import buggy_pypy_asyncgens
 from .._threads import (
     current_default_thread_limiter,
+    from_thread_check_cancelled,
     from_thread_run,
     from_thread_run_sync,
     to_thread_run_sync,
@@ -810,25 +809,28 @@ def test_from_thread_run_during_shutdown():
     save = []
     record = []
 
-    async def agen():
+    async def agen(token):
         try:
             yield
         finally:
             with pytest.raises(_core.RunFinishedError), _core.CancelScope(shield=True):
-                await to_thread_run_sync(from_thread_run, sleep, 0)
+                await to_thread_run_sync(
+                    partial(from_thread_run, sleep, 0, trio_token=token)
+                )
             record.append("ok")
 
-    async def main():
-        save.append(agen())
+    async def main(use_system_task):
+        save.append(agen(_core.current_trio_token() if use_system_task else None))
         await save[-1].asend(None)
 
-    _core.run(main)
+    _core.run(main, True)  # System nursery will be closed and raise RunFinishedError
+    _core.run(main, False)  # host task will not be rescheduled
     assert record == ["ok"]
 
 
 async def test_trio_token_weak_referenceable():
-    token = current_trio_token()
-    assert isinstance(token, TrioToken)
+    token = _core.current_trio_token()
+    assert isinstance(token, _core.TrioToken)
     weak_reference = weakref.ref(token)
     assert token is weak_reference()
 
@@ -842,3 +844,166 @@ def __bool__(self):
 
     with pytest.raises(NotImplementedError):
         await to_thread_run_sync(int, cancellable=BadBool())
+
+
+async def test_from_thread_reuses_task():
+    task = _core.current_task()
+
+    async def async_current_task():
+        return _core.current_task()
+
+    assert task is await to_thread_run_sync(from_thread_run_sync, _core.current_task)
+    assert task is await to_thread_run_sync(from_thread_run, async_current_task)
+
+
+async def test_recursive_to_thread():
+    tid = None
+
+    def get_tid_then_reenter():
+        nonlocal tid
+        tid = threading.get_ident()
+        return from_thread_run(to_thread_run_sync, threading.get_ident)
+
+    assert tid != await to_thread_run_sync(get_tid_then_reenter)
+
+
+async def test_from_thread_host_cancelled():
+    queue = stdlib_queue.Queue()
+
+    def sync_check():
+        from_thread_run_sync(cancel_scope.cancel)
+        try:
+            from_thread_run_sync(bool)
+        except _core.Cancelled:
+            queue.put(True)
+        else:
+            queue.put(False)
+
+    with _core.CancelScope() as cancel_scope:
+        await to_thread_run_sync(sync_check)
+
+    assert not cancel_scope.cancelled_caught
+    assert not queue.get_nowait()
+
+    with _core.CancelScope() as cancel_scope:
+        await to_thread_run_sync(sync_check, cancellable=True)
+
+    assert cancel_scope.cancelled_caught
+    assert await to_thread_run_sync(partial(queue.get, timeout=1))
+
+    async def no_checkpoint():
+        return True
+
+    def async_check():
+        from_thread_run_sync(cancel_scope.cancel)
+        try:
+            assert from_thread_run(no_checkpoint)
+        except _core.Cancelled:
+            queue.put(True)
+        else:
+            queue.put(False)
+
+    with _core.CancelScope() as cancel_scope:
+        await to_thread_run_sync(async_check)
+
+    assert not cancel_scope.cancelled_caught
+    assert not queue.get_nowait()
+
+    with _core.CancelScope() as cancel_scope:
+        await to_thread_run_sync(async_check, cancellable=True)
+
+    assert cancel_scope.cancelled_caught
+    assert await to_thread_run_sync(partial(queue.get, timeout=1))
+
+    async def async_time_bomb():
+        cancel_scope.cancel()
+        with fail_after(10):
+            await sleep_forever()
+
+    with _core.CancelScope() as cancel_scope:
+        await to_thread_run_sync(from_thread_run, async_time_bomb)
+
+    assert cancel_scope.cancelled_caught
+
+
+async def test_from_thread_check_cancelled():
+    q = stdlib_queue.Queue()
+
+    async def child(cancellable):
+        record.append("start")
+        try:
+            return await to_thread_run_sync(f, cancellable=cancellable)
+        except _core.Cancelled:
+            record.append("cancel")
+        finally:
+            record.append("exit")
+
+    def f():
+        try:
+            from_thread_check_cancelled()
+        except _core.Cancelled:  # pragma: no cover, test failure path
+            q.put("Cancelled")
+        else:
+            q.put("Not Cancelled")
+        ev.wait()
+        return from_thread_check_cancelled()
+
+    # Base case: nothing cancelled so we shouldn't see cancels anywhere
+    record = []
+    ev = threading.Event()
+    async with _core.open_nursery() as nursery:
+        nursery.start_soon(child, False)
+        await wait_all_tasks_blocked()
+        assert record[0] == "start"
+        assert q.get(timeout=1) == "Not Cancelled"
+        ev.set()
+    # implicit assertion, Cancelled not raised via nursery
+    assert record[1] == "exit"
+
+    # cancellable=False case: a cancel will pop out but be handled by
+    # the appropriate cancel scope
+    record = []
+    ev = threading.Event()
+    async with _core.open_nursery() as nursery:
+        nursery.start_soon(child, False)
+        await wait_all_tasks_blocked()
+        assert record[0] == "start"
+        assert q.get(timeout=1) == "Not Cancelled"
+        nursery.cancel_scope.cancel()
+        ev.set()
+    assert nursery.cancel_scope.cancelled_caught
+    assert "cancel" in record
+    assert record[-1] == "exit"
+
+    # cancellable=True case: slightly different thread behavior needed
+    # check thread is cancelled "soon" after abandonment
+    def f():  # noqa: F811
+        ev.wait()
+        try:
+            from_thread_check_cancelled()
+        except _core.Cancelled:
+            q.put("Cancelled")
+        else:  # pragma: no cover, test failure path
+            q.put("Not Cancelled")
+
+    record = []
+    ev = threading.Event()
+    async with _core.open_nursery() as nursery:
+        nursery.start_soon(child, True)
+        await wait_all_tasks_blocked()
+        assert record[0] == "start"
+        nursery.cancel_scope.cancel()
+        ev.set()
+    assert nursery.cancel_scope.cancelled_caught
+    assert "cancel" in record
+    assert record[-1] == "exit"
+    assert q.get(timeout=1) == "Cancelled"
+
+
+async def test_from_thread_check_cancelled_raises_in_foreign_threads():
+    with pytest.raises(RuntimeError):
+        from_thread_check_cancelled()
+    q = stdlib_queue.Queue()
+    _core.start_thread_soon(from_thread_check_cancelled, lambda _: q.put(_))
+    with pytest.raises(RuntimeError):
+        q.get(timeout=1).unwrap()

trio/_tests/verify_types_darwin.json

@@ -40,7 +40,7 @@
     ],
     "exportedSymbolCounts": {
       "withAmbiguousType": 0,
-      "withKnownType": 631,
+      "withKnownType": 632,
       "withUnknownType": 0
     },
     "ignoreUnknownTypesFromImports": true,

trio/_tests/verify_types_linux.json

@@ -28,7 +28,7 @@
     ],
     "exportedSymbolCounts": {
       "withAmbiguousType": 0,
-      "withKnownType": 628,
+      "withKnownType": 629,
       "withUnknownType": 0
     },
     "ignoreUnknownTypesFromImports": true,

trio/_tests/verify_types_windows.json

@@ -7,7 +7,7 @@
     "warningCount": 0
   },
   "typeCompleteness": {
-    "completenessScore": 0.9857369255150554,
+    "completenessScore": 0.9857594936708861,
     "diagnostics": [
       {
         "message": "Return type annotation is missing",
@@ -144,7 +144,7 @@
     ],
     "exportedSymbolCounts": {
       "withAmbiguousType": 0,
-      "withKnownType": 622,
+      "withKnownType": 623,
       "withUnknownType": 9
     },
     "ignoreUnknownTypesFromImports": true,