forked from localstack/localstack
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathasync_utils.py
More file actions
135 lines (103 loc) · 4.28 KB
/
async_utils.py
File metadata and controls
135 lines (103 loc) · 4.28 KB
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
import asyncio
import concurrent.futures
import functools
import time
from contextvars import copy_context
from localstack.utils import common
from localstack.utils.common import TMP_THREADS, start_worker_thread
from localstack.utils.run import FuncThread
# reference to named event loop instances
EVENT_LOOPS = {}
class AdaptiveThreadPool(concurrent.futures.ThreadPoolExecutor):
"""Thread pool executor that maintains a maximum of 'core_size' reusable threads in
the core pool, and creates new thread instances as needed (if the core pool is full)."""
DEFAULT_CORE_POOL_SIZE = 30
def __init__(self, core_size=None):
self.core_size = core_size or self.DEFAULT_CORE_POOL_SIZE
super(AdaptiveThreadPool, self).__init__(max_workers=self.core_size)
def submit(self, fn, *args, **kwargs):
# if idle threads are available, don't spin new threads
if self.has_idle_threads():
return super(AdaptiveThreadPool, self).submit(fn, *args, **kwargs)
def _run(*tmpargs):
return fn(*args, **kwargs)
thread = start_worker_thread(_run)
return thread.result_future
def has_idle_threads(self):
if hasattr(self, "_idle_semaphore"):
return self._idle_semaphore.acquire(timeout=0)
num_threads = len(self._threads)
return num_threads < self._max_workers
# Thread pool executor for running sync functions in async context.
# Note: For certain APIs like DynamoDB, we need 3x threads for each parallel request,
# as during request processing the API calls out to the DynamoDB API again (recursively).
# (TODO: This could potentially be improved if we move entirely to asyncio functions.)
THREAD_POOL = AdaptiveThreadPool()
TMP_THREADS.append(THREAD_POOL)
class AsyncThread(FuncThread):
def __init__(self, async_func_gen=None, loop=None):
"""Pass a function that receives an event loop instance and a shutdown event,
and returns an async function."""
FuncThread.__init__(self, self.run_func, None)
self.async_func_gen = async_func_gen
self.loop = loop
self.shutdown_event = None
def run_func(self, *args):
loop = self.loop or ensure_event_loop()
self.shutdown_event = asyncio.Event()
if self.async_func_gen:
self.async_func = async_func = self.async_func_gen(loop, self.shutdown_event)
if async_func:
loop.run_until_complete(async_func)
loop.run_forever()
def stop(self, quiet=None):
if self.shutdown_event:
self.shutdown_event.set()
self.shutdown_event = None
@classmethod
def run_async(cls, func=None, loop=None):
thread = cls(func, loop=loop)
thread.start()
TMP_THREADS.append(thread)
return thread
async def run_sync(func, *args, thread_pool=None, **kwargs):
loop = asyncio.get_running_loop()
thread_pool = thread_pool or THREAD_POOL
func_wrapped = functools.partial(func, *args, **kwargs)
return await loop.run_in_executor(thread_pool, copy_context().run, func_wrapped)
def run_coroutine(coroutine, loop=None):
"""Run an async coroutine in a threadsafe way in the main event loop"""
loop = loop or get_main_event_loop()
future = asyncio.run_coroutine_threadsafe(coroutine, loop)
return future.result()
def ensure_event_loop():
"""Ensure that an event loop is defined for the currently running thread"""
try:
return asyncio.get_event_loop()
except Exception:
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
return loop
def get_main_event_loop():
return get_named_event_loop("_main_")
def get_named_event_loop(name):
result = EVENT_LOOPS.get(name)
if result:
return result
def async_func_gen(loop, shutdown_event):
EVENT_LOOPS[name] = loop
AsyncThread.run_async(async_func_gen)
time.sleep(1)
return EVENT_LOOPS[name]
async def receive_from_queue(queue):
def get():
# run in a retry loop (instead of blocking forever) to allow for graceful shutdown
while True:
try:
if common.INFRA_STOPPED:
return
return queue.get(timeout=1)
except Exception:
pass
msg = await run_sync(get)
return msg