@enbox/dwn-server 0.0.2 → 0.0.4

package/src/connection/flow-controller.ts

@@ -0,0 +1,117 @@
+import type { SubscriptionMessage } from '@enbox/dwn-sdk-js';
+import type { JsonRpcId, JsonRpcSuccessResponse } from '@enbox/dwn-clients';
+
+import log from 'loglevel';
+
+import { createJsonRpcSuccessResponse } from '@enbox/dwn-clients';
+
+/** Default maximum number of unacknowledged events before pausing delivery. */
+export const DEFAULT_MAX_IN_FLIGHT = 32;
+
+/** Maximum buffer size before the subscription is force-closed to prevent OOM. */
+export const MAX_BUFFER_SIZE = 1000;
+
+/**
+ * Per-subscription flow controller that enforces a sliding window of
+ * unacknowledged events. When the window is full, incoming events are
+ * buffered. When the client sends `rpc.ack` with a cursor, events up
+ * to that cursor are acknowledged and buffered events are flushed.
+ *
+ * If the buffer exceeds {@link MAX_BUFFER_SIZE}, the subscription is
+ * closed via the provided `onOverflow` callback to prevent unbounded
+ * memory growth.
+ */
+export class FlowController {
+  /** Ordered list of cursors for events that have been sent but not yet acknowledged. */
+  private unacked: string[] = [];
+
+  /** Buffer of events waiting to be sent once the window opens. */
+  private buffer: SubscriptionMessage[] = [];
+
+  /** Whether the controller has been closed due to overflow. */
+  private closed = false;
+
+  constructor(
+    private readonly subscriptionId: JsonRpcId,
+    private readonly maxInFlight: number,
+    private readonly send: (response: JsonRpcSuccessResponse) => void,
+    private readonly onOverflow: () => void,
+  ) {}
+
+  /**
+   * Accept an incoming {@link SubscriptionMessage} from the EventLog listener.
+   * If the window has room, send immediately. Otherwise buffer.
+   */
+  public push(message: SubscriptionMessage): void {
+    if (this.closed) {
+      return;
+    }
+
+    if (this.unacked.length < this.maxInFlight) {
+      this.sendMessage(message);
+    } else {
+      this.buffer.push(message);
+
+      if (this.buffer.length > MAX_BUFFER_SIZE) {
+        log.warn(
+          `FlowController: buffer overflow for subscription ${String(this.subscriptionId)}, ` +
+          `closing subscription (buffer=${this.buffer.length}, unacked=${this.unacked.length})`
+        );
+        this.closed = true;
+        this.buffer = [];
+        this.unacked = [];
+        this.onOverflow();
+      }
+    }
+  }
+
+  /**
+   * Process an `rpc.ack` for this subscription. Acknowledges all events up
+   * to and including the given cursor, then flushes buffered events into the
+   * newly opened window slots.
+   */
+  public ack(cursor: string): void {
+    if (this.closed) {
+      return;
+    }
+
+    const idx = this.unacked.lastIndexOf(cursor);
+    if (idx === -1) {
+      // Unknown cursor — could be a stale or duplicate ack. Ignore silently.
+      log.debug(`FlowController: unknown cursor in ack for subscription ${String(this.subscriptionId)}: ${cursor}`);
+      return;
+    }
+
+    // Remove all entries up to and including the acked cursor.
+    this.unacked.splice(0, idx + 1);
+
+    // Flush buffered messages into the freed window slots.
+    while (this.buffer.length > 0 && this.unacked.length < this.maxInFlight) {
+      const buffered = this.buffer.shift()!;
+      this.sendMessage(buffered);
+    }
+  }
+
+  /**
+   * Returns the number of events currently in flight (sent but unacknowledged).
+   */
+  public get inFlightCount(): number {
+    return this.unacked.length;
+  }
+
+  /**
+   * Returns the number of events currently buffered (waiting to be sent).
+   */
+  public get bufferCount(): number {
+    return this.buffer.length;
+  }
+
+  /**
+   * Sends a single message over the wire and tracks its cursor.
+   */
+  private sendMessage(message: SubscriptionMessage): void {
+    const response = createJsonRpcSuccessResponse(this.subscriptionId, { subscription: message });
+    this.send(response);
+    this.unacked.push(message.cursor);
+  }
+}

package/src/connection/socket-connection.ts

