sonamu 0.9.4 → 0.9.6

package/src/stream/ws-delivery.ts

@@ -0,0 +1,138 @@
+import { randomUUID } from "node:crypto";
+
+import { type WebSocketAudience } from "./ws-audience";
+import { type WebSocketAudienceResolver } from "./ws-audience-resolver";
+import { type WebSocketClusterBus, type WebSocketClusterEnvelope } from "./ws-cluster-bus";
+import { type ManagedWebSocketConnection } from "./ws-core";
+import { type WebSocketLocalConnectionStore } from "./ws-local-connection-store";
+
+// fan-out을 event loop tick 단위로 나눠 한 번에 긴 동기 루프를 만들지 않게 함
+const FAN_OUT_BATCH_SIZE = 250;
+
+export class WebSocketDeliveryEngine {
+  private readonly pendingFanOutJobs: Array<{
+    targets: ManagedWebSocketConnection[];
+    event: string;
+    data: unknown;
+    cursor: number;
+  }> = [];
+  private fanOutFlushScheduled = false;
+  private readonly unsubscribe: () => void;
+
+  constructor(
+    private readonly options: {
+      nodeId: string;
+      localConnections: WebSocketLocalConnectionStore;
+      audienceResolver: WebSocketAudienceResolver;
+      clusterBus: WebSocketClusterBus;
+    },
+  ) {
+    this.unsubscribe = this.options.clusterBus.subscribe((envelope) => {
+      this.handleClusterEnvelope(envelope);
+    });
+  }
+
+  publishToAudience(audience: WebSocketAudience, event: string, data: unknown): void {
+    const routingPlan = this.options.audienceResolver.resolve(audience);
+    const localTargets = this.options.localConnections.getConnections(routingPlan.localSessionIds);
+    this.enqueueFanOut(localTargets, event, data);
+
+    if (routingPlan.remoteNodeIds.length === 0) {
+      return;
+    }
+
+    void this.options.clusterBus.publish({
+      id: randomUUID(),
+      sourceNodeId: this.options.nodeId,
+      targetNodeIds: routingPlan.remoteNodeIds,
+      audience,
+      event,
+      data,
+      emittedAt: Date.now(),
+    });
+  }
+
+  async shutdown(): Promise<void> {
+    this.unsubscribe();
+    await this.options.clusterBus.shutdown();
+  }
+
+  private handleClusterEnvelope(envelope: WebSocketClusterEnvelope): void {
+    if (envelope.sourceNodeId === this.options.nodeId) {
+      return;
+    }
+
+    if (
+      envelope.targetNodeIds &&
+      envelope.targetNodeIds.length > 0 &&
+      !envelope.targetNodeIds.includes(this.options.nodeId)
+    ) {
+      return;
+    }
+
+    const routingPlan = this.options.audienceResolver.resolve(envelope.audience);
+    this.enqueueFanOut(
+      this.options.localConnections.getConnections(routingPlan.localSessionIds),
+      envelope.event,
+      envelope.data,
+    );
+  }
+
+  private enqueueFanOut(targets: ManagedWebSocketConnection[], event: string, data: unknown): void {
+    if (targets.length === 0) {
+      return;
+    }
+
+    this.pendingFanOutJobs.push({
+      targets,
+      event,
+      data,
+      cursor: 0,
+    });
+    this.scheduleFanOutFlush();
+  }
+
+  private scheduleFanOutFlush(): void {
+    if (this.fanOutFlushScheduled) {
+      return;
+    }
+
+    this.fanOutFlushScheduled = true;
+    setImmediate(() => {
+      this.fanOutFlushScheduled = false;
+      this.flushFanOutJobs();
+    });
+  }
+
+  private flushFanOutJobs(): void {
+    let processed = 0;
+
+    while (this.pendingFanOutJobs.length > 0 && processed < FAN_OUT_BATCH_SIZE) {
+      const job = this.pendingFanOutJobs[0];
+
+      while (job.cursor < job.targets.length && processed < FAN_OUT_BATCH_SIZE) {
+        this.safePublish(job.targets[job.cursor], job.event, job.data);
+        job.cursor += 1;
+        processed += 1;
+      }
+
+      if (job.cursor >= job.targets.length) {
+        this.pendingFanOutJobs.shift();
+      }
+    }
+
+    if (this.pendingFanOutJobs.length > 0) {
+      this.scheduleFanOutFlush();
+    }
+  }
+
+  private safePublish(connection: ManagedWebSocketConnection, event: string, data: unknown): void {
+    try {
+      if (!connection.closed) {
+        connection.publishUntyped(event, data);
+      }
+    } catch {
+      connection.close(1011, "WebSocket publish failed");
+    }
+  }
+}

