openrvdas.logger.readers.websocket_reader

No module-level documentation available.
  1#!/usr/bin/env python3
  2
  3import asyncio
  4import logging
  5import queue
  6import ssl
  7import threading
  8import time
  9import websockets
 10
 11# Compatibility for websockets library versions < 10.0 and >= 10.0
 12# InvalidStatusCode was renamed to InvalidStatus in version 10.0
 13try:
 14    from websockets import InvalidStatus
 15except ImportError:
 16    from websockets.exceptions import InvalidStatus
 17
 18WS_InvalidStatus = InvalidStatus
 19
 20DEFAULT_SERVER_WEBSOCKET = 'localhost:8766'
 21
 22
 23################################################################################
 24class WebsocketReader():
 25    """Connect to a websocket served by a WebsocketWriter, and service read()
 26    requests from it.
 27    """
 28
 29    def __init__(self, uri, check_cert=False, **kwargs):
 30        """
 31        ```
 32        uri -      Hostname, port and protocol, (e.g. wss://localhost:8080) at which
 33                   to try to connect to a WebsocketWriter
 34
 35        check_cert  - If True, and uri protocol is 'wss', check the server's TLS certificate
 36                      for validity; if a str, use as local filepath location of .pem
 37                      file to check against.
 38        ```
 39        """
 40        super().__init__(**kwargs)
 41
 42        self.uri = uri
 43        self.check_cert = check_cert
 44
 45        # We won't initialize our websocket until the first read()
 46        # call. At that point we'll launch an async process in a separate
 47        # thread that will wait for data from the websocket and put it in
 48        # a queue that read() will pop from.
 49        self.websocket_thread = None
 50        self.queue = queue.Queue()
 51        self.quit_flag = False
 52
 53    ############################
 54    def _start_websocket(self):
 55        """We'll run this in a separate thread as soon as we get our first
 56        call to read()."""
 57
 58        ############################
 59        async def _websocket_loop(self):
 60            """Asynchronous inner function that will connect to websocket,
 61            iteratively try to read from it and put the result in our queue.
 62            """
 63            # Iterate if we lose the websocket for some reason other than a 'quit'
 64            while not self.quit_flag:
 65                try:
 66                    if self.uri.find('wss://') == 0:  # using wss
 67                        # If check_cert is a str, take it as the location of the
 68                        # .pem file we'll check for validity. Otherwise, if not
 69                        # False, take as a bool to verify by own means.
 70                        ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
 71                        if self.check_cert:
 72                            if isinstance(self.check_cert, str):
 73                                ssl_context.load_verify_locations(self.check_cert)
 74                            else:
 75                                ssl_context.verify_mode = ssl.CERT_REQUIRED
 76                        else:
 77                            ssl_context.check_hostname = False
 78                            ssl_context.verify_mode = ssl.CERT_NONE
 79
 80                    else:  # not using wss
 81                        ssl_context = None
 82
 83                    logging.debug(f'WebsocketReader connecting to {self.uri}')
 84                    async with websockets.connect(self.uri, ssl=ssl_context) as ws:  # type: ignore
 85                        logging.info(f'Connected to WebsocketWriter at {self.uri}')
 86
 87                        while not self.quit_flag:
 88                            record = await ws.recv()
 89                            logging.debug(f'WebsocketReader got record {record}')
 90                            self.queue.put(record)
 91
 92                except BrokenPipeError:
 93                    logging.info(f'WebsocketReader BrokenPipeError connecting to {self.uri}')
 94                    pass
 95                except AttributeError as e:
 96                    logging.warning(f'WebsocketReader websocket loop error: {e}')
 97                except websockets.exceptions.ConnectionClosed:  # type: ignore
 98                    logging.info('WebsocketReader lost websocket connection to '
 99                                 'data server; trying to reconnect.')
