From 83af12dc3e944a6e84df5254dc11743517d9d8ad Mon Sep 17 00:00:00 2001
From: Neil <4138956+nwparker@users.noreply.github.com>
Date: Thu, 30 Jul 2026 17:48:07 -0700
Subject: [PATCH] [P2] perf(relay): stop re-encoding frames to size PTY chunks
 (#11620)

* perf(relay): stop re-encoding frames to size PTY chunks

maxLegacyPtyDataChars now sizes the full candidate once (the common case)
and falls back to a binary search over an exact cheap byte formula instead
of fully encoding the frame at every probe. Publish paths thread their
already-computed frame estimate into enqueueFrame so each PTY publish
encodes the message once for admission instead of twice.

Co-authored-by: Orca <help@stably.ai>

* fix(relay): preserve dispatcher frame guards

---------

Co-authored-by: Orca <help@stably.ai>
---
 ...dispatcher-frame-guard-regressions.test.ts |  47 +++++
 src/relay/dispatcher.test.ts                  | 193 ++++++++++++++++++
 src/relay/dispatcher.ts                       |  63 ++++--
 3 files changed, 283 insertions(+), 20 deletions(-)
 create mode 100644 src/relay/dispatcher-frame-guard-regressions.test.ts

diff --git a/src/relay/dispatcher-frame-guard-regressions.test.ts b/src/relay/dispatcher-frame-guard-regressions.test.ts
new file mode 100644
index 000000000..71859e05c
--- /dev/null
+++ b/src/relay/dispatcher-frame-guard-regressions.test.ts
@@ -0,0 +1,47 @@
+import { describe, expect, it, vi } from 'vitest'
+import { RelayDispatcher } from './dispatcher'
+import type { JsonRpcNotification } from './protocol'
+
+type DispatcherInternals = {
+  primaryClient: object
+  estimateFrameBytes: (msg: JsonRpcNotification) => number
+  enqueueFrame: (client: object, msg: JsonRpcNotification, lane: string) => boolean
+}
+
+describe('RelayDispatcher frame guards', () => {
+  it('returns zero for invalid active-client limits without encoding', () => {
+    const dispatcher = new RelayDispatcher(() => true, {
+      writableHighWaterMark: () => 1024 * 1024,
+      writableLength: () => 0
+    })
+    try {
+      const spy = vi.spyOn(dispatcher as unknown as DispatcherInternals, 'estimateFrameBytes')
+      for (const limit of [0, -1, Number.NaN]) {
+        expect(dispatcher.maxLegacyPtyDataChars({ id: 'pty-1' }, 'hello', limit)).toBe(0)
+      }
+      expect(spy).not.toHaveBeenCalled()
+    } finally {
+      dispatcher.dispose()
+    }
+  })
+
+  it('does not estimate frames after disposal', () => {
+    const dispatcher = new RelayDispatcher(() => true)
+    const internals = dispatcher as unknown as DispatcherInternals
+    const spy = vi.spyOn(internals, 'estimateFrameBytes')
+    const msg: JsonRpcNotification = {
+      jsonrpc: '2.0',
+      method: 'pty.data',
+      params: {
+        data: {
+          toJSON: () => {
+            throw new Error('must not serialize')
+          }
+        }
+      }
+    }
+    dispatcher.dispose()
+    expect(internals.enqueueFrame(internals.primaryClient, msg, 'ordinary')).toBe(false)
+    expect(spy).not.toHaveBeenCalled()
+  })
+})
diff --git a/src/relay/dispatcher.test.ts b/src/relay/dispatcher.test.ts
index aaa083b77..3ec57dcf8 100644
--- a/src/relay/dispatcher.test.ts
+++ b/src/relay/dispatcher.test.ts
@@ -1,5 +1,6 @@
 import { describe, expect, it, vi, beforeEach, afterEach } from 'vitest'
 import { RelayDispatcher, type SinkWriteSettlement } from './dispatcher'
+import { relayWriterControlReserve } from './dispatcher-writer-admission'
 import {
   encodeJsonRpcFrame,
   encodeKeepAliveFrame,
@@ -716,4 +717,196 @@ describe('RelayDispatcher', () => {
       }
     })
   })
