hat.juggler

26class JugglerError(Exception):
27    """Juggler error"""
28
29    def __init__(self, data: json.Data):
30        self.__data = data
31
32    @property
33    def data(self) -> json.Data:
34        """Error data"""
35        return self.__data

38async def connect(address: str,
39                  notify_cb: NotifyCb | None = None,
40                  *,
41                  auth: aiohttp.BasicAuth | None = None,
42                  ssl_ctx: ssl.SSLContext | None = None,
43                  send_queue_size: int = 1024,
44                  max_segment_size: int = 64 * 1024,
45                  ping_delay: float = 30,
46                  ping_timeout: float = 30,
47                  **kwargs
48                  ) -> 'Client':
49    """Connect to remote server
50
51    `address` represents remote WebSocket URL formated as
52    ``<schema>://<host>:<port>/<path>`` where ``<schema>`` is ``ws`` or
53    ``wss``.
54
55    Additional arguments are passed directly to `aiohttp.ClientSession`.
56
57    """
58    client = Client()
59    client._notify_cb = notify_cb
60    client._loop = asyncio.get_running_loop()
61    client._async_group = aio.Group()
62    client._state = json.Storage()
63    client._res_futures = {}
64    client._next_req_ids = itertools.count(1)
65    client._session = aiohttp.ClientSession(**kwargs)
66
67    try:
68        ws = await client._session.ws_connect(address,
69                                              auth=auth,
70                                              ssl=ssl_ctx or False,
71                                              max_msg_size=0)
72
73    except BaseException:
74        await aio.uncancellable(client._session.close())
75        raise
76
77    client._transport = Transport(ws=ws,
78                                  msg_cb=client._on_msg,
79                                  send_queue_size=send_queue_size,
80                                  max_segment_size=max_segment_size,
81                                  ping_delay=ping_delay,
82                                  ping_timeout=ping_timeout)
83
84    client.async_group.spawn(aio.call_on_cancel, client._on_close)
85    client.async_group.spawn(aio.call_on_done,
86                             client._transport.wait_closing(), client.close)
87
88    return client

 91class Client(aio.Resource):
 92    """Client
 93
 94    For creating new client see `connect` coroutine.
 95
 96    """
 97
 98    @property
 99    def async_group(self) -> aio.Group:
