FastAPI WebSocket Stabilizer

A production-ready WebSocket stabilization layer for FastAPI applications that provides robust connection lifecycle management, automatic heartbeat/ping-pong detection, graceful shutdown, and reconnection support.

Motivation

Building reliable WebSocket applications in production requires handling many edge cases:

Connection timeouts: Proxy servers often have idle connection timeouts (60-90 seconds)
Dead connections: Network partitions can leave zombie connections that silently hang
Ungraceful restarts: Server restarts need clean connection closure to prevent client reconnection storms
Load balancer issues: Behind ALBs/NLBs, stale connections accumulate without proper heartbeat
Cloud platform quirks: Azure App Service, AWS Lambda, and other serverless platforms have specific WebSocket constraints
Session resumption: Clients may disconnect due to network switching and need to restore state

This library abstracts away these concerns with a simple, production-tested API.

Key Features

✅ Connection Management: Automatic tracking of active WebSocket connections
✅ Heartbeat/Ping-Pong: Configurable automatic heartbeat detection of dead connections
✅ Graceful Shutdown: Clean connection closure with proper async task cleanup
✅ Reconnection Tokens: Server-side session token support for seamless reconnection
✅ Structured Logging: Cloud-friendly logging for easy monitoring and debugging
✅ Broadcast Messaging: Send messages to individual clients, groups, or all clients
✅ Connection Limits: Configurable max connections with automatic enforcement
✅ Type Safety: Full type hints for all public APIs (Python 3.10+)
✅ Zero Dependencies: Only depends on FastAPI (which provides WebSocket support)
✅ Sharded Connection Pools: Reduce lock contention under high load with configurable sharding
✅ Memory Management: Track and limit memory usage per connection with eviction policies
✅ Concurrency Control: Semaphore-based limiting for cleanup and broadcast operations

Installation

Install via pip:

pip install fastapi-websocket-stabilizer

Or from source:

git clone https://github.com/fastapi-websocket-stabilizer/fastapi-websocket-stabilizer.git
cd fastapi-websocket-stabilizer
pip install -e ".[dev]"

Quick Start

Basic Chat Application

import json
from contextlib import asynccontextmanager
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi_websocket_stabilizer import WebSocketConnectionManager, WebSocketConfig

# Configure the manager
ws_config = WebSocketConfig(
    heartbeat_interval=30.0,
    heartbeat_timeout=60.0,
    max_connections=100,
)
ws_manager = WebSocketConnectionManager(ws_config)

# Setup lifespan events
@asynccontextmanager
async def lifespan(app: FastAPI):
    await ws_manager.start()
    yield
    await ws_manager.graceful_shutdown()

app = FastAPI(lifespan=lifespan)

@app.websocket("/ws/{client_id}")
async def websocket_endpoint(websocket: WebSocket, client_id: str):
    await websocket.accept()

    try:
        # Register connection
        await ws_manager.connect(client_id, websocket)

        # Handle messages
        while True:
            data = await websocket.receive_text()

            # Handle heartbeat responses
            if data == "pong":
                await ws_manager.handle_pong(client_id)
                continue

            # Broadcast to all clients
            await ws_manager.broadcast(json.dumps({
                "user": client_id,
                "message": data
            }))

    except WebSocketDisconnect:
        await ws_manager.disconnect(client_id)

Usage Patterns

Pattern 1: Simple Broadcast

# Send message to all clients
await ws_manager.broadcast("System message")

# Send to all except sender
await ws_manager.broadcast(
    json.dumps({"message": "chat"}),
    exclude_client=sender_id
)

Pattern 2: Direct Messaging

# Send to specific client
success = await ws_manager.send_to_client(client_id, "Direct message")

if not success:
    print(f"Client {client_id} is no longer connected")

Pattern 3: Reconnection Support

@app.post("/reconnect-token/{client_id}")
async def get_reconnect_token(client_id: str):
    """Get a token before intentional disconnection."""
    token = ws_manager.generate_reconnect_token(client_id)
    return {"token": token}

@app.websocket("/ws/{client_id}")
async def websocket_endpoint(websocket: WebSocket, client_id: str, token: str = None):
    await websocket.accept()

    # Restore session if token provided
    if token:
        try:
            restored_id = await ws_manager.validate_reconnect_token(token)
            # Continue with restored session
        except TokenExpiredError:
            await websocket.close(code=1008, reason="Token expired")
            return

    await ws_manager.connect(client_id, websocket)

Pattern 4: Statistics and Monitoring

@app.get("/stats")
async def get_stats():
    return {
        "active_connections": ws_manager.get_connection_count(),
        "client_ids": ws_manager.get_active_connections(),
        "memory": ws_manager.get_memory_usage(),  # v1.1.0+
    }

Pattern 5: Per-Connection User Data (v1.1.0+)

