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