100        """Async group"""
101        return self._async_group
102
103    @property
104    def state(self) -> json.Storage:
105        """Remote server state"""
106        return self._state
107
108    async def send(self,
109                   name: str,
110                   data: json.Data
111                   ) -> json.Data:
112        """Send request and wait for response
113
114        Args:
115            name: request name
116            data: request payload
117
118        Raises:
119            JugglerError
120            ConnectionError
121
122        """
123        if not self.is_open:
124            raise ConnectionError()
125
126        req_id = next(self._next_req_ids)
127        res_future = self._loop.create_future()
128        self._res_futures[req_id] = res_future
129
130        try:
131            await self._transport.send({'type': 'request',
132                                        'id': req_id,
133                                        'name': name,
134                                        'data': data})
135            return await res_future
136
137        finally:
138            self._res_futures.pop(req_id)
139
140    async def _on_close(self):
141        for f in self._res_futures.values():
142            if not f.done():
143                f.set_exception(ConnectionError())
144
145        await self._transport.async_close()
146        await self._session.close()
147
148    async def _on_msg(self, msg):
149        if msg['type'] == 'response':
150            res_future = self._res_futures.get(msg['id'])
151            if not res_future or res_future.done():
152                return
153
154            if msg['success']:
155                res_future.set_result(msg['data'])
156
157            else:
158                res_future.set_exception(JugglerError(msg['data']))
159
160        elif msg['type'] == 'state':
161            data = json.patch(self._state.data, msg['diff'])
162            self._state.set([], data)
163
164        elif msg['type'] == 'notify':
165            if not self._notify_cb:
166                return
167
168            await aio.call(self._notify_cb, self, msg['name'],
169                           msg['data'])
170
171        else:
172            raise Exception("invalid message type")

 32async def listen(host: str,
 33                 port: int,
 34                 connection_cb: ConnectionCb | None = None,
 35                 request_cb: RequestCb | None = None,
 36                 *,
 37                 ws_path: str | None = '/ws',
 38                 static_dir: pathlib.PurePath | None = None,
 39                 index_path: str | None = '/index.html',
 40                 htpasswd_file: pathlib.PurePath | None = None,
 41                 ssl_ctx: ssl.SSLContext | None = None,
 42                 autoflush_delay: float | None = 0.2,
 43                 shutdown_timeout: float = 0.1,
 44                 state: json.Storage | None = None,
 45                 parallel_requests: bool = False,
 46                 additional_routes: Iterable[aiohttp.web.RouteDef] = [],
 47                 send_queue_size: int = 1024,
 48                 max_segment_size: int = 64 * 1024,
 49                 ping_delay: float = 30,
 50                 ping_timeout: float = 30,
 51                 no_cache: bool = True
 52                 ) -> 'Server':
 53    """Create listening server
 54
 55    Each time server receives new incoming juggler connection, `connection_cb`
 56    is called with newly created connection.
 57
 58    For each connection, when server receives `request` message, `request_cb`
 59    is called with associated connection, request name and request data.
 60    If `request_cb` returns value, successful `response` message is sent
 61    with resulting value as data. If `request_cb` raises exception,
 62    unsuccessful `response` message is sent with raised exception as data.
 63    If `request_cb` is ``None``, each `request` message causes sending
 64    of unsuccessful `response` message.
 65
 66    If `static_dir` is set, server serves static files is addition to providing
 67    juggler communication.
 68
 69    If `index_path` is set, request for url path ``/`` are redirected to
 70    `index_path`.
 71
 72    If `htpasswd_file` is set, HTTP Basic Authentication is enabled.
 73    All requests are checked for ``Authorization`` header and only users
 74    specified by `htpassword_file` are accepted. `htpasswd_file` is read
 75    during initialization and changes to it's content, after initialization
 76    finishes, are not monitored.
 77
 78    If `ssl_ctx` is set, server provides `https/wss` communication instead
 79    of `http/ws` communication.
 80
 81    Argument `autoflush_delay` is associated with all connections associated
 82    with this server. `autoflush_delay` defines maximum time delay for
 83    automatic synchronization of `state` changes. If `autoflush_delay` is set
 84    to ``None``, automatic synchronization is disabled and user is responsible
 85    for calling :meth:`Connection.flush`. If `autoflush_delay` is set to ``0``,
 86    synchronization of `state` is performed on each change of `state` data.
 87
 88    `shutdown_timeout` defines maximum time duration server will wait for
 89    regular connection close procedures during server shutdown. All connections
 90    that are not closed during this period are forcefully closed.
 91
 92    If `state` is ``None``, each connection is initialized with it's own
 93    instance of server state. If `state` is set, provided state is shared
 94    between all connections.
 95
 96    If `parallel_requests` is set to ``True``, incoming requests will be
 97    processed in parallel - processing of subsequent requests can start (and
 98    finish) before prior responses are generated.
 99
100    Argument `additional_routes` can be used for providing addition aiohttp
101    route definitions handled by running web server.
102
103    `send_queue_size` limits number of messages that can be put in send queue.
104    This limit can impact blocking of :meth:`Connection.notify`.
105
106    `max_segment_size` limits maximum size of single segment
107    (transport payload size).
108
109    When connection doesn't receive incoming data,
110    `ping_delay` is time (in seconds) that connection waits before sending
111    ping request.
112
113    `ping_timeout` is time (in seconds), that connection waits for any kind
114    of incoming traffic before closed connection is assumed.
115
116    If `no_cache` is set to ``True``, server will include
117    ``Cache-Control: no-cache`` header in all responses.
118
119    Args:
120        host: listening hostname
121        port: listening TCP port
122        connection_cb: connection callback
123        request_cb: request callback
124        ws_path: WebSocket url path segment
125        static_dir: static files directory path
126        index_path: index path
127        htpasswd_file: htpasswd file path
128        ssl_ctx: SSL context
129        autoflush_delay: autoflush delay
130        shutdown_timeout: shutdown timeout
131        state: shared server state
132        parallel_requests: parallel request processing
133        additional_routes: additional route definitions
134        send_queue_size: send queue size
135        max_segment_size: maximum segment size
136        ping_delay: ping delay
137        ping_timeout: ping timeout
138        no_cache: no cache header
139
140    """
141    server = Server()
142    server._connection_cb = connection_cb
143    server._request_cb = request_cb
144    server._autoflush_delay = autoflush_delay
145    server._state = state
146    server._parallel_requests = parallel_requests
147    server._send_queue_size = send_queue_size
148    server._max_segment_size = max_segment_size
149    server._ping_delay = ping_delay
150    server._ping_timeout = ping_timeout
151    server._async_group = aio.Group()
152
153    middlewares = []
154
155    if htpasswd_file:
156        middlewares.append(BasicAuthMiddleware(htpasswd_file))
157
158    routes = []
159
160    if ws_path:
161        routes.append(aiohttp.web.get(ws_path, server._ws_handler))
162
163    if index_path:
164
165        async def root_handler(request):
166            raise aiohttp.web.HTTPFound(index_path)
167
168        routes.append(aiohttp.web.get('/', root_handler))
169
170    routes.extend(additional_routes)
171
172    if static_dir:
173        routes.append(aiohttp.web.static('/', static_dir))
174
175    app = aiohttp.web.Application(middlewares=middlewares)
176    app.add_routes(routes)
177
178    if no_cache:
179        app.on_response_prepare.append(_no_cache_prepare)
180
181    runner = aiohttp.web.AppRunner(app,
182                                   shutdown_timeout=shutdown_timeout)
183    await runner.setup()
184    server.async_group.spawn(aio.call_on_cancel, runner.cleanup)
185
186    try:
187        site = aiohttp.web.TCPSite(runner=runner,
188                                   host=host,
189                                   port=port,
190                                   ssl_context=ssl_ctx,
191                                   reuse_address=True)
192        await site.start()
193
194    except BaseException:
195        await aio.uncancellable(server.async_close())
196        raise
197
198    return server

