From b5b100c6fd04f548a8e2d07afe6f5a96dcd2e096 Mon Sep 17 00:00:00 2001 From: "nekomimi.pet" Date: Wed, 7 Oct 2026 17:30:21 -0400 Subject: [PATCH] perf(firehose): scan storage stats only on the leader, and bound the frame queue Standby workers no longer run the hourly full-bucket ListObjects scan; it starts with the protected work once this instance is the active worker and stops on step-down and shutdown. BunSubscription caps frames received but not yet consumed (16 MiB by default). A WebSocket cannot pause its reads, so past the cap it closes the connection, still delivers the frames already queued, in order, and reconnects at once from getParams()' cursor (or after the last delivered seq when there is none). Before, a slow consumer let the queue grow without bound. --- .../scripts/storage-stats-scans.ts | 113 +++++++++++++++ .../src/index.storage-stats.test.ts | 25 ++++ apps/firehose-service/src/index.ts | 5 +- .../bun-firehose/bench/backpressure.ts | 133 ++++++++++++++++++ .../bun-firehose/src/subscription.test.ts | 107 ++++++++++++++ .../bun-firehose/src/subscription.ts | 69 +++++++-- 6 files changed, 442 insertions(+), 10 deletions(-) create mode 100644 apps/firehose-service/scripts/storage-stats-scans.ts create mode 100644 apps/firehose-service/src/index.storage-stats.test.ts create mode 100644 packages/@wispplace/bun-firehose/bench/backpressure.ts diff --git a/apps/firehose-service/scripts/storage-stats-scans.ts b/apps/firehose-service/scripts/storage-stats-scans.ts new file mode 100644 index 0000000..c93cc80 --- /dev/null +++ b/apps/firehose-service/scripts/storage-stats-scans.ts @@ -0,0 +1,113 @@ +/** + * Counts the full-bucket storage-stat scans (S3 ListObjectsV2 requests) one + * firehose worker makes. The worker runs against a fake S3 and, as `standby`, + * against a fake leadership supervisor that never grants authority; as + * `single`, it runs without leader election, so it is the active worker. + * + * bun scripts/storage-stats-scans.ts [standby|single] [holdMs=3000] + * + * A refresh scans once when started and then hourly, so the scans in the + * first seconds are the scans per hour. Needs no database, Redis or relay. + */ +import { chmod, mkdtemp, rm, writeFile } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' + +export type WorkerRole = 'standby' | 'single' + +export interface StorageStatsScanCount { + role: WorkerRole + holdMs: number + listObjectsRequests: number + health: { readiness?: string; storage?: Record } +} + +const EMPTY_LISTING = + '' + + 'bench01000false' + +/** Answers on stdin's behalf like a supervisor which stays in standby until the worker exits. */ +const STANDBY_SUPERVISOR = `#!/usr/bin/env bun +process.stdout.write(JSON.stringify({ version: 1, type: 'state', state: 'standby', pid: process.pid }) + '\\n') +for await (const _ of process.stdin) {} +` + +function freePort(): number { + const probe = Bun.serve({ port: 0, fetch: () => new Response() }) + const port = probe.port + probe.stop(true) + if (port === undefined) throw new Error('no free port') + return port +} + +async function waitForHealth(url: string, timeoutMs: number): Promise { + const deadline = Date.now() + timeoutMs + while (Date.now() < deadline) { + const ok = await fetch(url).then( + () => true, + () => false, + ) + if (ok) return + await Bun.sleep(50) + } + throw new Error('firehose worker health endpoint did not come up') +} + +export async function countStorageStatsScans(role: WorkerRole, holdMs = 3_000): Promise { + let listObjectsRequests = 0 + const s3 = Bun.serve({ + port: 0, + fetch(req) { + if (new URL(req.url).searchParams.get('list-type') === '2') { + listObjectsRequests++ + return new Response(EMPTY_LISTING, { headers: { 'content-type': 'application/xml' } }) + } + return new Response(null, { status: 200 }) + }, + }) + const dir = await mkdtemp(join(tmpdir(), 'firehose-storage-scans-')) + const supervisorPath = join(dir, 'standby-supervisor') + await writeFile(supervisorPath, STANDBY_SUPERVISOR) + await chmod(supervisorPath, 0o755) + + const healthPort = freePort() + const worker = Bun.spawn(['bun', join(import.meta.dir, '../src/index.ts')], { + env: { + PATH: process.env.PATH, + HOME: process.env.HOME, + NODE_ENV: 'test', + HEALTH_PORT: String(healthPort), + S3_BUCKET: 'bench', + S3_ENDPOINT: `http://127.0.0.1:${s3.port}`, + AWS_ACCESS_KEY_ID: 'bench', + AWS_SECRET_ACCESS_KEY: 'bench', + // Nothing listens here: an active worker just keeps reconnecting. + FIREHOSE_SERVICE: 'ws://127.0.0.1:9', + LEADERSHIP_SUPERVISOR_ENABLED: role === 'standby' ? 'true' : 'false', + FIREHOSE_SUPERVISOR_PATH: supervisorPath, + }, + stdout: 'ignore', + stderr: 'ignore', + }) + try { + const healthUrl = `http://127.0.0.1:${healthPort}/health` + await waitForHealth(healthUrl, 15_000) + await Bun.sleep(holdMs) + const health = (await (await fetch(healthUrl)).json()) as StorageStatsScanCount['health'] + return { role, holdMs, listObjectsRequests, health } + } finally { + worker.kill('SIGKILL') + await worker.exited + s3.stop(true) + await rm(dir, { recursive: true, force: true }) + } +} + +if (import.meta.main) { + const role = (process.argv[2] ?? 'standby') as WorkerRole + if (role !== 'standby' && role !== 'single') throw new Error('role must be standby or single') + const result = await countStorageStatsScans(role, Number(process.argv[3] ?? 3_000)) + console.log( + JSON.stringify({ ...result, health: { readiness: result.health.readiness, storage: result.health.storage } }), + ) +} diff --git a/apps/firehose-service/src/index.storage-stats.test.ts b/apps/firehose-service/src/index.storage-stats.test.ts new file mode 100644 index 0000000..51cacac --- /dev/null +++ b/apps/firehose-service/src/index.storage-stats.test.ts @@ -0,0 +1,25 @@ +import { describe, expect, test } from 'bun:test' +import { countStorageStatsScans } from '../scripts/storage-stats-scans' + +// Spawns the real worker entry point against a fake S3 and a fake supervisor. +describe('storage-stat scans and leadership', () => { + test('a standby worker never lists the bucket', async () => { + const result = await countStorageStatsScans('standby', 1_500) + + expect(result.health.readiness).toBe('standby') + expect(result.listObjectsRequests).toBe(0) + expect(result.health.storage).toEqual({ + lastSuccessAgeMs: null, + lastErrorKind: null, + stale: true, + refreshing: false, + }) + }, 30_000) + + test('the active worker still refreshes storage stats', async () => { + const result = await countStorageStatsScans('single', 1_500) + + expect(result.listObjectsRequests).toBe(1) + expect(result.health.storage?.stale).toBe(false) + }, 30_000) +}) diff --git a/apps/firehose-service/src/index.ts b/apps/firehose-service/src/index.ts index c34c8d9..be43613 100644 --- a/apps/firehose-service/src/index.ts +++ b/apps/firehose-service/src/index.ts @@ -584,6 +584,9 @@ async function startProtectedWork( return false } startFirehose(cursor, onFailure) + // A full storage listing is only worth its cost on the one active worker; + // standbys report no capacity data (`stale: true`, `lastSuccessAgeMs: null`). + startStorageStatsRefresh() firehoseStartGate.open() return true } @@ -600,7 +603,6 @@ async function main() { fetch: app.fetch, port: config.healthPort, }) - startStorageStatsRefresh() logger.info('Health server listening', { port: config.healthPort }) if (config.leadershipSupervisorEnabled) { @@ -654,6 +656,7 @@ async function main() { ) }, async () => { + stopStorageStatsRefresh() const drain = await stopFirehose() const workerStop = await stopRevalidateWorker({ gracePeriodMs: config.firehoseDrainGraceMs }) await stopCursorSaving() diff --git a/packages/@wispplace/bun-firehose/bench/backpressure.ts b/packages/@wispplace/bun-firehose/bench/backpressure.ts new file mode 100644 index 0000000..e7332bd --- /dev/null +++ b/packages/@wispplace/bun-firehose/bench/backpressure.ts @@ -0,0 +1,133 @@ +/** + * Feeds frames to BunSubscription faster than its consumer handles them, like a + * relay replaying a backlog to a slow leader. Reports peak memory and proves + * every sequence number reaches the consumer once, in order. Each frame costs + * the consumer `workUs` of CPU plus one event-loop turn, as a handler awaiting + * I/O does. + * + * bun bench/backpressure.ts [frames=100000] [frameBytes=1024] [workUs=50] + * + * The relay runs in a child process so its send buffers are not measured. + */ +import { encode } from '@atproto/lex-cbor' +import { MessageQueue } from '../src/queue' +import { BunSubscription } from '../src/subscription' + +const [mode, ...rest] = process.argv.slice(2) + +if (mode === '--relay') serveRelay(Number(rest[0]), Number(rest[1])) +else await consume(Number(mode ?? 100_000), Number(rest[0] ?? 1024), Number(rest[1] ?? 50)) + +/** A relay that replays `frames` sequence numbers after the requested cursor as fast as the socket accepts them. */ +function serveRelay(frames: number, frameBytes: number): void { + const header = encode({ op: 1, t: '#frame' }) + const pad = new Uint8Array(frameBytes) + const frame = (seq: number) => { + const body = encode({ seq, pad }) + const bytes = new Uint8Array(header.length + body.length) + bytes.set(header) + bytes.set(body, header.length) + return bytes + } + type Stream = { next: number; open: boolean } + const pump = (ws: Bun.ServerWebSocket) => { + while (ws.data.open && ws.data.next <= frames) { + if (ws.send(frame(ws.data.next++)) === -1) return + } + } + const server = Bun.serve({ + port: 0, + fetch(req, srv) { + const cursor = Number(new URL(req.url).searchParams.get('cursor') ?? 0) + return srv.upgrade(req, { data: { next: cursor + 1, open: true } }) + ? undefined + : new Response('no', { status: 400 }) + }, + websocket: { + open: pump, + drain: pump, + close(ws) { + ws.data.open = false + }, + message() {}, + }, + }) + console.log(server.port) +} + +async function consume(frames: number, frameBytes: number, workUs: number): Promise { + const relay = Bun.spawn(['bun', import.meta.path, '--relay', String(frames), String(frameBytes)], { stdout: 'pipe' }) + const port = (await new Response(relay.stdout).body!.getReader().read()).value + const url = `ws://127.0.0.1:${new TextDecoder().decode(port).trim()}` + + Bun.gc(true) + const baseline = process.memoryUsage() + let peakRss = baseline.rss + let peakHeap = baseline.heapUsed + const sampler = setInterval(() => { + const usage = process.memoryUsage() + peakRss = Math.max(peakRss, usage.rss) + peakHeap = Math.max(peakHeap, usage.heapUsed) + }, 10) + + // Observe the subscription's frame queue without changing it. + let peakQueued = 0 + const push = MessageQueue.prototype.push + MessageQueue.prototype.push = function (this: MessageQueue, item: unknown) { + push.call(this, item) + peakQueued = Math.max(peakQueued, this.length) + } + + let last = 0 + let delivered = 0 + let duplicates = 0 + let gaps = 0 + let connects = 0 + const abort = new AbortController() + const subscription = new BunSubscription<{ seq: number }>({ + service: url, + method: 'bench.subscribe', + signal: abort.signal, + maxReconnectSeconds: 1, + validate: (body) => body as { seq: number }, + // The consumer's committed cursor: everything up to `last` was handled. + getParams: () => (last > 0 ? { cursor: last } : {}), + onConnect: () => connects++, + }) + + const started = performance.now() + for await (const { seq } of subscription) { + delivered++ + if (seq <= last) duplicates++ + else if (seq !== last + 1) gaps++ + last = Math.max(last, seq) + const until = performance.now() + workUs / 1000 + while (performance.now() < until) {} + // Like a handler awaiting I/O: lets the socket deliver more frames meanwhile. + await new Promise((resolve) => setImmediate(resolve)) + if (last === frames) break + } + const seconds = (performance.now() - started) / 1000 + clearInterval(sampler) + abort.abort() + relay.kill() + + const mib = (bytes: number) => `${(bytes / 1024 / 1024).toFixed(1)} MiB` + console.log( + JSON.stringify({ + frames, + frameBytes, + workUs, + delivered, + lastSeq: last, + gaps, + duplicates, + connects, + seconds: Number(seconds.toFixed(2)), + peakQueuedFrames: peakQueued, + peakRssGrowth: mib(peakRss - baseline.rss), + peakHeapGrowth: mib(peakHeap - baseline.heapUsed), + }), + ) + process.exit(gaps === 0 && last === frames ? 0 : 1) +} diff --git a/packages/@wispplace/bun-firehose/src/subscription.test.ts b/packages/@wispplace/bun-firehose/src/subscription.test.ts index 8d62680..8e49fa5 100644 --- a/packages/@wispplace/bun-firehose/src/subscription.test.ts +++ b/packages/@wispplace/bun-firehose/src/subscription.test.ts @@ -1,6 +1,9 @@ import { afterEach, describe, expect, test } from 'bun:test' import { createHash } from 'node:crypto' import { createServer, type Server } from 'node:net' +import { encode } from '@atproto/lex-cbor' +import type { ServerWebSocket } from 'bun' +import { MessageQueue } from './queue' import { BunSubscription } from './subscription' const WS_GUID = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11' @@ -106,3 +109,107 @@ describe('BunSubscription silence handling', () => { expect(harness.disconnects).toBeGreaterThanOrEqual(1) }) }) + +/** A relay replaying `frames` sequence numbers after the requested cursor as fast as the socket takes them. */ +function replayServer(frames: number, frameBytes: number): { url: string; cursors: Array } { + const header = encode({ op: 1, t: '#frame' }) + const pad = new Uint8Array(frameBytes) + const cursors: Array = [] + type Stream = { next: number; open: boolean } + const pump = (ws: ServerWebSocket) => { + while (ws.data.open && ws.data.next <= frames) { + const body = encode({ seq: ws.data.next++, pad }) + const bytes = new Uint8Array(header.length + body.length) + bytes.set(header) + bytes.set(body, header.length) + if (ws.send(bytes) === -1) return + } + } + const server = Bun.serve({ + port: 0, + fetch(req, srv) { + const cursor = new URL(req.url).searchParams.get('cursor') + cursors.push(cursor) + return srv.upgrade(req, { data: { next: Number(cursor ?? 0) + 1, open: true } }) + ? undefined + : new Response('no', { status: 400 }) + }, + websocket: { + open: pump, + drain: pump, + close(ws) { + ws.data.open = false + }, + message() {}, + }, + }) + cleanups.push(() => server.stop(true)) + return { url: `ws://127.0.0.1:${server.port}`, cursors } +} + +/** Consumes like a handler awaiting I/O per event, which lets the socket keep delivering meanwhile. */ +async function consumeSlowly( + url: string, + frames: number, + options: { maxQueuedBytes: number; committed: boolean }, +): Promise<{ seqs: number[]; peakQueued: number; connects: number }> { + let peakQueued = 0 + const push = MessageQueue.prototype.push + MessageQueue.prototype.push = function (this: MessageQueue, item: unknown) { + push.call(this, item) + peakQueued = Math.max(peakQueued, this.length) + } + cleanups.push(() => { + MessageQueue.prototype.push = push + }) + + const seqs: number[] = [] + let connects = 0 + const abort = new AbortController() + cleanups.push(() => abort.abort()) + const subscription = new BunSubscription<{ seq: number }>({ + service: url, + method: 'test.subscribe', + signal: abort.signal, + maxReconnectSeconds: 1, + maxQueuedBytes: options.maxQueuedBytes, + validate: (body) => body as { seq: number }, + getParams: () => (options.committed && seqs.length > 0 ? { cursor: seqs[seqs.length - 1] } : undefined), + onConnect: () => connects++, + }) + for await (const { seq } of subscription) { + seqs.push(seq) + await new Promise((resolve) => setImmediate(resolve)) + if (seq === frames) break + } + return { seqs, peakQueued, connects } +} + +describe('BunSubscription backpressure', () => { + const frames = 5_000 + const frameBytes = 512 + const maxQueuedBytes = 64 * 1024 + const expected = Array.from({ length: frames }, (_, i) => i + 1) + + test('bounds the frame queue and resumes from the committed cursor without losing or reordering', async () => { + const { url, cursors } = replayServer(frames, frameBytes) + + const result = await consumeSlowly(url, frames, { maxQueuedBytes, committed: true }) + + expect(result.seqs).toEqual(expected) + expect(result.peakQueued).toBeLessThanOrEqual(Math.ceil(maxQueuedBytes / frameBytes)) + expect(result.connects).toBeGreaterThan(1) + expect(cursors[0]).toBeNull() + expect(cursors.slice(1).every((cursor) => cursor !== null)).toBe(true) + }, 30_000) + + test('resumes after the last delivered frame when the consumer has no cursor yet', async () => { + const { url, cursors } = replayServer(frames, frameBytes) + + const result = await consumeSlowly(url, frames, { maxQueuedBytes, committed: false }) + + expect(result.seqs).toEqual(expected) + expect(result.connects).toBeGreaterThan(1) + expect(cursors.slice(1).every((cursor) => cursor !== null)).toBe(true) + }, 30_000) +}) diff --git a/packages/@wispplace/bun-firehose/src/subscription.ts b/packages/@wispplace/bun-firehose/src/subscription.ts index a63b14f..3de08ff 100644 --- a/packages/@wispplace/bun-firehose/src/subscription.ts +++ b/packages/@wispplace/bun-firehose/src/subscription.ts @@ -36,6 +36,9 @@ export function decodeFrame(bytes: Uint8Array): { header: FrameHeader; body: unk return { header: header.value as unknown as FrameHeader, body: body.value } } +/** Default for {@link BunSubscriptionOptions.maxQueuedBytes}. */ +export const DEFAULT_MAX_QUEUED_BYTES = 16 * 1024 * 1024 + export interface BunSubscriptionOptions { service: string method: string @@ -56,6 +59,14 @@ export interface BunSubscriptionOptions { heartbeatIntervalMs?: number /** How long a ping may go unanswered before the connection is dropped (default: 5000) */ heartbeatTimeoutMs?: number + /** + * Bytes of received frames the consumer may fall behind by (default: 16 MiB). A WebSocket + * cannot pause its reads, so past this the connection is closed instead: frames already + * queued are still delivered in order, then the subscription reconnects at once from + * `getParams()`, which should name the consumer's committed cursor. Without a cursor there, + * it resumes after the last sequence number it delivered, so no frame is skipped. + */ + maxQueuedBytes?: number } type PingWebSocket = WebSocket & { ping(data: string): unknown } @@ -64,6 +75,8 @@ export class BunSubscription { private ws: WebSocket | null = null private reconnectAttempts = 0 private aborted = false + /** `seq` of the last message handed to the consumer, the fallback cursor after an overflow. */ + private lastSeq: number | undefined constructor(public opts: BunSubscriptionOptions) { if (opts.signal) { @@ -74,8 +87,11 @@ export class BunSubscription { } } - private async getUrl(): Promise { - const params = (await this.opts.getParams?.()) ?? {} + private async getUrl(resumeAfterOverflow: boolean): Promise { + let params = (await this.opts.getParams?.()) ?? {} + if (resumeAfterOverflow && params.cursor === undefined && this.lastSeq !== undefined) { + params = { ...params, cursor: this.lastSeq } + } const query = encodeQueryParams(params) const base = this.opts.service.replace(/\/$/, '') return `${base}/xrpc/${this.opts.method}${query ? `?${query}` : ''}` @@ -89,19 +105,30 @@ export class BunSubscription { async *[Symbol.asyncIterator](): AsyncGenerator { const maxSilenceMs = this.opts.maxSilenceMs === undefined ? 15_000 : this.opts.maxSilenceMs + const maxQueuedBytes = this.opts.maxQueuedBytes ?? DEFAULT_MAX_QUEUED_BYTES + let resumeAfterOverflow = false while (!this.aborted) { let silenceTimer: ReturnType | null = null try { - const url = await this.getUrl() + const url = await this.getUrl(resumeAfterOverflow) + resumeAfterOverflow = false // Create a queue for messages const messageQueue = new MessageQueue() + let queuedBytes = 0 let resolveMessage: (() => void) | null = null let wsError: Error | null = null let wsOpen = false let wsClosed = false + let overflowed = false + let disconnectReported = false + const reportDisconnect = () => { + if (disconnectReported) return + disconnectReported = true + this.opts.onDisconnect?.() + } const dropSilentConnection = (reason: string) => { if (!wsClosed && !this.aborted) { @@ -132,7 +159,8 @@ export class BunSubscription { }, this.opts.heartbeatIntervalMs ?? 5_000) } - this.ws = new WebSocket(url) + const socket = new WebSocket(url) + this.ws = socket this.ws.binaryType = 'arraybuffer' this.ws.addEventListener('open', () => { @@ -151,7 +179,23 @@ export class BunSubscription { this.ws.addEventListener('message', (event) => { const data = event.data + if (overflowed) return if (data instanceof ArrayBuffer) { + if (messageQueue.length > 0 && queuedBytes + data.byteLength > maxQueuedBytes) { + // This frame and any still in flight are dropped unseen; the + // reconnect resumes from a cursor at or before the queued ones. + console.warn( + `[BunSubscription] ${messageQueue.length} frames (${queuedBytes} bytes) unconsumed, reconnecting from cursor`, + ) + overflowed = true + wsClosed = true + socket.close() + // Report now: the close event may land after the next connection opens. + reportDisconnect() + resolveMessage?.() + return + } + queuedBytes += data.byteLength messageQueue.push(new Uint8Array(data)) resetSilenceTimer() resolveMessage?.() @@ -164,7 +208,7 @@ export class BunSubscription { this.ws.addEventListener('close', () => { wsClosed = true - this.opts.onDisconnect?.() + reportDisconnect() resolveMessage?.() }) @@ -180,8 +224,8 @@ export class BunSubscription { throw wsError } - // Process messages - while (!this.aborted && !wsClosed) { + // Process messages, including those queued before the socket closed + while (!this.aborted) { // Wait for message if queue is empty while (messageQueue.length === 0 && !wsClosed && !this.aborted) { await new Promise((resolve) => { @@ -189,10 +233,11 @@ export class BunSubscription { }) } - if (wsClosed || this.aborted) break + if (this.aborted) break const bytes = messageQueue.shift() - if (!bytes) continue + if (!bytes) break + queuedBytes -= bytes.byteLength try { const { header, body } = decodeFrame(bytes) @@ -214,6 +259,7 @@ export class BunSubscription { if (result !== undefined) { yield result } + if (isPlainObject(body) && typeof body.seq === 'number') this.lastSeq = body.seq } } catch (err) { // Log decode errors but continue @@ -227,6 +273,11 @@ export class BunSubscription { this.ws = null if (this.aborted) break + // Falling behind is not a connection failure: resume without backoff. + if (overflowed) { + resumeAfterOverflow = true + continue + } // Reconnect this.reconnectAttempts++ -- 2.54.0