hat.juggler

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

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

 94class Client(aio.Resource):
 95    """Client
 96
 97    For creating new client see `connect` coroutine.
 98
 99    """
100
101    @property
102    def async_group(self) -> aio.Group:
103        """Async group"""
104        return self._async_group
105
106    @property
107    def state(self) -> json.Storage:
108        """Remote server state"""
109        return self._state
110
111    async def send(self,
112                   name: str,
113                   data: json.Data
114                   ) -> json.Data:
115        """Send request and wait for response
116
117        Args:
118            name: request name
119            data: request payload
120
121        Raises:
122            JugglerError
123            ConnectionError
124
125        """
126        if not self.is_open:
127            raise ConnectionError()
128
129        req_id = next(self._next_req_ids)
130        res_future = self._loop.create_future()
131        self._res_futures[req_id] = res_future
132
133        try:
134            await self._transport.send({'type': 'request',
135                                        'id': req_id,
136                                        'name': name,
137                                        'data': data})
138            return await res_future
139
140        finally:
141            self._res_futures.pop(req_id)
142
143    async def _on_close(self):
144        for f in self._res_futures.values():
145            if not f.done():
146                f.set_exception(ConnectionError())
147
148        await self._transport.async_close()
149        await self._session.close()
150
151    async def _on_msg(self, msg):
152        if msg['type'] == 'response':
153            res_future = self._res_futures.get(msg['id'])
154            if not res_future or res_future.done():
155                return
156
157            if msg['success']:
158                res_future.set_result(msg['data'])
159
160            else:
161                res_future.set_exception(JugglerError(msg['data']))
162
163        elif msg['type'] == 'state':
164            data = json.patch(self._state.data, msg['diff'])
165            self._state.set([], data)
166
167        elif msg['type'] == 'notify':
168            if not self._notify_cb:
169                return
170
171            await aio.call(self._notify_cb, self, msg['name'],
172                           msg['data'])
173
174        else:
175            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                                    remote=conn._remote)
239
240        conn.async_group.spawn(aio.call_on_cancel, conn._transport.async_close)
241        conn.async_group.spawn(aio.call_on_done,
242                               conn._transport.wait_closing(), conn.close)
243
244        conn.async_group.spawn(conn._sync_loop)
245
246        if self._connection_cb:
247            conn.async_group.spawn(aio.call, self._connection_cb, conn)
248
249        return conn
250
251    async def _ws_handler(self, request):
252        conn = await self.create_connection(request)
253
254        await conn.wait_closed()
255
256        return conn.ws

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