Building a Scalable WebSocket Server in Python With AsyncIO and Custom Routing

For production-ready real-time systems, a basic WebSocket echo server isn't enough. In this guide, we’ll build a more advanced WebSocket server in Python using asyncio and the websockets library, including support for multiple routes, event-based message handling, and client context tracking. This structure is useful for applications like collaborative editors, multiplayer games, or real-time dashboards. Why Use Native WebSockets + AsyncIO? Lightweight: No need for full frameworks like Django or Flask unless needed. Scalable: AsyncIO handles thousands of concurrent connections efficiently. Flexible: Full control over routing, protocol design, and client behavior. Project Structure websocket-server/ ├── main.py ├── router.py └── handlers/ ├── init.py ├── chat.py └── presence.py Step 1: WebSocket Router Create a router.py file to handle incoming routes. # router.py from handlers import chat, presence ROUTES = { "/chat": chat.handle_chat, "/presence": presence.handle_presence, } async def route(path, websocket, context): handler = ROUTES.get(path) if handler: await handler(websocket, context) else: await websocket.send("Route not found.") await websocket.close() Step 2: Define Event-Based Handlers In handlers/chat.py: # handlers/chat.py import json clients = set() async def handle_chat(ws, context): clients.add(ws) try: async for message in ws: data = json.loads(message) if data["event"] == "message": payload = json.dumps({ "event": "message", "user": context["user"], "text": data["text"] }) for client in clients: if client != ws: await client.send(payload) except: pass finally: clients.remove(ws) In handlers/presence.py: # handlers/presence.py import json connected = {} async def handle_presence(ws, context): connected[ws] = context["user"] try: async for msg in ws: pass # For now, just track connection except: pass finally: del connected[ws] Step 3: Main Entry Point In main.py: # main.py import asyncio import websockets from urllib.parse import urlparse, parse_qs from router import route async def handler(websocket, path): query = parse_qs(urlparse(path).query) user = query.get("user", ["anonymous"])[0] context = { "user": user } clean_path = urlparse(path).path await route(clean_path, websocket, context) start_server = websockets.serve(handler, "0.0.0.0", 8765) print("WebSocket server started on ws://localhost:8765") asyncio.get_event_loop().run_until_complete(start_server) asyncio.get_event_loop().run_forever() Testing It Use different paths for different services: ws://localhost:8765/chat?user=Alice — joins chat room ws://localhost:8765/presence?user=Bob — registers presence Advanced Add-ons Add Redis or PostgreSQL to persist messages or presence status. Introduce pub/sub for scaling with multiple server instances. Support JWT-based authentication to secure user identity. Conclusion This modular WebSocket server design sets you up for scalable, real-time applications in Python without bloated dependencies. You now have full control over routing, context handling, and connection management. If this post was helpful, please consider supporting me here: buymeacoffee.com/hexshift

Apr 17, 2025 - 03:50

Building a Scalable WebSocket Server in Python With AsyncIO and Custom Routing

For production-ready real-time systems, a basic WebSocket echo server isn't enough. In this guide, we’ll build a more advanced WebSocket server in Python using asyncio and the websockets library, including support for multiple routes, event-based message handling, and client context tracking. This structure is useful for applications like collaborative editors, multiplayer games, or real-time dashboards.

Why Use Native WebSockets + AsyncIO?

Lightweight: No need for full frameworks like Django or Flask unless needed.
Scalable: AsyncIO handles thousands of concurrent connections efficiently.
Flexible: Full control over routing, protocol design, and client behavior.

Project Structure

websocket-server/
├── main.py
├── router.py
└── handlers/
    ├── __init__.py
    ├── chat.py
    └── presence.py

Step 1: WebSocket Router

Create a router.py file to handle incoming routes.

# router.py
from handlers import chat, presence

ROUTES = {
    "/chat": chat.handle_chat,
    "/presence": presence.handle_presence,
}

async def route(path, websocket, context):
    handler = ROUTES.get(path)
    if handler:
        await handler(websocket, context)
    else:
        await websocket.send("Route not found.")
        await websocket.close()

Step 2: Define Event-Based Handlers

In handlers/chat.py:

# handlers/chat.py
import json

clients = set()

async def handle_chat(ws, context):
    clients.add(ws)
    try:
        async for message in ws:
            data = json.loads(message)
            if data["event"] == "message":
                payload = json.dumps({
                    "event": "message",
                    "user": context["user"],
                    "text": data["text"]
                })
                for client in clients:
                    if client != ws:
                        await client.send(payload)
    except:
        pass
    finally:
        clients.remove(ws)

In handlers/presence.py:

# handlers/presence.py
import json

connected = {}

async def handle_presence(ws, context):
    connected[ws] = context["user"]
    try:
        async for msg in ws:
            pass  # For now, just track connection
    except:
        pass
    finally:
        del connected[ws]

Step 3: Main Entry Point

In main.py:

# main.py
import asyncio
import websockets
from urllib.parse import urlparse, parse_qs
from router import route

async def handler(websocket, path):
    query = parse_qs(urlparse(path).query)
    user = query.get("user", ["anonymous"])[0]
    context = { "user": user }

    clean_path = urlparse(path).path
    await route(clean_path, websocket, context)

start_server = websockets.serve(handler, "0.0.0.0", 8765)

print("WebSocket server started on ws://localhost:8765")
asyncio.get_event_loop().run_until_complete(start_server)
asyncio.get_event_loop().run_forever()

Testing It

Use different paths for different services:

ws://localhost:8765/chat?user=Alice — joins chat room
ws://localhost:8765/presence?user=Bob — registers presence

Advanced Add-ons

Add Redis or PostgreSQL to persist messages or presence status.
Introduce pub/sub for scaling with multiple server instances.
Support JWT-based authentication to secure user identity.

Conclusion

This modular WebSocket server design sets you up for scalable, real-time applications in Python without bloated dependencies. You now have full control over routing, context handling, and connection management.

If this post was helpful, please consider supporting me here: buymeacoffee.com/hexshift