WebSocket Management

Build robust, scalable WebSocket connections with reconnection, heartbeat, and error handling

WebSockets enable real-time bidirectional communication. This guide covers production-ready WebSocket patterns.

🎯 Goals

What we need:
✅ Automatic reconnection
✅ Heartbeat/ping-pong
✅ Message queuing
✅ Error handling
✅ Connection state management

🔌 Pattern 1: Basic WebSocket Client

Simple WebSocket with reconnection.

class WebSocketClient {
  private ws: WebSocket | null = null;
  private url: string;
  private reconnectAttempts = 0;
  private maxReconnectAttempts = 10;
  private reconnectDelay = 1000;
  private listeners = new Map<string, Set<Function>>();
  
  constructor(url: string) {
    this.url = url;
    this.connect();
  }
  
  private connect() {
    console.log('Connecting to WebSocket...');
    
    this.ws = new WebSocket(this.url);
    
    this.ws.onopen = () => {
      console.log('WebSocket connected');
      this.reconnectAttempts = 0;
      this.reconnectDelay = 1000;
      this.emit('connected');
    };
    
    this.ws.onmessage = (event) => {
      try {
        const message = JSON.parse(event.data);
        this.emit('message', message);
        
        // Emit specific event types
        if (message.type) {
          this.emit(message.type, message.data);
        }
      } catch (error) {
        console.error('Failed to parse message:', error);
      }
    };
    
    this.ws.onerror = (error) => {
      console.error('WebSocket error:', error);
      this.emit('error', error);
    };
    
    this.ws.onclose = () => {
      console.log('WebSocket disconnected');
      this.emit('disconnected');
      this.reconnect();
    };
  }
  
  private reconnect() {
    if (this.reconnectAttempts >= this.maxReconnectAttempts) {
      console.error('Max reconnect attempts reached');
      this.emit('max-reconnect-attempts');
      return;
    }
    
    this.reconnectAttempts++;
    
    console.log(
      `Reconnecting in ${this.reconnectDelay}ms (attempt ${this.reconnectAttempts})`
    );
    
    setTimeout(() => {
      this.connect();
      // Exponential backoff
      this.reconnectDelay = Math.min(this.reconnectDelay * 2, 30000);
    }, this.reconnectDelay);
  }
  
  send(type: string, data: any) {
    if (this.ws?.readyState === WebSocket.OPEN) {
      this.ws.send(JSON.stringify({ type, data }));
    } else {
      console.warn('WebSocket not connected, message not sent');
    }
  }
  
  on(event: string, callback: Function) {
    if (!this.listeners.has(event)) {
      this.listeners.set(event, new Set());
    }
    this.listeners.get(event)!.add(callback);
  }
  
  off(event: string, callback: Function) {
    this.listeners.get(event)?.delete(callback);
  }
  
  private emit(event: string, data?: any) {
    this.listeners.get(event)?.forEach(callback => callback(data));
  }
  
  close() {
    this.maxReconnectAttempts = 0; // Prevent reconnection
    this.ws?.close();
  }
}

// Usage
const ws = new WebSocketClient('wss://api.example.com/ws');

ws.on('connected', () => {
  console.log('Successfully connected!');
});

ws.on('message', (message) => {
  console.log('Received:', message);
});

ws.on('notification', (data) => {
  console.log('New notification:', data);
});

ws.send('subscribe', { channel: 'notifications' });

💓 Pattern 2: Heartbeat (Ping/Pong)

Keep connection alive and detect broken connections.

class WebSocketWithHeartbeat {
  private ws: WebSocket | null = null;
  private url: string;
  private pingInterval: NodeJS.Timeout | null = null;
  private pongTimeout: NodeJS.Timeout | null = null;
  private pingIntervalMs = 30000; // 30 seconds
  private pongTimeoutMs = 5000; // 5 seconds
  
  constructor(url: string) {
    this.url = url;
    this.connect();
  }
  
