Fun with websockets and asyncio in Django Management Command

I've been playing around with WebSockets to replace an existing AJAX-heavy part of the front-end of an webapp. This has taken me into the realm of asyncio and websockets Python modules which I had not used much of before, and I thought I'd write this up.

From AJAX Polling to WebSocket-driven Updates

There was something that I've been meaning to rewrite for a while that required a background Django management command running constantly in the background, which does some processing every second or so, instead of every front-end client making an AJAX call to do the same processing every second, which becomes unscalable pretty quickly.

The idea was to notify the front-end clients with a WebSocket message whenever the centralised processing causes updates to the database which requires the front-end client to AJAX in the new data.

I'd never done much websocket or asyncio coding in Python, and it was quite fun to get into. Instead of handling event handles and message queues yourself in languages like C++, Python gives you pretty powerful asynchronous processing, using await/async, coroutine/Task/Futures.

An Example

Here's a slightly modified version of the Django management command, which:

• Registers/unregisters front-end clients,
• Does some processing periodically
• If any changes were made during the processing, notifies the currently registered clients

import logging

from optparse import make_option

from django.core.management.base import BaseCommand

import asyncio
import websockets

from.utils import configure_logging
from.models import Task
from.settings import (WEBSOCKET_BIND_PORT,
                      WEBSOCKET_BIND_ADDRESS)


clients = set()


async def notify_clients():
    if clients:  # asyncio.wait doesn't accept an empty list
        await asyncio.wait([client.send('updated') for client in clients])


async def main_loop():
    while True:
        modified = Task.objects.sync()
        if modified > 0:
            await notify_clients()
        await asyncio.sleep(0.5)


async def register(websocket):
    clients.add(websocket)
    logging.info(f'Registered new user: {websocket}')
    await notify_clients()


async def unregister(websocket):
    clients.remove(websocket)
    logging.info(f'Unregistered user: {websocket}')
    await notify_clients()


async def handle_client(websocket, path):
    await register(websocket)
    try:
        async for message in websocket:
            if message == 'success':
                pass
            else:
                await websocket.send('updated')
    finally:
        await unregister(websocket)


class Command(BaseCommand):
    """ Monitor tasks and notify websocket clients when they are updated """

    help = __doc__

    def handle(self, *args, **options):
        configure_logging(**options)

        start_server = websockets.serve(handle_client,
                                        WEBSOCKET_BIND_ADDRESS,
                                        WEBSOCKET_BIND_PORT)
        asyncio.get_event_loop().run_until_complete(start_server)
        asyncio.get_event_loop().run_until_complete(main_loop())
        asyncio.get_event_loop().run_forever()

On the front-end, we just wait for the 'updated' message and only trigger an AJAX call then:

var ws = new WebSocket("ws{{ websocket_secure | yesno:'s,' }}://" + window.location.hostname + ":{{ websocket_port }}/");
ws.onmessage = function (event) {
  var type = event.data;
  if (type == 'updated') {
    $.getJSON('/task-list', function(data) {
      // Populate some stuff
    });
  }
};