How to handle multiple websocket messages at the same time?

Question

If anyone could help me with Python and async/await, any help would be much appreciated!

I need to listen to a websocket for messages, so I set up the following code:

import websockets
import asyncio

my_socket = "ws://......."

# I set a "while True" here to reconnect websocket if it stop for any reason
while True:
    try:
        async with websockets.connect(my_socket) as ws:
            # I set a "while True" here to keep listening to messages forever
            while True:
                await on_message(await ws.recv())
    # If websocket gets closed for any reason, we catch exception and wait before new loop
    except Exception as e:
        print(e)
    # Wait 10 secs before new loop to avoid flooding server if it is unavailable for any reason
    await asyncio.sleep(10)

async def on_message(message):
    # Do what needs to be done with received message
    # This function is running for a few minutes, with a lot of sleep() time in it..
    # .. so it does no hold process for itself

What I would like to do is:

• Listen to messages
• As soon as a message arrives, apply various actions with on_message() function, for several minutes
• Keep listening to messages while previous messages are still in process with on_message()

What actually happens:

• Listen to messages
• Receive a message and start on_message() function
• And then program is waiting for on_message() function to end before receiving any new message, which takes a few minutes, and make the second message late and so on

I do understand why it does this, as await on_message() clearly says : wait for on_message() to end so it won't go back to listen for new message. The thing I don't know, is how I could handle messages without having to wait for this function to end.

My on_message() function has a lot of idle time with some await asyncio.sleep(1), so I know that I can run multiple task in the same time.

So, how could I keep be listening to new messages while running tasks for the first one?

Answer 1

In short, you need to change await on_message(await ws.recv()) to asyncio.create_task(on_message(await ws.recv())).

As you correctly pointed out, await doesn't work for you because it implies waiting for the task to finish. Although the code is async, in the sense that it's driven by the event loop and that you could start a number of such tasks in parallel, each individual loop is sequential.

The alternative to await is to spawn the job in the background using asyncio.create_task(). This will create a task that will execute the coroutine in pieces (each piece between two awaits that suspend) interspersed with equivalent pieces of other active coroutines. create_task() will return a handle to the task that you can (and possibly at some point should) await to wait for the task to finish and obtain its result or exception. Since in your case you don't care about the result, you don't even need to store the task.