ConnectionManager Class

Why a class, not just a dict

The minimal manager from Track 3 was a dict of sets. Real applications need three more dimensions: per-connection metadata (when did they connect, from what IP, with what session?), per-user lookups (which connection belongs to user X right now?), and observability (how many connections, by room, by user, by version?). A class consolidates all of this without scattering state across modules.

The three indexes

You almost always want three lookup directions: room → set of connections, user_id → connection, connection → metadata. Three small dicts; every operation updates all three coherently. Skipping one of the three is the source of every "stale connection" bug in production.

Code

Production-shape manager·python

from fastapi import WebSocket
from typing import Dict, Set, Optional
import time

class ConnectionManager:
    def __init__(self):
        self.rooms: Dict[str, Set[WebSocket]] = {}
        self.user_map: Dict[str, WebSocket] = {}
        self.ws_meta: Dict[WebSocket, dict] = {}

    async def connect(self, ws: WebSocket, user_id: str, room: str):
        await ws.accept()
        self.rooms.setdefault(room, set()).add(ws)
        self.user_map[user_id] = ws
        self.ws_meta[ws] = {
            'user_id': user_id,
            'room': room,
            'connected_at': time.time(),
            'ip': ws.client.host if ws.client else None,
        }

    def disconnect(self, ws: WebSocket):
        meta = self.ws_meta.pop(ws, None)
        if meta is None:
            return
        room = meta['room']
        members = self.rooms.get(room)
        if members:
            members.discard(ws)
            if not members:
                self.rooms.pop(room, None)
        existing = self.user_map.get(meta['user_id'])
        if existing is ws:
            self.user_map.pop(meta['user_id'], None)

    async def send_to_user(self, user_id: str, message: dict) -> bool:
        ws = self.user_map.get(user_id)
        if ws is None:
            return False
        try:
            await ws.send_json(message)
            return True
        except Exception:
            self.disconnect(ws)
            return False

    async def broadcast(self, room: str, message: dict, *, exclude: Optional[WebSocket] = None):
        dead = []
        for ws in list(self.rooms.get(room, ())):
            if ws is exclude:
                continue
            try:
                await ws.send_json(message)
            except Exception:
                dead.append(ws)
        for ws in dead:
            self.disconnect(ws)

    @property
    def total_connections(self) -> int:
        return len(self.ws_meta)

    def room_users(self, room: str) -> list[str]:
        return [
            self.ws_meta[ws]['user_id']
            for ws in self.rooms.get(room, ())
            if ws in self.ws_meta
        ]

Why a class, not just a dict

The three indexes

Code

External links

Exercise

Progress

Comments 0