@app.websocket("/ws/{client_id}")
async def websocket_endpoint(websocket: WebSocket, client_id: str):
    await websocket.accept()
    await ws_manager.connect(client_id, websocket)

    # Store user state per connection
    await ws_manager.set_user_data(client_id, "username", "Alice")
    await ws_manager.set_user_data(client_id, "room", "general")

    # Retrieve later
    username = await ws_manager.get_user_data(client_id, "username")
    print(f"{username} joined")  # "Alice joined"

Configuration Reference

WebSocketConfig

All configuration options with defaults:

from fastapi_websocket_stabilizer import WebSocketConfig
from datetime import timedelta

config = WebSocketConfig(
    # Heartbeat settings
    heartbeat_interval=30.0,      # Seconds between pings
    heartbeat_timeout=60.0,       # Seconds to wait for pong response
    heartbeat_type="ping",        # Type of heartbeat ("ping" or "custom")

    # Connection limits
    max_connections=None,         # Max concurrent connections (None = unlimited)
    max_message_size=1048576,     # Max message size in bytes (1MB default)

    # Reconnection
    reconnect_token_ttl=timedelta(hours=1),  # Token lifetime
    enable_reconnect=True,        # Enable token support

    # Compression
    compression="deflate",        # "deflate" or "none"

    # Logging
    log_level="INFO",             # DEBUG, INFO, WARNING, ERROR, CRITICAL
    enable_metrics=False,         # Collect metrics (experimental)

    # Sharding (reduces lock contention under high load)
    num_shards=16,                # Number of connection pool shards

    # Concurrency control (prevents cleanup/broadcast storms)
    max_concurrent_cleanup=100,   # Max concurrent connection closures
    max_concurrent_broadcast=1000, # Max concurrent broadcast sends

    # Memory management
    max_memory_per_connection=None,  # Max bytes per connection (None = unlimited)
    max_total_memory=None,           # Max total bytes (None = unlimited)
    memory_eviction_policy="lru",    # "lru", "fifo", or "oldest"
    memory_check_interval=60.0,      # Seconds between memory checks
)

Example Configurations

High-frequency trading/gaming (low latency)

WebSocketConfig(
    heartbeat_interval=5.0,       # More frequent heartbeats
    heartbeat_timeout=10.0,       # Detect dead connections quickly
    max_connections=10000,        # Expect many connections
)

Behind slow proxy (e.g., old reverse proxy)

WebSocketConfig(
    heartbeat_interval=45.0,      # Slower heartbeat (adapt to proxy timeout)
    heartbeat_timeout=90.0,
    compression="none",           # Reduce overhead
)

Serverless (Azure Functions, AWS Lambda)

WebSocketConfig(
    heartbeat_interval=20.0,      # Serverless platforms have stricter timeouts
    max_connections=500,          # Limited resource per instance
)

High-load production (10k+ connections)

from fastapi_websocket_stabilizer import WebSocketConfig, MemoryEvictionPolicy

WebSocketConfig(
    heartbeat_interval=30.0,
    max_connections=50000,
    num_shards=32,                      # More shards for reduced lock contention
    max_concurrent_cleanup=50,          # Prevent cleanup storms
    max_concurrent_broadcast=500,       # Control broadcast load
    max_memory_per_connection=1024*1024,  # 1MB per connection
    max_total_memory=1024*1024*1024,      # 1GB total
    memory_eviction_policy=MemoryEvictionPolicy.LRU,  # Evict least recently used
)

Deployment Guide

Local Development

# Install dependencies
pip install "fastapi[standard]" uvicorn

# Run example
uvicorn examples.basic_app:app --reload

# Test with curl or WebSocket client
# ws://localhost:8000/ws/user123

Behind Reverse Proxy (nginx, Apache)

nginx configuration:

upstream app {
    server localhost:8000;
}

server {
    listen 80;

    location /ws/ {
        proxy_pass http://app;
        proxy_http_version 1.1;
        proxy_set_header Upgrade $http_upgrade;
        proxy_set_header Connection "upgrade";
        proxy_read_timeout 3600s;
        proxy_send_timeout 3600s;
        # Set read timeout higher than heartbeat timeout + margin
    }
}

Key considerations:

Ensure proxy read_timeout and send_timeout exceed your heartbeat_interval
Set to at least 2x your heartbeat_interval for safety
Default is often 60s, which may be too low; increase to 300s+ for long-lived connections

AWS Application Load Balancer (ALB)

Target Group Configuration:

Deregistration delay: 30-60 seconds
Connection idle timeout: 60+ seconds (ALB default is 60s)

Application Configuration:

config = WebSocketConfig(
    heartbeat_interval=25.0,      # Less than ALB's 60s default
    heartbeat_timeout=50.0,
)

Azure App Service

Azure App Service has specific WebSocket constraints:

