@trevrpc/trevrpc-js 1.0.0 → 1.2.0

package/src/datagram.js

@@ -0,0 +1,357 @@
+// Unreliable RPC messages over WebTransport datagrams. One DatagramEndpoint per session owns the
+// session's single datagram reader loop, its single shared writer, the flow ID allocator, and the
+// demultiplexer from flow IDs to calls. Calls that stream an unreliable direction open a
+// DatagramFlow on the endpoint of the session that carries them. QMux over WebSocket has no
+// datagrams, so those calls never get a flow and everything stays on the stream.
+import { byteView } from "./bytes.js";
+import { FrameType } from "./framing.js";
+import { readVarint, varintLength, writeVarint } from "./varint.js";
+
+/**
+ * Largest datagram assumed when the browser does not report `maxDatagramSize`. Browsers silently
+ * drop datagrams above the real limit, so this stays below a typical path MTU.
+ */
+export const FallbackMaxDatagramSize = 1100;
+
+/**
+ * How much a flow buffers for a call that is not ready to take datagrams yet: at most this many
+ * messages or bytes, then further datagrams are dropped.
+ */
+export const FlowQueueMessages = 32;
+export const FlowQueueBytes = 1 << 20;
+
+/**
+ * Encoded size of one outgoing message (`{ type, payload }` from encodeMessage) as a MESSAGE or
+ * ZSTD_MESSAGE frame inside a datagram.
+ */
+export function messageFrameSize({ payload }) {
+  return 1 + varintLength(payload.byteLength) + payload.byteLength;
+}
+
+/**
+ * Encodes a datagram payload: the flow ID as a QUIC varint, then the frame of each outgoing
+ * message, in order.
+ */
+export function encodeDatagram(flow, messages) {
+  let size = varintLength(flow);
+  for (const message of messages) {
+    size += messageFrameSize(message);
+  }
+  const datagram = new Uint8Array(size);
+  let offset = writeVarint(datagram, 0, flow);
+  for (const { type, payload } of messages) {
+    datagram[offset] = type;
+    offset = writeVarint(datagram, offset + 1, payload.byteLength);
+    datagram.set(payload, offset);
+    offset += payload.byteLength;
+  }
+  return datagram;
+}
+
+/**
+ * Decodes a datagram payload as `{ flow, messages }`, or null if it has no complete flow ID. The
+ * messages are its MESSAGE and ZSTD_MESSAGE frames as `{ type, payload }`; every other frame is
+ * skipped. A datagram whose frames do not exactly fill it, or that holds a ZSTD_STREAM_MESSAGE,
+ * which needs a stream's ordered delivery, decodes to no messages at all.
+ */
+export function decodeDatagram(datagram) {
+  const bytes = byteView(datagram);
+  const flow = readVarint(bytes, 0);
+  if (flow === null) {
+    return null;
+  }
+  const messages = [];
+  let offset = flow.next;
+  while (offset < bytes.byteLength) {
+    const type = readVarint(bytes, offset);
+    const length = type === null ? null : readVarint(bytes, type.next);
+    if (length === null || typeof length.value !== "number") {
+      return { flow: flow.value, messages: [] };
+    }
+    const end = length.next + length.value;
+    if (end > bytes.byteLength) {
+      return { flow: flow.value, messages: [] };
+    }
+    if (type.value === FrameType.ZstdStreamMessage) {
+      return { flow: flow.value, messages: [] };
+    }
+    if (type.value === FrameType.Message || type.value === FrameType.ZstdMessage) {
+      messages.push({ type: type.value, payload: bytes.subarray(length.next, end) });
+    }
+    offset = end;
+  }
+  return { flow: flow.value, messages };
+}
+
+const endpoints = new WeakMap();
+
+/**
+ * The shared endpoint of a session's `datagrams` object, or undefined if the session has none or
+ * its datagram reader is unusable. The endpoint lives as long as the object does.
+ */
+export function datagramEndpoint(datagrams) {
+  if (datagrams == null || typeof datagrams !== "object") {
+    return undefined;
+  }
+  let endpoint = endpoints.get(datagrams);
+  if (endpoint === undefined) {
+    endpoint = new DatagramEndpoint(datagrams);
+    endpoints.set(datagrams, endpoint);
+  }
+  return endpoint.usable ? endpoint : undefined;
+}
+
+class DatagramEndpoint {
+  #datagrams;
+  /** Flows that accept datagrams, by ID. */
+  #flows = new Map();
+  #nextFlow = 1;
+  #reading = false;
+  #closed;
+  /** The shared writer: undefined until first needed, null once unavailable. */
+  #writer;
+
+  constructor(datagrams) {
+    this.#datagrams = datagrams;
+    this.#closed = typeof datagrams.readable?.getReader !== "function";
+  }
+
+  get usable() {
+    return !this.#closed;
+  }
+
+  /**
+   * Allocates the next flow of this session. A flow that `receive`s is registered with the
+   * reader, so datagrams for its ID reach it; `inflate` maps each received `{ type, payload }`
+   * message to its payload, decompressed, and throws if it is invalid. Returns undefined if the
+   * session cannot read datagrams.
+   */
+  openFlow({ receive, inflate }) {
+    if (!this.#start()) {
+      return undefined;
+    }
+    const flow = new DatagramFlow(this, this.#nextFlow, inflate);
+    this.#nextFlow += 1;
+    if (receive) {
+      this.#flows.set(flow.id, flow);
+    }
+    return flow;
+  }
+
+  /** Stops routing datagrams to a flow; its ID is never reused. */
+  closeFlow(id) {
+    this.#flows.delete(id);
+  }
+
+  /**
+   * How many bytes of MESSAGE frames fit in a datagram of `flow` right now, or 0 if this session
+   * cannot send datagrams.
+   */
+  frameRoom(flow) {
+    return this.#writerOrNull() === null
+      ? 0
+      : Math.max(0, this.maxDatagramSize - varintLength(flow));
+  }
+
+  /**
+   * Sends messages as one datagram of `flow`. Resolves true once the datagram was handed to the
+   * transport, and false if the messages have to go on the stream instead: they do not fit, there
+   * is no writer, or the write failed. Waits for send-buffer space until `signal` aborts, then
+   * rejects with the abort reason.
+   */
+  async send(flow, messages, signal) {
+    const writer = this.#writerOrNull();
+    const datagram = writer === null ? null : encodeDatagram(flow, messages);
+    if (datagram === null || datagram.byteLength > this.maxDatagramSize) {
+      return false;
+    }
+    try {
+      if (!(writer.desiredSize > 0)) {
+        await abortable(writer.ready, signal);
+      }
+      await abortable(writer.write(datagram), signal);
+      return true;
+    } catch {
+      if (signal.aborted) {
+        throw signal.reason;
+      }
+      return false;
+    }
+  }
+
+  /** The largest datagram payload the session accepts right now, flow ID included. */
+  get maxDatagramSize() {
+    const size = this.#datagrams.maxDatagramSize;
+    return Number.isFinite(size) && size > 0 ? size : FallbackMaxDatagramSize;
+  }
+
+  #start() {
+    if (this.#closed) {
+      return false;
+    }
+    if (!this.#reading) {
+      let reader;
+      try {
+        reader = this.#datagrams.readable.getReader();
+      } catch {
+        this.#closed = true;
+        return false;
+      }
+      this.#reading = true;
+      void this.#read(reader);
+    }
+    return true;
+  }
+
+  /** The reader loop. It never blocks on a call: delivery is a synchronous, bounded push. */
+  async #read(reader) {
+    try {
+      for (;;) {
+        const { value, done } = await reader.read();
+        if (done) {
+          break;
+        }
+        const datagram = decodeDatagram(value);
+        const flow = datagram === null ? undefined : this.#flows.get(datagram.flow);
+        // Unknown and ended flows drop the datagram, and so do malformed ones (no messages).
+        flow?.deliverAll(datagram.messages);
+      }
+    } catch {
+      // A failed session ends its calls through their streams.
+    } finally {
+      this.#closed = true;
+      this.#flows.clear();
+    }
+  }
+
+  #writerOrNull() {
+    if (this.#writer === undefined) {
+      this.#writer = this.#openWriter();
+    }
+    return this.#writer;
+  }
+
+  /** Locks the session's datagram writable once; a lost writer disables sending datagrams. */
+  #openWriter() {
+    try {
+      const datagrams = this.#datagrams;
+      const writable =
+        typeof datagrams.createWritable === "function"
+          ? datagrams.createWritable()
+          : datagrams.writable;
+      const writer = writable.getWriter();
+      writer.closed.catch(() => {
+        this.#writer = null;
+      });
+      return writer;
+    } catch {
+      return null;
+    }
+  }
+}
+
+/**
+ * One call's datagram flow. Datagrams for the flow wait in a bounded buffer until the call
+ * connects, then go straight to its handler. The call releases the flow when it ends.
+ */
+export class DatagramFlow {
+  #endpoint;
+  #inflate;
+  #handler;
+  #buffered = [];
+  #bufferedBytes = 0;
+  #released = false;
+
+  constructor(endpoint, id, inflate) {
+    this.#endpoint = endpoint;
+    this.#inflate = inflate;
+    this.id = id;
+  }
+
+  /**
+   * Delivers the messages of one received datagram, decompressed. A datagram with an invalid
+   * message is dropped whole.
+   */
+  deliverAll(messages) {
+    let payloads;
+    try {
+      payloads = messages.map(this.#inflate);
+    } catch {
+      return;
+    }
+    for (const payload of payloads) {
+      this.#deliver(payload);
+    }
+  }
+
+  /**
+   * Delivers one received message. It must not block or throw. Without a handler it is buffered,
+   * and dropped when the buffer is full.
+   */
+  #deliver(message) {
+    if (this.#handler !== undefined) {
+      this.#handler(message);
+    } else if (
+      this.#buffered.length < FlowQueueMessages &&
+      this.#bufferedBytes + message.byteLength <= FlowQueueBytes
+    ) {
+      this.#buffered.push(message);
+      this.#bufferedBytes += message.byteLength;
+    }
+  }
+
+  /** Hands the buffered messages, then all later ones, to `handler`. Only the first call counts. */
+  connect(handler) {
+    if (this.#handler !== undefined || this.#released) {
+      return;
+    }
+    this.#handler = handler;
+    const buffered = this.#buffered;
+    this.#buffered = [];
+    this.#bufferedBytes = 0;
+    for (const message of buffered) {
+      handler(message);
+    }
+  }
+
+  /** See DatagramEndpoint.frameRoom; 0 once the flow is released. */
+  get frameRoom() {
+    return this.#released ? 0 : this.#endpoint.frameRoom(this.id);
+  }
+
+  /** See DatagramEndpoint.send. */
+  send(messages, signal) {
+    return this.#released ? Promise.resolve(false) : this.#endpoint.send(this.id, messages, signal);
+  }
+
+  /** Ends the flow: later datagrams for it are dropped. */
+  release() {
+    if (!this.#released) {
+      this.#released = true;
+      this.#handler = undefined;
+      this.#buffered = [];
+      this.#endpoint.closeFlow(this.id);
+    }
+  }
+}
+
+/** Settles like `promise`, or rejects with the abort reason as soon as `signal` aborts. */
+function abortable(promise, signal) {
+  if (signal.aborted) {
+    return Promise.reject(signal.reason);
+  }
+  return new Promise((resolve, reject) => {
+    const onAbort = () => reject(signal.reason);
+    signal.addEventListener("abort", onAbort, { once: true });
+    promise.then(
+      (value) => {
+        signal.removeEventListener("abort", onAbort);
+        resolve(value);
+      },
+      (error) => {
+        signal.removeEventListener("abort", onAbort);
+        reject(error);
+      },
+    );
+  });
+}

