sonamu 0.11.16 → 0.12.0

package/src/stream/ws-registry.ts

@@ -1,15 +1,21 @@
+import { type z } from "zod";
+
+import { runtimeTypeOf } from "../runtime-type";
 import { isFunctionValue } from "../utils/runtime-value";
 import { WebSocketAudience, type WebSocketAudience as WebSocketAudienceSpec } from "./ws-audience";
 import { WebSocketAudienceResolver } from "./ws-audience-resolver";
 import { NoopWebSocketClusterBus, type WebSocketClusterBus } from "./ws-cluster-bus";
 import {
   type ManagedWebSocketConnection,
+  type MaybePromise,
+  type WebSocketFrameConnection,
   type WebSocketRegistryStats,
   type WebSocketRoomId,
   type WebSocketUserId,
 } from "./ws-core";
-import { WebSocketDeliveryEngine } from "./ws-delivery";
+import { WebSocketDeliveryEngine, type WebSocketRouteDefinition } from "./ws-delivery";
 import { WebSocketLocalConnectionStore } from "./ws-local-connection-store";
+import { WebSocketPresenceIndex } from "./ws-presence-index";
 import {
   InMemoryWebSocketPresenceStore,
   type WebSocketPresenceStore,
@@ -29,6 +35,22 @@ export type {
   WebSocketUserId,
 };
 export type WebSocketConnectionMeta = WebSocketSessionPresence;
+export type WebSocketRoomMemberCount = {
+  count: number;
+  // store 조회에 실패해 노드 로컬 수로 대체했으면 true
+  partial: boolean;
+};
+
+const DEFAULT_DRAIN_TIMEOUT_MS = 5000;
+
+type StoreMutationOperation =
+  | "register"
+  | "activate"
+  | "unregister"
+  | "setUserId"
+  | "clearUserId"
+  | "join"
+  | "leave";
 
 export type WebSocketRegistryOptions = {
   // 분산 환경에서 노드 간 라우팅 식별자로 쓰이므로 반드시 호출 측(WebSocketRuntime)에서 결정해 넘긴다
@@ -36,38 +58,52 @@ export type WebSocketRegistryOptions = {
   presenceStore?: WebSocketPresenceStore;
   clusterBus?: WebSocketClusterBus;
   telemetryController?: WebSocketTelemetryController;
+  // shutdown 시 진행 중인 store mutation과 publish를 기다리는 상한
+  drainTimeoutMs?: number;
 };
 
+/**
+ * 노드 로컬 presence를 메모리에 소유하고, 설정된 store에는 세션별 순서대로 replica를 반영한다.
+ * mutation과 publish는 호출 시점에 메모리를 동기로 갱신하고, store 실패는 telemetry로만 기록하며 reject하지 않는다.
+ */
 export class WebSocketRegistry {
   readonly nodeId: string;
   readonly localConnections = new WebSocketLocalConnectionStore();
   readonly presenceStore: WebSocketPresenceStore;
   readonly clusterBus: WebSocketClusterBus;
-  readonly audienceResolver: WebSocketAudienceResolver;
-  readonly deliveryEngine: WebSocketDeliveryEngine;
   readonly telemetryController: WebSocketTelemetryController;
+  private readonly memory = new WebSocketPresenceIndex();
+  private readonly deliveryEngine: WebSocketDeliveryEngine;
+  private readonly drainTimeoutMs: number;
+  // 세션별 store mutation 큐의 꼬리. 같은 세션의 호출은 이 promise 뒤에 이어 붙어 호출 순서대로 store에 도달한다
+  private readonly mutationQueues = new Map<string, Promise<void>>();
 
   constructor(options: WebSocketRegistryOptions) {
     this.nodeId = options.nodeId;
     this.presenceStore = options.presenceStore ?? new InMemoryWebSocketPresenceStore();
     this.clusterBus = options.clusterBus ?? new NoopWebSocketClusterBus();
+    this.drainTimeoutMs = options.drainTimeoutMs ?? DEFAULT_DRAIN_TIMEOUT_MS;
+    assertMatchingNodeId("presenceStore", this.presenceStore.nodeId, this.nodeId);
+    assertMatchingNodeId("clusterBus", this.clusterBus.nodeId, this.nodeId);
     this.telemetryController =
       options.telemetryController ?? new NoopWebSocketTelemetryController();
-    this.audienceResolver = new WebSocketAudienceResolver({
-      nodeId: this.nodeId,
-      presenceStore: this.presenceStore,
-    });
     this.deliveryEngine = new WebSocketDeliveryEngine({
       nodeId: this.nodeId,
       localConnections: this.localConnections,
-      audienceResolver: this.audienceResolver,
+      resolveLocalSessionIds: (audience) =>
+        this.memory.queryAudience(audience).map((session) => session.sessionId),
+      audienceResolver: new WebSocketAudienceResolver({
+        nodeId: this.nodeId,
+        presenceStore: this.presenceStore,
+      }),
       clusterBus: this.clusterBus,
       telemetryController: this.telemetryController,
     });
 
+    // telemetry는 노드 로컬 메모리만 집계한다. 클러스터 합계를 모든 노드가 보고하면 집계 시 중복되기 때문이다
     this.telemetryController.registerMetricSource({
       collect: (_now) => {
-        const stats = this.presenceStore.getStats();
+        const stats = this.memory.getStats();
         const localSnapshot = this.localConnections.getTelemetrySnapshot();
         return {
           activeConnections: stats.totalConnections,
@@ -81,17 +117,31 @@ export class WebSocketRegistry {
     });
   }
 
-  register(
-    connection: ManagedWebSocketConnection,
+  registerRoute(routeKey: string, outEvents: z.ZodObject): void {
+    this.deliveryEngine.registerRoute(routeKey, outEvents);
+  }
+
+  replaceRoutes(routes: WebSocketRouteDefinition[]): void {
+    this.deliveryEngine.replaceRoutes(routes);
+  }
+
+  /**
+   * 연결을 메모리에 즉시 등록하고 store.register가 끝나면 resolve한다.
+   * store가 실패해도 메모리 전용 등록으로 resolve한다.
+   */
+  async register(
+    connection: WebSocketFrameConnection,
     active: boolean = true,
-  ): WebSocketConnectionMeta {
-    this.localConnections.register(connection);
-    const meta = this.presenceStore.register({
+  ): Promise<WebSocketConnectionMeta> {
+    const registration = {
       sessionId: connection.id,
       nodeId: this.nodeId,
       namespace: connection.namespace,
       active,
-    });
+      connectedAt: new Date(),
+    };
+    this.localConnections.register(connection);
+    this.memory.register(registration);
     this.telemetryController.emit({
       name: "ws.connection.registered",
       level: "info",
@@ -108,16 +158,27 @@ export class WebSocketRegistry {
       tags: { outcome: "registered", namespace: connection.namespace },
       userId: connection.userId,
     });
-    return meta;
+
+    await this.enqueueStoreMutation(connection.id, "register", () =>
+      this.presenceStore.register(registration),
+    );
+    // 등록 도중 닫혔으면 메모리에서 이미 빠졌으므로 등록 시점 값을 돌려준다
+    return (
+      this.memory.getConnection(connection.id) ?? {
+        ...registration,
+        rooms: new Set(),
+        lastSeenAt: registration.connectedAt,
+      }
+    );
   }
 
-  activate(connectionId: string): void {
-    if (!this.presenceStore.getConnection(connectionId)) {
+  activate(connectionId: string): Promise<void> {
+    if (!this.memory.has(connectionId)) {
       this.emitMutationSkipped(connectionId, { operation: "activate" });
-      return;
+      return Promise.resolve();
     }
 
-    this.presenceStore.activate(connectionId);
+    this.memory.activate(connectionId);
     const connection = this.localConnections.getConnection(connectionId);
     this.telemetryController.emit({
       name: "ws.connection.activated",
@@ -126,14 +187,16 @@ export class WebSocketRegistry {
       userId: connection?.userId,
       ...getConnectionTelemetryContext(connection),
     });
+    return this.enqueueStoreMutation(connectionId, "activate", () =>
+      this.presenceStore.activate(connectionId),
+    );
   }
 
-  unregister(connectionId: string): void {
+  unregister(connectionId: string): Promise<void> {
     const connection = this.localConnections.getConnection(connectionId);
     const telemetryContext = getConnectionTelemetryContext(connection);
-    const userId =
-      connection?.userId ?? formatUserId(this.presenceStore.getConnection(connectionId)?.userId);
-    const meta = this.presenceStore.unregister(connectionId);
+    const meta = this.memory.unregister(connectionId);
+    const userId = connection?.userId ?? formatUserId(meta?.userId);
     this.localConnections.unregister(connectionId);
     this.telemetryController.emit({
       name: "ws.connection.unregistered",
@@ -143,50 +206,60 @@ export class WebSocketRegistry {
       userId,
       ...telemetryContext,
     });
-    if (meta) {
-      this.telemetryController.recordMetric({
-        name: "sonamu.ws.connections",
-        kind: "counter",
-        value: 1,
-        unit: "1",
-        tags: { outcome: "unregistered", namespace: meta.namespace },
-        userId,
-      });
+    if (!meta) {
+      return Promise.resolve();
     }
+
+    this.telemetryController.recordMetric({
+      name: "sonamu.ws.connections",
+      kind: "counter",
+      value: 1,
+      unit: "1",
+      tags: { outcome: "unregistered", namespace: meta.namespace },
+      userId,
+    });
+    // 같은 세션 큐에 이어 붙으므로 진행 중인 register보다 먼저 store에 도달하지 않는다
+    return this.enqueueStoreMutation(connectionId, "unregister", () =>
+      this.presenceStore.unregister(connectionId),
+    );
   }
 
+  // lastSeenAt은 메모리만 갱신한다. store에는 마지막 mutation 시점 값이 남는다
   touch(connectionId: string): void {
-    this.presenceStore.touch(connectionId);
+    this.memory.touch(connectionId);
   }
 
-  setUserId(connectionId: string, userId: WebSocketUserId): void {
-    if (!this.presenceStore.getConnection(connectionId)) {
+  setUserId(connectionId: string, userId: WebSocketUserId): Promise<void> {
+    const meta = this.memory.getConnection(connectionId);
+    if (!meta) {
       this.emitMutationSkipped(connectionId, { operation: "setUserId", userId });
-      return;
+      return Promise.resolve();
     }
 
-    this.presenceStore.setUserId(connectionId, userId);
-    const meta = this.presenceStore.getConnection(connectionId);
+    this.memory.setUserId(connectionId, userId);
     this.telemetryController.emit({
       name: "ws.user.bound",
       level: "debug",
       connectionId,
-      namespace: meta?.namespace,
+      namespace: meta.namespace,
       userId: String(userId),
       detail: { userId },
       ...getConnectionTelemetryContext(this.localConnections.getConnection(connectionId)),
     });
+    return this.enqueueStoreMutation(connectionId, "setUserId", () =>
+      this.presenceStore.setUserId(connectionId, userId),
+    );
   }
 
-  clearUserId(connectionId: string): void {
-    const meta = this.presenceStore.getConnection(connectionId);
+  clearUserId(connectionId: string): Promise<void> {
+    const meta = this.memory.getConnection(connectionId);
     if (!meta) {
       this.emitMutationSkipped(connectionId, { operation: "clearUserId" });
-      return;
+      return Promise.resolve();
     }
 
     const userId = meta.userId;
-    this.presenceStore.clearUserId(connectionId);
+    this.memory.clearUserId(connectionId);
     this.telemetryController.emit({
       name: "ws.user.cleared",
       level: "debug",
@@ -196,36 +269,42 @@ export class WebSocketRegistry {
       detail: userId !== undefined ? { userId } : undefined,
       ...getConnectionTelemetryContext(this.localConnections.getConnection(connectionId)),
     });
+    return this.enqueueStoreMutation(connectionId, "clearUserId", () =>
+      this.presenceStore.clearUserId(connectionId),
+    );
   }
 
-  join(connectionId: string, roomId: WebSocketRoomId): void {
-    if (!this.presenceStore.getConnection(connectionId)) {
+  join(connectionId: string, roomId: WebSocketRoomId): Promise<void> {
+    const meta = this.memory.getConnection(connectionId);
+    if (!meta) {
       this.emitMutationSkipped(connectionId, { operation: "join", roomId });
-      return;
+      return Promise.resolve();
     }
 
-    this.presenceStore.join(connectionId, roomId);
-    const meta = this.presenceStore.getConnection(connectionId);
+    this.memory.join(connectionId, roomId);
     const connection = this.localConnections.getConnection(connectionId);
     this.telemetryController.emit({
       name: "ws.room.joined",
       level: "debug",
       connectionId,
-      namespace: meta?.namespace,
-      userId: connection?.userId ?? formatUserId(meta?.userId),
+      namespace: meta.namespace,
+      userId: connection?.userId ?? formatUserId(meta.userId),
       detail: { roomId },
       ...getConnectionTelemetryContext(connection),
     });
+    return this.enqueueStoreMutation(connectionId, "join", () =>
+      this.presenceStore.join(connectionId, roomId),
+    );
   }
 
-  leave(connectionId: string, roomId: WebSocketRoomId): void {
-    const meta = this.presenceStore.getConnection(connectionId);
+  leave(connectionId: string, roomId: WebSocketRoomId): Promise<void> {
+    const meta = this.memory.getConnection(connectionId);
     if (!meta) {
       this.emitMutationSkipped(connectionId, { operation: "leave", roomId });
-      return;
+      return Promise.resolve();
     }
 
-    this.presenceStore.leave(connectionId, roomId);
+    this.memory.leave(connectionId, roomId);
     const connection = this.localConnections.getConnection(connectionId);
     this.telemetryController.emit({
       name: "ws.room.left",
@@ -236,26 +315,13 @@ export class WebSocketRegistry {
       detail: { roomId },
       ...getConnectionTelemetryContext(connection),
     });
+    return this.enqueueStoreMutation(connectionId, "leave", () =>
+      this.presenceStore.leave(connectionId, roomId),
+    );
   }
 
-  private emitMutationSkipped(
-    connectionId: string,
-    detail: {
-      operation: "activate" | "setUserId" | "clearUserId" | "join" | "leave";
-      roomId?: WebSocketRoomId;
-      userId?: WebSocketUserId;
-    },
-  ): void {
-    this.telemetryController.emit({
-      name: "ws.registry.mutation.skipped",
-      level: "debug",
-      connectionId,
-      detail: { ...detail, reason: "connectionMissing" },
-    });
-  }
-
-  broadcast<Data>(event: string, data: Data, namespace?: string): void {
-    this.publishToAudience(WebSocketAudience.all(namespace), event, data);
+  broadcast<Data>(event: string, data: Data, namespace?: string): Promise<void> {
+    return this.publishToAudience(WebSocketAudience.all(namespace), event, data);
   }
 
   publishToRoom<Data>(
@@ -263,8 +329,8 @@ export class WebSocketRegistry {
     event: string,
     data: Data,
     namespace?: string,
-  ): void {
-    this.publishToAudience(WebSocketAudience.room(roomId, namespace), event, data);
+  ): Promise<void> {
+    return this.publishToAudience(WebSocketAudience.room(roomId, namespace), event, data);
   }
 
   publishToUser<Data>(
@@ -272,40 +338,161 @@ export class WebSocketRegistry {
     event: string,
     data: Data,
     namespace?: string,
-  ): void {
-    this.publishToAudience(WebSocketAudience.user(userId, namespace), event, data);
+  ): Promise<void> {
+    return this.publishToAudience(WebSocketAudience.user(userId, namespace), event, data);
   }
 
-  publishToAudience<Data>(audience: WebSocketAudienceSpec, event: string, data: Data): void {
-    this.deliveryEngine.publishToAudience(audience, event, data);
+  publishToAudience<Data>(
+    audience: WebSocketAudienceSpec,
+    event: string,
+    data: Data,
+  ): Promise<void> {
+    return this.deliveryEngine.publishToAudience(audience, event, data);
   }
 
+  // 아래 읽기 API는 노드 로컬 snapshot(복사본)이다
   getConnection(connectionId: string): WebSocketConnectionMeta | undefined {
-    return this.presenceStore.getConnection(connectionId);
+    return this.memory.getConnection(connectionId);
   }
 
   getConnectionCount(namespace?: string): number {
-    return this.presenceStore.getConnectionCount(namespace);
+    return this.memory.getConnectionCount(namespace);
   }
 
   getRoomMembers(roomId: WebSocketRoomId, namespace?: string): WebSocketConnectionMeta[] {
-    return this.presenceStore.getRoomMembers(roomId, namespace);
+    return this.memory.getRoomMembers(roomId, namespace);
   }
 
   getStats(): WebSocketRegistryStats {
-    return this.presenceStore.getStats();
+    return this.memory.getStats();
+  }
+
+  /**
+   * 클러스터 전체의 활성 방 인원수. store가 실패하면 노드 로컬 수에 partial: true를 붙여 반환하며 reject하지 않는다.
+   */
+  async getRoomMemberCount(
+    roomId: WebSocketRoomId,
+    namespace?: string,
+  ): Promise<WebSocketRoomMemberCount> {
+    try {
+      const count = await callStore(() => this.presenceStore.getRoomMemberCount(roomId, namespace));
+      return { count, partial: false };
+    } catch (cause) {
+      this.telemetryController.emit({
+        name: "ws.registry.read.failed",
+        level: "warn",
+        detail: { operation: "getRoomMemberCount", roomId, errorType: describeError(cause) },
+      });
+      return { count: this.memory.getRoomMemberCount(roomId, namespace), partial: true };
+    }
   }
 
   closeAll(code?: number, reason?: string): void {
     this.localConnections.closeAll(code, reason);
   }
 
+  /**
+   * 로컬 연결을 닫고, 진행 중인 store mutation과 publish를 drainTimeoutMs까지 기다린 뒤 store와 bus를 종료한다.
+   */
   async shutdown(code?: number, reason?: string): Promise<void> {
     this.closeAll(code, reason);
-    await this.deliveryEngine.shutdown();
+
+    let timer: ReturnType<typeof setTimeout> | undefined;
+    const deadline = new Promise<void>((resolve) => {
+      timer = setTimeout(resolve, this.drainTimeoutMs);
+    });
+    try {
+      await Promise.race([
+        Promise.all([...this.mutationQueues.values(), this.deliveryEngine.drainPublishes()]),
+        deadline,
+      ]);
+    } finally {
+      clearTimeout(timer);
+    }
+
+    try {
+      await callStore(() => this.presenceStore.shutdown?.());
+    } catch (cause) {
+      this.telemetryController.emit({
+        name: "ws.registry.shutdown.failed",
+        level: "error",
+        detail: { errorType: describeError(cause) },
+      });
+    }
+    // 구독 해제는 drain이 끝났으므로 즉시 포기할 수 있게 이미 끝난 deadline을 넘긴다
+    await this.deliveryEngine.shutdown(Promise.resolve());
+  }
+
+  /**
+   * 세션별 큐 뒤에 store 호출을 이어 붙인다. 실패는 telemetry로만 기록하고, 반환한 promise는 reject하지 않는다.
+   */
+  private enqueueStoreMutation(
+    sessionId: string,
+    operation: StoreMutationOperation,
+    call: () => MaybePromise<void>,
+  ): Promise<void> {
+    const previous = this.mutationQueues.get(sessionId) ?? Promise.resolve();
+    const next: Promise<void> = previous
+      .then(call)
+      .then(
+        () => undefined,
+        (cause: unknown) => {
+          this.telemetryController.emit({
+            name: "ws.registry.mutation.failed",
+            level: "warn",
+            connectionId: sessionId,
+            detail: { operation, errorType: describeError(cause) },
+          });
+        },
+      )
+      .then(() => {
+        // 마지막 mutation이 끝나면 큐를 비워 세션 수만큼 promise가 쌓이지 않게 한다
+        if (this.mutationQueues.get(sessionId) === next) {
+          this.mutationQueues.delete(sessionId);
+        }
+      });
+    this.mutationQueues.set(sessionId, next);
+    return next;
+  }
+
+  private emitMutationSkipped(
+    connectionId: string,
+    detail: {
+      operation: "activate" | "setUserId" | "clearUserId" | "join" | "leave";
+      roomId?: WebSocketRoomId;
+      userId?: WebSocketUserId;
+    },
+  ): void {
+    this.telemetryController.emit({
+      name: "ws.registry.mutation.skipped",
+      level: "debug",
+      connectionId,
+      detail: { ...detail, reason: "connectionMissing" },
+    });
   }
 }
 
+// store 호출의 동기 throw를 rejection으로 바꾼다
+function callStore<T>(call: () => MaybePromise<T>): Promise<T> {
+  return Promise.resolve().then(call);
+}
+
+function assertMatchingNodeId(
+  source: "presenceStore" | "clusterBus",
+  sourceNodeId: string | undefined,
+  nodeId: string,
+): void {
+  if (sourceNodeId !== undefined && sourceNodeId !== nodeId) {
+    throw new Error(
+      `WebSocket ${source}.nodeId (${sourceNodeId}) must match the runtime nodeId (${nodeId}).`,
+    );
+  }
+}
+
+function describeError(cause: unknown): string {
+  return cause instanceof Error ? cause.name : runtimeTypeOf(cause);
+}
+
 function getConnectionTelemetryContext(
   connection: ManagedWebSocketConnection | undefined,
 ): WebSocketTelemetryConnectionContext {
@@ -316,7 +503,7 @@ function getConnectionTelemetryContext(
   return connection.getTelemetryContext();
 }
 
-// telemetry record의 userId는 string 정규화 — presence store는 number | string 그대로 보관함
+// telemetry record의 userId는 string 정규화 — presence는 number | string 그대로 보관함
 function formatUserId(userId: WebSocketUserId | undefined): string | undefined {
   return userId === undefined ? undefined : String(userId);
 }

package/src/stream/ws-session-id.ts

@@ -0,0 +1,39 @@
+import { randomBytes, randomUUID } from "node:crypto";
+import { hostname } from "node:os";
+
+// 세션 ID는 `${nodeId}:${uuid}` 형식이다. uuid에는 ':'가 없으므로 마지막 ':'로 나누면
+// 노드 ID에 ':' 같은 구분자가 들어 있어도 그대로 복원된다
+const SESSION_ID_SEPARATOR = ":";
+const SESSION_SUFFIX_PATTERN = /^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/;
+
+/**
+ * 기본 노드 ID를 만든다. 형식은 `${hostname()}-${process.pid}-${짧은 랜덤}`이다.
+ * 컨테이너 재시작 뒤 hostname과 pid가 같아도 이전 노드의 stale 상태를 물려받지 않도록 랜덤 접미사를 붙인다.
+ */
+export function createWebSocketNodeId(): string {
+  return `${hostname()}-${process.pid}-${randomBytes(6).toString("base64url")}`;
+}
+
+/**
+ * 노드 ID를 포함한 세션 ID를 만든다. `connections` audience는 이 ID만으로 대상 노드를 찾으므로 store 조회가 필요 없다.
+ */
+export function createWebSocketSessionId(nodeId: string): string {
+  return `${nodeId}${SESSION_ID_SEPARATOR}${randomUUID()}`;
+}
+
+/**
+ * 세션 ID에서 노드 ID를 복원한다. Sonamu가 만든 형식이 아니면 undefined를 반환한다.
+ */
+export function parseWebSocketSessionNodeId(sessionId: string): string | undefined {
+  const separatorIndex = sessionId.lastIndexOf(SESSION_ID_SEPARATOR);
+  if (separatorIndex < 0) {
+    return undefined;
+  }
+
+  const suffix = sessionId.slice(separatorIndex + SESSION_ID_SEPARATOR.length);
+  if (!SESSION_SUFFIX_PATTERN.test(suffix)) {
+    return undefined;
+  }
+
+  return sessionId.slice(0, separatorIndex);
+}