Settings in Azure Portal:

Enable WebSocket: ON
Connection timeout (Web socket idle timeout): 600 seconds

Application Configuration:

config = WebSocketConfig(
    heartbeat_interval=30.0,      # Well within Azure's limits
    max_connections=1000,         # App Service tier dependent
)

Deployment via Docker:

FROM python:3.11-slim
WORKDIR /app
COPY requirements.txt .
RUN pip install -r requirements.txt
COPY . .
CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8000"]

Kubernetes

ConfigMap for configuration:

apiVersion: v1
kind: ConfigMap
metadata:
  name: websocket-config
data:
  heartbeat_interval: "30"
  heartbeat_timeout: "60"
  max_connections: "1000"

Deployment:

apiVersion: apps/v1
kind: Deployment
metadata:
  name: websocket-app
spec:
  replicas: 3
  template:
    spec:
      containers:
      - name: app
        image: websocket-app:1.0
        ports:
        - containerPort: 8000
        livenessProbe:
          httpGet:
            path: /health
            port: 8000
          initialDelaySeconds: 10
          periodSeconds: 10

Docker Compose

version: '3.8'

services:
  app:
    build: .
    ports:
      - "8000:8000"
    environment:
      LOG_LEVEL: INFO
    command: uvicorn app:app --host 0.0.0.0 --port 8000 --workers 1

Troubleshooting

Connections Disconnecting Unexpectedly

Problem: Clients disconnect after ~60 seconds of inactivity

Solution: Reduce heartbeat_interval below your proxy/firewall timeout

# Adjust based on your environment's idle timeout
WebSocketConfig(heartbeat_interval=25.0)  # < 60s proxy timeout

High CPU Usage

Problem: Manager consuming excessive CPU

Solution: Reduce heartbeat frequency or use larger batches

WebSocketConfig(
    heartbeat_interval=60.0,  # Increase from default 30s
    max_connections=500,      # Reduce if possible
)

"Connection Limit Exceeded" Errors

Problem: New connections rejected with limit errors

Solutions:

Increase max_connections:
```
WebSocketConfig(max_connections=5000)
```
Implement graceful connection rejection with backoff
Scale horizontally (multiple app instances)

Reconnection Tokens Not Working

Problem: Token validation fails repeatedly

Solutions:

Check token hasn't expired: tokenExpiredError means TTL passed
Verify same WebSocketConnectionManager instance used for both generation and validation
Ensure client sends token in query parameter or auth header

Memory Leaks

Problem: Memory usage increases over time

Solutions:

Ensure graceful_shutdown() is called during app shutdown
Check that disconnected clients are properly cleaned up
Monitor token storage (should be auto-cleaned every 5 minutes)

API Reference

WebSocketConnectionManager

class WebSocketConnectionManager:
    def __init__(self, config: Optional[WebSocketConfig] = None) -> None:
        """Initialize manager with optional config."""

    async def start(self) -> None:
        """Start heartbeat and cleanup tasks."""

    async def connect(
        self,
        client_id: str,
        websocket: WebSocket,
        metadata: Optional[dict] = None
    ) -> None:
        """Register a new WebSocket connection."""

    async def disconnect(self, client_id: str) -> None:
        """Gracefully disconnect a client."""

    async def send_to_client(
        self,
        client_id: str,
        message: str | dict
    ) -> bool:
        """Send message to specific client."""

    async def broadcast(
        self,
        message: str | dict,
        exclude_client: Optional[str] = None
    ) -> BroadcastResult:
        """Send message to all clients."""

    def get_active_connections(self) -> list[str]:
        """Get list of connected client IDs."""

    def get_connection_count(self) -> int:
        """Get number of active connections."""

    def generate_reconnect_token(self, client_id: str) -> str:
        """Generate a reconnection token."""

    async def validate_reconnect_token(self, token: str) -> Optional[str]:
        """Validate token and return client_id."""

    async def handle_pong(self, client_id: str) -> None:
        """Handle pong response from client."""

    async def graceful_shutdown(self, timeout: float = 30.0) -> ShutdownReport:
        """Shut down manager and close all connections."""

    # Memory management methods (v1.1.0+)
    def get_memory_usage(self) -> dict[str, int | float]:
        """Get memory usage statistics.
        Returns: {total_memory, connection_count, per_connection_average}"""

    async def set_user_data(self, client_id: str, key: str, value: Any) -> None:
        """Set user data for a connection (with memory limit checking)."""

    async def get_user_data(self, client_id: str, key: str) -> Any | None:
        """Get user data for a connection."""

Exceptions