package/src/stream/ws-local-connection-store.ts

@@ -0,0 +1,44 @@
+import { type ManagedWebSocketConnection } from "./ws-core";
+
+export class WebSocketLocalConnectionStore {
+  private readonly connections = new Map<string, ManagedWebSocketConnection>();
+
+  register(connection: ManagedWebSocketConnection): void {
+    this.connections.set(connection.id, connection);
+  }
+
+  unregister(connectionId: string): void {
+    this.connections.delete(connectionId);
+  }
+
+  getConnection(connectionId: string): ManagedWebSocketConnection | undefined {
+    return this.connections.get(connectionId);
+  }
+
+  getConnections(connectionIds: string[]): ManagedWebSocketConnection[] {
+    const targets: ManagedWebSocketConnection[] = [];
+    const seen = new Set<string>();
+
+    for (const connectionId of connectionIds) {
+      if (seen.has(connectionId)) {
+        continue;
+      }
+      seen.add(connectionId);
+
+      const connection = this.connections.get(connectionId);
+      if (!connection || connection.closed) {
+        continue;
+      }
+
+      targets.push(connection);
+    }
+
+    return targets;
+  }
+
+  closeAll(code?: number, reason?: string): void {
+    for (const connection of this.connections.values()) {
+      connection.close(code, reason);
+    }
+  }
+}

package/src/stream/ws-presence-store.ts