100                    await asyncio.sleep(0.2)
101                # FIX: Use the compatibility alias defined at module level
102                except WS_InvalidStatus:  # type: ignore
103                    logging.info('WebsocketReader InvalidStatus/Code connecting to '
104                                 'data server; trying to reconnect.')
105                    await asyncio.sleep(0.2)
106                except OSError as e:
107                    logging.info('Unable to connect to websocket. Sleeping to try again...')
108                    logging.info(f'Connection error: {e}')
109                    await asyncio.sleep(2)
110
111        # In the outer function, get a new event loop and fire up the
112        # inner, async routine.
113        # Could we also use asyncio.ensure_future(_websocket_loop(self)) ?
114        websocket_event_loop = asyncio.new_event_loop()
115        websocket_event_loop.run_until_complete(_websocket_loop(self))
116        websocket_event_loop.close()
117
118    ############################
119    def quit(self, seconds=0):
120        """Sleep N seconds, then signal quit."""
121        time.sleep(seconds)
122        self.quit_flag = True
123
124    ############################
125    def read(self):
126        """Read/wait for data from the websocket."""
127
128        # If we've not yet fired up the websocket thread, do that now.
129        if not self.websocket_thread:
130            self.websocket_thread = threading.Thread(
131                name='websocket_thread',
132                target=self._start_websocket,
133                daemon=True)
134            self.websocket_thread.start()
135
136        # Use a timeout in our queue get() so we can periodically check if
137        # we've gotten a 'quit'
138        while not self.quit_flag:
139            try:
140                result = self.queue.get(block=True, timeout=2)
141                logging.debug('Got result from queue: %s', result)
142                return result
143            except queue.Empty:
144                logging.debug('get() timed out - trying again')
145                pass
146
147        # If we've fallen out because of a quit...
148        return None
WS_InvalidStatus = <class 'websockets.exceptions.InvalidStatus'>
DEFAULT_SERVER_WEBSOCKET = 'localhost:8766'
class WebsocketReader:
 25class WebsocketReader():
 26    """Connect to a websocket served by a WebsocketWriter, and service read()
 27    requests from it.
 28    """
 29
 30    def __init__(self, uri, check_cert=False, **kwargs):
 31        """
 32        ```
 33        uri -      Hostname, port and protocol, (e.g. wss://localhost:8080) at which
 34                   to try to connect to a WebsocketWriter
 35
 36        check_cert  - If True, and uri protocol is 'wss', check the server's TLS certificate
 37                      for validity; if a str, use as local filepath location of .pem
 38                      file to check against.
 39        ```
 40        """
 41        super().__init__(**kwargs)
 42
 43        self.uri = uri
 44        self.check_cert = check_cert
 45
 46        # We won't initialize our websocket until the first read()
 47        # call. At that point we'll launch an async process in a separate
 48        # thread that will wait for data from the websocket and put it in
 49        # a queue that read() will pop from.
 50        self.websocket_thread = None
 51        self.queue = queue.Queue()
 52        self.quit_flag = False
 53
 54    ############################
 55    def _start_websocket(self):
 56        """We'll run this in a separate thread as soon as we get our first
 57        call to read()."""
 58
 59        ############################
 60        async def _websocket_loop(self):
 61            """Asynchronous inner function that will connect to websocket,
 62            iteratively try to read from it and put the result in our queue.
 63            """
 64            # Iterate if we lose the websocket for some reason other than a 'quit'
 65            while not self.quit_flag:
 66                try:
 67                    if self.uri.find('wss://') == 0:  # using wss
 68                        # If check_cert is a str, take it as the location of the
 69                        # .pem file we'll check for validity. Otherwise, if not
 70                        # False, take as a bool to verify by own means.
 71                        ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
 72                        if self.check_cert:
 73                            if isinstance(self.check_cert, str):
 74                                ssl_context.load_verify_locations(self.check_cert)
 75                            else:
 76                                ssl_context.verify_mode = ssl.CERT_REQUIRED
 77                        else:
 78                            ssl_context.check_hostname = False
 79                            ssl_context.verify_mode = ssl.CERT_NONE
 80
 81                    else:  # not using wss
 82                        ssl_context = None
 83
 84                    logging.debug(f'WebsocketReader connecting to {self.uri}')
 85                    async with websockets.connect(self.uri, ssl=ssl_context) as ws:  # type: ignore
 86                        logging.info(f'Connected to WebsocketWriter at {self.uri}')
 87
 88                        while not self.quit_flag:
 89                            record = await ws.recv()
 90                            logging.debug(f'WebsocketReader got record {record}')
 91                            self.queue.put(record)
 92
 93                except BrokenPipeError:
 94                    logging.info(f'WebsocketReader BrokenPipeError connecting to {self.uri}')
 95                    pass
 96                except AttributeError as e:
 97                    logging.warning(f'WebsocketReader websocket loop error: {e}')
 98                except websockets.exceptions.ConnectionClosed:  # type: ignore
 99                    logging.info('WebsocketReader lost websocket connection to '
100                                 'data server; trying to reconnect.')
101                    await asyncio.sleep(0.2)
102                # FIX: Use the compatibility alias defined at module level
103                except WS_InvalidStatus:  # type: ignore
104                    logging.info('WebsocketReader InvalidStatus/Code connecting to '
105                                 'data server; trying to reconnect.')
106                    await asyncio.sleep(0.2)
107                except OSError as e:
108                    logging.info('Unable to connect to websocket. Sleeping to try again...')
109                    logging.info(f'Connection error: {e}')
110                    await asyncio.sleep(2)
111
112        # In the outer function, get a new event loop and fire up the
113        # inner, async routine.
114        # Could we also use asyncio.ensure_future(_websocket_loop(self)) ?
115        websocket_event_loop = asyncio.new_event_loop()
116        websocket_event_loop.run_until_complete(_websocket_loop(self))
117        websocket_event_loop.close()
118
119    ############################
120    def quit(self, seconds=0):
121        """Sleep N seconds, then signal quit."""
122        time.sleep(seconds)
123        self.quit_flag = True
124
125    ############################
126    def read(self):
127        """Read/wait for data from the websocket."""
128
129        # If we've not yet fired up the websocket thread, do that now.
130        if not self.websocket_thread:
131            self.websocket_thread = threading.Thread(
132                name='websocket_thread',
133                target=self._start_websocket,
134                daemon=True)
135            self.websocket_thread.start()
136
137        # Use a timeout in our queue get() so we can periodically check if
138        # we've gotten a 'quit'
139        while not self.quit_flag:
140            try:
141                result = self.queue.get(block=True, timeout=2)
142                logging.debug('Got result from queue: %s', result)
143                return result
144            except queue.Empty:
145                logging.debug('get() timed out - trying again')
146                pass
147
148        # If we've fallen out because of a quit...
149        return None