201class Server(aio.Resource):
202    """Server
203
204    For creating new server see `listen` coroutine.
205
206    When server is closed, all incoming connections are also closed.
207
208    """
209
210    @property
211    def async_group(self) -> aio.Group:
212        """Async group"""
213        return self._async_group
214
215    async def create_connection(self,
216                                request: aiohttp.web.Request
217                                ) -> 'Connection':
218        """Create connection"""
219        conn = Connection()
220
221        conn._ws = aiohttp.web.WebSocketResponse()
222        await conn._ws.prepare(request)
223
224        conn._remote = _get_remote(request)
225        conn._async_group = self.async_group.create_subgroup()
226        conn._request_cb = self._request_cb
227        conn._autoflush_delay = self._autoflush_delay
228        conn._state = self._state or json.Storage()
229        conn._parallel_requests = self._parallel_requests
230        conn._flush_queue = aio.Queue()
231
232        conn._transport = Transport(ws=conn._ws,
233                                    msg_cb=conn._on_msg,
234                                    send_queue_size=self._send_queue_size,
235                                    max_segment_size=self._max_segment_size,
236                                    ping_delay=self._ping_delay,
237                                    ping_timeout=self._ping_timeout)
238
239        conn.async_group.spawn(aio.call_on_cancel, conn._transport.async_close)
240        conn.async_group.spawn(aio.call_on_done,
241                               conn._transport.wait_closing(), conn.close)
242
243        conn.async_group.spawn(conn._sync_loop)
244
245        if self._connection_cb:
246            conn.async_group.spawn(aio.call, self._connection_cb, conn)
247
248        return conn
249
250    async def _ws_handler(self, request):
251        conn = await self.create_connection(request)
252
253        await conn.wait_closed()
254
255        return conn.ws

