python websockets authentication reconnect and broadcast: Pr

python websockets authentication reconnect and broadcast requires a clear understanding of the core syntax, runtime behavior, and practical implementation patterns demonstrated in the examples below.

Handling authentication, reconnection, and broadcast together is the most common set of requirements for a Python WebSocket service. The websockets library provides the primitives, but combining them correctly requires care about connection lifecycle, shared state, and error handling. This article walks through a complete implementation that covers all three, with code you can adapt to your own protocol.

A Minimal WebSocket Server with the `websockets` Library

The websockets library is the de facto standard for WebSocket communication in Python. A basic server is built around an async handler that receives a connection and exchanges messages. The following example starts a server on localhost:8765 and echoes every message back to the sender:

import asyncio
import websockets

async def echo(websocket):
    async for message in websocket:
        await websocket.send(message)

async def main():
    async with websockets.serve(echo, "localhost", 8765):
        await asyncio.Future()  # run forever

asyncio.run(main())

The serve function takes a handler that receives a WebSocketServerProtocol object. The handler runs for the lifetime of the connection. This is the foundation on which authentication, reconnection, and broadcast are built.

Authenticating Clients Before the Connection Is Accepted

Authentication must happen before the application starts exchanging data. The websockets library allows you to inspect the initial HTTP request via the path and headers attributes of the WebSocketServerProtocol. A common pattern is to require a token in the query string or in a custom header.

The following handler checks a token passed as a query parameter. If the token is invalid, it closes the connection with a policy violation code (1008) and returns without processing any messages:

async def authenticated_handler(websocket):
    token = websocket.query_params.get("token")
    if not is_valid_token(token):
        await websocket.close(code=1008, reason="Invalid token")
        return

    async for message in websocket:
        # process authenticated messages
        pass

In practice, you should validate tokens against a database or an external auth service. Avoid embedding secrets in URLs if they might appear in logs; a header like Authorization: Bearer <token> is safer. The websockets library exposes headers via websocket.request_headers.

async def auth_handler(websocket):
    auth_header = websocket.request_headers.get("Authorization", "")
    token = auth_header.replace("Bearer ", "")
    if not is_valid_token(token):
        await websocket.close(code=1008, reason="Unauthorized")
        return
    # ...

Implementing Reconnection with Exponential Backoff

Reconnection is a client-side concern. When a connection drops, the client must attempt to reconnect, but doing so immediately can overload the server. Exponential backoff with jitter is a standard approach. The websockets client uses websockets.connect, which raises ConnectionClosed when the connection is lost.

Here is a client that reconnects with exponential backoff, starting at 1 second and doubling up to 30 seconds, with a random jitter to prevent thundering herds:

import asyncio
import random
import websockets

async def connect_with_backoff(uri):
    delay = 1
    max_delay = 30
    while True:
        try:
            async with websockets.connect(uri) as websocket:
                print("Connected")
                await handle_connection(websocket)
        except websockets.ConnectionClosed:
            pass
        except OSError:
            pass

        await asyncio.sleep(delay + random.uniform(0, 1))
        delay = min(delay * 2, max_delay)

When the connection is re-established, the client should restore any necessary state, such as resubscribing to channels or replaying missed messages. This is often handled by a higher-level protocol message, like a subscribe or resume command.

Broadcasting Messages to All Connected Clients

Broadcasting means sending a message to every currently connected client. The server must maintain a registry of active connections. A simple set is sufficient, but because the server is async, you must protect it with a lock to avoid concurrent modification during iteration.

connected = set()
lock = asyncio.Lock()

async def broadcast(message):
    async with lock:
        for websocket in list(connected):
            try:
                await websocket.send(message)
            except websockets.ConnectionClosed:
                connected.remove(websocket)

The list(connected) creates a snapshot so you can safely remove connections while iterating. The lock ensures that add and remove operations are atomic. In the connection handler, you add the socket to the set when the connection is established and remove it when it closes:

async def handler(websocket):
    async with lock:
        connected.add(websocket)
    try:
        async for message in websocket:
            # handle incoming messages
            pass
    finally:
        async with lock:
            connected.remove(websocket)

This pattern is straightforward and works well for a single server process. For horizontal scaling, you would need a pub/sub backend like Redis, but that is a separate concern.

Managing Shared State and Concurrency

The broadcast example above illustrates a key concurrency issue: the connected set is shared across all connection handlers. Without the lock, two handlers could add or remove simultaneously, leading to RuntimeError: Set changed size during iteration. The lock serializes access, but it also means that sending a message to a slow client can block the broadcast for everyone.

An alternative is to use a asyncio.Queue per client or to send messages concurrently with asyncio.gather. However, for most applications, the lock-based approach is acceptable because send is usually fast and the lock is held only for the iteration, not for the network I/O. If you need to broadcast to thousands of clients, consider using asyncio.wait with a timeout to avoid blocking on a single dead connection.

Another concurrency concern is the lifecycle of the connection. When a client disconnects, the async for loop in the handler exits. You must ensure that the socket is removed from the set in a finally block, as shown above. Failing to do so will cause memory leaks and attempts to send to closed connections.

Handling Token Expiry and Heartbeats

Authentication tokens often expire. If a token expires while the connection is open, you have two options: close the connection and force the client to reconnect, or allow the connection to continue until the next reconnection. The first is simpler and more secure. You can implement a periodic check inside the handler:

async def handler(websocket):
    token = get_token(websocket)
    while True:
        if not is_token_valid(token):
            await websocket.close(code=1008, reason="Token expired")
            break
        try:
            message = await asyncio.wait_for(websocket.recv(), timeout=60)
            # process message
        except asyncio.TimeoutError:
            continue
        except websockets.ConnectionClosed:
            break

Heartbeats are essential for detecting dead connections. The websockets library has a built-in ping/pong mechanism, but you may want to implement your own application-level heartbeat to verify that the client is responsive. A common pattern is to send a ping every 30 seconds and close the connection if no pong is received within 10 seconds. The library's ping() method returns a future you can await with a timeout.

Security and Operational Considerations

Authentication and broadcast introduce security and operational concerns that go beyond the basic echo server. Always validate input from clients; never trust the content of a message just because the connection is authenticated. Use TLS in production by passing an ssl context to websockets.serve. For reconnection, be mindful of the load on your server; exponential backoff with jitter prevents a thundering herd when a server restarts.

Monitoring is also important. Log connection events, authentication failures, and broadcast errors. Use metrics to track the number of active connections and the rate of reconnections. This data helps you tune the backoff parameters and detect issues like token validation latency.

Finally, consider the lifecycle of the server itself. When you shut down, you should close all connections gracefully. The websockets library provides a close method on the server object, but you may also need to notify clients that the server is going away. A simple approach is to broadcast a shutdown message before closing the server.

Choosing between a single server and a load-balanced deployment depends on your scale. The broadcast implementation shown here works only within one process. If you need to broadcast across multiple servers, you must use an external message broker. The core principles of authentication, reconnection, and broadcast remain the same, but the shared state moves to a distributed system.

A Minimal WebSocket Server with the websockets Library