Best approach to multiple websocket client connections in Python?

Viewed 3176

I appreciate that the question I am about to ask is rather broad but, as a newcomer to Python, I am struggling to find the [best] way of doing something which would be trivial in, say, Node.js, and pretty trivial in other environments such as C#.

Let's say that there is a warehouse full of stuff. And let's say that there is a websocket interface onto that warehouse with two characteristics: on client connection it pumps out a full list of the warehouse's current inventory, and it then follows that up with further streaming updates when the inventory changes.

The web is full of examples of how, in Python, you connect to the warehouse and respond to changes in its state. But...

What if I want to connect to two warehouses and do something based on the combined information retrieved separately from each one? And what if I want to do things based on factors such as time, rather than solely being driven by inventory changes and incoming websocket messages?

In all the examples I've seen - and it's beginning to feel like hundreds - there is, somewhere, in some form, a run() or a run_forever() or a run_until_complete() etc. In other words, the I/O may be asynchronous, but there is always a massive blocking operation in the code, and always two fundamental assumptions which don't fit my case: that there will only be one websocket connection, and that all processing will be driven by events sent out by the [single] websocket server.

It's very unclear to me whether the answer to my question is some sort of use of multiple event loops, or of multiple threads, or something else.

To date, experimenting with Python has felt rather like being on the penthouse floor, admiring the quirky but undeniably elegant decor. But then you get in the elevator, press the button marked "parallelism" or "concurrency", and the evelator goes into freefall, eventually depositing you in a basement filled with some pretty ugly and steaming pipes.

... Returning from flowery metaphors back to the technical, the key thing I'm struggling with is the Python equivalent of, say, Node.js code which could be as trivially simple as the following example [left inelegant for simplicity]:

var aggregateState = { ... some sort of representation of combined state ... };

var socket1 = new WebSocket("wss://warehouse1");
socket1.on("message", OnUpdateFromWarehouse);

var socket2 = new WebSocket("wss://warehouse2");
socket2.on("message", OnUpdateFromWarehouse);

function OnUpdateFromWarehouse(message)
{
  ... Take the information and use it to update aggregate state from both warehouses ...
}


2 Answers

Answering my own question, in the hope that it may help other Python newcomers... asyncio seems to be the way to go (though there are gotchas such as the alarming ease with which you can deadlock the event loop).

Assuming the use of an asyncio-friendly websocket module such as websockets, what seems to work is a framework along the following lines - shorn, for simplicity, of logic such as reconnects. (The premise remains a warehouse which sends an initial list of its full inventory, and then sends updates to that initial state.)

class Warehouse:
    def __init__(self, warehouse_url):
        self.warehouse_url = warehouse_url
        self.inventory = {}  # Some description of the warehouse's inventory
    
    async def destroy():
        if (self.websocket.open):
            self.websocket.close()  # Terminates any recv() in wait_for_incoming() 
            await self.incoming_message_task  # keep asyncio happy by awaiting the "background" task

    async def start(self):
        try:
            # Connect to the warehouse
            self.websocket = await connect(self.warehouse_url)          
            # Get its initial message which describes its full state
            initial_inventory = await self.websocket.recv()
            # Store the initial inventory
            process_initial_inventory(initial_inventory)
            # Set up a "background" task for further streaming reads of the web socket
            self.incoming_message_task = asyncio.create_task(self.wait_for_incoming())
            # Done
            return True
        except:
            # Connection failed (or some unexpected error)
            return False

    async def wait_for_incoming(self):
        while self.websocket.open:
            try:
                update_message = await self.websocket.recv()
                asyncio.create_task(self.process_update_message(update_message))
            except:
                # Presumably, socket closure
                pass

    def process_initial_inventory(self, initial_inventory_message):
        ... Process initial_inventory_message into self.inventory ...
    
    async def process_update_message(self, update_message):
        ... Merge update_message into self.inventory ...
        ... And fire some sort of event so that the object's 
        ... creator can detect the change. There seems to be no ...
        ... consensus about what is a pythonic way of implementing events, ... 
        ... so I'll declare that - potentially trivial - element as out-of-scope ...

