Optimize asyncio.to_thread to avoid contextvars.copy_context() overhead for empty contexts

Author: heliang666s

Lib/asyncio/threads.py

@@ -21,5 +21,8 @@ async def to_thread(func, /, *args, **kwargs):
     """
     loop = events.get_running_loop()
     ctx = contextvars.copy_context()
-    func_call = functools.partial(ctx.run, func, *args, **kwargs)
-    return await loop.run_in_executor(None, func_call)
+    if len(ctx) == 0:
+        callback = functools.partial(func, *args, **kwargs)
+    else:
+        callback = functools.partial(ctx.run, func, *args, **kwargs)
+    return await loop.run_in_executor(None, callback)

Lib/test/test_asyncio/test_threads.py

@@ -2,13 +2,14 @@
 
 import asyncio
 import unittest
+import functools
 
 from contextvars import ContextVar
 from unittest import mock
 
 
 def tearDownModule():
-    asyncio._set_event_loop_policy(None)
+    asyncio.set_event_loop_policy(None)
 
 
 class ToThreadTests(unittest.IsolatedAsyncioTestCase):
@@ -61,6 +62,41 @@ def get_ctx():
 
         self.assertEqual(result, 'parrot')
 
+    @mock.patch('asyncio.base_events.BaseEventLoop.run_in_executor')
+    async def test_to_thread_optimization_path(self, run_in_executor):
+        # This test ensures that `to_thread` uses the correct execution path
+        # based on whether the context is empty or not.
+
+        # `to_thread` awaits the future returned by `run_in_executor`.
+        # We need to provide a completed future as a return value for the mock.
+        fut = asyncio.Future()
+        fut.set_result(None)
+        run_in_executor.return_value = fut
+
+        def myfunc():
+            pass
+
+        # Test with an empty context (optimized path)
+        await asyncio.to_thread(myfunc)
+        run_in_executor.assert_called_once()
+
+        callback = run_in_executor.call_args.args[1]
+        self.assertIsInstance(callback, functools.partial)
+        self.assertIs(callback.func, myfunc)
+        run_in_executor.reset_mock()
+
+        # Test with a non-empty context (standard path)
+        var = ContextVar('var')
+        var.set('value')
+
+        await asyncio.to_thread(myfunc)
+        run_in_executor.assert_called_once()
+
+        callback = run_in_executor.call_args.args[1]
+        self.assertIsInstance(callback, functools.partial)
+        self.assertIsNot(callback.func, myfunc)  # Should be ctx.run
+        self.assertIs(callback.args[0], myfunc)
+
 
 if __name__ == "__main__":
     unittest.main()