@@ -0,0 +1,326 @@
+import { type WebSocketAudience } from "./ws-audience";
+import { type WebSocketRegistryStats, type WebSocketRoomId, type WebSocketUserId } from "./ws-core";
+
+export type WebSocketSessionPresence = {
+  sessionId: string;
+  nodeId: string;
+  namespace: string;
+  active: boolean;
+  rooms: Set<WebSocketRoomId>;
+  connectedAt: Date;
+  lastSeenAt: Date;
+  userId?: WebSocketUserId;
+};
+
+export interface WebSocketPresenceStore {
+  register(input: {
+    sessionId: string;
+    nodeId: string;
+    namespace: string;
+    active?: boolean;
+  }): WebSocketSessionPresence;
+  activate(sessionId: string): void;
+  unregister(sessionId: string): WebSocketSessionPresence | undefined;
+  touch(sessionId: string): void;
+  setUserId(sessionId: string, userId: WebSocketUserId): void;
+  clearUserId(sessionId: string): void;
+  join(sessionId: string, roomId: WebSocketRoomId): void;
+  leave(sessionId: string, roomId: WebSocketRoomId): void;
+  getConnection(sessionId: string): WebSocketSessionPresence | undefined;
+  getConnectionCount(namespace?: string): number;
+  getRoomMembers(roomId: WebSocketRoomId, namespace?: string): WebSocketSessionPresence[];
+  getStats(): WebSocketRegistryStats;
+  queryAudience(audience: WebSocketAudience): WebSocketSessionPresence[];
+}
+
+export class InMemoryWebSocketPresenceStore implements WebSocketPresenceStore {
+  private readonly connections = new Map<string, WebSocketSessionPresence>();
+  private readonly rooms = new Map<string, Set<string>>();
+  private readonly users = new Map<string, Set<string>>();
+
+  register(input: {
+    sessionId: string;
+    nodeId: string;
+    namespace: string;
+    active?: boolean;
+  }): WebSocketSessionPresence {
+    const meta: WebSocketSessionPresence = {
+      sessionId: input.sessionId,
+      nodeId: input.nodeId,
+      namespace: input.namespace,
+      active: input.active ?? true,
+      rooms: new Set(),
+      connectedAt: new Date(),
+      lastSeenAt: new Date(),
+    };
+    this.connections.set(input.sessionId, meta);
+    return meta;
+  }
+
+  activate(sessionId: string): void {
+    const meta = this.connections.get(sessionId);
+    if (!meta) {
+      return;
+    }
+
+    meta.active = true;
+  }
+
+  unregister(sessionId: string): WebSocketSessionPresence | undefined {
+    const meta = this.connections.get(sessionId);
+    if (!meta) {
+      return undefined;
+    }
+
+    if (meta.userId !== undefined) {
+      this.removeUserBinding(meta.namespace, meta.userId, sessionId);
+    }
+
+    for (const roomId of meta.rooms) {
+      this.removeRoomBinding(meta.namespace, roomId, sessionId);
+    }
+
+    this.connections.delete(sessionId);
+    return meta;
+  }
+
+  touch(sessionId: string): void {
+    const meta = this.connections.get(sessionId);
+    if (!meta) {
+      return;
+    }
+
+    meta.lastSeenAt = new Date();
+  }
+
+  setUserId(sessionId: string, userId: WebSocketUserId): void {
+    const meta = this.connections.get(sessionId);
+    if (!meta) {
+      return;
+    }
+
+    if (meta.userId !== undefined) {
+      this.removeUserBinding(meta.namespace, meta.userId, sessionId);
+    }
+
+    meta.userId = userId;
+    const key = getBindingKey(meta.namespace, String(userId));
+    const bound = this.users.get(key) ?? new Set<string>();
+    bound.add(sessionId);
+    this.users.set(key, bound);
+  }
+
+  clearUserId(sessionId: string): void {
+    const meta = this.connections.get(sessionId);
+    if (!meta || meta.userId === undefined) {
+      return;
+    }
+
+    this.removeUserBinding(meta.namespace, meta.userId, sessionId);
+    delete meta.userId;
+  }
+
+  join(sessionId: string, roomId: WebSocketRoomId): void {
+    const meta = this.connections.get(sessionId);
+    if (!meta) {
+      return;
+    }
+
+    meta.rooms.add(roomId);
+    const key = getBindingKey(meta.namespace, roomId);
+    const members = this.rooms.get(key) ?? new Set<string>();
+    members.add(sessionId);
+    this.rooms.set(key, members);
+  }
+
+  leave(sessionId: string, roomId: WebSocketRoomId): void {
+    const meta = this.connections.get(sessionId);
+    if (!meta) {
+      return;
+    }
+
+    meta.rooms.delete(roomId);
+    this.removeRoomBinding(meta.namespace, roomId, sessionId);
+  }
+
+  getConnection(sessionId: string): WebSocketSessionPresence | undefined {
+    return this.connections.get(sessionId);
+  }
+
+  getConnectionCount(namespace?: string): number {
+    let count = 0;
+
+    for (const meta of this.connections.values()) {
+      if (!meta.active || (namespace && meta.namespace !== namespace)) {
+        continue;
+      }
+
+      count += 1;
+    }
+
+    return count;
+  }
+
+  getRoomMembers(roomId: WebSocketRoomId, namespace?: string): WebSocketSessionPresence[] {
+    return this.getBoundConnections(this.rooms, roomId, namespace);
+  }
+
+  getStats(): WebSocketRegistryStats {
+    const byNamespace: Record<string, number> = {};
+
+    for (const meta of this.connections.values()) {
+      if (!meta.active) {
+        continue;
+      }
+      byNamespace[meta.namespace] = (byNamespace[meta.namespace] ?? 0) + 1;
+    }
+
+    return {
+      totalConnections: Object.values(byNamespace).reduce((sum, count) => sum + count, 0),
+      totalRooms: this.rooms.size,
+      byNamespace,
+    };
+  }
+
+  queryAudience(audience: WebSocketAudience): WebSocketSessionPresence[] {
+    switch (audience.type) {
+      case "all":
+        return this.getActiveConnections(audience.namespace);
+      case "room":
+        return this.getBoundConnections(this.rooms, audience.roomId, audience.namespace);
+      case "user":
+        return this.getBoundConnections(this.users, String(audience.userId), audience.namespace);
+      case "connections": {
+        const metas: WebSocketSessionPresence[] = [];
+
+        for (const connectionId of audience.connectionIds) {
+          const meta = this.connections.get(connectionId);
+          if (!meta?.active) {
+            continue;
+          }
+          if (audience.namespace && meta.namespace !== audience.namespace) {
+            continue;
+          }
+          metas.push(meta);
+        }
+
+        return metas;
+      }
+      case "union": {
+        const seen = new Set<string>();
+        const metas: WebSocketSessionPresence[] = [];
+
+        for (const item of audience.audiences) {
+          for (const meta of this.queryAudience(item)) {
+            if (seen.has(meta.sessionId)) {
+              continue;
+            }
+
+            seen.add(meta.sessionId);
+            metas.push(meta);
+          }
+        }
+
+        return metas;
+      }
+    }
+  }
+
+  private getActiveConnections(namespace?: string): WebSocketSessionPresence[] {
+    const metas: WebSocketSessionPresence[] = [];
+
+    for (const meta of this.connections.values()) {
+      if (!meta.active || (namespace && meta.namespace !== namespace)) {
+        continue;
+      }
+
+      metas.push(meta);
+    }
+
+    return metas;
+  }
+
+  private getBoundConnections(
+    bindings: Map<string, Set<string>>,
+    id: string,
+    namespace?: string,
+  ): WebSocketSessionPresence[] {
+    const metas: WebSocketSessionPresence[] = [];
+    const seen = new Set<string>();
+
+    for (const members of this.getScopedBindings(bindings, id, namespace)) {
+      for (const connectionId of members) {
+        if (seen.has(connectionId)) {
+          continue;
+        }
+        seen.add(connectionId);
+
+        const meta = this.connections.get(connectionId);
+        if (meta?.active) {
+          metas.push(meta);
+        }
+      }
+    }
+
+    return metas;
+  }
+
+  private getScopedBindings(
+    bindings: Map<string, Set<string>>,
+    id: string,
+    namespace?: string,
+  ): Set<string>[] {
+    if (namespace) {
+      const exact = bindings.get(getBindingKey(namespace, id));
+      return exact ? [exact] : [];
+    }
+
+    const matches: Set<string>[] = [];
+    for (const [key, members] of bindings) {
+      const parsed = parseBindingKey(key);
+      if (parsed.id === id) {
+        matches.push(members);
+      }
+    }
+
+    return matches;
+  }
+
+  private removeRoomBinding(namespace: string, roomId: WebSocketRoomId, sessionId: string): void {
+    const key = getBindingKey(namespace, roomId);
+    const members = this.rooms.get(key);
+    if (!members) {
+      return;
+    }
+
+    members.delete(sessionId);
+    if (members.size === 0) {
+      this.rooms.delete(key);
+    }
+  }
+
+  private removeUserBinding(namespace: string, userId: WebSocketUserId, sessionId: string): void {
+    const key = getBindingKey(namespace, String(userId));
+    const members = this.users.get(key);
+    if (!members) {
+      return;
+    }
+
+    members.delete(sessionId);
+    if (members.size === 0) {
+      this.users.delete(key);
+    }
+  }
+}
+
+function getBindingKey(namespace: string, id: string): string {
+  return `${namespace}::${id}`;
+}
+
+function parseBindingKey(key: string): { namespace: string; id: string } {
+  const [namespace, ...rest] = key.split("::");
+  return {
+    namespace,
+    id: rest.join("::"),
+  };
+}

