[trio.from_thread]

Rework thread local storage to use the canonical threading.local()
accessor, update from_thread.* unit tests to better reflect use cases,
updated no token error msg to give a specific reason for failure.

Author: epellis

trio/_threads.py

@@ -16,6 +16,9 @@
     "BlockingTrioPortal",
 ]
 
+# Global due to Threading API, thread local storage for trio token
+TOKEN_LOCAL = threading.local()
+
 
 class BlockingTrioPortal:
     """A portal that synchronous threads can reach through to run code in the
@@ -375,23 +378,31 @@ def do_release_then_return_result():
 
     # This is the function that runs in the worker thread to do the actual
     # work and then schedule the call to report_back_in_trio_thread_fn
-    def worker_thread_fn():
+    # Since this is spawned in a new thread, the trio token needs to be passed
+    # explicitly to it so it can inject it into thread local storage
+    def worker_thread_fn(trio_token):
+        TOKEN_LOCAL.token = trio_token
         result = outcome.capture(sync_fn, *args)
         try:
             token.run_sync_soon(report_back_in_trio_thread_fn, result)
         except trio.RunFinishedError:
             # The entire run finished, so our particular task is certainly
             # long gone -- it must have cancelled.
             pass
+        finally:
+            del TOKEN_LOCAL.token
 
     await limiter.acquire_on_behalf_of(placeholder)
     try:
         # daemon=True because it might get left behind if we cancel, and in
         # this case shouldn't block process exit.
+        current_trio_token = trio.hazmat.current_trio_token()
         thread = threading.Thread(
-            target=worker_thread_fn, name=name, daemon=True
+            target=worker_thread_fn,
+            args=(current_trio_token,),
+            name=name,
+            daemon=True
         )
-        setattr(thread, 'current_trio_token', trio.hazmat.current_trio_token())
         thread.start()
     except:
         limiter.release_on_behalf_of(placeholder)
@@ -412,11 +423,15 @@ def _run_fn_as_system_task(cb, fn, *args, trio_token=None):
 
     Since this internally uses TrioToken.run_sync_soon, all warnings about
     raised exceptions canceling all tasks should be noted.
-
     """
+
     if not trio_token:
-        current_thread = threading.current_thread()
-        trio_token = getattr(current_thread, 'current_trio_token')
+        try:
+            trio_token = TOKEN_LOCAL.token
+        except AttributeError:
+            raise RuntimeError(
+                "this thread wasn't created by Trio, pass kwarg trio_token=..."
+            )
 
     try:
         trio.hazmat.current_task()

trio/tests/test_threads.py

@@ -463,10 +463,8 @@ def bad_start(self):
 async def test_trio_run_sync_in_thread_token():
     # Test that run_sync_in_thread automatically injects the current trio token
     # into a spawned thread
-
     def thread_fn():
-        current_thread = threading.current_thread()
-        callee_token = getattr(current_thread, 'current_trio_token')
+        callee_token = run_sync(_core.current_trio_token)
         return callee_token
 
     caller_token = _core.current_trio_token()
@@ -478,25 +476,28 @@ async def test_trio_from_thread_run_sync():
     # Test that run_sync_in_thread correctly "hands off" the trio token to
     # trio.from_thread.run_sync()
     def thread_fn():
-        start = run_sync(_core.current_time)
-        end = run_sync(_core.current_time)
-        return end - start
+        trio_time = run_sync(_core.current_time)
+        return trio_time
 
-    duration = await run_sync_in_thread(thread_fn)
-    assert duration > 0
+    trio_time = await run_sync_in_thread(thread_fn)
+    assert isinstance(trio_time, float)
 
 
 async def test_trio_from_thread_run():
     # Test that run_sync_in_thread correctly "hands off" the trio token to
     # trio.from_thread.run()
+    record = []
+
+    async def back_in_trio_fn():
+        _core.current_time()  # implicitly checks that we're in trio
+        record.append("back in trio")
+
     def thread_fn():
-        start = time.perf_counter()
-        run(sleep, 0.05)
-        end = time.perf_counter()
-        return end - start
+        record.append("in thread")
+        run(back_in_trio_fn)
 
-    duration = await run_sync_in_thread(thread_fn)
-    assert duration > 0
+    await run_sync_in_thread(thread_fn)
+    assert record == ["in thread", "back in trio"]
 
 
 async def test_trio_from_thread_token():
@@ -523,23 +524,9 @@ def thread_fn(token):
     assert callee_token == caller_token
 
 
-async def test_trio_from_thread_both_run():
-    # Test that trio.from_thread.run() and from_thread.run_sync() can run in
-    # the same thread together
-
-    def thread_fn():
-        start = run_sync(_core.current_time)
-        run(sleep, 0.05)
-        end = run_sync(_core.current_time)
-        return end - start
-
-    duration = await run_sync_in_thread(thread_fn)
-    assert duration > 0
-
-
-async def test_trio_from_thread_raw_call():
+async def test_from_thread_no_token():
     # Test that a "raw call" to trio.from_thread.run() fails because no token
     # has been provided
 
-    with pytest.raises(AttributeError):
+    with pytest.raises(RuntimeError):
         run_sync(_core.current_time)