diff --git a/sitemap/src/index.ts b/sitemap/src/index.ts index a2c692d48..195790262 100644 --- a/sitemap/src/index.ts +++ b/sitemap/src/index.ts @@ -15,11 +15,13 @@ function hydrantUrl(env: Env): string { export interface SyncDeps { subscribe: Subscribe; idleMs: number; + closeMs: number; } const defaultDeps: SyncDeps = { subscribe: subscribeJetstream, idleMs: 3000, + closeMs: 10_000, }; export function parseConfigLimits(env: Env): { maxUrls: number; maxBytes: number } { @@ -136,6 +138,7 @@ export async function syncFromHydrant( collections: Object.keys(COLLECTION_SPECS), maxEvents: batchSize, idleMs: deps.idleMs, + closeMs: deps.closeMs, untilUs: now.getTime() * 1000, maxReconnects: 3, seen, diff --git a/sitemap/src/jetstream.ts b/sitemap/src/jetstream.ts index 167bf5fb7..2eb44b3f3 100644 --- a/sitemap/src/jetstream.ts +++ b/sitemap/src/jetstream.ts @@ -39,47 +39,50 @@ export const subscribeJetstream: Subscribe = (url, cursor, collections, onClose) for (const collection of collections) endpoint.searchParams.append('wantedCollections', collection); const queue: JetstreamEvent[] = []; - let closed = false; + let ended = false; let wake: (() => void) | undefined; let socket: WebSocket | undefined; - const finish = () => { - if (closed) return; - closed = true; - onClose(); + const end = () => { + ended = true; wake?.(); }; + // onClose means the socket is really gone, so it fires after return() too + const gone = () => { + end(); + onClose(); + }; const opened = connect(endpoint).then( (ws) => { socket = ws; ws.addEventListener('message', (msg) => { + if (ended) return; const evt = parseEvent(msg.data); if (!evt) return; queue.push(evt); wake?.(); }); - ws.addEventListener('close', finish); - ws.addEventListener('error', finish); + ws.addEventListener('close', gone); + ws.addEventListener('error', gone); }, (err) => { console.warn('[sitemap] jetstream connect failed:', err); - finish(); + gone(); }, ); return { async next(): Promise> { await opened; - while (queue.length === 0 && !closed) { + while (queue.length === 0 && !ended) { await new Promise((resolve) => (wake = resolve)); } const evt = queue.shift(); return evt ? { value: evt, done: false } : { value: undefined, done: true }; }, async return(): Promise> { - closed = true; - wake?.(); + end(); socket?.close(1000, 'done'); return { value: undefined, done: true }; }, @@ -96,6 +99,8 @@ export interface ReadOptions { idleMs: number; // anything newer happened after the run started, so the backlog is done untilUs: number; + // how long to wait for a closed socket to actually go + closeMs: number; maxReconnects: number; seen: Set; } @@ -134,6 +139,10 @@ export async function readJetstream(opts: ReadOptions, subscribe = subscribeJets } if (outcome !== 'dropped') { + // a worker gets six open connections, and a socket we closed holds one + // until the server's close frame lands behind all the backlog it already + // sent. so wait, or the d1 queries after a few batches queue up forever + await Promise.race([drop, new Promise((resolve) => setTimeout(resolve, opts.closeMs))]); return { events, cursor, full: outcome === 'full' }; } if (++drops > opts.maxReconnects) { diff --git a/sitemap/test/jetstream.test.ts b/sitemap/test/jetstream.test.ts index 543bcf1b8..0010f44c6 100644 --- a/sitemap/test/jetstream.test.ts +++ b/sitemap/test/jetstream.test.ts @@ -11,6 +11,7 @@ function options(overrides: Partial = {}): ReadOptions { collections: ['sh.tangled.repo'], maxEvents: 1000, idleMs: 5, + closeMs: 5, untilUs: Number.MAX_SAFE_INTEGER, maxReconnects: 3, seen: new Set(), diff --git a/sitemap/test/sync.test.ts b/sitemap/test/sync.test.ts index d292eba85..403918cce 100644 --- a/sitemap/test/sync.test.ts +++ b/sitemap/test/sync.test.ts @@ -50,7 +50,7 @@ class MockR2Bucket { } function deps(subscribe: SyncDeps['subscribe']): SyncDeps { - return { subscribe, idleMs: 5 }; + return { subscribe, idleMs: 5, closeMs: 5 }; } function shardOf(db: TestD1, url: string): number | undefined {