+
+  describe('legacy PTY chunk sizing', () => {
+    type DispatcherInternals = {
+      primaryClient: object
+      estimateFrameBytes: (msg: JsonRpcNotification) => number
+      enqueueFrame: (
+        client: object,
+        msg: JsonRpcNotification,
+        lane: string,
+        onSettled?: (result: SinkWriteSettlement) => void,
+        estimatedBytes?: number
+      ) => boolean
+    }
+
+    // Pre-optimization sizing loop, kept verbatim for differential verification.
+    function referenceMaxChars(
+      capacities: number[],
+      params: Record<string, unknown>,
+      data: string,
+      limit: number
+    ): number {
+      if (capacities.length === 0) {
+        return Math.min(data.length, limit)
+      }
+      let low = 0
+      let high = Math.min(data.length, limit)
+      while (low < high) {
+        const mid = Math.ceil((low + high) / 2)
+        const msg: JsonRpcNotification = {
+          jsonrpc: '2.0',
+          method: 'pty.data',
+          params: { ...params, data: data.slice(0, mid) }
+        }
+        const bytes = encodeJsonRpcFrame(msg, 0, 0).length
+        if (capacities.every((capacity) => bytes <= capacity)) {
+          low = mid
+        } else {
+          high = mid - 1
+        }
+      }
+      return low
+    }
+
+    function makeDispatcher(highWaterMarks: number[]): {
+      sized: RelayDispatcher
+      capacities: number[]
+    } {
+      const [primaryHwm, ...rest] = highWaterMarks
+      const sized = new RelayDispatcher(() => true, {
+        writableHighWaterMark: () => primaryHwm,
+        writableLength: () => 0
+      })
+      for (const hwm of rest) {
+        sized.attachClient(() => true, {
+          writableHighWaterMark: () => hwm,
+          writableLength: () => 0
+        })
+      }
+      return {
+        sized,
+        capacities: highWaterMarks.map((hwm) => Math.max(0, hwm - relayWriterControlReserve(hwm)))
+      }
+    }
+
+    function mulberry32(seed: number): () => number {
+      let a = seed
+      return () => {
+        a = (a + 0x6d2b79f5) | 0
+        let t = Math.imul(a ^ (a >>> 15), 1 | a)
+        t = (t + Math.imul(t ^ (t >>> 7), 61 | t)) ^ t
+        return ((t ^ (t >>> 14)) >>> 0) / 4294967296
+      }
+    }
+
+    // Multi-byte UTF-8, astral pairs, lone surrogates, controls, quotes, and backslashes.
+    const alphabet = 'aZ9 "\\\n\r\t éß€中𝄞😀𐀀�'
+
+    it('matches the pre-optimization sizing loop across randomized inputs', () => {
+      const random = mulberry32(0xc0ffee)
+      const hwmChoices = [1030, 1100, 1250, 1400, 2048, 1024 * 1024]
+      for (let trial = 0; trial < 300; trial++) {
+        const length = Math.floor(random() * 240)
+        let data = ''
+        for (let i = 0; i < length; i++) {
+          data += alphabet[Math.floor(random() * alphabet.length)]
+        }
+        const clientCount = 1 + Math.floor(random() * 2)
+        const hwms = Array.from(
+          { length: clientCount },
+          () => hwmChoices[Math.floor(random() * hwmChoices.length)]
+        )
+        const params = { id: `pty-${trial}`, seq: trial }
+        const { sized, capacities } = makeDispatcher(hwms)
+        try {
+          for (const limit of [0, 1, Math.floor(length / 2), length, length + 17]) {
+            expect(sized.maxLegacyPtyDataChars(params, data, limit)).toBe(
+              referenceMaxChars(capacities, params, data, limit)
+            )
+          }
+          expect(sized.maxLegacyPtyDataChars(params, data)).toBe(
+            referenceMaxChars(capacities, params, data, data.length)
+          )
+        } finally {
+          sized.dispose()
+        }
+      }
+    })
+
+    it('sizes a chunk that fits every client with a single frame encode', () => {
+      const { sized } = makeDispatcher([1024 * 1024])
+      try {
+        const spy = vi.spyOn(sized as unknown as DispatcherInternals, 'estimateFrameBytes')
+        const data = 'x'.repeat(16 * 1024)
+        expect(sized.maxLegacyPtyDataChars({ id: 'pty-1' }, data)).toBe(data.length)
+        expect(spy).toHaveBeenCalledTimes(1)
+      } finally {
+        sized.dispose()
+      }
+    })
+
+    it('publishes PTY data with a single frame estimate', () => {
+      const frames: Buffer[] = []
+      const publisher = new RelayDispatcher((data) => {
+        frames.push(Buffer.from(data))
+        return true
+      })
+      try {
+        const spy = vi.spyOn(publisher as unknown as DispatcherInternals, 'estimateFrameBytes')
+        expect(publisher.tryNotifyPtyData({ id: 'pty-1', data: 'hello' })).toBe(true)
+        expect(frames).toHaveLength(1)
+        expect(spy).toHaveBeenCalledTimes(1)
+      } finally {
+        publisher.dispose()
+      }
+    })
+
+    it('enqueueFrame with a caller-supplied estimate matches the computed default', () => {
+      const frames: Buffer[] = []
+      const publisher = new RelayDispatcher((data) => {
+        frames.push(Buffer.from(data))
+        return true
+      })
+      try {
+        const internals = publisher as unknown as DispatcherInternals
+        const msg: JsonRpcNotification = {
+          jsonrpc: '2.0',
+          method: 'pty.data',
+          params: { id: 'pty-1', data: 'héllo "𝄞"\\\n\uD800' }
+        }
+        expect(internals.enqueueFrame(internals.primaryClient, msg, 'ordinary')).toBe(true)
+        expect(
+          internals.enqueueFrame(
+            internals.primaryClient,
+            msg,
+            'ordinary',
+            undefined,
+            internals.estimateFrameBytes(msg)
+          )
+        ).toBe(true)
+        expect(frames).toHaveLength(2)
+        expect(
+          decodeFirstFrame(frames[1]).payload.equals(decodeFirstFrame(frames[0]).payload)
+        ).toBe(true)
+      } finally {
+        publisher.dispose()
+      }
+    })
+
+    it('enqueueFrame rejects identically with and without a caller-supplied estimate', () => {
+      const { sized } = makeDispatcher([1030])
+      try {
+        const internals = sized as unknown as DispatcherInternals
+        const msg: JsonRpcNotification = {
+          jsonrpc: '2.0',
+          method: 'pty.data',
+          params: { id: 'pty-1', data: 'x'.repeat(512) }
+        }
+        expect(internals.enqueueFrame(internals.primaryClient, msg, 'ordinary')).toBe(false)
+        expect(
+          internals.enqueueFrame(
+            internals.primaryClient,
+            msg,
+            'ordinary',
+            undefined,
+            internals.estimateFrameBytes(msg)
+          )
+        ).toBe(false)
+      } finally {
+        sized.dispose()
+      }
+    })
+  })
 })
