feat(peer): backpressure support (#195)

Author: pi0x

docs/1.guide/2.hooks.md

@@ -41,6 +41,12 @@ const hooks = defineHooks({
   error(peer, error) {
     console.log("[ws] error", peer, error);
   },
+
+  drain(peer) {
+    // Send buffer drained after backpressure — safe to resume sending.
+    // Pair with `peer.bufferedAmount`. Not all adapters emit this.
+    console.log("[ws] drain", peer);
+  },
 });
 ```
 

docs/1.guide/3.peer.md

@@ -60,12 +60,40 @@ All topics, this peer has been subscribed to.
 
 Peer's pubsub namespace.
 
+### `peer.bufferedAmount`
+
+Number of bytes queued for transmission but not yet flushed to the client.
+
+Use this to apply **backpressure**: pause sending while it grows past a high watermark and resume once it drops. This prevents a fast producer (for example a streaming async generator) from filling the adapter's send buffer faster than the client can drain it. The easiest way to wait is [`peer.waitForDrain()`](#peerwaitfordrain).
+
+> [!NOTE]
+> Not all adapters expose a buffer signal; those return `0`. Refer to the [compatibility table](#compatibility) for more info.
+
 ## Instance methods
 
 ### `peer.send(message, { compress? })`
 
 Send a message to the connected client.
 
+### `peer.waitForDrain({ threshold?, pollInterval?, signal? })`
+
+Returns a promise that resolves once [`peer.bufferedAmount`](#peerbufferedamount) drops to or below `threshold` bytes (default `0`), letting you `await` backpressure relief in a send loop:
+
+```ts
+// Throttle a fast producer to the client's pace
+for (const chunk of stream) {
+  peer.send(chunk);
+  if (peer.bufferedAmount > 1024 * 1024 /* 1MB */) {
+    await peer.waitForDrain({ threshold: 256 * 1024 /* 256KB */ });
+  }
+}
+```
+
+It resolves immediately when there is no backpressure (or on adapters that do not report `bufferedAmount`). Otherwise it polls every `pollInterval` ms (default `100`) until the buffer drains, and also resolves early if the connection is no longer open so a send loop never hangs on a dropped client. Pass a `signal` (e.g. `AbortSignal.timeout(ms)`) to abort the wait — the promise then rejects with the signal's reason.
+
+> [!TIP]
+> For event-driven backpressure (instead of `await`ing in a loop), use the [`drain` hook](/guide/hooks) together with `peer.bufferedAmount`.
+
 ### `peer.subscribe(channel)`
 
 Join a broadcast channel.
@@ -113,6 +141,8 @@ To gracefully close the connection, use `peer.close()`.
 | `publish()` / `subscribe()` | ✓          | ⨉                 | ✓ [^1]                      | ✓ [^1]       | ✓ [^1]              | ✓                             | ✓ [^1]     |
 | `close()`                   | ✓          | ✓                 | ✓                           | ✓            | ✓                   | ✓                             | ✓          |
 | `terminate()`               | ✓          | ✓ [^2]            | ✓                           | ✓            | ✓                   | ✓                             | ✓ [^2]     |
+| `bufferedAmount`            | ✓          | ⨉ [^8]            | ⨉ [^8]                      | ✓            | ✓                   | ✓                             | ⨉ [^8]     |
+| `drain` hook                | ✓          | ⨉ [^9]            | ⨉ [^9]                      | ⨉ [^9]       | ✓                   | ✓                             | ⨉ [^9]     |
 | `request`                   | ✓          | ✓                 | ✓ [^30]                     | ✓            | ✓ [^31]             | ✓ [^31]                       | ✓          |
 | `remoteAddress`             | ✓          | ⨉                 | ⨉                           | ✓            | ✓                   | ✓                             | ⨉          |
 | `websocket.url`             | ✓          | ✓                 | ✓                           | ✓            | ✓                   | ✓                             | ✓          |
@@ -145,3 +175,7 @@ To gracefully close the connection, use `peer.close()`.
 [^6]: [`websocket.readyState`](https://developer.mozilla.org/en-US/docs/Web/API/WebSocket/readyState) is polyfilled by tracking open/close events.
 
 [^7]: Some runtimes have non standard values including `"nodebuffer"` and `"uint8array"`. crossws auto converts them for [`message.data`](/guide/message#messagedata).
+
+[^8]: The runtime exposes no send-buffer signal, so `peer.bufferedAmount` reports `0`. (Cloudflare buffers and applies backpressure internally; SSE has no equivalent.)
+
+[^9]: The runtime emits no drain signal. Poll [`peer.bufferedAmount`](#peerbufferedamount) instead where it is available.

src/adapters/bun.ts

@@ -73,6 +73,11 @@ const bunAdapter: Adapter<BunAdapter, BunOptions> = (options = {}) => {
         peers.delete(peer);
         hooks.callHook("close", peer, { code, reason });
       },
+      drain: (ws) => {
+        const peers = getPeers(globalPeers, ws.data.namespace);
+        const peer = getPeer(ws, peers);
+        hooks.callHook("drain", peer);
+      },
     },
   };
 };
@@ -109,6 +114,10 @@ class BunPeer extends Peer<{
     return this._internal.ws.data.context;
   }
 
+  override get bufferedAmount(): number {
+    return this._internal.ws.getBufferedAmount();
+  }
+
   send(data: unknown, options?: { compress?: boolean }): number {
     return this._internal.ws.send(toBufferLike(data), options?.compress);
   }

src/adapters/node.ts

@@ -81,8 +81,17 @@ const nodeAdapter: Adapter<NodeAdapter, NodeOptions> = (options = {}) => {
       peers.delete(peer);
       hooks.callHook("error", peer, new WSError(error));
     });
+    // `ws` has no drain event of its own; the underlying TCP socket emits
+    // `drain` after a backpressured write flushes. Note this tracks the OS
+    // socket buffer, which is an approximation of `peer.bufferedAmount` (the
+    // latter also includes ws's internal sender queue) — treat it as a resume
+    // nudge, not an exact "bufferedAmount reached 0" signal.
+    const socket = (ws as WebSocketT & { _socket?: Duplex })._socket;
+    const onDrain = () => hooks.callHook("drain", peer);
+    socket?.on("drain", onDrain);
     ws.on("close", (code: number, reason: Buffer) => {
       peers.delete(peer);
+      socket?.off("drain", onDrain);
       hooks.callHook("close", peer, {
         code,
         reason: reason?.toString(),
@@ -166,7 +175,7 @@ class NodePeer extends Peer<{
       binary: isBinary,
       ...options,
     });
-    return 0;
+    return this._internal.ws.bufferedAmount;
   }
 
   publish(topic: string, data: unknown, options?: { compress?: boolean }): void {

src/adapters/uws.ts

@@ -61,6 +61,11 @@ const uwsAdapter: Adapter<UWSAdapter, UWSOptions> = (options = {}) => {
         const peer = getPeer(ws, peers);
         hooks.callHook("message", peer, new Message(message, peer));
       },
+      drain(ws) {
+        const peers = getPeers(globalPeers, ws.getUserData().namespace);
+        const peer = getPeer(ws, peers);
+        hooks.callHook("drain", peer);
+      },
       open(ws) {
         const peers = getPeers(globalPeers, ws.getUserData().namespace);
         const peer = getPeer(ws, peers);

src/hooks.ts

@@ -146,6 +146,15 @@ export interface Hooks {
   /** A socket is closed */
   close: (peer: Peer, details: { code?: number; reason?: string }) => MaybePromise<void>;
 
+  /**
+   * The send buffer has drained after backpressure, so it is safe to resume
+   * sending. Pair with {@link Peer.bufferedAmount} to throttle senders.
+   *
+   * **Note:** Only emitted by adapters that expose a drain signal. Refer to the
+   * [compatibility table](https://crossws.h3.dev/guide/peer#compatibility).
+   */
+  drain: (peer: Peer) => MaybePromise<void>;
+
   /** An error occurs */
   error: (peer: Peer, error: WSError) => MaybePromise<void>;
 }

src/index.ts

@@ -10,7 +10,7 @@ export type { Adapter, AdapterInstance, AdapterOptions } from "./adapter.ts";
 export type { Message } from "./message.ts";
 
 // Peer
-export type { Peer, PeerContext, AdapterInternal } from "./peer.ts";
+export type { Peer, PeerContext, AdapterInternal, WaitForDrainOptions } from "./peer.ts";
 
 // Error
 export type { WSError } from "./error.ts";

src/peer.ts

@@ -3,6 +3,28 @@ import { kNodeInspect } from "./utils.ts";
 
 export interface PeerContext extends Record<string, unknown> {}
 
+export interface WaitForDrainOptions {
+  /**
+   * Resolve once {@link Peer.bufferedAmount} drops to or below this many bytes.
+   *
+   * @default 0
+   */
+  threshold?: number;
+
+  /**
+   * Polling interval (in milliseconds) used to re-check {@link Peer.bufferedAmount}.
+   *
+   * @default 100
+   */
+  pollInterval?: number;
+
+  /**
+   * Abort the wait (e.g. `AbortSignal.timeout(ms)`). The returned promise
+   * rejects with the signal's `reason`.
+   */
+  signal?: AbortSignal;
+}
+
 export interface AdapterInternal {
   ws: unknown;
   request: Request;
@@ -78,6 +100,69 @@ export abstract class Peer<Internal extends AdapterInternal = AdapterInternal> {
     return this._topics;
   }
 
+  /**
+   * Number of bytes queued for transmission but not yet flushed to the client.
+   *
+   * Use this to apply backpressure: pause sending while it grows past a high
+   * watermark and resume once it drops (or on the `drain` hook). Returns `0` on
+   * adapters that do not expose a buffer signal. Refer to the
+   * [compatibility table](https://crossws.h3.dev/guide/peer#compatibility).
+   */
+  get bufferedAmount(): number {
+    return (this._internal.ws as Partial<web.WebSocket>)?.bufferedAmount ?? 0;
+  }
+
+  /**
+   * Wait until the send buffer drains to `threshold` bytes (default `0`).
+   *
+   * Resolves immediately when there is no backpressure (or on adapters that do
+   * not expose {@link Peer.bufferedAmount}). Otherwise it polls every
+   * `pollInterval` milliseconds until the buffer drains, also resolving early if
+   * the connection is no longer open so a send loop never hangs on a dropped
+   * client.
+   *
+   * ```ts
+   * for (const chunk of stream) {
+   *   peer.send(chunk);
+   *   if (peer.bufferedAmount > 1024 * 1024) {
+   *     await peer.waitForDrain({ threshold: 256 * 1024 });
+   *   }
+   * }
+   * ```
+   */
+  waitForDrain(opts: WaitForDrainOptions = {}): Promise<void> {
+    const threshold = opts.threshold ?? 0;
+    if (this.bufferedAmount <= threshold) {
+      return Promise.resolve();
+    }
+    const signal = opts.signal;
+    if (signal?.aborted) {
+      return Promise.reject(signal.reason);
+    }
+    return new Promise<void>((resolve, reject) => {
+      const check = () => {
+        // Resolve once drained, or if the socket left the OPEN (1) state — a
+        // closed peer never drains, so this prevents a permanent hang + leak.
+        if (this.bufferedAmount <= threshold || (this.websocket.readyState ?? 1) > 1) {
+          cleanup();
+          resolve();
+        }
+      };
+      const onAbort = () => {
+        cleanup();
+        reject(signal!.reason);
+      };
+      const timer = setInterval(check, opts.pollInterval ?? 100);
+      // Don't keep the event loop alive just for a pending drain check.
+      timer.unref?.();
+      const cleanup = () => {
+        clearInterval(timer);
+        signal?.removeEventListener("abort", onAbort);
+      };
+      signal?.addEventListener("abort", onAbort, { once: true });
+    });
+  }
+
   abstract close(code?: number, reason?: string): void;
 
   /** Abruptly close the connection */

test/fixture/_shared.ts

@@ -28,6 +28,7 @@ export function createDemo<T extends Adapter<any, any>>(
           peer.send({
             id: peer.id,
             remoteAddress: peer.remoteAddress,
+            bufferedAmount: peer.bufferedAmount,
             context: peer.context,
             request: {
               url: peer.request?.url,
@@ -44,6 +45,10 @@ export function createDemo<T extends Adapter<any, any>>(
           });
           break;
         }
+        case "waitForDrain": {
+          peer.waitForDrain({ pollInterval: 10 }).then(() => peer.send("drained"));
+          break;
+        }
         case "peers": {
           peer.send({
             peers: [...peer.peers].map((p) => p.id),

test/tests.ts

@@ -119,6 +119,21 @@ export function wsTests(getURL: () => string, opts: WSTestOpts): void {
     }
   });
 
+  test("peer.bufferedAmount", async () => {
+    const ws = await wsConnect(getURL(), { skip: 1 });
+    await ws.send("debug");
+    const { bufferedAmount } = await ws.next();
+    // Adapters without a buffer signal report 0; capable ones report >= 0.
+    expect(typeof bufferedAmount).toBe("number");
+    expect(bufferedAmount).toBeGreaterThanOrEqual(0);
+  });
+
+  test("peer.waitForDrain", async () => {
+    const ws = await wsConnect(getURL(), { skip: 1 });
+    await ws.send("waitForDrain");
+    expect(await ws.next()).toBe("drained");
+  });
+
   test("peer.websocket", async () => {
     const ws = await wsConnect(getURL() + "?foo=bar", {
       skip: 1,