@@ -1,54 +1,80 @@
-import type { Dwn, GenericMessage, MessageEvent } from "@enbox/dwn-sdk-js";
-import { DwnMethodName } from "@enbox/dwn-sdk-js";
+import type { ActivityLog } from '../admin/activity-log.js';
+import type { AdminConnectionSnapshot } from '../admin/types.js';
+import type { AdminStore } from '../admin/admin-store.js';
+import type { DwnServerConfig } from '../config.js';
+import type { MessageProcessedHook } from '../message-processed-hook.js';
+import type { RateLimiter } from '../rate-limiter.js';
+import type { RegistrationStore } from '../registration/registration-store.js';
+import type { RequestContext } from '../lib/json-rpc-router.js';
+import type { ServerWebSocket } from 'bun';
+import type { WsData } from '../http-api.js';
+import type { Dwn, GenericMessage, SubscriptionMessage } from '@enbox/dwn-sdk-js';
+import type { JsonRpcErrorResponse, JsonRpcId, JsonRpcRequest, JsonRpcResponse, JsonRpcSubscription } from '@enbox/dwn-clients';
 
-import type { WebSocket } from "ws";
 import log from 'loglevel';
-import { v4 as uuidv4 } from 'uuid';
-
-import type { RequestContext } from "../lib/json-rpc-router.js";
-import type { JsonRpcErrorResponse, JsonRpcId, JsonRpcRequest, JsonRpcResponse, JsonRpcSubscription } from "../lib/json-rpc.js";
 
-import { requestCounter } from "../metrics.js";
-import { jsonRpcRouter } from "../json-rpc-api.js";
-import { JsonRpcErrorCodes, createJsonRpcErrorResponse, createJsonRpcSuccessResponse } from "../lib/json-rpc.js";
-import { DwnServerError, DwnServerErrorCode } from "../dwn-error.js";
+import { DwnMethodName } from '@enbox/dwn-sdk-js';
+import { jsonRpcRouter } from '../json-rpc-api.js';
+import { requestCounter } from '../metrics.js';
+import { v4 as uuidv4 } from 'uuid';
+import { createJsonRpcErrorResponse, JsonRpcErrorCodes } from '@enbox/dwn-clients';
+import { DEFAULT_MAX_IN_FLIGHT, FlowController } from './flow-controller.js';
+import { DwnServerError, DwnServerErrorCode } from '../dwn-error.js';
 
 const HEARTBEAT_INTERVAL = 30_000;
 
 /**
  * SocketConnection handles a WebSocket connection to a DWN using JSON RPC.
  * It also manages references to the long running RPC subscriptions for the connection.
+ *
+ * With Bun's native WebSocket, the message/close/error events are dispatched by the
+ * Bun.serve() websocket handlers in http-api.ts, which delegate to the public `message()`
+ * and `close()` methods on this class.
  */
 export class SocketConnection {
-  private heartbeatInterval: NodeJS.Timer;
+  /** Unique identifier for this connection (for admin introspection). */
+  public readonly id: string = uuidv4();
+
+  /** Timestamp when the connection was established (for admin introspection). */
+  public readonly connectedAt: number = Date.now();
+
+  private heartbeatInterval: ReturnType<typeof setInterval>;
   private subscriptions: Map<JsonRpcId, JsonRpcSubscription> = new Map();
+  private flowControllers: Map<JsonRpcId, FlowController> = new Map();
   private isAlive: boolean;
 
   constructor(
-    private socket: WebSocket,
+    private socket: ServerWebSocket<WsData>,
     private dwn: Dwn,
-    private onClose?: () => void
+    private onCloseCallback?: () => void,
+    private maxInFlight: number = DEFAULT_MAX_IN_FLIGHT,
+    private activityLog?: ActivityLog,
+    private adminStore?: AdminStore,
+    private registrationStore?: RegistrationStore,
+    private serverConfig?: DwnServerConfig,
+    private tenantRateLimiter?: RateLimiter,
+    private messageProcessedHooks?: MessageProcessedHook[],
   ){
-    socket.on('message', this.message.bind(this));
-    socket.on('close', this.close.bind(this));
-    socket.on('error', this.error.bind(this));
-    socket.on('pong', this.pong.bind(this));
-
-    // Sometimes connections between client <-> server can get borked in such a way that
-    // leaves both unaware of the borkage. ping messages can be used as a means to verify
-    // that the remote endpoint is still responsive. Server will ping each socket every 30s
-    // if a pong hasn't received from a socket by the next ping, the server will terminate
-    // the socket connection
+    // Bun handles ping/pong automatically at the protocol level, but we still
+    // want an application-level heartbeat to detect dead connections.
     this.isAlive = true;
     this.heartbeatInterval = setInterval(() => {
       if (this.isAlive === false) {
         this.close();
+        return;
       }
       this.isAlive = false;
       this.socket.ping();
     }, HEARTBEAT_INTERVAL);
   }
 
+  /**
+   * Called when a pong is received (triggered by Bun's built-in ping/pong handling).
+   */
+  pong(): void {
+    this.isAlive = true;
+  }
+
   /**
    * Checks to see if the incoming `JsonRpcId` is already in use for a subscription.
    */
@@ -65,7 +91,7 @@ export class SocketConnection {
       throw new DwnServerError(
         DwnServerErrorCode.ConnectionSubscriptionJsonRpcIdExists,
         `the subscription with id ${subscription.id} already exists`
-      )
+      );
     }
 
     this.subscriptions.set(subscription.id, subscription);
