|
1 |
| -import abc |
2 |
| -import asyncio |
| 1 | +from __future__ import annotations |
| 2 | + |
| 3 | +import sys |
| 4 | +from asyncio import Future, Queue |
| 5 | +from asyncio.tasks import FIRST_COMPLETED, ensure_future, gather, wait |
3 | 6 | from logging import getLogger
|
4 |
| -from typing import Any, AsyncIterator, Awaitable, Callable, Dict |
| 7 | +from typing import Any, AsyncIterator, Awaitable, Callable, List, Sequence, Tuple |
| 8 | +from weakref import WeakSet |
5 | 9 |
|
6 | 10 | from anyio import create_task_group
|
7 |
| -from anyio.abc import TaskGroup |
| 11 | + |
| 12 | +from idom.utils import Ref |
8 | 13 |
|
9 | 14 | from .layout import Layout, LayoutEvent, LayoutUpdate
|
10 |
| -from .utils import HasAsyncResources, async_resource |
| 15 | + |
| 16 | + |
| 17 | +if sys.version_info >= (3, 7): # pragma: no cover |
| 18 | + from contextlib import asynccontextmanager # noqa |
| 19 | +else: # pragma: no cover |
| 20 | + from async_generator import asynccontextmanager |
11 | 21 |
|
12 | 22 |
|
13 | 23 | logger = getLogger(__name__)
|
|
16 | 26 | RecvCoroutine = Callable[[], Awaitable[LayoutEvent]]
|
17 | 27 |
|
18 | 28 |
|
19 |
| -class AbstractDispatcher(HasAsyncResources, abc.ABC): |
20 |
| - """A base class for implementing :class:`~idom.core.layout.Layout` dispatchers.""" |
| 29 | +async def dispatch_single_view( |
| 30 | + layout: Layout, |
| 31 | + send: SendCoroutine, |
| 32 | + recv: RecvCoroutine, |
| 33 | +) -> None: |
| 34 | + with layout: |
| 35 | + async with create_task_group() as task_group: |
| 36 | + task_group.start_soon(_single_outgoing_loop, layout, send) |
| 37 | + task_group.start_soon(_single_incoming_loop, layout, recv) |
21 | 38 |
|
22 |
| - __slots__ = "_layout" |
23 | 39 |
|
24 |
| - def __init__(self, layout: Layout) -> None: |
25 |
| - super().__init__() |
26 |
| - self._layout = layout |
| 40 | +_SharedDispatchFuture = Callable[[SendCoroutine, RecvCoroutine], Future] |
27 | 41 |
|
28 |
| - async def start(self) -> None: |
29 |
| - await self.__aenter__() |
30 | 42 |
|
31 |
| - async def stop(self) -> None: |
32 |
| - await self.task_group.cancel_scope.cancel() |
33 |
| - await self.__aexit__(None, None, None) |
| 43 | +@asynccontextmanager |
| 44 | +async def create_shared_view_dispatcher( |
| 45 | + layout: Layout, run_forever: bool = False |
| 46 | +) -> AsyncIterator[_SharedDispatchFuture]: |
| 47 | + with layout: |
| 48 | + ( |
| 49 | + dispatch_shared_view, |
| 50 | + model_state, |
| 51 | + all_update_queues, |
| 52 | + ) = await _make_shared_view_dispatcher(layout) |
34 | 53 |
|
35 |
| - @async_resource |
36 |
| - async def layout(self) -> AsyncIterator[Layout]: |
37 |
| - async with self._layout as layout: |
38 |
| - yield layout |
| 54 | + dispatch_tasks: List[Future] = [] |
39 | 55 |
|
40 |
| - @async_resource |
41 |
| - async def task_group(self) -> AsyncIterator[TaskGroup]: |
42 |
| - async with create_task_group() as group: |
43 |
| - yield group |
| 56 | + def dispatch_shared_view_soon( |
| 57 | + send: SendCoroutine, recv: RecvCoroutine |
| 58 | + ) -> Future: |
| 59 | + future = ensure_future(dispatch_shared_view(send, recv)) |
| 60 | + dispatch_tasks.append(future) |
| 61 | + return future |
44 | 62 |
|
45 |
| - async def run(self, send: SendCoroutine, recv: RecvCoroutine, context: Any) -> None: |
46 |
| - """Start an unending loop which will drive the layout. |
| 63 | + yield dispatch_shared_view_soon |
47 | 64 |
|
48 |
| - This will call :meth:`AbstractLayout.render` and :meth:`Layout.dispatch` |
49 |
| - to render new models and execute events respectively. |
50 |
| - """ |
51 |
| - await self.task_group.spawn(self._outgoing_loop, send, context) |
52 |
| - await self.task_group.spawn(self._incoming_loop, recv, context) |
53 |
| - return None |
| 65 | + gathered_dispatch_tasks = gather(*dispatch_tasks, return_exceptions=True) |
54 | 66 |
|
55 |
| - async def _outgoing_loop(self, send: SendCoroutine, context: Any) -> None: |
56 |
| - try: |
57 |
| - while True: |
58 |
| - await send(await self._outgoing(self.layout, context)) |
59 |
| - except Exception: |
60 |
| - logger.info("Failed to send outgoing update", exc_info=True) |
61 |
| - raise |
| 67 | + while True: |
| 68 | + ( |
| 69 | + update_future, |
| 70 | + dispatchers_completed_future, |
| 71 | + ) = await _wait_until_first_complete( |
| 72 | + layout.render(), |
| 73 | + gathered_dispatch_tasks, |
| 74 | + ) |
| 75 | + |
| 76 | + if dispatchers_completed_future.done(): |
| 77 | + update_future.cancel() |
| 78 | + break |
| 79 | + else: |
| 80 | + update: LayoutUpdate = update_future.result() |
| 81 | + |
| 82 | + model_state.current = update.apply_to(model_state.current) |
| 83 | + # push updates to all dispatcher callbacks |
| 84 | + for queue in all_update_queues: |
| 85 | + queue.put_nowait(update) |
| 86 | + |
| 87 | + |
| 88 | +def ensure_shared_view_dispatcher_future( |
| 89 | + layout: Layout, |
| 90 | +) -> Tuple[Future, _SharedDispatchFuture]: |
| 91 | + dispatcher_future = Future() |
| 92 | + |
| 93 | + async def dispatch_shared_view_forever(): |
| 94 | + with layout: |
| 95 | + ( |
| 96 | + dispatch_shared_view, |
| 97 | + model_state, |
| 98 | + all_update_queues, |
| 99 | + ) = await _make_shared_view_dispatcher(layout) |
| 100 | + |
| 101 | + dispatcher_future.set_result(dispatch_shared_view) |
62 | 102 |
|
63 |
| - async def _incoming_loop(self, recv: RecvCoroutine, context: Any) -> None: |
64 |
| - try: |
65 | 103 | while True:
|
66 |
| - await self._incoming(self.layout, context, await recv()) |
67 |
| - except Exception: |
68 |
| - logger.info("Failed to receive incoming event", exc_info=True) |
69 |
| - raise |
70 |
| - |
71 |
| - @abc.abstractmethod |
72 |
| - async def _outgoing(self, layout: Layout, context: Any) -> Any: |
73 |
| - ... |
| 104 | + update = await layout.render() |
| 105 | + model_state.current = update.apply_to(model_state.current) |
| 106 | + # push updates to all dispatcher callbacks |
| 107 | + for queue in all_update_queues: |
| 108 | + queue.put_nowait(update) |
74 | 109 |
|
75 |
| - @abc.abstractmethod |
76 |
| - async def _incoming(self, layout: Layout, context: Any, message: Any) -> None: |
77 |
| - ... |
| 110 | + async def dispatch(send: SendCoroutine, recv: RecvCoroutine) -> None: |
| 111 | + await (await dispatcher_future)(send, recv) |
78 | 112 |
|
| 113 | + return ensure_future(dispatch_shared_view_forever()), dispatch |
79 | 114 |
|
80 |
| -class SingleViewDispatcher(AbstractDispatcher): |
81 |
| - """Each client of the dispatcher will get its own model. |
82 | 115 |
|
83 |
| - ..note:: |
84 |
| - The ``context`` parameter of :meth:`SingleViewDispatcher.run` should just |
85 |
| - be ``None`` since it's not used. |
86 |
| - """ |
| 116 | +_SharedDispatchCoroutine = Callable[[SendCoroutine, RecvCoroutine], Awaitable[None]] |
87 | 117 |
|
88 |
| - __slots__ = "_current_model_as_json" |
89 | 118 |
|
90 |
| - def __init__(self, layout: Layout) -> None: |
91 |
| - super().__init__(layout) |
92 |
| - self._current_model_as_json = "" |
| 119 | +async def _make_shared_view_dispatcher( |
| 120 | + layout: Layout, |
| 121 | +) -> Tuple[_SharedDispatchCoroutine, Ref[Any], WeakSet[Queue[LayoutUpdate]]]: |
| 122 | + initial_update = await layout.render() |
| 123 | + model_state = Ref(initial_update.apply_to({})) |
93 | 124 |
|
94 |
| - async def _outgoing(self, layout: Layout, context: Any) -> LayoutUpdate: |
95 |
| - return await layout.render() |
| 125 | + # We push updates to queues instead of pushing directly to send() callbacks in |
| 126 | + # order to isolate the render loop from any errors dispatch callbacks might |
| 127 | + # raise. |
| 128 | + all_update_queues: WeakSet[Queue[LayoutUpdate]] = WeakSet() |
96 | 129 |
|
97 |
| - async def _incoming(self, layout: Layout, context: Any, event: LayoutEvent) -> None: |
98 |
| - await layout.dispatch(event) |
| 130 | + async def dispatch_shared_view(send: SendCoroutine, recv: RecvCoroutine) -> None: |
| 131 | + update_queue: Queue[LayoutUpdate] = Queue() |
| 132 | + async with create_task_group() as inner_task_group: |
| 133 | + all_update_queues.add(update_queue) |
| 134 | + await send(LayoutUpdate.create_from({}, model_state.current)) |
| 135 | + inner_task_group.start_soon(_single_incoming_loop, layout, recv) |
| 136 | + inner_task_group.start_soon(_shared_outgoing_loop, send, update_queue) |
99 | 137 | return None
|
100 | 138 |
|
| 139 | + return dispatch_shared_view, model_state, all_update_queues |
101 | 140 |
|
102 |
| -class SharedViewDispatcher(SingleViewDispatcher): |
103 |
| - """Each client of the dispatcher shares the same model. |
104 | 141 |
|
105 |
| - The client's ID is indicated by the ``context`` argument of |
106 |
| - :meth:`SharedViewDispatcher.run` |
107 |
| - """ |
| 142 | +async def _single_outgoing_loop(layout: Layout, send: SendCoroutine) -> None: |
| 143 | + while True: |
| 144 | + await send(await layout.render()) |
108 | 145 |
|
109 |
| - __slots__ = "_update_queues", "_model_state" |
110 | 146 |
|
111 |
| - def __init__(self, layout: Layout) -> None: |
112 |
| - super().__init__(layout) |
113 |
| - self._model_state: Any = {} |
114 |
| - self._update_queues: Dict[str, asyncio.Queue[LayoutUpdate]] = {} |
| 147 | +async def _single_incoming_loop(layout: Layout, recv: RecvCoroutine) -> None: |
| 148 | + while True: |
| 149 | + await layout.dispatch(await recv()) |
115 | 150 |
|
116 |
| - @async_resource |
117 |
| - async def task_group(self) -> AsyncIterator[TaskGroup]: |
118 |
| - async with create_task_group() as group: |
119 |
| - await group.spawn(self._render_loop) |
120 |
| - yield group |
121 | 151 |
|
122 |
| - async def run( |
123 |
| - self, send: SendCoroutine, recv: RecvCoroutine, context: str, join: bool = False |
124 |
| - ) -> None: |
125 |
| - await super().run(send, recv, context) |
126 |
| - if join: |
127 |
| - await self._join_event.wait() |
| 152 | +async def _shared_outgoing_loop( |
| 153 | + send: SendCoroutine, queue: Queue[LayoutUpdate] |
| 154 | +) -> None: |
| 155 | + while True: |
| 156 | + await send(await queue.get()) |
128 | 157 |
|
129 |
| - async def _render_loop(self) -> None: |
130 |
| - while True: |
131 |
| - update = await super()._outgoing(self.layout, None) |
132 |
| - self._model_state = update.apply_to(self._model_state) |
133 |
| - # append updates to all other contexts |
134 |
| - for queue in self._update_queues.values(): |
135 |
| - await queue.put(update) |
136 |
| - |
137 |
| - async def _outgoing_loop(self, send: SendCoroutine, context: Any) -> None: |
138 |
| - self._update_queues[context] = asyncio.Queue() |
139 |
| - await send(LayoutUpdate.create_from({}, self._model_state)) |
140 |
| - await super()._outgoing_loop(send, context) |
141 |
| - |
142 |
| - async def _outgoing(self, layout: Layout, context: str) -> LayoutUpdate: |
143 |
| - return await self._update_queues[context].get() |
144 |
| - |
145 |
| - @async_resource |
146 |
| - async def _join_event(self) -> AsyncIterator[asyncio.Event]: |
147 |
| - event = asyncio.Event() |
148 |
| - try: |
149 |
| - yield event |
150 |
| - finally: |
151 |
| - event.set() |
| 158 | + |
| 159 | +async def _wait_until_first_complete( |
| 160 | + *tasks: Awaitable[Any], |
| 161 | +) -> Sequence[Future]: |
| 162 | + futures = [ensure_future(t) for t in tasks] |
| 163 | + await wait(futures, return_when=FIRST_COMPLETED) |
| 164 | + return futures |
0 commit comments