sonamu 0.9.4 → 0.9.6

package/src/stream/ws.ts

@@ -0,0 +1,591 @@
+import { randomUUID } from "node:crypto";
+
+import { type WebSocket } from "ws";
+import { z } from "zod";
+
+import { type WebSocketAudience } from "./ws-audience";
+import { type WebSocketClusterBus } from "./ws-cluster-bus";
+import { type WebSocketPresenceStore } from "./ws-presence-store";
+import {
+  type ManagedWebSocketConnection,
+  WebSocketRegistry,
+  type WebSocketRegistryOptions,
+  type WebSocketRoomId,
+  type WebSocketUserId,
+} from "./ws-registry";
+
+// transport-level 상수와 queue threshold를 한 파일에 모아 lifecycle/backpressure 정책을 중앙화함
+const WS_CONNECTING = 0;
+const WS_OPEN = 1;
+const WS_CLOSED = 3;
+// RFC 6455 close codes used by Sonamu's WebSocket runtime.
+const WS_CLOSE_CODE_GOING_AWAY = 1001;
+const WS_CLOSE_CODE_INVALID_FRAME_PAYLOAD_DATA = 1007;
+const WS_CLOSE_CODE_POLICY_VIOLATION = 1008;
+const WS_CLOSE_CODE_MESSAGE_TOO_BIG = 1009;
+const WS_CLOSE_CODE_INTERNAL_ERROR = 1011;
+const WS_CLOSE_CODE_TRY_AGAIN_LATER = 1013;
+const MAX_PENDING_MESSAGES = 100;
+const MAX_PENDING_OUTBOUND_MESSAGES = 1_000;
+const MAX_SOCKET_BUFFERED_AMOUNT = 1_048_576;
+const OUTBOUND_BATCH_SIZE = 50;
+const OUTBOUND_RETRY_DELAY_MS = 5;
+
+// envelope을 `{event, data}` 형태로 고정해 server handler와 generated client가 같은 framing contract를 쓰게 함
+const WebSocketEnvelopeSchema = z.object({
+  event: z.string(),
+  data: z.unknown(),
+});
+
+type MessageHandler<T> = (data: T) => void | Promise<void>;
+type CloseHandler = () => void | Promise<void>;
+
+export type WebSocketEventMap = Record<string, unknown>;
+
+type InferWebSocketEventMap<TSchema extends z.ZodRawShape> = z.infer<z.ZodObject<TSchema>>;
+
+export type WebSocketOutEvents<TOut extends WebSocketEventMap = WebSocketEventMap> = TOut;
+
+export type WebSocketInEvents<TIn extends WebSocketEventMap = WebSocketEventMap> = TIn;
+
+export interface WebSocketConnection<
+  TOut extends WebSocketEventMap = WebSocketEventMap,
+  TIn extends WebSocketEventMap = WebSocketEventMap,
+> extends ManagedWebSocketConnection {
+  transport: "ws";
+  onClose(callback: CloseHandler): void;
+  onMessage<K extends keyof WebSocketInEvents<TIn>>(
+    event: K,
+    handler: MessageHandler<WebSocketInEvents<TIn>[K]>,
+  ): void;
+  publish<K extends keyof WebSocketOutEvents<TOut>>(
+    event: K,
+    data: WebSocketOutEvents<TOut>[K],
+  ): void;
+  waitForClose(): Promise<void>;
+  join(roomId: WebSocketRoomId): void;
+  leave(roomId: WebSocketRoomId): void;
+  setUserId(userId: WebSocketUserId): void;
+  clearUserId(): void;
+}
+
+export type AnyWebSocketConnection = WebSocketConnection<WebSocketEventMap, WebSocketEventMap>;
+
+type ParsedEnvelope = z.infer<typeof WebSocketEnvelopeSchema>;
+
+type WebSocketConnectionOptions<TOut extends z.ZodRawShape, TIn extends z.ZodRawShape> = {
+  namespace?: string;
+  heartbeat?: number;
+  maxPayload?: number;
+  active?: boolean;
+  outEvents: z.ZodObject<TOut>;
+  inEvents: z.ZodObject<TIn>;
+  registry: WebSocketRegistry;
+};
+
+export type WebSocketRuntimeOptions = {
+  nodeId?: string;
+  presenceStore?: WebSocketPresenceStore;
+  clusterBus?: WebSocketClusterBus;
+};
+
+// registry를 소유하고 connection 생성/shutdown을 담당함. Sonamu 애플리케이션 수명주기와 같이 움직이도록 설계함
+export class WebSocketRuntime {
+  readonly registry: WebSocketRegistry;
+
+  constructor(options: WebSocketRuntimeOptions = {}) {
+    const registryOptions: WebSocketRegistryOptions = {
+      nodeId: options.nodeId,
+      presenceStore: options.presenceStore,
+      clusterBus: options.clusterBus,
+    };
+    this.registry = new WebSocketRegistry(registryOptions);
+  }
+
+  registerConnection<TOutSchema extends z.ZodRawShape, TInSchema extends z.ZodRawShape>(
+    socket: WebSocket,
+    options: Omit<WebSocketConnectionOptions<TOutSchema, TInSchema>, "registry">,
+  ): WebSocketConnection<InferWebSocketEventMap<TOutSchema>, InferWebSocketEventMap<TInSchema>> {
+    return new WebSocketConnectionImpl(socket, {
+      ...options,
+      registry: this.registry,
+    });
+  }
+
+  activateConnection(connectionId: string): void {
+    this.registry.activate(connectionId);
+  }
+
+  broadcast(event: string, data: unknown, namespace?: string): void {
+    this.registry.broadcast(event, data, namespace);
+  }
+
+  publishToRoom(roomId: WebSocketRoomId, event: string, data: unknown, namespace?: string): void {
+    this.registry.publishToRoom(roomId, event, data, namespace);
+  }
+
+  publishToUser(userId: WebSocketUserId, event: string, data: unknown, namespace?: string): void {
+    this.registry.publishToUser(userId, event, data, namespace);
+  }
+
+  publishToAudience(audience: WebSocketAudience, event: string, data: unknown): void {
+    this.registry.publishToAudience(audience, event, data);
+  }
+
+  // 프로세스 종료 시 살아있는 연결이 남지 않도록 registry를 순회해 일괄 종료함
+  async shutdown(
+    code: number = WS_CLOSE_CODE_GOING_AWAY,
+    reason = "Server shutting down",
+  ): Promise<void> {
+    await this.registry.shutdown(code, reason);
+  }
+}
+
+export function createWebSocketRuntime(options: WebSocketRuntimeOptions = {}): WebSocketRuntime {
+  return new WebSocketRuntime(options);
+}
+
+class WebSocketConnectionImpl<
+  TOutSchema extends z.ZodRawShape,
+  TInSchema extends z.ZodRawShape,
+> implements WebSocketConnection<
+  InferWebSocketEventMap<TOutSchema>,
+  InferWebSocketEventMap<TInSchema>
+> {
+  readonly id = randomUUID();
+  readonly transport = "ws";
+  readonly namespace: string;
+
+  private readonly closeCallbacks: CloseHandler[] = [];
+  private readonly messageHandlers = new Map<string, Array<MessageHandler<unknown>>>();
+  private readonly pendingMessages: ParsedEnvelope[] = [];
+  private readonly pendingOutboundMessages: string[] = [];
+  private readonly closePromise: Promise<void>;
+  private readonly resolveClosePromise: () => void;
+  private readonly heartbeatMs: number;
+  private readonly maxPayload?: number;
+  private readonly eventSchemasIn: Record<string, z.ZodTypeAny>;
+  private readonly eventSchemasOut: Record<string, z.ZodTypeAny>;
+
+  private closedInternal = false;
+  private closeStarted = false;
+  private awaitingPong = false;
+  private heartbeatTimer: ReturnType<typeof setInterval> | null = null;
+  private messageQueue: Promise<void> = Promise.resolve();
+  private outboundFlushScheduled = false;
+
+  constructor(
+    private readonly socket: WebSocket,
+    private readonly options: WebSocketConnectionOptions<TOutSchema, TInSchema>,
+  ) {
+    this.namespace = options.namespace ?? "default";
+    this.heartbeatMs = options.heartbeat ?? 30000;
+    this.maxPayload = options.maxPayload;
+    this.eventSchemasIn = options.inEvents.shape as unknown as Record<string, z.ZodTypeAny>;
+    this.eventSchemasOut = options.outEvents.shape as unknown as Record<string, z.ZodTypeAny>;
+
+    let resolveClosePromise!: () => void;
+    this.closePromise = new Promise<void>((resolve) => {
+      resolveClosePromise = resolve;
+    });
+    this.resolveClosePromise = resolveClosePromise;
+
+    this.options.registry.register(this, options.active ?? true);
+    this.socket.on("message", this.handleMessage);
+    this.socket.on("close", this.handleClose);
+    this.socket.on("error", this.handleError);
+    this.socket.on("pong", this.handlePong);
+    this.startHeartbeat();
+  }
+
+  get closed(): boolean {
+    return this.closedInternal;
+  }
+
+  onClose(callback: CloseHandler): void {
+    this.closeCallbacks.push(callback);
+  }
+
+  onMessage<K extends keyof InferWebSocketEventMap<TInSchema>>(
+    event: K,
+    handler: MessageHandler<InferWebSocketEventMap<TInSchema>[K]>,
+  ): void {
+    const eventKey = String(event);
+    const handlers = this.messageHandlers.get(eventKey) ?? [];
+    handlers.push(handler as MessageHandler<unknown>);
+    this.messageHandlers.set(eventKey, handlers);
+    this.flushPendingMessages(eventKey);
+  }
+
+  publish<K extends keyof InferWebSocketEventMap<TOutSchema>>(
+    event: K,
+    data: InferWebSocketEventMap<TOutSchema>[K],
+  ): void {
+    this.publishValidated(String(event), data);
+  }
+
+  publishUntyped(event: string, data: unknown): void {
+    this.publishValidated(event, data);
+  }
+
+  waitForClose(): Promise<void> {
+    return this.closePromise;
+  }
+
+  join(roomId: WebSocketRoomId): void {
+    this.options.registry.join(this.id, roomId);
+  }
+
+  leave(roomId: WebSocketRoomId): void {
+    this.options.registry.leave(this.id, roomId);
+  }
+
+  setUserId(userId: WebSocketUserId): void {
+    this.options.registry.setUserId(this.id, userId);
+  }
+
+  clearUserId(): void {
+    this.options.registry.clearUserId(this.id);
+  }
+
+  // transport 종료 도중 예외가 나도 markClosed가 반드시 실행되도록 try/finally로 감쌈
+  close(code?: number, reason?: string): void {
+    if (this.closedInternal || this.closeStarted || this.socket.readyState === WS_CLOSED) {
+      return;
+    }
+
+    this.closeStarted = true;
+    try {
+      this.closeTransport(code, reason);
+    } finally {
+      this.markClosed();
+    }
+  }
+
+  // 인바운드 메시지를 순차 처리 큐에 올림. payload size → envelope 파싱 순으로 transport 레벨 검증을 우선 수행함
+  private readonly handleMessage = (raw: unknown) => {
+    this.enqueueMessageTask(async () => {
+      const text = normalizeMessage(raw);
+      if (this.maxPayload !== undefined && Buffer.byteLength(text) > this.maxPayload) {
+        this.close(WS_CLOSE_CODE_MESSAGE_TOO_BIG, "Message too large");
+        return;
+      }
+
+      const parsedEnvelope = safeParseEnvelope(text);
+      if (!parsedEnvelope) {
+        this.close(WS_CLOSE_CODE_INVALID_FRAME_PAYLOAD_DATA, "Invalid message payload");
+        return;
+      }
+
+      this.options.registry.touch(this.id);
+      await this.dispatchEnvelope(parsedEnvelope);
+    });
+  };
+
+  private readonly handleClose = () => {
+    this.markClosed();
+  };
+
+  // 소켓이 transport error를 emit하면 즉시 1011 close로 수렴시켜 상태 누락을 막음
+  private readonly handleError = () => {
+    this.close(WS_CLOSE_CODE_INTERNAL_ERROR, "WebSocket transport error");
+  };
+
+  private readonly handlePong = () => {
+    this.awaitingPong = false;
+    this.options.registry.touch(this.id);
+  };
+
+  // event 존재 여부 → schema 검증 → handler 실행 순으로 분기함
+  // handler가 아직 등록되지 않은 초기 메시지는 버퍼에 보관했다가 onMessage 등록 시 flush함
+  // handler는 `await`으로 순차 실행해 한 connection 안의 메시지 순서를 보장함
+  private async dispatchEnvelope(envelope: ParsedEnvelope): Promise<void> {
+    const handlers = this.messageHandlers.get(envelope.event);
+    const schema = this.eventSchemasIn[envelope.event];
+
+    if (!schema) {
+      this.close(WS_CLOSE_CODE_POLICY_VIOLATION, "Unknown event");
+      return;
+    }
+
+    const parsed = schema.safeParse(envelope.data);
+    if (!parsed.success) {
+      this.close(WS_CLOSE_CODE_INVALID_FRAME_PAYLOAD_DATA, "Invalid event data");
+      return;
+    }
+
+    if (!handlers || handlers.length === 0) {
+      if (this.pendingMessages.length >= MAX_PENDING_MESSAGES) {
+        this.pendingMessages.shift();
+      }
+      this.pendingMessages.push(envelope);
+      return;
+    }
+
+    for (const handler of handlers) {
+      await handler(parsed.data);
+    }
+  }
+
+  private flushPendingMessages(event: string): void {
+    const remaining: ParsedEnvelope[] = [];
+    const toFlush: ParsedEnvelope[] = [];
+
+    for (const message of this.pendingMessages) {
+      if (message.event !== event) {
+        remaining.push(message);
+        continue;
+      }
+
+      toFlush.push(message);
+    }
+
+    this.pendingMessages.length = 0;
+    this.pendingMessages.push(...remaining);
+
+    for (const message of toFlush) {
+      this.enqueueMessageTask(async () => {
+        await this.dispatchEnvelope(message);
+      });
+    }
+  }
+
+  private publishValidated(event: string, data: unknown): void {
+    const schema = this.eventSchemasOut[event];
+    if (!schema) {
+      throw new Error(`Unknown websocket event: ${event}`);
+    }
+
+    const parsed = schema.safeParse(data);
+    if (!parsed.success) {
+      throw new Error(`Invalid websocket event payload: ${event}`);
+    }
+
+    if (this.closedInternal || this.socket.readyState !== WS_OPEN) {
+      return;
+    }
+
+    this.enqueueOutboundMessage(
+      JSON.stringify({
+        event,
+        data: parsed.data,
+      }),
+    );
+  }
+
+  // listener 해제 / heartbeat 중단 / pending queue 비움 / registry unregister / onClose 실행 / waitForClose resolve 을 한 곳에 모아 원자적으로 처리함
+  // async onClose가 reject해도 unhandled rejection으로 새지 않도록 catch로 격리함
+  private markClosed(): void {
+    if (this.closedInternal) {
+      return;
+    }
+
+    this.closedInternal = true;
+    this.closeStarted = false;
+    this.stopHeartbeat();
+    this.socket.off("message", this.handleMessage);
+    this.socket.off("close", this.handleClose);
+    this.socket.off("error", this.handleError);
+    this.socket.off("pong", this.handlePong);
+    this.awaitingPong = false;
+    this.pendingMessages.length = 0;
+    this.pendingOutboundMessages.length = 0;
+    this.options.registry.unregister(this.id);
+    for (const callback of this.closeCallbacks.splice(0)) {
+      try {
+        const result = callback();
+        if (isPromiseLike(result)) {
+          void result.catch(() => {
+            // async close callbacks must not escape as unhandled rejections
+          });
+        }
+      } catch {
+        // close callbacks must not block transport cleanup
+      }
+    }
+    this.resolveClosePromise();
+  }
+
+  // pong이 오지 않은 상태에서 다음 tick이 오면 timeout close로 처리해 zombie connection을 정리함
+  private startHeartbeat(): void {
+    if (this.heartbeatMs <= 0) {
+      return;
+    }
+
+    this.heartbeatTimer = setInterval(() => {
+      if (this.closedInternal || this.socket.readyState !== WS_OPEN) {
+        return;
+      }
+
+      if (this.awaitingPong) {
+        this.close(WS_CLOSE_CODE_GOING_AWAY, "Heartbeat timeout");
+        return;
+      }
+
+      this.awaitingPong = true;
+      this.socket.ping();
+    }, this.heartbeatMs);
+  }
+
+  private stopHeartbeat(): void {
+    if (!this.heartbeatTimer) {
+      return;
+    }
+
+    clearInterval(this.heartbeatTimer);
+    this.heartbeatTimer = null;
+  }
+
+  // close가 실패해도 terminate 폴백까지 시도하고, 끝내 실패하면 markClosed에 상태 정리를 위임함
+  private closeTransport(code?: number, reason?: string): void {
+    try {
+      if (this.socket.readyState === WS_OPEN || this.socket.readyState === WS_CONNECTING) {
+        this.socket.close(code, truncateCloseReason(reason));
+        return;
+      }
+
+      this.socket.terminate();
+    } catch {
+      try {
+        this.socket.terminate();
+      } catch {
+        // transport is already broken; state cleanup is handled by markClosed()
+      }
+    }
+  }
+
+  // 인바운드 handler를 promise chain으로 serialize 하고, handler 예외는 connection-local 1011 close로 축소함
+  private enqueueMessageTask(task: () => Promise<void>): void {
+    this.messageQueue = this.messageQueue
+      .then(async () => {
+        if (this.closedInternal) {
+          return;
+        }
+
+        await task();
+      })
+      .catch(() => {
+        this.close(WS_CLOSE_CODE_INTERNAL_ERROR, "Message handling failed");
+      });
+  }
+
+  // queue가 한계에 도달하면 1013으로 닫아 느린 소비자가 메모리를 끝없이 잡아먹지 못하게 함
+  private enqueueOutboundMessage(payload: string): void {
+    if (this.pendingOutboundMessages.length >= MAX_PENDING_OUTBOUND_MESSAGES) {
+      this.close(WS_CLOSE_CODE_TRY_AGAIN_LATER, "WebSocket backpressure overflow");
+      return;
+    }
+
+    this.pendingOutboundMessages.push(payload);
+    this.scheduleOutboundFlush();
+  }
+
+  private scheduleOutboundFlush(delayMs: number = 0): void {
+    if (this.outboundFlushScheduled || this.closedInternal) {
+      return;
+    }
+
+    this.outboundFlushScheduled = true;
+    const flush = () => {
+      this.outboundFlushScheduled = false;
+      this.flushOutboundMessages();
+    };
+
+    if (delayMs > 0) {
+      setTimeout(flush, delayMs);
+      return;
+    }
+
+    setImmediate(flush);
+  }
+
+  // bufferedAmount가 임계치를 넘으면 flush를 미뤄 socket 내부 큐가 터지지 않도록 backpressure를 존중함
+  // 한 번에 배치 단위로만 send해 동기 루프가 이벤트 루프를 장시간 점유하지 않게 함
+  private flushOutboundMessages(): void {
+    if (this.closedInternal || this.socket.readyState !== WS_OPEN) {
+      return;
+    }
+
+    if (this.socket.bufferedAmount > MAX_SOCKET_BUFFERED_AMOUNT) {
+      this.scheduleOutboundFlush(OUTBOUND_RETRY_DELAY_MS);
+      return;
+    }
+
+    let sent = 0;
+    while (
+      sent < OUTBOUND_BATCH_SIZE &&
+      this.pendingOutboundMessages.length > 0 &&
+      this.socket.readyState === WS_OPEN
+    ) {
+      const payload = this.pendingOutboundMessages.shift();
+      if (!payload) {
+        break;
+      }
+
+      try {
+        this.socket.send(payload);
+      } catch {
+        this.close(WS_CLOSE_CODE_INTERNAL_ERROR, "Outbound publish failed");
+        return;
+      }
+
+      sent += 1;
+      if (this.socket.bufferedAmount > MAX_SOCKET_BUFFERED_AMOUNT) {
+        break;
+      }
+    }
+
+    if (this.pendingOutboundMessages.length > 0) {
+      this.scheduleOutboundFlush(
+        this.socket.bufferedAmount > MAX_SOCKET_BUFFERED_AMOUNT ? OUTBOUND_RETRY_DELAY_MS : 0,
+      );
+    }
+  }
+}
+
+function normalizeMessage(raw: unknown): string {
+  if (typeof raw === "string") {
+    return raw;
+  }
+
+  if (raw instanceof Buffer) {
+    return raw.toString("utf-8");
+  }
+
+  if (raw instanceof ArrayBuffer) {
+    return Buffer.from(raw).toString("utf-8");
+  }
+
+  if (Array.isArray(raw)) {
+    return Buffer.concat(raw.filter((chunk): chunk is Buffer => chunk instanceof Buffer)).toString(
+      "utf-8",
+    );
+  }
+
+  return JSON.stringify(raw);
+}
+
+function safeParseEnvelope(raw: string): ParsedEnvelope | null {
+  try {
+    const parsed = JSON.parse(raw) as unknown;
+    const validated = WebSocketEnvelopeSchema.safeParse(parsed);
+    return validated.success ? validated.data : null;
+  } catch {
+    return null;
+  }
+}
+
+// RFC 6455가 close frame reason을 123 byte로 제한하므로 초과분은 잘라 전송 실패를 방지함
+function truncateCloseReason(reason?: string): string | undefined {
+  if (!reason) {
+    return undefined;
+  }
+
+  return Buffer.byteLength(reason, "utf-8") <= 123
+    ? reason
+    : Buffer.from(reason).subarray(0, 123).toString("utf-8");
+}
+
+function isPromiseLike(value: unknown): value is Promise<void> {
+  return typeof value === "object" && value !== null && "then" in value && "catch" in value;
+}