  private connect() {
    this.ws = new WebSocket(this.url);
    
    this.ws.onopen = () => {
      console.log('Connected');
      this.startHeartbeat();
    };
    
    this.ws.onmessage = (event) => {
      const message = JSON.parse(event.data);
      
      // Handle pong response
      if (message.type === 'pong') {
        this.handlePong();
        return;
      }
      
      // Handle other messages
      this.handleMessage(message);
    };
    
    this.ws.onclose = () => {
      console.log('Disconnected');
      this.stopHeartbeat();
      this.reconnect();
    };
  }
  
  private startHeartbeat() {
    // Send ping every 30 seconds
    this.pingInterval = setInterval(() => {
      this.sendPing();
    }, this.pingIntervalMs);
  }
  
  private stopHeartbeat() {
    if (this.pingInterval) {
      clearInterval(this.pingInterval);
      this.pingInterval = null;
    }
    
    if (this.pongTimeout) {
      clearTimeout(this.pongTimeout);
      this.pongTimeout = null;
    }
  }
  
  private sendPing() {
    if (this.ws?.readyState === WebSocket.OPEN) {
      this.ws.send(JSON.stringify({ type: 'ping' }));
      
      // Expect pong within 5 seconds
      this.pongTimeout = setTimeout(() => {
        console.error('Pong timeout, reconnecting...');
        this.ws?.close();
      }, this.pongTimeoutMs);
    }
  }
  
  private handlePong() {
    // Clear pong timeout
    if (this.pongTimeout) {
      clearTimeout(this.pongTimeout);
      this.pongTimeout = null;
    }
  }
  
  private handleMessage(message: any) {
    console.log('Message:', message);
  }
  
  private reconnect() {
    setTimeout(() => this.connect(), 2000);
  }
  
  send(data: any) {
    if (this.ws?.readyState === WebSocket.OPEN) {
      this.ws.send(JSON.stringify(data));
    }
  }
  
  close() {
    this.stopHeartbeat();
    this.ws?.close();
  }
}

Backend (Node.js)

const WebSocket = require('ws');

const wss = new WebSocket.Server({ port: 8080 });

wss.on('connection', (ws) => {
  console.log('Client connected');
  
  ws.on('message', (data) => {
    const message = JSON.parse(data);
    
    // Handle ping
    if (message.type === 'ping') {
      ws.send(JSON.stringify({ type: 'pong' }));
      return;
    }
    
    // Handle other messages
    console.log('Received:', message);
  });
  
  ws.on('close', () => {
    console.log('Client disconnected');
  });
});

📬 Pattern 3: Message Queue

Queue messages when disconnected.

interface QueuedMessage {
  type: string;
  data: any;
  timestamp: number;
  retries: number;
}

class WebSocketWithQueue {
  private ws: WebSocket | null = null;
  private url: string;
  private messageQueue: QueuedMessage[] = [];
  private maxQueueSize = 100;
  private maxRetries = 3;
  
  constructor(url: string) {
    this.url = url;
    this.connect();
  }
  
  private connect() {
    this.ws = new WebSocket(this.url);
    
    this.ws.onopen = () => {
      console.log('Connected, flushing queue...');
      this.flushQueue();
    };
    
    this.ws.onclose = () => {
      console.log('Disconnected');
      this.reconnect();
    };
  }
  
  send(type: string, data: any) {
    const message: QueuedMessage = {
      type,
      data,
      timestamp: Date.now(),
      retries: 0
    };
    
    if (this.ws?.readyState === WebSocket.OPEN) {
      this.sendMessage(message);
    } else {
      this.queueMessage(message);
    }
  }
  
  private sendMessage(message: QueuedMessage) {
    try {
      this.ws!.send(JSON.stringify({
        type: message.type,
        data: message.data
      }));
    } catch (error) {
      console.error('Failed to send message:', error);
      this.queueMessage(message);
    }
  }
  