from fastapi_websocket_stabilizer import (
    WebSocketStabilizerError,          # Base exception
    ConnectionNotFoundError,            # Client not found
    ConnectionLimitExceededError,       # Max connections reached
    ConnectionTimeoutError,             # Heartbeat timeout
    TokenExpiredError,                  # Token TTL exceeded
    InvalidTokenError,                  # Token invalid/tampered
    BroadcastFailedError,               # Broadcast failed
    ShutdownError,                      # Shutdown error
    MemoryLimitExceededError,           # Memory limit exceeded (v1.1.0+)
)

Roadmap

Completed (v1.1.0)

Sharded connection pools for high-load scenarios
Memory management with eviction policies
Concurrency control for cleanup/broadcast storms

Planned Features

Metrics collection (Prometheus-compatible)
Connection lifecycle hooks (on_connect, on_disconnect callbacks)
Message compression strategies
Distributed mode (Redis-backed connection registry for multi-instance deployments)
Built-in rate limiting
Automatic reconnection client library (JavaScript/TypeScript)

Possible Enhancements

OpenTelemetry instrumentation
Support for WebSocket subprotocols
Custom heartbeat message types
Connection grouping/rooms functionality
Pub/Sub integration

Contributing

We welcome contributions! Please:

Fork the repository
Create a feature branch (git checkout -b feature/amazing-feature)
Write tests for new functionality
Ensure all tests pass: pytest tests/
Format code with black: black src/ tests/
Lint with ruff: ruff check src/ tests/
Commit with clear messages
Push and create a Pull Request

Development Setup

# Clone repository
git clone https://github.com/fastapi-websocket-stabilizer/fastapi-websocket-stabilizer.git
cd fastapi-websocket-stabilizer

# Create virtual environment
python -m venv venv
source venv/bin/activate  # On Windows: venv\Scripts\activate

# Install in development mode
pip install -e ".[dev]"

# Run tests
pytest tests/ -v

# Run specific test
pytest tests/test_manager.py::TestConnectionManager::test_should_track_connection_when_connect_called -v

# Check types
mypy src/

# Format code
black src/ tests/

# Lint
ruff check src/ tests/

License

This project is licensed under the MIT License - see the LICENSE file for details.

Support

Acknowledgments

Inspired by production WebSocket patterns in large-scale FastAPI deployments, with particular consideration for:

FastAPI's excellent WebSocket support
Real-world challenges in cloud environments
Community feedback on WebSocket reliability

Version History

1.1.0

Sharded Connection Pools: Reduce lock contention with configurable num_shards
Concurrency Control: Semaphore-based max_concurrent_cleanup and max_concurrent_broadcast
Memory Management: Per-connection and total memory limits with eviction policies (LRU, FIFO, OLDEST)
New methods: get_memory_usage(), set_user_data(), get_user_data()
New exception: MemoryLimitExceededError
Bug fixes: hmac.hexdigest() for Python 3.14 compatibility, reconnection within limits

1.0.0

Stable release

0.1.1

Fix GitHub URLs and setuptools compatibility

0.1.0 (Initial Release)

Core connection management
Heartbeat/ping-pong mechanism
Reconnection token support
Graceful shutdown
Comprehensive tests and documentation

Name		Name	Last commit message	Last commit date
examples		examples
src/fastapi_websocket_stabilizer		src/fastapi_websocket_stabilizer
tests		tests
.gitignore		.gitignore
ARCHITECTURE.md		ARCHITECTURE.md
LICENSE		LICENSE
PROJECT_SUMMARY.md		PROJECT_SUMMARY.md
PUBLIC_API.md		PUBLIC_API.md
QUICKSTART.md		QUICKSTART.md
README.md		README.md
pyproject.toml		pyproject.toml

Latest commit

History

Folders and files

Repository files navigation

FastAPI WebSocket Stabilizer

Motivation

Key Features

Installation

Quick Start

Basic Chat Application

Usage Patterns

Pattern 1: Simple Broadcast

Pattern 2: Direct Messaging

Pattern 3: Reconnection Support

Pattern 4: Statistics and Monitoring

Pattern 5: Per-Connection User Data (v1.1.0+)

Configuration Reference

WebSocketConfig

Example Configurations

High-frequency trading/gaming (low latency)

Behind slow proxy (e.g., old reverse proxy)

Serverless (Azure Functions, AWS Lambda)

High-load production (10k+ connections)

Deployment Guide

Local Development

Behind Reverse Proxy (nginx, Apache)

AWS Application Load Balancer (ALB)

Azure App Service

Kubernetes

Docker Compose

Troubleshooting

Connections Disconnecting Unexpectedly

High CPU Usage

"Connection Limit Exceeded" Errors

Reconnection Tokens Not Working

Memory Leaks

API Reference

WebSocketConnectionManager

Exceptions

Roadmap

Completed (v1.1.0)

Planned Features

Possible Enhancements

Contributing

Development Setup

License

Support

Acknowledgments

Version History

1.1.0

1.0.0

0.1.1

0.1.0 (Initial Release)

About

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages