openrvdas.logger.readers.cached_data_reader

  1#!/usr/bin/env python3
  2
  3import asyncio
  4import json
  5import logging
  6import queue
  7import ssl
  8import threading
  9import time
 10
 11try:
 12    import websockets
 13    WEBSOCKETS_ENABLED = True
 14except ModuleNotFoundError:
 15    WEBSOCKETS_ENABLED = False
 16
 17from logger.readers.reader import Reader  # noqa: E402
 18from logger.utils.das_record import to_das_record_list  # noqa: E402
 19
 20DEFAULT_SERVER_WEBSOCKET = 'localhost:8766'
 21
 22
 23################################################################################
 24class CachedDataReader(Reader):
 25    """Subscribe to and read field values from a CachedDataServer via
 26    websocket connection.
 27    """
 28
 29    def __init__(self, subscription, data_server=DEFAULT_SERVER_WEBSOCKET,
 30                 bundle_seconds=0, return_das_record=False, data_id=None,
 31                 use_wss=False, check_cert=False, **kwargs):
 32        """
 33        ```
 34        subscription - a dictionary corresponding to the full
 35            fields/seconds, etc that the reader wishes, following the
 36            conventions described in logger/utils/cached_data_server.py
 37            e.g:
 38
 39            subscription = {'fields':{'S330CourseTrue':{seconds:0},
 40                                      'S330HeadingTrue':{seconds:0}}}
 41
 42            If the value of 'fields' is a list instead of a dict, it will be
 43            interpreted as a list of field names to be subscribed to with a
 44            value of seconds = 0. e.g.:
 45
 46            subscription = {'fields':['S330CourseTrue', 'S330HeadingTrue']}
 47
 48        data_server - the host and port at which to try to connect to a
 49            CachedDataServer
 50
 51        bundle_seconds - If specified, aggregate this many seconds worth of records
 52                      and return as a list of records. Note that as implemented, it
 53                      bundles by system clock time, and not by DASRecord timestamp.
 54
 55        return_das_record - If True, return results as DASRecords.
 56
 57        data_id - If return_das_record, use this as the records' data_id
 58
 59        use_wss -     If True, use secure websockets
 60
 61        check_cert  - If True and use_wss is True, check the server's TLS certificate
 62                      for validity; if a str, use as local filepath location of .pem
 63                      file to check against.
 64        ```
 65        When invoked in a config file, this would be:
 66        ```
 67          readers:
 68            class: CachedDataServer
 69            kwargs:
 70              data_server: localhost:8766
 71              subscription:
 72                fields:
 73                  S330CourseTrue:
 74                    seconds: 0
 75                  S330HeadingTrue:
 76                    seconds: 0
 77        ```
 78        """
 79        super().__init__(**kwargs)
 80
 81        if not WEBSOCKETS_ENABLED:
 82            raise ModuleNotFoundError('CachedDataReader(): websockets module is not '
 83                                      'installed. Please try "pip3 install '
 84                                      'websockets" prior to use.')
 85        if not (isinstance(bundle_seconds, int) or isinstance(bundle_seconds, float)) \
 86                or bundle_seconds < 0:
 87            raise ValueError('CachedDataReader parameter "bundle_seconds" must be a number '
 88                             f'greater than or equal to zero. Found "{bundle_seconds}"')
 89
 90        # To simplify templating, subscription may be a list of fields instead of a dict.
 91        # If so, convert it to a dict here.from
 92        subscription_fields = subscription.get('fields')
 93        if not subscription_fields:
 94            raise ValueError('CachedDataReader subscription - no "fields" found!')
 95        elif isinstance(subscription_fields, list):
 96            new_fields = {field: {'seconds': 0} for field in subscription_fields}
 97            subscription['fields'] = new_fields
 98
 99        self.subscription = subscription
