perf(relay): index relay table for O(1) hot-path lookups

scripts/bench.relay.ts

@@ -0,0 +1,129 @@
+/**
+ * Microbenchmark for the UDP relay hot path (`UDPRelayHandler`).
+ *
+ * Drives the real relay methods — same table lookups, prom-client metrics and
+ * event emit as production — with a no-op socket injected in place of a real
+ * UDP socket, so we measure the CPU cost of the relay logic itself rather than
+ * loopback syscall throughput.
+ *
+ * Compares two entry points, apples-to-apples in one process:
+ *   - relayRaw(...)   the shipped hot path: nested-map lookup, zero allocation
+ *   - relay(NetAddr)  same lookup, but a NetAddress is allocated per packet
+ * The gap between them is the per-packet allocation cost that the raw path saves.
+ *
+ * Run: bun scripts/bench.relay.ts
+ */
+import { UDPRelayHandler } from "../src/relay/udp.relay.handler.ts";
+import { RelayEntry } from "../src/relay/relay.entry.ts";
+import { NetAddress } from "../src/relay/net.address.ts";
+
+const TABLE_SIZES = [10, 100, 1_000, 5_000, 10_000];
+const ITERATIONS = 2_000_000;
+const PACKET = Buffer.alloc(128, 0x61);
+
+function buildHandler(n: number): UDPRelayHandler {
+  const handler = new UDPRelayHandler();
+  const pool = handler.socketPool as any;
+
+  for (let i = 0; i < n; i++) {
+    const port = 10_000 + i;
+    pool.sockets.set(port, { port, send() {} });
+    pool.freePorts.push(port);
+  }
+
+  for (let i = 0; i < n; i++) {
+    handler.createRelay(
+      new RelayEntry({
+        address: new NetAddress({
+          address: `10.${(i >> 16) & 255}.${(i >> 8) & 255}.${i & 255}`,
+          port: 30_000 + i,
+        }),
+        port: -1,
+      }),
+    );
+  }
+
+  return handler;
+}
+
+interface Pairs {
+  addr: string[];
+  port: number[];
+  target: number[];
+}
+
+function buildPairs(handler: UDPRelayHandler, n: number): Pairs {
+  const table = handler.relayTable;
+  const P = 4096;
+  const addr = new Array<string>(P);
+  const port = new Array<number>(P);
+  const target = new Array<number>(P);
+  for (let i = 0; i < P; i++) {
+    const s = (Math.random() * n) | 0;
+    let t = (Math.random() * n) | 0;
+    if (t === s) t = (t + 1) % n;
+    addr[i] = table[s].address.address;
+    port[i] = table[s].address.port;
+    target[i] = table[t].port;
+  }
+  return { addr, port, target };
+}
+
+function benchRaw(handler: UDPRelayHandler, p: Pairs): number {
+  for (let i = 0; i < 100_000; i++) {
+    const k = i & 4095;
+    handler.relayRaw(PACKET, p.addr[k], p.port[k], p.target[k]);
+  }
+  const start = Bun.nanoseconds();
+  for (let i = 0; i < ITERATIONS; i++) {
+    const k = i & 4095;
+    handler.relayRaw(PACKET, p.addr[k], p.port[k], p.target[k]);
+  }
+  return 1e9 / ((Bun.nanoseconds() - start) / ITERATIONS);
+}
+
+function benchAlloc(handler: UDPRelayHandler, p: Pairs): number {
+  for (let i = 0; i < 100_000; i++) {
+    const k = i & 4095;
+    handler.relay(
+      PACKET,
+      new NetAddress({ address: p.addr[k], port: p.port[k] }),
+      p.target[k],
+    );
+  }
+  const start = Bun.nanoseconds();
+  for (let i = 0; i < ITERATIONS; i++) {
+    const k = i & 4095;
+    // Mirrors the OLD relay.ts: a fresh NetAddress per packet.
+    handler.relay(
+      PACKET,
+      new NetAddress({ address: p.addr[k], port: p.port[k] }),
+      p.target[k],
+    );
+  }
+  return 1e9 / ((Bun.nanoseconds() - start) / ITERATIONS);
+}
+
+console.log(
+  `relay() microbenchmark — ${ITERATIONS.toLocaleString()} calls per table size\n`,
+);
+console.log(
+  "table size | relayRaw (no alloc) | relay (per-pkt alloc) | alloc cost",
+);
+console.log(
+  "-----------+---------------------+-----------------------+-----------",
+);
+
+for (const n of TABLE_SIZES) {
+  const handler = buildHandler(n);
+  const pairs = buildPairs(handler, n);
+  const raw = benchRaw(handler, pairs);
+  const alloc = benchAlloc(handler, pairs);
+  handler.clear();
+
+  const nsRaw = 1e9 / raw;
+  const nsAlloc = 1e9 / alloc;
+  console.log(
+    `${String(n).padStart(10)} | ${(Math.round(raw).toLocaleString() + " /s").padStart(19)} | ${(Math.round(alloc).toLocaleString() + " /s").padStart(21)} | ${(nsAlloc - nsRaw).toFixed(0).padStart(6)} ns`,
+  );
+}

src/relay/relay.ts

@@ -14,7 +14,6 @@ import { UDPRemoteRegistrar } from "./udp.remote.registrar.ts";
 import { hostRepository } from "../hosts/host.ts";
 import { useDynamicRelay } from "./dynamic.relaying.ts";
 import { UDPSocketPool } from "./udp.socket.pool.ts";
