|
14 | 14 |
|
15 | 15 | from asyncio import AbstractEventLoop, new_event_loop, run_coroutine_threadsafe
|
16 | 16 | from concurrent.futures import Future
|
17 |
| -from threading import Thread |
18 |
| -from typing import ContextManager |
| 17 | +from threading import Thread, Lock |
| 18 | +from typing import ContextManager, Generic, TypeVar, Optional, Callable |
19 | 19 |
|
| 20 | +_T = TypeVar("_T") |
20 | 21 |
|
21 |
| -class ManagedEventLoop(ContextManager): |
| 22 | + |
| 23 | +class _Lazy(Generic[_T]): |
| 24 | + _Factory = Callable[[], _T] |
| 25 | + |
| 26 | + _lock: Lock |
| 27 | + _factory: _Factory |
| 28 | + _impl: Optional[_T] |
| 29 | + |
| 30 | + def __init__(self, factory: _Factory): |
| 31 | + self._lock = Lock() |
| 32 | + self._factory = factory |
| 33 | + self._impl = None |
| 34 | + |
| 35 | + def get(self) -> _T: |
| 36 | + with self._lock: |
| 37 | + if self._impl is None: |
| 38 | + self._impl = self._factory() |
| 39 | + return self._impl |
| 40 | + |
| 41 | + |
| 42 | +class _ManagedEventLoopImpl(ContextManager): |
22 | 43 | _loop: AbstractEventLoop
|
23 | 44 | _thread: Thread
|
24 | 45 |
|
25 | 46 | def __init__(self, name=None):
|
26 | 47 | self._loop = new_event_loop()
|
27 |
| - self._thread = Thread(target=lambda: self._loop.run_forever(), name=name) |
| 48 | + self._thread = Thread( |
| 49 | + target=lambda: self._loop.run_forever(), name=name, daemon=True |
| 50 | + ) |
28 | 51 |
|
29 | 52 | def __enter__(self):
|
30 | 53 | self._thread.start()
|
| 54 | + return self |
31 | 55 |
|
32 | 56 | def __exit__(self, exc_type, exc_value, traceback):
|
33 | 57 | self._loop.call_soon_threadsafe(self._loop.stop)
|
34 | 58 | self._thread.join()
|
35 | 59 |
|
36 | 60 | def submit(self, coro) -> Future:
|
37 | 61 | return run_coroutine_threadsafe(coro, self._loop)
|
| 62 | + |
| 63 | + |
| 64 | +# TODO(dpcollins): Remove when underlying issue is fixed. |
| 65 | +# This is a workaround for https://github.com/grpc/grpc/issues/25364, a grpc |
| 66 | +# issue which prevents grpc-asyncio working with multiple event loops in the |
| 67 | +# same process. This workaround enables multiple topic publishing as well as |
| 68 | +# publish/subscribe from the same process, but does not enable use with other |
| 69 | +# grpc-asyncio clients. Once this issue is fixed, roll back the PR which |
| 70 | +# introduced this to return to a single event loop per client for isolation. |
| 71 | +_global_event_loop: _Lazy[_ManagedEventLoopImpl] = _Lazy( |
| 72 | + lambda: _ManagedEventLoopImpl(name="PubSubLiteEventLoopThread").__enter__() |
| 73 | +) |
| 74 | + |
| 75 | + |
| 76 | +class ManagedEventLoop(ContextManager): |
| 77 | + _loop: _ManagedEventLoopImpl |
| 78 | + |
| 79 | + def __init__(self, name=None): |
| 80 | + self._loop = _global_event_loop.get() |
| 81 | + |
| 82 | + def __enter__(self): |
| 83 | + return self |
| 84 | + |
| 85 | + def __exit__(self, exc_type, exc_value, traceback): |
| 86 | + pass |
| 87 | + |
| 88 | + def submit(self, coro) -> Future: |
| 89 | + return self._loop.submit(coro) |
0 commit comments