package/src/syncer/__tests__/api-parser.websocket-type-ref.test.ts

@@ -0,0 +1,78 @@
+import { mkdtemp, rm, writeFile } from "fs/promises";
+import { tmpdir } from "os";
+import path from "path";
+
+import { afterEach, beforeEach, describe, expect, it } from "vitest";
+import { z } from "zod";
+
+import { registeredApis } from "../../api/decorators";
+import { readApisFromFile } from "../api-parser";
+
+describe("readApisFromFile websocket type refs", () => {
+  let originalApis: Array<(typeof registeredApis)[number]> = [];
+  let tempDir: string | null = null;
+
+  beforeEach(() => {
+    originalApis = [...registeredApis];
+  });
+
+  afterEach(async () => {
+    registeredApis.splice(0, registeredApis.length, ...originalApis);
+    if (tempDir) {
+      await rm(tempDir, { recursive: true, force: true });
+      tempDir = null;
+    }
+  });
+
+  it("captures websocket event identifiers for generated service imports", async () => {
+    registeredApis.push({
+      modelName: "ChatFrame",
+      methodName: "subscribeChat",
+      path: "/chat/subscribeChat",
+      options: {
+        httpMethod: "GET",
+      },
+      websocketOptions: {
+        outEvents: z.object({
+          ready: z.object({
+            ok: z.boolean(),
+          }),
+        }),
+        inEvents: z.object({
+          ping: z.object({
+            at: z.string(),
+          }),
+        }),
+      },
+    });
+
+    tempDir = await mkdtemp(path.join(tmpdir(), "sonamu-api-parser-"));
+    const filePath = path.join(tempDir, "chat.frame.ts");
+    await writeFile(
+      filePath,
+      `
+class ChatFrameClass {
+  @api()
+  @websocket({
+    outEvents: ChatOutEvents,
+    inEvents: ChatInEvents,
+  })
+  async subscribeChat(): Promise<void> {
+    return;
+  }
+}
+      `.trim(),
+    );
+
+    const [api] = await readApisFromFile(filePath as `${string}.ts`);
+
+    expect(api?.websocketOptions?.outEventsTypeRef).toEqual({
+      t: "ref",
+      id: "ChatOutEvents",
+    });
+    expect(api?.websocketOptions?.inEventsTypeRef).toEqual({
+      t: "ref",
+      id: "ChatInEvents",
+    });
+  });
+});