Connect to a websocket served by a WebsocketWriter, and service read() requests from it.

WebsocketReader(uri, check_cert=False, **kwargs)
30    def __init__(self, uri, check_cert=False, **kwargs):
31        """
32        ```
33        uri -      Hostname, port and protocol, (e.g. wss://localhost:8080) at which
34                   to try to connect to a WebsocketWriter
35
36        check_cert  - If True, and uri protocol is 'wss', check the server's TLS certificate
37                      for validity; if a str, use as local filepath location of .pem
38                      file to check against.
39        ```
40        """
41        super().__init__(**kwargs)
42
43        self.uri = uri
44        self.check_cert = check_cert
45
46        # We won't initialize our websocket until the first read()
47        # call. At that point we'll launch an async process in a separate
48        # thread that will wait for data from the websocket and put it in
49        # a queue that read() will pop from.
50        self.websocket_thread = None
51        self.queue = queue.Queue()
52        self.quit_flag = False
uri -      Hostname, port and protocol, (e.g. wss://localhost:8080) at which
           to try to connect to a WebsocketWriter

check_cert  - If True, and uri protocol is 'wss', check the server's TLS certificate
              for validity; if a str, use as local filepath location of .pem
              file to check against.
uri
check_cert
websocket_thread
queue
quit_flag
def quit(self, seconds=0):
120    def quit(self, seconds=0):
121        """Sleep N seconds, then signal quit."""
122        time.sleep(seconds)
123        self.quit_flag = True

Sleep N seconds, then signal quit.

def read(self):
126    def read(self):
127        """Read/wait for data from the websocket."""
128
129        # If we've not yet fired up the websocket thread, do that now.
130        if not self.websocket_thread:
131            self.websocket_thread = threading.Thread(
132                name='websocket_thread',
133                target=self._start_websocket,
134                daemon=True)
135            self.websocket_thread.start()
136
137        # Use a timeout in our queue get() so we can periodically check if
138        # we've gotten a 'quit'
139        while not self.quit_flag:
140            try:
141                result = self.queue.get(block=True, timeout=2)
142                logging.debug('Got result from queue: %s', result)
143                return result
144            except queue.Empty:
145                logging.debug('get() timed out - trying again')
146                pass
147
148        # If we've fallen out because of a quit...
149        return None

Read/wait for data from the websocket.