src/websocket/wsServer.ts

import http from 'http';
import { WebSocketServer as WSServer, WebSocket } from 'ws';
import { WebSocketHandler } from './wsHandler';
import { CONFIG } from '../config/app.config';
import { logger } from '../utils/logger';

export class TaskBrokerWebSocketServer {
  private wss: WSServer;
  private handler: WebSocketHandler;

  constructor(server: http.Server, handler: WebSocketHandler) {
    this.handler = handler;
    this.wss = new WSServer({ server, path: CONFIG.WS.PATH });
    this.init();
  }

  private init(): void {
    logger.info(`WebSocket Server initialized on endpoint path [${CONFIG.WS.PATH}]`);

    this.wss.on('connection', (ws: WebSocket) => {
      this.handler.handleConnection(ws);
    });

    // Send periodic pings to keep connections alive through cloud proxies/load balancers
    const pingInterval = setInterval(() => {
      this.wss.clients.forEach((client: any) => {
        if (client.readyState === WebSocket.OPEN) {
          client.ping();
        }
      });
    }, CONFIG.WS.PING_INTERVAL_MS);

    this.wss.on('close', () => {
      clearInterval(pingInterval);
    });
  }

  public sendToWorker(workerId: string, message: any): boolean {
    const ws = this.handler.getSocketForWorker(workerId);
    if (ws && ws.readyState === WebSocket.OPEN) {
      ws.send(JSON.stringify(message));
      return true;
    }
    return false;
  }
}