diff --git a/src/relay/dispatcher.ts b/src/relay/dispatcher.ts
index 833081c65..3809b9f5f 100644
--- a/src/relay/dispatcher.ts
+++ b/src/relay/dispatcher.ts
@@ -193,20 +193,34 @@ export class RelayDispatcher {
     limit = data.length
   ): number {
     const clients = this.activeClients()
+    const max = Math.min(data.length, limit)
     if (clients.length === 0) {
-      return Math.min(data.length, limit)
+      return max
     }
-    let low = 0
-    let high = Math.min(data.length, limit)
-    while (low < high) {
-      const mid = Math.ceil((low + high) / 2)
-      const msg: JsonRpcNotification = {
+    if (!(max > 0)) {
+      return 0
+    }
+    const fitsAll = (bytes: number): boolean =>
+      clients.every((client) => bytes <= client.writer.producerFrameCapacity)
+    const sizeFrame = (chunk: string): number =>
+      this.estimateFrameBytes({
         jsonrpc: '2.0',
         method: 'pty.data',
-        params: { ...params, data: data.slice(0, mid) }
-      }
-      const bytes = this.estimateFrameBytes(msg)
-      if (clients.every((client) => bytes <= client.writer.producerFrameCapacity)) {
+        params: { ...params, data: chunk }
+      })
+    // Fast path: the whole chunk usually fits — one encode instead of log2(n).
+    if (fitsAll(sizeFrame(data.slice(0, max)))) {
+      return max
+    }
+    // Exact per-step size: only the escaped data string varies; its quotes are in baseBytes.
+    const baseBytes = sizeFrame('')
+    const bytesFor = (chars: number): number =>
+      baseBytes + Buffer.byteLength(JSON.stringify(data.slice(0, chars))) - 2
+    let low = 0
+    let high = max
+    while (low < high) {
+      const mid = Math.ceil((low + high) / 2)
+      if (fitsAll(bytesFor(mid))) {
         low = mid
       } else {
         high = mid - 1
@@ -822,19 +836,21 @@ export class RelayDispatcher {
     client: RelayClient,
     msg: JsonRpcRequest | JsonRpcResponse | JsonRpcNotification,
     lane: DispatcherWriterLane,
-    onSettled: (result: SinkWriteSettlement) => void = () => {}
+    onSettled: (result: SinkWriteSettlement) => void = () => {},
+    // Why: publish paths already sized the frame; avoid a redundant encode.
+    estimatedBytes?: number
   ): boolean {
     if (this.disposed || client.closed) {
       return false
     }
-    const estimatedBytes = this.estimateFrameBytes(msg)
+    const frameBytes = estimatedBytes ?? this.estimateFrameBytes(msg)
     return client.writer.enqueue(
       lane,
       () => {
         const seq = client.nextOutgoingSeq++
         return encodeJsonRpcFrame(msg, seq, client.highestReceivedSeq)
       },
-      estimatedBytes,
+      frameBytes,
       onSettled
     )
   }
@@ -898,7 +914,7 @@ export class RelayDispatcher {
         return false
       }
       for (let index = 0; index < clients.length; index++) {
-        if (!this.enqueueLeasedFrame(clients[index], msg, lane, leases[index])) {
+        if (!this.enqueueLeasedFrame(clients[index], msg, lane, leases[index], bytes)) {
           if (this.disposed || clients[index].closed) {
             continue
           }
@@ -949,7 +965,7 @@ export class RelayDispatcher {
     if (!leases) {
       return false
     }
-    return this.enqueueLeasedFrame(client, msg, lane, leases[0], onSettled)
+    return this.enqueueLeasedFrame(client, msg, lane, leases[0], bytes, onSettled)
   }
 
   private publishBulkWhenAvailable(client: RelayClient, msg: JsonRpcNotification): Promise<void> {
@@ -998,13 +1014,20 @@ export class RelayDispatcher {
     msg: JsonRpcNotification,
     lane: 'interactive' | 'ordinary' | 'fixed-bulk' | 'bulk',
     lease: LegacyPublicationLease,
+    estimatedBytes: number,
     onSettled: (result: SinkWriteSettlement) => void = () => {}
   ): boolean {
-    const accepted = this.enqueueFrame(client, msg, lane, (result) => {
-      lease.release()
-      onSettled(result)
-      this.notifyLegacyCapacityIfLow()
-    })
+    const accepted = this.enqueueFrame(
+      client,
+      msg,
+      lane,
+      (result) => {
+        lease.release()
+        onSettled(result)
+        this.notifyLegacyCapacityIfLow()
+      },
+      estimatedBytes
+    )
     if (!accepted) {
       lease.release()
       this.notifyLegacyCapacityIfLow()