After completing the initial connection logic, one key thing is setting up a "background" task which repeatedly reads further update messages coming in over the websocket. The code above doesn't include any firing of events, but there are all sorts of ways in which process_update_message() can/could do this (many of them trivially simple), allowing the object's creator to deal with notifications whenever and however it sees fit. The streaming messages will continue to be received, and any events will be continued to be fired, for as long as the object's creator continues to play nicely with asyncio and to participate in co-operative multitasking.

With that in place, a connection can be established along the following lines:

async def main():
    warehouse1 = Warehouse("wss://warehouse1")
    if await warehouse1.start():
        ... Connection succeeded. Update messages will now be processed 
        in the "background" provided that other users of the event loop 
        yield in some way ...
    else:
        ... Connection failed ...

asyncio.run(main())

Multiple warehouses can be initiated in several ways, including doing a create_task(warehouse.start()) on each one and then doing a gather on the tasks to ensure/check that they're all okay.

When it's time to quit, to keep asyncio happy, and to stop it complaining about orphaned tasks, and to allow everything to shut down nicely, it's necessary to call destroy() on each warehouse.

But there's one common element which this doesn't cover. Extending the original premise above, let's say that the warehouse also accepts requests from our websocket client, such as "ship X to Y". The success/failure responses to these requests will come in alongside the general update messages; it generally won't be possible to guarantee that the first recv() after the send() of a request will be the response to that request. This complicates process_update_message().

The best answer I've found may or may not be considered "pythonic" because it uses a Future in a way which is strongly analogous to a TaskCompletionSource in .NET.

Let's invent a couple of implementation details; any real-world scenario is likely to look something like this:

  • We can supply a request_id when submitting an instruction to the warehouse
  • The success/failure response from the warehouse repeats the request_id back to us (and thus also distinguishing between command-response messages versus inventory-update messages)

The first step is to have a dictionary which maps the ID of pending, in-progress requests to Future objects:

    def __init__(self, warehouse_url):
        ...
        self.pending_requests = {}

The definition of a coroutine which sends a request then looks something like this:

    async def send_request(self, some_request_definition)
        # Allocate a unique ID for the  request
        request_id = <some unique request id>
        # Create a Future for the pending request
        request_future = asyncio.Future()
        # Store the map of the ID -> Future in the dictionary of pending requests
        self.pending_requests[request_id] = request_future
        # Build a request message to send to the server, somehow including the request_id
        request_msg = <some request definition, including the request_id>
        # Send the message 
        await self.websocket.send(request_msg) 
        # Wait for the future to complete - we're now asynchronously awaiting
        # activity in a separate function
        await asyncio.wait_for(command_future, timeout = None)
        # Return the result of the Future as the return value of send_request()
        return request_future.result()

A caller can create a request and wait for its asynchronous response using something like the following:

     some_result = await warehouse.send_request(<some request def>)

The key to making this all work is then to modify and extend process_update_message() to do the following:

  • Distinguish between request responses versus inventory updates
  • For the former, extract the request ID (which our invented scenario says gets repeated back to us)
  • Look up the pending Future for the request
  • Do a set_result() on it (whose value can be anything depending on what the server's response says). This releases send_request() and causes the await from it to be resolved.

For example:

    async def process_update_message(self, update_message):
        if <some test that update_message is a request response>:
            request_id = <extract the request ID repeated back in update_message>
            # Get the Future for this request ID
            request_future = self.pending_requests[request_id]
            # Create some sort of return value for send_request() based on the response
            return_value = <some result of the request>
            # Complete the Future, causing send_request() to return
            request_future.set_result(return_value)
        else:
            ... handle inventory updates as before ...

I've not used sockets with asyncio, but you're likely just looking for asyncio's open_connection

async def socket_activity(address, callback):
    reader, _ = await asyncio.open_connection(address)
    while True:
        message = await reader.read()
        if not message:  # empty bytes on EOF
            break  # connection was closed
        await callback(message)

Then add these to the event loop

    tasks = []  # keeping a reference prevents these from being garbage collected
    for address in ["wss://warehouse1", "wss://warehouse2"]:
        tasks.append(asyncio.create_task(
            socket_activity(address, callback)
        ))
    # return tasks  # or work with them

If you want to wait in a coroutine until N operations are complete, you can use .gather()

Alternatively, you may find Tornado does everything you want and more (I based my Answer off this one)
Tornado websocket client: how to async on_message? (coroutine was never awaited)

Related