  private queueMessage(message: QueuedMessage) {
    if (this.messageQueue.length >= this.maxQueueSize) {
      // Remove oldest message
      this.messageQueue.shift();
      console.warn('Message queue full, dropping oldest message');
    }
    
    this.messageQueue.push(message);
    console.log(`Queued message (queue size: ${this.messageQueue.length})`);
  }
  
  private flushQueue() {
    console.log(`Flushing ${this.messageQueue.length} queued messages`);
    
    while (this.messageQueue.length > 0) {
      const message = this.messageQueue.shift()!;
      
      if (message.retries >= this.maxRetries) {
        console.error('Message exceeded max retries, dropping:', message);
        continue;
      }
      
      message.retries++;
      this.sendMessage(message);
    }
  }
  
  private reconnect() {
    setTimeout(() => this.connect(), 2000);
  }
  
  close() {
    this.ws?.close();
  }
}

🔐 Pattern 4: Authenticated WebSocket

Handle authentication and token refresh.

class AuthenticatedWebSocket {
  private ws: WebSocket | null = null;
  private url: string;
  private getToken: () => Promise<string>;
  
  constructor(url: string, getToken: () => Promise<string>) {
    this.url = url;
    this.getToken = getToken;
    this.connect();
  }
  
  private async connect() {
    // Get fresh token
    const token = await this.getToken();
    
    // Add token to WebSocket URL
    const wsUrl = `${this.url}?token=${token}`;
    
    this.ws = new WebSocket(wsUrl);
    
    this.ws.onopen = () => {
      console.log('Authenticated WebSocket connected');
    };
    
    this.ws.onerror = (error) => {
      console.error('WebSocket error:', error);
    };
    
    this.ws.onclose = (event) => {
      // Check if closed due to auth error
      if (event.code === 4001) {
        console.log('Auth error, refreshing token...');
        this.reconnectWithNewToken();
      } else {
        console.log('Disconnected');
        this.reconnect();
      }
    };
    
    this.ws.onmessage = (event) => {
      const message = JSON.parse(event.data);
      
      // Handle auth errors
      if (message.type === 'auth_error') {
        console.log('Auth error, reconnecting...');
        this.reconnectWithNewToken();
        return;
      }
      
      this.handleMessage(message);
    };
  }
  
  private async reconnectWithNewToken() {
    // Close old connection
    this.ws?.close();
    
    // Wait a bit before reconnecting
    await new Promise(resolve => setTimeout(resolve, 1000));
    
    // Connect with fresh token
    this.connect();
  }
  
  private reconnect() {
    setTimeout(() => this.connect(), 2000);
  }
  
  private handleMessage(message: any) {
    console.log('Message:', message);
  }
  
  send(data: any) {
    if (this.ws?.readyState === WebSocket.OPEN) {
      this.ws.send(JSON.stringify(data));
    }
  }
  
  close() {
    this.ws?.close();
  }
}

// Usage
const ws = new AuthenticatedWebSocket(
  'wss://api.example.com/ws',
  async () => {
    // Fetch or refresh token
    const response = await fetch('/api/auth/token');
    const { token } = await response.json();
    return token;
  }
);

⚛️ Pattern 5: React Hook

Use WebSocket in React components.

import { useEffect, useRef, useState } from 'react';

interface UseWebSocketOptions {
  onMessage?: (data: any) => void;
  onConnected?: () => void;
  onDisconnected?: () => void;
  reconnect?: boolean;
}

