hat.monitor.observer.client

  1"""Observer Client"""
  2
  3import logging
  4import typing
  5
  6from hat import aio
  7from hat import json
  8from hat import util
  9from hat.drivers import chatter
 10from hat.drivers import tcp
 11
 12from hat.monitor.observer import common
 13
 14
 15mlog: logging.Logger = logging.getLogger(__name__)
 16"""Module logger"""
 17
 18StateCb: typing.TypeAlias = aio.AsyncCallable[['Client', 'State'], None]
 19"""State callback"""
 20
 21CloseReqCb: typing.TypeAlias = aio.AsyncCallable[['Client'], None]
 22"""Close request callback"""
 23
 24
 25class State(typing.NamedTuple):
 26    """Client state"""
 27    info: common.ComponentInfo | None
 28    components: list[common.ComponentInfo]
 29
 30
 31async def connect(addr: tcp.Address,
 32                  name: str,
 33                  group: str,
 34                  *,
 35                  data: json.Data = None,
 36                  state_cb: StateCb | None = None,
 37                  close_req_cb: CloseReqCb | None = None,
 38                  **kwargs
 39                  ) -> 'Client':
 40    """Connect to Observer Server
 41
 42    Additional arguments are passed directly to `hat.drivers.chatter.connect`.
 43
 44    """
 45    conn = await chatter.connect(addr, **kwargs)
 46
 47    try:
 48        return Client(conn=conn,
 49                      name=name,
 50                      group=group,
 51                      data=data,
 52                      state_cb=state_cb,
 53                      close_req_cb=close_req_cb)
 54
 55    except Exception:
 56        await aio.uncancellable(conn.async_close())
 57        raise
 58
 59
 60class Client(aio.Resource):
 61    """Observer Client
 62
 63    For creating new client see `connect` coroutine.
 64
 65    """
 66
 67    def __init__(self,
 68                 conn: chatter.Connection,
 69                 name: str,
 70                 group: str,
 71                 data: json.Data,
 72                 state_cb: StateCb | None,
 73                 close_req_cb: CloseReqCb | None):
 74        self._conn = conn
 75        self._name = name
 76        self._group = group
 77        self._data = json.encode(data)
 78        self._state_cb = state_cb
 79        self._close_req_cb = close_req_cb
 80        self._state = State(info=None,
 81                            components=[])
 82        self._blessing_res = common.BlessingRes(token=None,
 83                                                ready=False)
 84
 85        self.async_group.spawn(self._receive_loop)
 86
 87    @property
 88    def async_group(self) -> aio.Group:
 89        """Async group"""
 90        return self._conn.async_group
 91
 92    @property
 93    def state(self) -> State:
 94        """Client's state"""
 95        return self._state
 96
 97    async def set_blessing_res(self, res: common.BlessingRes):
 98        """Set blessing response"""
 99        if res == self._blessing_res:
100            return
101
102        self._blessing_res = res
103        await self._send_msg_client(res)
104
105    async def _receive_loop(self):
106        mlog.debug("starting receive loop")
107        try:
108            await self._send_msg_client(self._blessing_res)
109
110            while True:
111                msg_type, msg_data = await common.receive_msg(self._conn)
112
113                if msg_type == 'HatObserver.MsgServer':
114                    mlog.debug("received msg server")
115                    components = [common.component_info_from_sbs(i)
116                                  for i in msg_data['components']]
117                    await self._process_msg_server(cid=msg_data['cid'],
118                                                   mid=msg_data['mid'],
119                                                   components=components)
120
121                elif msg_type == 'HatObserver.MsgClose':
122                    mlog.debug("received msg close")
123                    if self._close_req_cb:
124                        await aio.call(self._close_req_cb, self)
125                    break
126
127                else:
128                    raise Exception('unsupported message type')
129
130        except ConnectionError:
131            mlog.debug("connection closed")
132
133        except Exception as e:
134            mlog.warning("monitor client error: %s", e, exc_info=e)
135
136        finally:
137            mlog.debug("stopping receive loop")
138            self.close()
139
140    async def _send_msg_client(self, blessing_res):
141        await common.send_msg(self._conn, 'HatObserver.MsgClient', {
142            'name': self._name,
143            'group': self._group,
144            'data': self._data,
145            'blessingRes': common.blessing_res_to_sbs(blessing_res)})
146
147    async def _process_msg_server(self, cid, mid, components):
148        info = util.first(components, lambda i: (i.cid == cid and
149                                                 i.mid == mid))
150        state = State(info=info,
151                      components=components)
152        if self._state == state:
153            return
154
155        self._state = state
156        if self._state_cb:
157            await aio.call(self._state_cb, self, state)