package/src/framing.js

@@ -1,12 +1,22 @@
 import { byteView } from "./bytes.js";
-import { Code, invalidArgument, resourceExhausted, unavailable } from "./status.js";
-import { RpcStreamFrame, RpcStreamFrameKind, preflightWireMessage } from "./wire.js";
+import { internal, invalidArgument, resourceExhausted, unavailable } from "./status.js";
+import { readVarint, varintLength, writeVarint } from "./varint.js";
+import { RpcAnswer, RpcHeaders, RpcRequest, RpcStatus, preflightWireMessage } from "./wire.js";
 
 export const DefaultMaxFrameSize = 4 * 1024 * 1024;
 
-const HeaderLength = 4;
-/** Tag of RpcStreamFrame.body: field 4, length-delimited. */
-const StreamFrameBodyTag = (4 << 3) | 2;
+/** Frame types; every other type is unknown and skipped by receivers. */
+export const FrameType = Object.freeze({
+  Message: 0,
+  Request: 1,
+  Headers: 2,
+  Status: 3,
+  ZstdMessage: 4,
+  ZstdStreamMessage: 5,
+  Answer: 6,
+});
+
+const KnownFrameTypes = new Set(Object.values(FrameType));
 const EmptyBytes = new Uint8Array(0);
 
 /** Encodes a protobuf message body. */