function useWebSocket(url: string, options: UseWebSocketOptions = {}) {
  const wsRef = useRef<WebSocket | null>(null);
  const [isConnected, setIsConnected] = useState(false);
  const reconnectTimeoutRef = useRef<NodeJS.Timeout>();
  
  useEffect(() => {
    function connect() {
      const ws = new WebSocket(url);
      
      ws.onopen = () => {
        console.log('WebSocket connected');
        setIsConnected(true);
        options.onConnected?.();
      };
      
      ws.onmessage = (event) => {
        try {
          const data = JSON.parse(event.data);
          options.onMessage?.(data);
        } catch (error) {
          console.error('Failed to parse message:', error);
        }
      };
      
      ws.onerror = (error) => {
        console.error('WebSocket error:', error);
      };
      
      ws.onclose = () => {
        console.log('WebSocket disconnected');
        setIsConnected(false);
        options.onDisconnected?.();
        
        // Reconnect if enabled
        if (options.reconnect !== false) {
          reconnectTimeoutRef.current = setTimeout(() => {
            connect();
          }, 2000);
        }
      };
      
      wsRef.current = ws;
    }
    
    connect();
    
    return () => {
      if (reconnectTimeoutRef.current) {
        clearTimeout(reconnectTimeoutRef.current);
      }
      wsRef.current?.close();
    };
  }, [url]);
  
  const send = (data: any) => {
    if (wsRef.current?.readyState === WebSocket.OPEN) {
      wsRef.current.send(JSON.stringify(data));
    } else {
      console.warn('WebSocket not connected');
    }
  };
  
  return { send, isConnected };
}

// Usage
function ChatComponent() {
  const [messages, setMessages] = useState<any[]>([]);
  
  const { send, isConnected } = useWebSocket('wss://api.example.com/chat', {
    onMessage: (message) => {
      setMessages(prev => [...prev, message]);
    },
    onConnected: () => {
      console.log('Chat connected!');
    }
  });
  
  const sendMessage = (text: string) => {
    send({ type: 'message', text });
  };
  
  return (
    <div>
      <div>Status: {isConnected ? '🟢 Connected' : '🔴 Disconnected'}</div>
      
      {messages.map((msg, i) => (
        <div key={i}>{msg.text}</div>
      ))}
      
      <input
        onKeyPress={(e) => {
          if (e.key === 'Enter') {
            sendMessage(e.currentTarget.value);
            e.currentTarget.value = '';
          }
        }}
      />
    </div>
  );
}

📊 Connection State Machine

Manage connection states properly.

type ConnectionState = 
  | 'disconnected'
  | 'connecting'
  | 'connected'
  | 'reconnecting'
  | 'failed';

class StatefulWebSocket {
  private ws: WebSocket | null = null;
  private state: ConnectionState = 'disconnected';
  private listeners = new Map<string, Set<Function>>();
  
  constructor(private url: string) {
    this.connect();
  }
  
  private setState(newState: ConnectionState) {
    const oldState = this.state;
    this.state = newState;
    console.log(`State: ${oldState} → ${newState}`);
    this.emit('state-change', { old: oldState, new: newState });
  }
  
  private connect() {
    if (this.state === 'connecting' || this.state === 'connected') {
      return;
    }
    
    this.setState('connecting');
    
    this.ws = new WebSocket(this.url);
    
    this.ws.onopen = () => {
      this.setState('connected');
    };
    
    this.ws.onclose = () => {
      if (this.state === 'connected') {
        this.setState('reconnecting');
        setTimeout(() => this.connect(), 2000);
      } else {
        this.setState('disconnected');
      }
    };
    
    this.ws.onerror = () => {
      this.setState('failed');
    };
  }
  
  getState(): ConnectionState {
    return this.state;
  }
  
  on(event: string, callback: Function) {
    if (!this.listeners.has(event)) {
      this.listeners.set(event, new Set());
    }
    this.listeners.get(event)!.add(callback);
  }
  
  private emit(event: string, data?: any) {
    this.listeners.get(event)?.forEach(callback => callback(data));
  }
}

📚 Key Takeaways

Always implement reconnection - Networks fail
Use heartbeat - Detect dead connections
Queue messages when offline
Handle authentication - Refresh tokens
Manage state - Know your connection status
Exponential backoff - Don't hammer server
Test offline - Toggle airplane mode

WebSockets are powerful but need careful management. Don't wing it in production!

WebSocket Management

On this page