About explicitly trans websocket

[vanker-x]

Hello,
I have a problem when using emmett when I tried to add websocket object to the list,
the new connected websocket object would overwrite the original websocket object.
Have you considered explicitly transferring websocket object to avoid this problem?

from emmett import websocket
from emmett import App

app = App(__name__)


class Connections:
    def __init__(self):
        self.connections = []

    def add(self, connection):
        self.connections.append(connection)

    async def broadcast(self, value):
        for connection in self.connections:
            await connection.send(f"broadcast_message:{value}")


connections = Connections()


@app.websocket("/")
async def echo():
    connections.add(websocket)
    print(connections.connections)
    while True:
        message = await websocket.receive()
        if message == "broadcast":
            await connections.broadcast("broadcast")
            continue
        await websocket.send(f"received:{message}")


"""
The std output:

[INFO] Starting granian (main PID: 4404)
[INFO] Listening at: http://127.0.0.1:8000
[INFO] Spawning worker-1 with pid: 4406
[INFO] Started worker-1
[INFO] Started worker-1 runtime-1
[<emmett.rsgi.wrappers.Websocket object at 0x1048115f0>]
[<emmett.rsgi.wrappers.Websocket object at 0x104e0e9e0>, <emmett.rsgi.wrappers.Websocket object at 0x104e0e9e0>]
"""

[gi0baro]

@vanker-x the websocket object in Emmett is a wrapper for a contextvar, so theoretically using the copy module might work, but I suggest you to avoid that path.

Here is how I would personally implement the above code (added json just for the sake of routing messages):

import asyncio
import uuid

from emmett import App, Pipe, websocket
from emmett.tools import service

app = App(__name__)


class Broadcast:
    def __init__(self):
        self.consumers = {}

    def new_consumer(self):
        consumer = Consumer()
        self.consumers[consumer.uuid] = consumer
        return consumer

    def del_consumer(self, consumer):
        self.consumers.pop(consumer.uuid, None)

    def send(self, message):
        for consumer in self.consumers.values():
            consumer.push(message)


class Consumer:
    def __init__(self):
        self.uuid = uuid.uuid4()
        self.queue = asyncio.Queue()

    def push(self, message):
        self.queue.put_nowait(message)

    async def consume(self):
        while message := await self.queue.get():
            yield message


class ConsumerPipe(Pipe):
    def __init__(self) -> None:
        self.broadcast = Broadcast()

    async def open(self):
        websocket.broadcast = self.broadcast
        websocket.consumer = self.broadcast.new_consumer()

    async def close(self):
        self.broadcast.del_consumer(websocket.consumer)


@app.websocket("/", pipeline=[ConsumerPipe()])
@service.json
async def echo():
    await websocket.accept()

    async def channel():
        async for message in websocket.consumer.consume():
            await websocket.send({"type": "broadcast", "data": message})

    async def echo_loop():
        while True:
            message = await websocket.receive()
            if message["type"] == "broadcast":
                websocket.broadcast.send(message["data"])
                continue
            await websocket.send(message)

    await asyncio.gather(echo_loop(), channel())

and this is a quick test code for it:

import asyncio
import json
import websockets


async def loop(cli_no):
    async with websockets.connect(
        'ws://localhost:8000/'
    ) as websocket:
        async def recv():
            async for message in websocket:
                print(f"recv @ cli {cli_no}: {message}")

        async def send():
            while True:
                await websocket.send(json.dumps({"type": "message", "data": "foo"}))
                await websocket.send(json.dumps({"type": "broadcast", "data": f"foo from {cli_no}"}))
                await asyncio.sleep(0.2)

        await asyncio.gather(send(), recv())


if __name__ == "__main__":
    async def main():
        await asyncio.gather(loop(1), loop(2))

    asyncio.run(main())

Please mind that this implementation only works with a single Emmett worker; for production usage you should add some external service like redis to broadcast messages between clients connected to different workers.

[vanker-x]

thanks

[vanker-x]

The example you gave is very good, but I think this method is somewhat different from emmett's 'easy to use' feature. Maybe I can discuss with you to make it better.

[gi0baro]

Which part you find hard?
Even if context variables like websocket would have been transferable between tasks, you still need multiple asyncio tasks in order to handle a 4-way communication. In your original code the socket would block sending and receiving from clients while you deal with broadcasting..

[vanker-x]

Yes, the broadcast I wrote is a block code, but the point is not here, can we have a similar or better way to deal with Fastapi? I will try to find a better way.Thank you for such a wonderful reply