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++