100        subscription['type'] = 'subscribe'
101        self.data_server = data_server
102        self.bundle_seconds = bundle_seconds
103        self.return_das_record = return_das_record
104        self.data_id = data_id
105        self.use_wss = use_wss
106        self.check_cert = check_cert
107
108        # We won't initialize our websocket until the first read()
109        # call. At that point we'll launch an async process in a separate
110        # thread that will wait for data from the websocket and put it in
111        # a queue that read() will pop from.
112        self.websocket_thread = None
113        self.queue = queue.Queue()
114        self.quit_flag = False
115
116    ############################
117    def _parse_response(self, response):
118        """Parse a CachedDataServer response and enqueue the resulting data."""
119        if not response.get('type') == 'data':
120            logging.info('Non-"data" response received from data '
121                         'server: %s', response)
122            return
123        if not response.get('status') == 200:
124            logging.warning('Non-"200" status received from data '
125                            'server: %s', response)
126            return
127        data = response.get('data')
128        if not data:
129            logging.debug('No data found in data server response?: %s', response)
130            return
131
132        # If we've gotten a list, assume/hope it's a list of
133        # DASRecord-like dicts; that means it's already collated for us by
134        # timestamp.
135        if type(data) is list:
136            for entry in data:
137                self.queue.put(entry)
138            return
139
140        # Otherwise we expect it to be a field dict, and need to collate
141        # by timestamp manually.
142        if not type(data) is dict:
143            logging.warning('Data from data server not a dict?!?: %s', response)
144            return
145
146        # Collate the fields/values by timestamp
147        timestamp_dict = {}
148        for field, values in data.items():
149            for timestamp, value in values:  # should be list of [ts, value] pairs
150                if timestamp not in timestamp_dict:
151                    timestamp_dict[timestamp] = {}
152                timestamp_dict[timestamp][field] = value
153
154        # Enqueue entries by timestamp
155        for timestamp in sorted(timestamp_dict.keys()):
156            entry = {'timestamp': timestamp, 'fields': timestamp_dict[timestamp]}
157            logging.debug('Enqueuing from CDS: %s', entry)
158            self.queue.put(entry)
159
160    ############################
161    def _start_websocket(self):
162        """We'll run this in a separate thread as soon as we get our first
163        call to read()."""
164
165        ############################
166        async def _websocket_loop(self):
167            """Asynchronous inner function that will read from websocket and put
168            the result in our queue.
169            """
170            # Iterate if we lose the websocket for some reason other than a 'quit'
171            while not self.quit_flag:
172                try:
173                    if self.use_wss:
174                        # If check_cert is a str, take it as the location of the
175                        # .pem file we'll check for validity. Otherwise, if not
176                        # False, take as a bool to verify by own means.
177                        ws_data_server = 'wss://' + self.data_server
178                        ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS)
179                        if self.check_cert:
180                            if isinstance(self.check_cert, str):
181                                ssl_context.load_verify_locations(self.check_cert)
182                            else:
183                                ssl_context.verify_mode = ssl.CERT_REQUIRED
184                        else:
185                            ssl_context.verify_mode = ssl.CERT_NONE
186
187                    else:  # not using wss
188                        ws_data_server = 'ws://' + self.data_server
189                        ssl_context = None
190
191                    logging.info(f'CachedDataReader connecting to {ws_data_server}')
192                    async with websockets.connect(ws_data_server, ssl=ssl_context) as ws:
193                        logging.info(f'Connected to data server {ws_data_server}')
194                        # Send our subscription request
195                        await ws.send(json.dumps(self.subscription))
196                        result = await ws.recv()
197                        response = json.loads(result)
198
199                        while not self.quit_flag:
200                            await ws.send(json.dumps({'type': 'ready'}))
201                            result = await ws.recv()
202                            response = json.loads(result)
203                            logging.debug('Got CachedDataServer response: %s', response)
204                            self._parse_response(response)
205
206                except BrokenPipeError:
207                    pass
208                except AttributeError as e:
209                    logging.info('CachedDataReader websocket loop error: %s', e)
210                except websockets.exceptions.ConnectionClosed:
211                    logging.warning('CachedDataReader lost websocket connection to '
212                                    'data server; trying to reconnect.')
213                    await asyncio.sleep(0.2)
214
215                except websockets.exceptions.InvalidStatusCode:
216                    logging.warning('CachedDataWriter InvalidStatusCode connecting to '
217                                    'data server; trying to reconnect.')
218                    await asyncio.sleep(0.2)
219
220                except OSError as e:
221                    logging.info('Unable to connect to data server. '
222                                 'Sleeping to try again...')
223                    logging.info('Connection error: %s', str(e))
224                    await asyncio.sleep(5)
225
226        # In the outer function, get a new event loop and fire up the
227        # inner, async routine.
228        self.websocket_initialized = True
229
230        # Could we also use asyncio.ensure_future(_websocket_loop(self)) ?
231
232        websocket_event_loop = asyncio.new_event_loop()
233        websocket_event_loop.run_until_complete(_websocket_loop(self))
234        websocket_event_loop.close()
235
236    ############################
237    def quit(self, seconds=0):
238        """Sleep N seconds, then signal quit."""
239        time.sleep(seconds)
240        self.quit_flag = True
241
242    ############################
243    def read(self):
244        """Read/wait for data from the websocket."""
245
246        # If we've not yet fired up the websocket thread, do that now.
247        if not self.websocket_thread:
248            self.websocket_thread = threading.Thread(
249                name='websocket_thread',
250                target=self._start_websocket,
251                daemon=True)
252            self.websocket_thread.start()
253
254        start_time = time.time()
255        result_list = []
256
257        # Use a timeout in our queue get() so we can periodically check if
258        # we've gotten a 'quit'
259        while not self.quit_flag:
260            try:
261                result = self.queue.get(timeout=1)
262                logging.debug('Got result from queue: %s', result)
263
264                # If we're not bundling results, just return the result
265                if not self.bundle_seconds:
266                    if self.return_das_record:
267                        result = to_das_record_list(result, data_id=self.data_id)
268                    return result
269                else:
270                    result_list.append(result)
271            except queue.Empty:
272                logging.debug('get() timed out - trying again')
273                pass
274
275            # If we've been bundling long enough, return the list of records,
276            # or None if the list is empty.
277            now = time.time()
278            if now > start_time + self.bundle_seconds:
279                if self.return_das_record:
280                    result_list = to_das_record_list(result_list, data_id=self.data_id)
281                return result_list
282
283        # If we've fallen out because of a quit...
284        return None