How to process WebSocket messages in parallel using Django Channels?

user_1234438 0
Viewed 784

We're getting started with Django Channels and are struggling with the following use case:

Our app receives multiple requests from a single client (another server) in a short time. Creating each response takes a long time. The order in which responses are sent to the client doesn't matter.

We want to keep an open WebSocket connection to reduce connection overhead for sending many requests and responses from and to the same client.

Django Channels seems to process messages on the same WebSocket connection strictly in order, and won't start processing the next frame before the previous one has been responded to.

Consider the following example:

Example

Server-side

import asyncio
from channels.generic.websocket import AsyncWebsocketConsumer

class QuestionConsumer(AsyncWebsocketConsumer):
    async def websocket_connect(self, event):
        await self.accept()

    async def complicated_answer(self, question):
        await asyncio.sleep(3)
        return {
             "What is the Answer to Life, The Universe and Everything?": "42",
             "Why?": "Because.",
        }.get(question, "Don't know")


    async def receive(self, text_data=None, bytes_data=None):
        # while awaiting below, we should start processing the next WS frame
        answer = await self.complicated_answer(text_data)
        await self.send(answer)

asgi.py:

from django.urls import re_path
from channels.routing import ProtocolTypeRouter, URLRouter
application = ProtocolTypeRouter(
    {"websocket": URLRouter([
        re_path(r"^questions", QuestionConsumer.as_asgi(), name="questions",)
    ]}
  )
)

Client-side

import asyncio
import websockets
from time import time

async def main():
    async with websockets.connect("ws://0.0.0.0:8000/questions") as ws:
        tasks = []
        for m in [
                "What is the Answer to Life, The Universe and Everything?",
                "Why?"
        ]:
            tasks.append(ws.send(m))
        # send all requests (without waiting for response)
        time_before = time()
        await asyncio.gather(*tasks)
        # wait for responses
        for t in tasks:
            print(await ws.recv())
            print("{:.1f} seconds since first request".format(time() - time_before))
asyncio.get_event_loop().run_until_complete(main())

Result

Actual

42
3.0 seconds since first request
Because.
6.0 seconds since first request

Desired

42
3.0 seconds since first request
Because.
3.0 seconds since first request

In other words, we would like the event loop to switch between async tasks not only for multiple consumers, but also for all tasks handled by the same consumer. Is this possible or is there a workaround we are overlooking? Have you used Django Channels for similar challenges and how did you solve them?

django-channels websocket python
1 Answers

The consumer's receive function is called sequentially for each incoming WebSocket message, and when the await of the first receive is reached, the receive method wasn't called for the second message and hence switching context to the second co-routine is not yet possible. I couldn't find a source for this, but I'm guessing that this is part of the ASGI protocol itself. For many use-cases, handling WebSocket messages stricty in the order of receiving is probably desired.

The solution to handle messages asynchronously is to not send the response from the receive method, but instead send the response from a coroutine scheduled through loop.create_task.

Scheduling the long-running coroutine which generates response allows receive to complete, and for the next receive to begin. Once the second message's response generation has been scheduled, two coroutines will have been scheduled, and the interpreter can switch contexts to execute them asynchronously.

For the example in the question, this is the solution I found:

class QuestionConsumer(AsyncWebsocketConsumer):

    async def complicated_answer(self, question):
        await asyncio.sleep(3)
        answer = {
             "What is the Answer to Life, The Universe and Everything?": "42",
             "Why?": "Because.",
        }.get(question, "Don't know")
        # instead of returning the answer, send it directly to client as a response
        await self.send(answer)


    async def receive(self, text_data=None, bytes_data=None):
        # instead of awaiting, schedule the coroutine
        loop = asyncio.get_running_loop()
        loop.create_task(
            self.complicated_answer(text_data)
        )

The output of this altered consumer matches the desired output given by the question. Note that responses may be returned out of order, and clients are responsible for matching requests to responses.

Note that for Python versions <3.7, get_event_loop should be used instead of get_running_loop.

user_12344380
Related