@@ -34,63 +44,65 @@ export function unmarshalMessage(messageType, body) {
   }
 }
 
-/** Encodes a protobuf message into a length-prefixed frame. */
-export function encodeFrame(messageType, message, maxFrameSize = DefaultMaxFrameSize) {
-  const body = marshalMessage(messageType, message);
-  checkFrameLength(body.byteLength, maxFrameSize);
-  const frame = new Uint8Array(HeaderLength + body.byteLength);
-  new DataView(frame.buffer).setUint32(0, body.byteLength, false);
-  frame.set(body, HeaderLength);
+/** Encodes one frame carrying an already-encoded payload. */
+export function encodeFrame(type, payload, maxFrameSize = DefaultMaxFrameSize) {
+  checkFrameLength(payload.byteLength, maxFrameSize);
+  const frame = new Uint8Array(frameSize(type, payload.byteLength));
+  writeFrame(frame, 0, type, payload);
   return frame;
 }
 
+/** Encodes a MESSAGE frame; the payload is the raw protobuf message. */
+export function encodeMessageFrame(body, maxFrameSize = DefaultMaxFrameSize) {
+  return encodeFrame(FrameType.Message, body, maxFrameSize);
+}
+
 /**
- * Encodes a streaming message frame carrying an already-encoded protobuf body.
- * This is byte-identical to encodeFrame(RpcStreamFrame, { body }) without building the
- * intermediate message.
+ * Encodes an outgoing message, `{ type, payload, size }` from encodeMessage, as its MESSAGE or
+ * ZSTD_MESSAGE frame. The message size counts against the limit, not just the payload's.
  */
-export function encodeMessageFrame(body, maxFrameSize = DefaultMaxFrameSize) {
-  const frameLength =
-    body.byteLength === 0 ? 0 : 1 + varintLength(body.byteLength) + body.byteLength;
-  checkFrameLength(frameLength, maxFrameSize);
-
-  const frame = new Uint8Array(HeaderLength + frameLength);
-  new DataView(frame.buffer).setUint32(0, frameLength, false);
-  if (body.byteLength > 0) {
-    let offset = HeaderLength;
-    frame[offset++] = StreamFrameBodyTag;
-    offset = writeVarint(frame, offset, body.byteLength);
-    frame.set(body, offset);
-  }
+export function encodeOutgoingFrame(message, maxFrameSize = DefaultMaxFrameSize) {
+  checkFrameLength(message.size, maxFrameSize);
+  return encodeFrame(message.type, message.payload, maxFrameSize);
+}
+
+/**
+ * Encodes a REQUEST frame, followed by the frame of message, an outgoing message from
+ * encodeMessage, when it is given, in one buffer so both go out in a single write.
+ */
+export function encodeRequestFrame(request, message, maxFrameSize = DefaultMaxFrameSize) {
+  const payload = marshalMessage(RpcRequest, request);
+  checkFrameLength(payload.byteLength, maxFrameSize);
+  if (message === undefined) {
+    return encodeFrame(FrameType.Request, payload, maxFrameSize);
+  }
+  checkFrameLength(message.size, maxFrameSize);
+  const requestSize = frameSize(FrameType.Request, payload.byteLength);
+  const frame = new Uint8Array(requestSize + frameSize(message.type, message.payload.byteLength));
+  writeFrame(frame, requestSize, message.type, message.payload);
+  writeFrame(frame, 0, FrameType.Request, payload);
   return frame;
 }
 
-/** Decodes the body of a response stream frame, validating its kind. */
-export function decodeStreamFrame(body) {
-  const bytes = byteView(body);
-  const messageBody = bareMessageBody(bytes);
-  if (messageBody !== undefined) {
-    return { kind: RpcStreamFrameKind.Message, status: Code.Ok, message: "", body: messageBody };
-  }
+/** Encodes a HEADERS frame from RpcHeaders fields. */
+export function encodeHeadersFrame(headers, maxFrameSize = DefaultMaxFrameSize) {
+  return encodeFrame(FrameType.Headers, marshalMessage(RpcHeaders, headers), maxFrameSize);
+}
 
-  const frame = unmarshalMessage(RpcStreamFrame, bytes);
-  if (frame.kind !== RpcStreamFrameKind.Message && frame.kind !== RpcStreamFrameKind.Status) {
-    throw invalidArgument(`unsupported stream frame kind ${frame.kind}`);
-  }
-  return frame;
+/** Encodes an ANSWER frame from RpcAnswer fields. */
+export function encodeAnswerFrame(answer, maxFrameSize = DefaultMaxFrameSize) {
+  return encodeFrame(FrameType.Answer, marshalMessage(RpcAnswer, answer), maxFrameSize);
 }
 
-/** Validates the body length stored in a frame header. */
-export function frameBodyLength(header, maxFrameSize = DefaultMaxFrameSize) {
-  const length = new DataView(header.buffer, header.byteOffset, header.byteLength).getUint32(
-    0,
-    false,
-  );
-  checkFrameLength(length, maxFrameSize);
-  return length;
+/** Encodes a STATUS frame from RpcStatus fields. */
+export function encodeStatusFrame(status, maxFrameSize = DefaultMaxFrameSize) {
+  return encodeFrame(FrameType.Status, marshalMessage(RpcStatus, status), maxFrameSize);
 }
 
-/** Reads length-prefixed frames from a byte stream reader. */
+/**
+ * Reads frames from a byte stream reader. Frame headers may be split across chunks; payloads
+ * inside one chunk are returned as views without copying. Unknown frame types are skipped.
+ */
 export class FrameReader {
   #reader;
   /** Bytes received from the reader but not yet consumed. */
@@ -101,44 +113,61 @@ export class FrameReader {
     this.#reader = reader;
   }
 
-  /** Reads and decodes one frame; a stream that ends first is an error. */
-  async readFrame(messageType, maxFrameSize = DefaultMaxFrameSize) {
-    return unmarshalMessage(messageType, await this.#readBody(maxFrameSize, false));
+  /**
+   * Reads the next known frame as `{ type, payload }`, or returns null when the stream ended
+   * cleanly between frames. The declared length of every frame, unknown ones included, is checked
+   * against maxFrameSize before any payload is read or allocated.
+   */
+  async readFrame(maxFrameSize = DefaultMaxFrameSize) {
+    for (;;) {
+      const type = await this.#readVarint(true);
+      if (type === null) {
+        return null;
+      }
+      const length = await this.#readVarint(false);
+      checkFrameLength(length, maxFrameSize);
+      if (!KnownFrameTypes.has(type)) {
+        // Types of 2^53 and above are bigints, which are never known.
+        await this.#skip(length);
+        continue;
+      }
+      return { type, payload: await this.#readExact(length) };
+    }
   }
 
-  /** Reads one streaming RPC frame, or returns null when the stream ended cleanly. */
-  async readStreamFrame(maxFrameSize = DefaultMaxFrameSize) {
-    const body = await this.#readBody(maxFrameSize, true);
-    return body === null ? null : decodeStreamFrame(body);
+  /**
+   * Resolves when the rest of the stream consists of unknown frames and a clean end; rejects if
+   * a known frame or a malformed frame arrives first.
+   */
+  async expectEnd(maxFrameSize = DefaultMaxFrameSize) {
+    const frame = await this.readFrame(maxFrameSize);
+    if (frame !== null) {
+      throw internal(`frame of type ${frame.type} after the final STATUS frame`);
+    }
   }
 
-  /** Resolves when the stream ends cleanly; rejects if more data arrives first. */
-  async expectEnd() {
-    if (this.#pending.byteLength === 0) {
+  /** Reads a varint; null if allowEnd and the stream ended before its first byte. */
+  async #readVarint(allowEnd) {
+    while (this.#pending.byteLength === 0) {
       const { value, done } = await this.#nextChunk();
       if (done) {
-        return;
+        if (allowEnd) {
+          return null;
+        }
+        throw unavailable("stream ended before a complete frame");
       }
       this.#pending = value;
     }
-    throw unavailable("stream carried data after the final frame");
-  }
-
-  async #readBody(maxFrameSize, allowEnd) {
-    const header = await this.#readExact(HeaderLength, allowEnd);
-    if (header === null) {
-      return null;
-    }
-    return await this.#readExact(frameBodyLength(header, maxFrameSize), false);
+    const bytes = await this.#readExact(1 << (this.#pending[0] >> 6));
+    return readVarint(bytes, 0).value;
   }
 
-  /** Returns exactly size bytes, or null if allowEnd and the stream ended before the first byte. */
-  async #readExact(size, allowEnd) {
+  /** Returns exactly size bytes; copies only when the read spans chunks. */
+  async #readExact(size) {
     if (this.#pending.byteLength >= size) {
       return this.#take(size);
     }
 
-    // Copies are needed only when a read spans chunks; whole frames inside one chunk are views.
     const out = new Uint8Array(size);
     let filled = this.#pending.byteLength;
     out.set(this.#pending);
@@ -146,9 +175,6 @@ export class FrameReader {
     while (filled < size) {
       const { value, done } = await this.#nextChunk();
       if (done) {
-        if (allowEnd && filled === 0) {
-          return null;
-        }
         throw unavailable("stream ended before a complete frame");
       }
       const count = Math.min(value.byteLength, size - filled);
@@ -159,6 +185,23 @@ export class FrameReader {
     return out;
   }
 
+  /** Discards size bytes without buffering them. */
+  async #skip(size) {
+    let left = size;
+    while (left > 0) {
+      if (this.#pending.byteLength === 0) {
+        const { value, done } = await this.#nextChunk();
+        if (done) {
+          throw unavailable("stream ended before a complete frame");
+        }
+        this.#pending = value;
+      }
+      const count = Math.min(this.#pending.byteLength, left);
+      this.#pending = this.#pending.subarray(count);
+      left -= count;
+    }
+  }
+
   async #nextChunk() {
     for (;;) {
       const { value, done } = await this.#reader.read();
@@ -178,47 +221,18 @@ export class FrameReader {
   }
 }
 
+/** Rejects payload lengths above the limit; length may be a bigint for huge varints. */
 function checkFrameLength(length, maxFrameSize) {
   if (length > maxFrameSize) {
     throw resourceExhausted(`frame length ${length} exceeds maximum ${maxFrameSize}`);
   }
 }
 
-/** Returns the body of a message-only stream frame, or undefined if it carries anything else. */
-function bareMessageBody(bytes) {
-  if (bytes.byteLength === 0) {
-    return EmptyBytes;
-  }
-  if (bytes[0] !== StreamFrameBodyTag) {
-    return undefined;
-  }
-
-  let length = 0;
-  let offset = 1;
-  for (let shift = 0; shift < 35 && offset < bytes.byteLength; shift += 7) {
-    const byte = bytes[offset++];
-    length += (byte & 0x7f) * 2 ** shift;
-    if (byte < 0x80) {
-      return offset + length === bytes.byteLength ? bytes.subarray(offset) : undefined;
-    }
-  }
-  return undefined;
-}
-
-function writeVarint(output, offset, value) {
-  let remaining = value;
-  while (remaining >= 0x80) {
-    output[offset++] = (remaining % 0x80) | 0x80;
-    remaining = Math.floor(remaining / 0x80);
-  }
-  output[offset++] = remaining;
-  return offset;
+function frameSize(type, length) {
+  return varintLength(type) + varintLength(length) + length;
 }
 
-function varintLength(value) {
-  let length = 1;
-  for (let remaining = value; remaining >= 0x80; remaining = Math.floor(remaining / 0x80)) {
-    length += 1;
-  }
-  return length;
+function writeFrame(output, offset, type, payload) {
+  const next = writeVarint(output, writeVarint(output, offset, type), payload.byteLength);
+  output.set(payload, next);
 }