Python websockets send to client and keep connection alive

Question

I'm using python websockets: https://websockets.readthedocs.io/

They have a simple client/server example, where the server echoes the client's input back once. The code looks like this:

Client side:

# WS client example

import asyncio
import websockets

async def hello():
    async with websockets.connect(
            'ws://localhost:8765') as websocket:
        name = input("What's your name? ")

        await websocket.send(name)
        print(f"> {name}")

        greeting = await websocket.recv()
        print(f"< {greeting}")

asyncio.get_event_loop().run_until_complete(hello())

Server side:

# WS server example

import asyncio
import websockets

async def hello(websocket, path):
    name = await websocket.recv()
    print(f"< {name}")

    greeting = f"Hello {name}!"

    await websocket.send(greeting)
    print(f"> {greeting}")

start_server = websockets.serve(hello, 'localhost', 8765)

asyncio.get_event_loop().run_until_complete(start_server)
asyncio.get_event_loop().run_forever()

I want to adapt just the server side only, so that it does the following upon a socket connection:

1. Send an acknowledgement message to the client. e.g. Hello Client! Please wait for your data.
2. Keep the connection alive.
3. Process some data that takes some time.
4. After the data has completed processing, notify the client on the existing websocket connection. e.g. Your data is here!

The python websockets documentation doesn't have a code sample which does this.

Answer 1

To keep the connection open do not terminate the handler after processing the first message. For example, you can have an endless-loop that will keep processing the incoming messages until the connection is closed by the client:

async def hello(websocket, path):
    while True:
        try:
            name = await websocket.recv()
        except websockets.ConnectionClosed:
            print(f"Terminated")
            break

        print(f"< {name}")
        greeting = f"Hello {name}!"

        await websocket.send(greeting)
        print(f"> {greeting}")

In the async fun you can then await any long running operation as suggested here.

You will however need to adapt both server and client side in the similar way. Your client also terminates after receiving the first message.

Answer 2

Presumably your function that processes the data is blocking, otherwise you'd simply await it inside the coroutine. The straightforward approach is to use run_in_executor to run it in another thread, and await it in your handler coroutine:

async def hello(websocket, path):
    loop = asyncio.get_event_loop()
    await websocket.send("Hello Client! Please wait for your data.")
    data = await loop.run_in_executor(None, get_data)
    await websocket.send("Your data is here!")
    await websocket.send(data)

def get_data():
    # something that takes a long time to calculate
    x = 19134702400093278081449423917**300000 % 256
    return bytes([x])