diff --git a/scripts/bench.relay.ts b/scripts/bench.relay.ts new file mode 100644 index 0000000..f517847 --- /dev/null +++ b/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(P); + const port = new Array(P); + const target = new Array(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`, + ); +} diff --git a/src/relay/relay.ts b/src/relay/relay.ts index 85eec1b..55c5473 100644 --- a/src/relay/relay.ts +++ b/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!"); diff --git a/src/relay/udp.relay.handler.ts b/src/relay/udp.relay.handler.ts index b09ebdf..b56dccc 100644 --- a/src/relay/udp.relay.handler.ts +++ b/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>(); + private _byPort = new Map(); + + 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(); + 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? 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(); diff --git a/test/spec/relay/udp.relay.handler.test.ts b/test/spec/relay/udp.relay.handler.test.ts index 9cb881a..63d9e98 100644 --- a/test/spec/relay/udp.relay.handler.test.ts +++ b/test/spec/relay/udp.relay.handler.test.ts @@ -242,6 +242,19 @@ describe("UDPRelayHandler", () => { assert(socketPool.getSocket.notCalled, "Socket queried!"); assert(socket.send.notCalled, "Message sent?"); assert(dropHandler.calledOnce, "Drop event not emitted!"); + + // The drop event carries a NetAddress reconstructed from the raw sender + // address/port. It need not be the caller's instance, but must be + // value-equal — dynamic relaying keys off this address by value. + const droppedSender = dropHandler.firstCall.args[2]; + assert( + droppedSender instanceof NetAddress, + "Drop sender is not a NetAddress!", + ); + assert( + droppedSender.equals(new NetAddress({ address: "88.59.62.107", port: 65227 })), + "Drop sender address does not match the packet's origin!", + ); }); it("should ignore on missing socket", async () => { // Given @@ -295,6 +308,143 @@ describe("UDPRelayHandler", () => { assert(socket.send.notCalled, "Message sent?"); }); }); + describe("relayRaw", () => { + it("should relay through the raw hot path", async () => { + // Given + const message = Buffer.from("Hello!", "utf-8"); + const socket = sinon.createStubInstance(dgram.Socket); + const socketPool = sinon.createStubInstance(UDPSocketPool); + socketPool.getPort.onFirstCall().returns(10001); + socketPool.getPort.onSecondCall().returns(10002); + socketPool.getSocket.returns(socket); + socket.removeAllListeners.returnsThis(); + const handler = sinon.stub(); + + const relayHandler = new UDPRelayHandler({ socketPool }); + relayHandler.on("transmit", handler); + + // Sender, allocated local port 10001 + relayHandler.createRelay( + new RelayEntry({ + port: 57789, + address: new NetAddress({ address: "10.0.0.1", port: 1111 }), + }), + ); + // Target, allocated local port 10002 + relayHandler.createRelay( + new RelayEntry({ + port: 57789, + address: new NetAddress({ address: "10.0.0.2", port: 2222 }), + }), + ); + socketPool.getSocket.resetHistory(); + + // When — call the shipped hot path directly, no NetAddress allocation + const success = relayHandler.relayRaw(message, "10.0.0.1", 1111, 10002); + + // Then + assert(success, "Relay failed!"); + assert( + socketPool.getSocket.calledOnceWith(10001), + "Sender socket not queried!", + ); + assert( + socket.send.calledWith(message, 2222, "10.0.0.2"), + "Message not sent to target address!", + ); + assert(handler.calledOnce, "Transmit event not emitted!"); + }); + + it("should drop after the target is freed (deindexed from _byPort)", async () => { + // Given + const message = Buffer.from("Hello!", "utf-8"); + const socket = sinon.createStubInstance(dgram.Socket); + const socketPool = sinon.createStubInstance(UDPSocketPool); + socketPool.getPort.onFirstCall().returns(10001); + socketPool.getPort.onSecondCall().returns(10002); + socketPool.getSocket.returns(socket); + socket.removeAllListeners.returnsThis(); + const dropHandler = sinon.spy(); + + const relayHandler = new UDPRelayHandler({ socketPool }); + + relayHandler.createRelay( + new RelayEntry({ + port: 57789, + address: new NetAddress({ address: "10.0.0.1", port: 1111 }), + }), + ); + const target = new RelayEntry({ + port: 57789, + address: new NetAddress({ address: "10.0.0.2", port: 2222 }), + }); + relayHandler.createRelay(target); + + // Sanity: routing to the target works while it exists + assert( + relayHandler.relayRaw(message, "10.0.0.1", 1111, 10002), + "Precondition failed: target not routable", + ); + + relayHandler.freeRelay(target); + relayHandler.on("drop", dropHandler); + socketPool.getSocket.resetHistory(); + socket.send.resetHistory(); + + // When — the target port is now gone from the index + const success = relayHandler.relayRaw(message, "10.0.0.1", 1111, 10002); + + // Then + assert(!success, "Relay succeeded after target freed?"); + assert(socket.send.notCalled, "Message sent to freed target?"); + assert(dropHandler.calledOnce, "Drop event not emitted!"); + }); + + it("should keep a shared address routable after freeing one of its ports", async () => { + // Given — two relays on the SAME address, different source ports + const socket = sinon.createStubInstance(dgram.Socket); + const socketPool = sinon.createStubInstance(UDPSocketPool); + socketPool.getPort.onFirstCall().returns(10001); + socketPool.getPort.onSecondCall().returns(10002); + socketPool.getSocket.returns(socket); + socket.removeAllListeners.returnsThis(); + + const relayHandler = new UDPRelayHandler({ socketPool }); + + const relayA = new RelayEntry({ + port: 57789, + address: new NetAddress({ address: "10.0.0.1", port: 1111 }), + }); + const relayB = new RelayEntry({ + port: 57789, + address: new NetAddress({ address: "10.0.0.1", port: 2222 }), + }); + relayHandler.createRelay(relayA); + relayHandler.createRelay(relayB); + + // When — free only one of the two entries on this address + relayHandler.freeRelay(relayA); + + // Then — the sibling on the same address is still indexed... + assert(!relayHandler.hasRelay(relayA), "Freed entry still present!"); + assert(relayHandler.hasRelay(relayB), "Sibling entry was dropped!"); + // ...and the address bucket is retained while any port remains + assert( + (relayHandler as any)._byAddress.has("10.0.0.1"), + "Address bucket removed while a port still exists!", + ); + + // When — free the last entry on the address + relayHandler.freeRelay(relayB); + + // Then — the now-empty address bucket is cleaned up + assert(!relayHandler.hasRelay(relayB), "Freed entry still present!"); + assert( + !(relayHandler as any)._byAddress.has("10.0.0.1"), + "Empty address bucket not cleaned up!", + ); + }); + }); describe("hasRelay", () => { it("should not have relay", async () => { // Given