forked from django/asgiref
-
Notifications
You must be signed in to change notification settings - Fork 0
/
sync.py
201 lines (175 loc) · 6.45 KB
/
sync.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
import asyncio
import functools
import os
import sys
import threading
from concurrent.futures import Future, ThreadPoolExecutor
try:
import contextvars # Python 3.7+ only.
except ImportError:
contextvars = None
class AsyncToSync:
"""
Utility class which turns an awaitable that only works on the thread with
the event loop into a synchronous callable that works in a subthread.
Must be initialised from the main thread.
"""
# Maps launched Tasks to the threads that launched them
launch_map = {}
def __init__(self, awaitable, force_new_loop=False):
self.awaitable = awaitable
if force_new_loop:
# They have asked that we always run in a new sub-loop.
self.main_event_loop = None
else:
try:
self.main_event_loop = asyncio.get_event_loop()
except RuntimeError:
# There's no event loop in this thread. Look for the threadlocal if
# we're inside SyncToAsync
self.main_event_loop = getattr(
SyncToAsync.threadlocal, "main_event_loop", None
)
def __call__(self, *args, **kwargs):
# You can't call AsyncToSync from a thread with a running event loop
try:
event_loop = asyncio.get_event_loop()
except RuntimeError:
pass
else:
if event_loop.is_running():
raise RuntimeError(
"You cannot use AsyncToSync in the same thread as an async event loop - "
"just await the async function directly."
)
# Make a future for the return information
call_result = Future()
# Get the source thread
source_thread = threading.current_thread()
# Use call_soon_threadsafe to schedule a synchronous callback on the
# main event loop's thread if it's there, otherwise make a new loop
# in this thread.
if not (self.main_event_loop and self.main_event_loop.is_running()):
# Make our own event loop and run inside that.
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
try:
loop.run_until_complete(
self.main_wrap(args, kwargs, call_result, source_thread)
)
finally:
try:
if hasattr(loop, "shutdown_asyncgens"):
loop.run_until_complete(loop.shutdown_asyncgens())
finally:
loop.close()
asyncio.set_event_loop(self.main_event_loop)
else:
self.main_event_loop.call_soon_threadsafe(
self.main_event_loop.create_task,
self.main_wrap(args, kwargs, call_result, source_thread),
)
# Wait for results from the future.
return call_result.result()
def __get__(self, parent, objtype):
"""
Include self for methods
"""
return functools.partial(self.__call__, parent)
async def main_wrap(self, args, kwargs, call_result, source_thread):
"""
Wraps the awaitable with something that puts the result into the
result/exception future.
"""
current_task = SyncToAsync.get_current_task()
self.launch_map[current_task] = source_thread
try:
result = await self.awaitable(*args, **kwargs)
except Exception as e:
call_result.set_exception(e)
else:
call_result.set_result(result)
finally:
del self.launch_map[current_task]
class SyncToAsync:
"""
Utility class which turns a synchronous callable into an awaitable that
runs in a threadpool. It also sets a threadlocal inside the thread so
calls to AsyncToSync can escape it.
"""
executor = None
# Maps launched threads to the coroutines that spawned them
launch_map = {}
# Storage for main event loop references
threadlocal = threading.local()
def __init__(self, func):
self.func = func
if self.executor is None:
if "ASGI_THREADS" in os.environ:
executor_kwargs = {
"max_workers": int(os.environ["ASGI_THREADS"]),
}
else:
executor_kwargs = {}
if sys.version_info >= (3, 6):
executor_kwargs["thread_name_prefix"] = "sync-to-async-"
self.executor = ThreadPoolExecutor(**executor_kwargs)
async def __call__(self, *args, **kwargs):
loop = asyncio.get_event_loop()
if contextvars is not None:
context = contextvars.copy_context()
child = functools.partial(self.func, *args, **kwargs)
func = context.run
args = (child,)
kwargs = {}
else:
func = self.func
future = loop.run_in_executor(
self.executor,
functools.partial(
self.thread_handler,
loop,
self.get_current_task(),
func,
*args,
**kwargs
),
)
return await asyncio.wait_for(future, timeout=None)
def __get__(self, parent, objtype):
"""
Include self for methods
"""
return functools.partial(self.__call__, parent)
def thread_handler(self, loop, source_task, func, *args, **kwargs):
"""
Wraps the sync application with exception handling.
"""
# Set the threadlocal for AsyncToSync
self.threadlocal.main_event_loop = loop
# Set the task mapping (used for the locals module)
current_thread = threading.current_thread()
self.launch_map[current_thread] = source_task
# Run the function
try:
return func(*args, **kwargs)
finally:
del self.launch_map[current_thread]
@staticmethod
def get_current_task():
"""
Cross-version implementation of asyncio.current_task()
Returns None if there is no task.
"""
try:
if hasattr(asyncio, "current_task"):
# Python 3.7 and up
return asyncio.current_task()
else:
# Python 3.6
return asyncio.Task.current_task()
except RuntimeError:
return None
# Lowercase is more sensible for most things
sync_to_async = SyncToAsync
async_to_sync = AsyncToSync