-import { NetAddress } from "./net.address.ts";
 
 export const udpSocketPool = new UDPSocketPool();
 
@@ -120,11 +119,7 @@ async function bindPortForRelaying(
   await udpSocketPool.allocatePort(port, {
     socket: {
       data(socket, data, port, address) {
-        udpRelayHandler.relay(
-          data,
-          new NetAddress({ address, port }),
-          socket.port,
-        );
+        udpRelayHandler.relayRaw(data, address, port, socket.port);
       },
       error(_socket, error) {
         log.error(error, "UDP relay socket encountered an error!");

src/relay/udp.relay.handler.ts

@@ -61,6 +61,35 @@ export class UDPRelayHandler extends EventEmitter {
   public socketPool: UDPSocketPool;
   private _relayTable: RelayEntry[] = [];
 
+  // Hot-path indexes: avoid O(N) linear scans in relay() on every packet.
+  // Nested map (address -> port -> entry) so the hot path can look up by the
+  // raw address string + port with no per-packet string concatenation.
+  private _byAddress = new Map<string, Map<number, RelayEntry>>();
+  private _byPort = new Map<number, RelayEntry>();
+
+  private _lookupAddr(address: string, port: number): RelayEntry | undefined {
+    return this._byAddress.get(address)?.get(port);
+  }
+
+  private _index(entry: RelayEntry): void {
+    let ports = this._byAddress.get(entry.address.address);
+    if (!ports) {
+      ports = new Map<number, RelayEntry>();
+      this._byAddress.set(entry.address.address, ports);
+    }
+    ports.set(entry.address.port, entry);
+    this._byPort.set(entry.port, entry);
+  }
+
+  private _deindex(entry: RelayEntry): void {
+    const ports = this._byAddress.get(entry.address.address);
+    if (ports) {
+      ports.delete(entry.address.port);
+      if (ports.size === 0) this._byAddress.delete(entry.address.address);
+    }
+    this._byPort.delete(entry.port);
+  }
+
   /**
    * Relay table used for relaying.
    */
@@ -91,7 +120,7 @@ export class UDPRelayHandler extends EventEmitter {
     if (this.hasRelay(relay)) {
       // We already have this relay entry
       log.trace({ relay }, "Relay already exists, ignoring");
-      return this._relayTable.find((e) => e.equals(relay))!;
+      return this._lookupAddr(relay.address.address, relay.address.port)!;
     }
 
     relay.port = this.socketPool.getPort();
@@ -100,6 +129,7 @@ export class UDPRelayHandler extends EventEmitter {
     relay.lastReceived = time();
     relay.created = time();
     this._relayTable.push(relay);
+    this._index(relay);
     log.trace({ relay }, "Relay created");
 
     activeRelayGauge.inc();
@@ -113,7 +143,9 @@ export class UDPRelayHandler extends EventEmitter {
    * NOTE: This only compares the addresses, not the allocated port.
    */
   hasRelay(relay: RelayEntry): boolean {
-    return this._relayTable.find((e) => e.equals(relay)) !== undefined;
+    return (
+      this._lookupAddr(relay.address.address, relay.address.port) !== undefined
+    );
   }
 
   /**
@@ -126,10 +158,13 @@ export class UDPRelayHandler extends EventEmitter {
       return false;
     }
 
+    const stored = this._relayTable[idx];
+
     this.emit("destroy", relay);
 
-    this.socketPool.returnPort(relay.port);
+    this.socketPool.returnPort(stored.port);
     this._relayTable = this.relayTable.filter((_, i) => i !== idx);
+    this._deindex(stored);
 
     activeRelayGauge.dec();
 
@@ -141,6 +176,8 @@ export class UDPRelayHandler extends EventEmitter {
    */
   clear() {
     this._relayTable.forEach((entry) => this.freeRelay(entry));
+    this._byAddress.clear();
+    this._byPort.clear();
 
     activeRelayGauge.reset();
   }
@@ -153,17 +190,40 @@ export class UDPRelayHandler extends EventEmitter {
    */
   // TODO: Why was the return type documented as Promise<boolean>?
   relay(msg: Buffer, sender: NetAddress, target: number): boolean {
+    return this.relayRaw(msg, sender.address, sender.port, target);
+  }
+
+  /**
+   * Relay a message, addressing the sender by its raw address + port.
+   *
+   * This is the allocation-free hot path: it avoids constructing a NetAddress
+   * per packet. A NetAddress is only built on the (cold) drop path, to preserve
+   * the `drop` event contract.
+   *
+   * @fires UDPRelayHandler#transmit
+   * @fires UDPRelayHandler#drop
+   */
+  relayRaw(
+    msg: Buffer,
+    senderAddress: string,
+    senderPort: number,
+    target: number,
+  ): boolean {
     const measure = relayDurationHistogram.startTimer();
 
-    const senderRelay = this._relayTable.find(
-      (r) =>
-        r.address.port === sender.port && r.address.address === sender.address,
-    );
-    const targetRelay = this._relayTable.find((r) => r.port === target);
+    const senderRelay = this._lookupAddr(senderAddress, senderPort);
+    const targetRelay = this._byPort.get(target);
 
     if (!senderRelay || !targetRelay) {
       // We don't have a relay for the sender, target, or both
-      this.emit("drop", senderRelay, targetRelay, sender, target, msg);
+      this.emit(
+        "drop",
+        senderRelay,
+        targetRelay,
+        new NetAddress({ address: senderAddress, port: senderPort }),
+        target,
+        msg,
+      );
 
       relayDropCounter.inc();
       measure();