@@ -81,12 +107,24 @@ export class SocketConnection {
       throw new DwnServerError(
         DwnServerErrorCode.ConnectionSubscriptionJsonRpcIdNotFound,
         `the subscription with id ${id} was not found`
-      )
+      );
     }
 
     const connection = this.subscriptions.get(id);
     await connection.close();
     this.subscriptions.delete(id);
+    this.flowControllers.delete(id);
+  }
+
+  /**
+   * Acknowledges subscription events up to the given cursor, advancing the
+   * flow-control window for the subscription.
+   */
+  ackSubscription(id: JsonRpcId, cursor: string): void {
+    const fc = this.flowControllers.get(id);
+    if (fc) {
+      fc.ack(cursor);
+    }
   }
 
   /**
@@ -95,10 +133,7 @@ export class SocketConnection {
   async close(): Promise<void> {
     clearInterval(this.heartbeatInterval);
 
-    // clean up all socket event listeners
-    this.socket.removeAllListeners();
-
-    const closePromises = [];
+    const closePromises: Promise<void>[] = [];
     for (const [id, subscription] of this.subscriptions) {
       closePromises.push(subscription.close());
       this.subscriptions.delete(id);
@@ -107,56 +142,52 @@ export class SocketConnection {
     // close all of the associated subscriptions
     await Promise.all(closePromises);
 
+    // clear all flow controllers
+    this.flowControllers.clear();
+
     // close the socket.
     this.socket.close();
 
     // if there was a close handler passed call it after the connection has been closed
-    if (this.onClose !== undefined) {
-      this.onClose();
+    if (this.onCloseCallback !== undefined) {
+      this.onCloseCallback();
     }
   }
 
-  /**
-   * Pong messages are automatically sent in response to ping messages as required by
-   * the websocket spec. So, no need to send explicit pongs.
-   */
-  private pong(): void {
-    this.isAlive = true;
-  }
-
   /**
    * Log the error and close the connection.
    */
-  private async error(error:Error): Promise<void>{
+  async error(error: Error): Promise<void> {
     log.error(`SocketConnection error, terminating connection`, error);
-    this.socket.terminate();
+    this.socket.close();
     await this.close();
   }
 
   /**
    * Handles a `JSON RPC 2.0` encoded message.
+   * This is called by Bun's websocket message handler via http-api.ts.
    */
-  private async message(dataBuffer: Buffer): Promise<void> {
+  async message(dataBuffer: Buffer): Promise<void> {
     const requestData = dataBuffer.toString();
     if (!requestData) {
       return this.send(createJsonRpcErrorResponse(
         uuidv4(),
         JsonRpcErrorCodes.BadRequest,
         'request payload required.'
-      ))
+      ));
     }
 
     let jsonRequest: JsonRpcRequest;
     try {
       jsonRequest = JSON.parse(requestData);
-    } catch(error) {
+    } catch (error) {
       const errorResponse = createJsonRpcErrorResponse(
         uuidv4(),
         JsonRpcErrorCodes.BadRequest,
         (error as Error).message
       );
       return this.send(errorResponse);
-    };
+    }
 
     const requestContext = await this.buildRequestContext(jsonRequest);
     const { jsonRpcResponse } = await jsonRpcRouter.handle(jsonRequest, requestContext);