258class Connection(aio.Resource):
259    """Connection
260
261    For creating new connection see `listen` coroutine.
262
263    """
264
265    @property
266    def async_group(self) -> aio.Group:
267        """Async group"""
268        return self._async_group
269
270    @property
271    def remote(self) -> str:
272        """Remote IP address
273
274        Address is obtained from Forwarded or X-Forwarded-For headers. If
275        these headers are not available, socket's remote address is used.
276
277        """
278        return self._remote
279
280    @property
281    def state(self) -> json.Storage:
282        """Server state"""
283        return self._state
284
285    @property
286    def ws(self) -> aiohttp.web.WebSocketResponse:
287        """Associated WebSocket"""
288        return self._ws
289
290    async def flush(self):
291        """Force synchronization of state data
292
293        Raises:
294            ConnectionError
295
296        """
297        try:
298            flush_future = asyncio.Future()
299            self._flush_queue.put_nowait(flush_future)
300            await flush_future
301
302        except aio.QueueClosedError:
303            raise ConnectionError()
304
305    async def notify(self,
306                     name: str,
307                     data: json.Data):
308        """Send notification
309
310        Raises:
311            ConnectionError
312
313        """
314        if not self.is_open:
315            raise ConnectionError()
316
317        await self._transport.send({'type': 'notify',
318                                    'name': name,
319                                    'data': data})
320
321    async def _on_msg(self, msg):
322        if msg['type'] != 'request':
323            raise Exception("invalid message type")
324
325        if self._parallel_requests:
326            self.async_group.spawn(self._process_request, msg)
327
328        else:
329            await self._process_request(msg)
330
331    async def _process_request(self, req):
332        try:
333            res = {'type': 'response',
334                   'id': req['id']}
335
336            if req['name']:
337                try:
338                    if not self._request_cb:
339                        raise Exception('request handler not implemented')
340
341                    res['data'] = await aio.call(self._request_cb, self,
342                                                 req['name'], req['data'])
343                    res['success'] = True
344
345                except Exception as e:
346                    res['data'] = str(e)
347                    res['success'] = False
348
349            else:
350                res['data'] = req['data']
351                res['success'] = True
352
353            await self._transport.send(res)
354
355        except ConnectionError:
356            self.close()
357
358        except Exception as e:
359            self.close()
360            mlog.error("process request error: %s", e, exc_info=e)
361
362    async def _sync_loop(self):
363        flush_future = None
364        data = None
365        synced_data = None
366        data_queue = aio.Queue()
367
368        try:
369            with self._state.register_change_cb(data_queue.put_nowait):
370                data_queue.put_nowait(self._state.data)
371
372                if not self.is_open:
373                    return
374
375                get_data_future = self.async_group.spawn(data_queue.get)
376                get_flush_future = self.async_group.spawn(
377                    self._flush_queue.get)
378
379                while True:
380                    await asyncio.wait([get_data_future, get_flush_future],
381                                       return_when=asyncio.FIRST_COMPLETED)
382
383                    if get_flush_future.done():
384                        flush_future = get_flush_future.result()
385                        get_flush_future = self.async_group.spawn(
386                            self._flush_queue.get)
387
388                    else:
389                        await asyncio.wait([get_flush_future],
390                                           timeout=self._autoflush_delay)
391
392                        if get_flush_future.done():
393                            flush_future = get_flush_future.result()
394                            get_flush_future = self.async_group.spawn(
395                                self._flush_queue.get)
396
397                        else:
398                            flush_future = None
399
400                    if get_data_future.done():
401                        data = get_data_future.result()
402                        get_data_future = self.async_group.spawn(
403                            data_queue.get)
404
405                    if self._autoflush_delay != 0:
406                        if not data_queue.empty():
407                            data = data_queue.get_nowait_until_empty()
408
409                    if synced_data is not data:
410                        diff = json.diff(synced_data, data)
411                        synced_data = data
412
413                        if diff:
414                            await self._transport.send({'type': 'state',
415                                                        'diff': diff})
416
417                    if flush_future and not flush_future.done():
418                        flush_future.set_result(True)
419
420        except Exception as e:
421            mlog.error("sync loop error: %s", e, exc_info=e)
422
423        finally:
424            self.close()
425
426            self._flush_queue.close()
427            while True:
428                if flush_future and not flush_future.done():
429                    flush_future.set_exception(ConnectionError())
430
431                if self._flush_queue.empty():
432                    break
433
434                flush_future = self._flush_queue.get_nowait()