openrvdas.server.websocket_server

  1#!/usr/bin/env python3
  2"""A simple websocket server. Producer and consumer routines may be
  3passed in that will generate messages to be sent to the websocket
  4(producer) and will take messages retrieved from the websocket and
  5process them (consumer).
  6
  7The attached executable script uses default routines that look for
  8messages from a producer_queue to send, and add received messages to a
  9consumer_queue. When given a command line, the script's behavior is to
 10send it out to *all* attached clients. Retrieved messages are queued,
 11but nothing is done with them.
 12
 13To run, try:
 14
 15    server/websocket_server.py --websocket localhost:8765 -v
 16
 17In a second window, start up a LoggerRunner as a websocket client:
 18
 19    server/logger_runner.py --websocket localhost:8765 \
 20            --host_id client.host -v
 21
 22You should see messages from the LoggerRunner appear on the
 23websocket_server's console, beginning with an identifying message,
 24then repeated status messages. If you type a command into the
 25websocket console, it will be relayed to the LoggerRunner, which (most
 26likely) will complain that it is an unrecognized command. (The only
 27commands the LoggerRunner recognizes over the websocket are 'quit' and
 28'set_configs' followed by a JSON encoding of a complete set of logger
 29configurations).
 30
 31To explore the WebsocketServer in context, please look at the
 32documentation for LoggerManager, which uses a WebsocketServer to
 33dispatch configurations to client LoggerRunners.
 34"""
 35import asyncio
 36import logging
 37import queue
 38import threading
 39import time
 40import websockets
 41
 42
 43################################################################################
 44class WebsocketServer:
 45    ############################
 46    def __init__(self, host, port, consumer, producer,
 47                 on_connect=None, on_disconnect=None):
 48        """host, port - host and port to open as websocket
 49
 50        consumer - async routine that takes a str argument and a
 51           client_id and does something with it. Strings retrieved from
 52           the websocket will be passed to this routine.
 53
 54        producer - async routine that takes a client_id and produces the
 55           strings we want to send out to that client on the websocket.
 56
 57        on_connect - optional routine to be called when a new client connects.
 58           Should take three parameters:
 59               websocket - the websocket itself, in case someone wants to store it
 60               client_id - an integer client id
 61               path - the path with which the client connected
 62
 63        on_disconnect - optional routine to be called when a client
 64           disconnects. Should take a single integer representing client's
 65           unique client_id.
 66        """
 67        self.host = host
 68        self.port = port
 69        self.consumer = consumer
 70        self.producer = producer
 71
 72        self.on_connect = on_connect
 73        self.on_disconnect = on_disconnect
 74
 75        self.client_lock = threading.Lock()
 76        self.num_clients = 0
 77        self.client_map = {}
 78
 79        self.quit_requested = False
 80
 81    ############################
 82    async def _consumer_handler(self, websocket, client_id):
 83        try:
 84            async for message in websocket:
 85                logging.debug('WebsocketServer received message: ' + message)
 86                await self.consumer(message, client_id)
 87
 88                if self.quit_requested:
 89                    return
 90        except:  # noqa E722
 91            logging.info('Websocket connection lost')
 92
 93    ############################
 94    async def _producer_handler(self, websocket, client_id):
 95        """Here, we could either await some other producer to give us a command
 96        or poll."""
 97        # If we're waiting for an external command producer:
 98        try:
 99            while not self.quit_requested:
100                message = await self.producer(client_id)
101                if message:
102                    await websocket.send(message)
103                    logging.debug('WebsocketServer sent message: ' + message)
104                else:
105                    await asyncio.sleep(1)
106        except websockets.exceptions.ConnectionClosed:
107            logging.info('Websocket connection lost')
108
109    ############################
110    async def _handler(self, websocket, path):
111        with self.client_lock:
112            client_id = self.num_clients
113            logging.warning('New client #%d attached', client_id)
114            self.client_map[client_id] = websocket
115            self.num_clients += 1
116            if self.on_connect:
117                self.on_connect(websocket, client_id, path)
118
119        tasks = []
120
121        # Task that receives data from websocket and does something with it
122        if self.consumer:
123            tasks.append(asyncio.ensure_future(self._consumer_handler(websocket,
124                                                                      client_id)))
125        # Task that produces data that we're going to send out on websocket
126        if self.producer:
127            tasks.append(asyncio.ensure_future(self._producer_handler(websocket,
128                                                                      client_id)))
129        done, pending = await asyncio.wait(tasks,
130                                           return_when=asyncio.FIRST_COMPLETED)
131        for task in pending:
132            task.cancel()
133
134        # When client disconnects, delete the queues it was using
135        with self.client_lock:
136            logging.info('WebsocketServer client #%d completed', client_id)
137            del self.client_map[client_id]
138            if self.on_disconnect:
139                self.on_disconnect(client_id)
140
141    ############################
142    def clients(self):
143        """Return a dict mapping client_id->websocket."""
144        return self.client_map
145
146    ############################
147    def run(self):
148        start_server = websockets.serve(self._handler, self.host, self.port)
149        asyncio.get_event_loop().run_until_complete(start_server)
150        asyncio.get_event_loop().run_forever()
151
152    ############################
153    def quit(self):
154        """NOTE: This doesn't really shut things down because the event loop
155        keeps running."""
156        self.quit_requested = True
157
158
159################################################################################
160"""Below are tools for a standalone executable that takes messages
161from the command line and sends them to every connected client, and
162receives messages from the client websockets and prints them to the
163console.
164
165To remain non-blocking, it operates via a pair of queues:
166
167  send_queue = {}     # client_id->queue for messages to send to ws
168  receive_queue = {}  # client_id->queue for messages from ws
169
170The websocket server is initialized with a pair of routines that use
171these queues:
172
173  queued_consumer - a consumer that takes a message and "consumes"
174                    them, in this case by pushing them on the
175                    receive_queue for others to process
176
177  queued_producer - a routine that "produces" a message (by pulling
178                    one off the send_queue) for the ws server to send.
179"""
180
181# Queues for our queued consumers/producers
182send_queue = {}     # client_id->queue for messages to send to ws
183receive_queue = {}  # client_id->queue for messages from ws
184
185# A lock to make sure only one thread is messing with the above
186# maps at any given time.
187websocket_map_lock = threading.Lock()
188
189##########################
190# This is a consumer - it takes a message and does something with it
191
192
193async def queued_consumer(message, client_id):
194    global receive_queue
195    logging.debug('Received message from client #%d: %s', client_id, message)
196    receive_queue[client_id].put(message)
197
198############################
199# This is a producer - it produces a message (from the queue) to send
200
201
202async def queued_producer(client_id):
203    global send_queue
204    while True:
205        try:
206            message = send_queue[client_id].get_nowait()
207            if message.strip():
208                logging.debug('Sending message to client #%d: %s', client_id, message)
209                return message
210        except queue.Empty:
211            await asyncio.sleep(0.1)
212
213############################
214
215
216def register_websocket_client(websocket, client_id, path):
217    """We've been alerted that a websocket client has connected.
218    Register it properly."""
219    global send_queue, receive_queue, websocket_map_lock
220    with websocket_map_lock:
221        if client_id not in send_queue:
222            send_queue[client_id] = queue.Queue()
223            receive_queue[client_id] = queue.Queue()
224        logging.warning('Websocket client #%d has connected', client_id)
225
226############################
227
228
229def unregister_websocket_client(client_id):
230    """We've been alerted that a websocket client has disconnected.
231    Unegister it properly."""
232    global send_queue, receive_queue, websocket_map_lock
233    with websocket_map_lock:
234        if client_id in send_queue:
235            del send_queue[client_id]
236        if client_id in receive_queue:
237            del receive_queue[client_id]
238        logging.warning('Websocket client #%d has disconnected', client_id)
239
240############################
241
242
243def read_commands():
244    while not server.quit_requested:
245        command = input('Command? ')
246        for client_id, sender in send_queue.items():
247            logging.warning('Pushing command to client %d: %s', client_id, command)
248            send_queue[client_id].put(command)
249
250        if command == 'quit':
251            logging.warning('Quitting!')
252            server.quit()
253
254############################
255
256
257def process_results():
258    SHOW_LEN = 30
259    while not server.quit_requested:
260        with server.client_lock:
261            for client_id, receiver in receive_queue.items():
262                try:
263                    message = receive_queue[client_id].get_nowait()
264                    logging.info('#%d: %s%s', client_id, message[:SHOW_LEN],
265                                 '' if len(message) < SHOW_LEN else '...')
266                except queue.Empty:
267                    pass
268            time.sleep(0.1)
269
270
271################################################################################
272if __name__ == '__main__':
273    import argparse
274
275    parser = argparse.ArgumentParser()
276    parser.add_argument('--websocket', dest='websocket', action='store',
277                        required=True, type=str,
278                        help='Attempt to open specified host:port as websocket '
279                        'and begin reading/writing data on it.')
280
281    parser.add_argument('-v', '--verbosity', dest='verbosity',
282                        default=0, action='count',
283                        help='Increase output verbosity')
284    args = parser.parse_args()
285
286    # Set logger format and verbosity
287    LOGGING_FORMAT = '%(asctime)-15s %(filename)s:%(lineno)d %(message)s'
288    logging.basicConfig(format=LOGGING_FORMAT)
289    LOG_LEVELS = {0: logging.WARNING, 1: logging.INFO, 2: logging.DEBUG}
290    args.verbosity = min(args.verbosity, max(LOG_LEVELS))
291    logging.getLogger().setLevel(LOG_LEVELS[args.verbosity])
292
293    try:
294        host, port_str = args.websocket.split(':')
295        port = int(port_str)
296    except ValueError:
297        logging.error('--websocket arg "%s" not in host:port format', args.websocket)
298        exit(1)
299
300    # Create the websocket server, setting up queued senders/receivers
301    server = WebsocketServer(host=host, port=port,
302                             consumer=queued_consumer, producer=queued_producer,
303                             on_connect=register_websocket_client,
304                             on_disconnect=unregister_websocket_client)
305
306    read_command_thread = threading.Thread(target=read_commands)
307    read_command_thread.start()
308    process_results_thread = threading.Thread(target=process_results)
309    process_results_thread.start()
310
311    # Start websocket server
312    try:
313        server.run()
314    except KeyboardInterrupt:
315        logging.warning('Got interrupt')
316
317    read_command_thread.join()
318    process_results_thread.join()