@@ -164,55 +195,101 @@ export class SocketConnection {
       requestCounter.inc({ method: jsonRequest.method, error: 1 });
     } else {
       requestCounter.inc({
-        method: jsonRequest.method,
-        status: jsonRpcResponse?.result?.reply?.status?.code || 0,
+        method : jsonRequest.method,
+        status : jsonRpcResponse?.result?.reply?.status?.code || 0,
       });
     }
     this.send(jsonRpcResponse);
   }
 
   /**
-   * Sends a JSON encoded Buffer through the Websocket.
+   * Returns the number of active subscriptions on this connection.
+   */
+  get subscriptionCount(): number {
+    return this.subscriptions.size;
+  }
+
+  /**
+   * Returns a serializable snapshot of this connection for the admin inspector.
+   */
+  toSnapshot(): AdminConnectionSnapshot {
+    const subscriptions = Array.from(this.flowControllers.entries()).map(
+      ([id, fc]): AdminConnectionSnapshot['subscriptions'][number] => ({
+        id       : id as string | number,
+        inflight : fc.inFlightCount,
+        buffered : fc.bufferCount,
+      }),
+    );
+
+    return {
+      id                : this.id,
+      connectedAt       : new Date(this.connectedAt).toISOString(),
+      subscriptionCount : this.subscriptions.size,
+      subscriptions,
+    };
+  }
+
+  /**
+   * Sends a JSON encoded string through the WebSocket.
    */
   private send(response: JsonRpcResponse | JsonRpcErrorResponse): void {
     this.socket.send(JSON.stringify(response));
   }
 
   /**
-   * Creates a subscription handler to send messages matching the subscription requested.
-   *
-   * Wraps the incoming `message` in a `JSON RPC Success Response` using the original subscription`JSON RPC Id` to send through the WebSocket.
+   * Creates a flow-controlled subscription handler that enforces the
+   * `maxInFlight` window. Returns a `SubscriptionListener` to be passed
+   * to the EventLog, and stores the `FlowController` for later `rpc.ack`
+   * processing.
    */
-  private createSubscriptionHandler(id: JsonRpcId): (message: MessageEvent) => void {
-    return (event) => {
-      const response = createJsonRpcSuccessResponse(id, { event });
-      this.send(response);
-    }
+  private createSubscriptionHandler(id: JsonRpcId): (message: SubscriptionMessage) => void {
+    const fc = new FlowController(
+      id,
+      this.maxInFlight,
+      (response) => {
+        this.send(response);
+      },
+      () => {
+        // overflow: close the subscription to prevent OOM
+        this.closeSubscription(id).catch((err) => {
+          log.error(`FlowController: error closing subscription ${String(id)} on overflow`, err);
+        });
+      },
+    );
+
+    this.flowControllers.set(id, fc);
+
+    return (message) => {
+      fc.push(message);
+    };
   }
 
   /**
    * Builds a `RequestContext` object to use with the `JSON RPC API`.
-   *
-   * Adds a `subscriptionHandler` for `Subscribe` messages.
    */
   private async buildRequestContext(request: JsonRpcRequest): Promise<RequestContext> {
     const { params, method, subscription } = request;
 
     const requestContext: RequestContext = {
-      transport        : 'ws',
-      dwn              : this.dwn,
-      socketConnection : this,
-    }
+      transport             : 'ws',
+      dwn                   : this.dwn,
+      socketConnection      : this,
+      activityLog           : this.activityLog,
+      adminStore            : this.adminStore,
+      registrationStore     : this.registrationStore,
+      config                : this.serverConfig,
+      tenantRateLimiter     : this.tenantRateLimiter,
+      messageProcessedHooks : this.messageProcessedHooks,
+    };
 
     // methods that expect a long-running subscription begin with `rpc.subscribe.`
     if (method.startsWith('rpc.subscribe.') && subscription) {
       const { message } = params as { message?: GenericMessage };
       if (message?.descriptor.method === DwnMethodName.Subscribe) {
-        const handlerFunc = this.createSubscriptionHandler(subscription.id);
         requestContext.subscriptionRequest = {
-          id: subscription.id,
-          subscriptionHandler: (message): void => handlerFunc(message),
-        }
+          id                  : subscription.id,
+          subscriptionHandler : this.createSubscriptionHandler(subscription.id),
+        };
       }
     }