3장: WebSocket — 양방향 실시간 통신

[클라이언트 요청]
GET /ws/chat HTTP/1.1
Host: api.example.com
Upgrade: websocket
Connection: Upgrade
Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==
Sec-WebSocket-Version: 13
 
[서버 응답]
HTTP/1.1 101 Switching Protocols
Upgrade: websocket
Connection: Upgrade
Sec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=

 0                   1                   2                   3
 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
+-+-+-+-+-------+-+-------------+-------------------------------+
|F|R|R|R| opcode|M| Payload len |    Extended payload length    |
|I|S|S|S|  (4)  |A|     (7)     |          (16/64)              |
|N|V|V|V|       |S|             |   (if payload len==126/127)   |
| |1|2|3|       |K|             |                               |
+-+-+-+-+-------+-+-------------+-------------------------------+
|     Masking-key (0 or 4 bytes)                                |
+-------------------------------+-------------------------------+
|                    Payload Data                                |
+---------------------------------------------------------------+

interface WSMessage {
  type: "chat" | "cancel" | "token" | "done" | "error";
  payload: string;
  requestId: string;
}
 
class AIChatClient {
  private ws: WebSocket;
  private currentRequestId: string | null = null;
 
  constructor(url: string) {
    this.ws = new WebSocket(url);
    this.ws.onmessage = this.handleMessage.bind(this);
  }
 
  sendMessage(content: string) {
    this.currentRequestId = crypto.randomUUID();
    const msg: WSMessage = {
      type: "chat",
      payload: content,
      requestId: this.currentRequestId,
    };
    this.ws.send(JSON.stringify(msg));
  }
 
  cancelGeneration() {
    if (this.currentRequestId) {
      const msg: WSMessage = {
        type: "cancel",
        payload: "",
        requestId: this.currentRequestId,
      };
      this.ws.send(JSON.stringify(msg));
    }
  }
 
  private handleMessage(event: MessageEvent) {
    const msg: WSMessage = JSON.parse(event.data);
 
    switch (msg.type) {
      case "token":
        appendTokenToUI(msg.payload);
        break;
      case "done":
        finishResponse();
        break;
      case "error":
        showError(msg.payload);
        break;
    }
  }
}

import { WebSocketServer, WebSocket } from "ws";
 
const wss = new WebSocketServer({ port: 8080 });
 
// 연결 상태 저장소
const connections = new Map<
  string,
  {
    ws: WebSocket;
    userId: string;
    sessionId: string;
    lastPing: number;
  }
>();
 
wss.on("connection", (ws, request) => {
  const connectionId = crypto.randomUUID();
 
  connections.set(connectionId, {
    ws,
    userId: extractUserId(request),
    sessionId: extractSessionId(request),
    lastPing: Date.now(),
  });
 
  // 하트비트 설정
  const heartbeat = setInterval(() => {
    if (ws.readyState === WebSocket.OPEN) {
      ws.ping();
    }
  }, 30000);
 
  ws.on("pong", () => {
    const conn = connections.get(connectionId);
    if (conn) conn.lastPing = Date.now();
  });
 
  ws.on("close", () => {
    clearInterval(heartbeat);
    connections.delete(connectionId);
    cleanupSession(connectionId);
  });
 
  ws.on("error", (error) => {
    console.error(`Connection ${connectionId} error:`, error);
    clearInterval(heartbeat);
    connections.delete(connectionId);
  });
});
 
// 좀비 연결 정리 (60초 이상 무응답)
setInterval(() => {
  const now = Date.now();
  for (const [id, conn] of connections) {
    if (now - conn.lastPing > 60000) {
      conn.ws.terminate();
      connections.delete(id);
    }
  }
}, 30000);

import { Server } from "socket.io";
 
const io = new Server(server, {
  cors: { origin: "*" },
  transports: ["websocket", "polling"],
});
 
// 네임스페이스로 AI 채팅 분리
const aiChat = io.of("/ai-chat");
 
aiChat.on("connection", (socket) => {
  // 룸을 활용한 세션 관리
  socket.on("join-session", (sessionId: string) => {
    socket.join(sessionId);
  });
 
  socket.on("chat-message", async (data) => {
    const { sessionId, message } = data;
 
    // 같은 세션의 모든 참여자에게 사용자 메시지 브로드캐스트
    aiChat.to(sessionId).emit("user-message", {
      userId: socket.id,
      message,
    });
 
    // AI 응답 스트리밍
    for await (const token of streamLLMResponse(message)) {
      aiChat.to(sessionId).emit("ai-token", { token });
    }
 
    aiChat.to(sessionId).emit("ai-done");
  });
 
  socket.on("cancel-generation", (data) => {
    cancelLLMRequest(data.sessionId);
    aiChat.to(data.sessionId).emit("generation-cancelled");
  });
});

import { createClient } from "redis";
import { WebSocketServer } from "ws";
 
const publisher = createClient();
const subscriber = createClient();
 
await publisher.connect();
await subscriber.connect();
 
const wss = new WebSocketServer({ port: 8080 });
const localConnections = new Map<string, WebSocket>();
 
// 로컬 WebSocket 메시지를 Redis로 발행
wss.on("connection", (ws) => {
  const id = crypto.randomUUID();
  localConnections.set(id, ws);
 
  ws.on("message", async (data) => {
    const msg = JSON.parse(data.toString());
    // 다른 서버의 클라이언트에게도 전달
    await publisher.publish(
      `session:${msg.sessionId}`,
      JSON.stringify(msg)
    );
  });
 
  ws.on("close", () => localConnections.delete(id));
});
 
// Redis에서 메시지를 수신하여 로컬 클라이언트에 전달
await subscriber.pSubscribe("session:*", (message) => {
  const msg = JSON.parse(message);
  for (const [, ws] of localConnections) {
    if (ws.readyState === WebSocket.OPEN) {
      ws.send(message);
    }
  }
});

// 그레이스풀 셧다운
process.on("SIGTERM", async () => {
  // 1. 로드밸런서에서 제거 (새 연결 차단)
  await deregisterFromLoadBalancer();
 
  // 2. 기존 연결에 종료 예고
  for (const [id, conn] of connections) {
    conn.ws.send(
      JSON.stringify({
        type: "server-shutdown",
        reconnectAfter: 5000,
      })
    );
  }
 
  // 3. 진행 중인 스트리밍이 완료될 때까지 대기 (최대 30초)
  const deadline = Date.now() + 30000;
  while (connections.size > 0 && Date.now() < deadline) {
    await new Promise((resolve) => setTimeout(resolve, 1000));
  }
 
  // 4. 남은 연결 강제 종료
  for (const [, conn] of connections) {
    conn.ws.close(1001, "Server shutting down");
  }
 
  process.exit(0);
});