package/src/stream/ws-registry.ts

@@ -0,0 +1,138 @@
+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 WebSocketRegistryStats,
+  type WebSocketRoomId,
+  type WebSocketUserId,
+} from "./ws-core";
+import { WebSocketDeliveryEngine } from "./ws-delivery";
+import { WebSocketLocalConnectionStore } from "./ws-local-connection-store";
+import {
+  InMemoryWebSocketPresenceStore,
+  type WebSocketPresenceStore,
+  type WebSocketSessionPresence,
+} from "./ws-presence-store";
+
+export type {
+  ManagedWebSocketConnection,
+  WebSocketRegistryStats,
+  WebSocketRoomId,
+  WebSocketUserId,
+};
+export type WebSocketConnectionMeta = WebSocketSessionPresence;
+
+export type WebSocketRegistryOptions = {
+  nodeId?: string;
+  presenceStore?: WebSocketPresenceStore;
+  clusterBus?: WebSocketClusterBus;
+};
+
+export class WebSocketRegistry {
+  readonly nodeId: string;
+  readonly localConnections = new WebSocketLocalConnectionStore();
+  readonly presenceStore: WebSocketPresenceStore;
+  readonly clusterBus: WebSocketClusterBus;
+  readonly audienceResolver: WebSocketAudienceResolver;
+  readonly deliveryEngine: WebSocketDeliveryEngine;
+
+  constructor(options: WebSocketRegistryOptions = {}) {
+    this.nodeId = options.nodeId ?? "local";
+    this.presenceStore = options.presenceStore ?? new InMemoryWebSocketPresenceStore();
+    this.clusterBus = options.clusterBus ?? new NoopWebSocketClusterBus();
+    this.audienceResolver = new WebSocketAudienceResolver({
+      nodeId: this.nodeId,
+      presenceStore: this.presenceStore,
+    });
+    this.deliveryEngine = new WebSocketDeliveryEngine({
+      nodeId: this.nodeId,
+      localConnections: this.localConnections,
+      audienceResolver: this.audienceResolver,
+      clusterBus: this.clusterBus,
+    });
+  }
+
+  register(
+    connection: ManagedWebSocketConnection,
+    active: boolean = true,
+  ): WebSocketConnectionMeta {
+    this.localConnections.register(connection);
+    return this.presenceStore.register({
+      sessionId: connection.id,
+      nodeId: this.nodeId,
+      namespace: connection.namespace,
+      active,
+    });
+  }
+
+  activate(connectionId: string): void {
+    this.presenceStore.activate(connectionId);
+  }
+
+  unregister(connectionId: string): void {
+    this.presenceStore.unregister(connectionId);
+    this.localConnections.unregister(connectionId);
+  }
+
+  touch(connectionId: string): void {
+    this.presenceStore.touch(connectionId);
+  }
+
+  setUserId(connectionId: string, userId: WebSocketUserId): void {
+    this.presenceStore.setUserId(connectionId, userId);
+  }
+
+  clearUserId(connectionId: string): void {
+    this.presenceStore.clearUserId(connectionId);
+  }
+
+  join(connectionId: string, roomId: WebSocketRoomId): void {
+    this.presenceStore.join(connectionId, roomId);
+  }
+
+  leave(connectionId: string, roomId: WebSocketRoomId): void {
+    this.presenceStore.leave(connectionId, roomId);
+  }
+
+  broadcast(event: string, data: unknown, namespace?: string): void {
+    this.publishToAudience(WebSocketAudience.all(namespace), event, data);
+  }
+
+  publishToRoom(roomId: WebSocketRoomId, event: string, data: unknown, namespace?: string): void {
+    this.publishToAudience(WebSocketAudience.room(roomId, namespace), event, data);
+  }
+
+  publishToUser(userId: WebSocketUserId, event: string, data: unknown, namespace?: string): void {
+    this.publishToAudience(WebSocketAudience.user(userId, namespace), event, data);
+  }
+
+  publishToAudience(audience: WebSocketAudienceSpec, event: string, data: unknown): void {
+    this.deliveryEngine.publishToAudience(audience, event, data);
+  }
+
+  getConnection(connectionId: string): WebSocketConnectionMeta | undefined {
+    return this.presenceStore.getConnection(connectionId);
+  }
+
+  getConnectionCount(namespace?: string): number {
+    return this.presenceStore.getConnectionCount(namespace);
+  }
+
+  getRoomMembers(roomId: WebSocketRoomId, namespace?: string): WebSocketConnectionMeta[] {
+    return this.presenceStore.getRoomMembers(roomId, namespace);
+  }
+
+  getStats(): WebSocketRegistryStats {
+    return this.presenceStore.getStats();
+  }
+
+  closeAll(code?: number, reason?: string): void {
+    this.localConnections.closeAll(code, reason);
+  }
+
+  async shutdown(code?: number, reason?: string): Promise<void> {
+    this.closeAll(code, reason);
+    await this.deliveryEngine.shutdown();
+  }
+}