From 4f0ce1adaff0e8b6996f0a5abc5624da0a932ea9 Mon Sep 17 00:00:00 2001 From: Florian <45694132+flo-bit@users.noreply.github.com> Date: Thu, 6 Aug 2026 03:55:28 +0200 Subject: [PATCH] Add resumable bootstrap source adapters --- .changeset/bootstrap-sources.md | 2 +- packages/contrail/src/core/backfill.ts | 210 ++---- packages/contrail/src/core/bootstrap.ts | 150 ++++- .../contrail/src/core/jetstream-source.ts | 318 ++++++++++ packages/contrail/src/core/pds-snapshot.ts | 596 ++++++++++++++++++ packages/contrail/src/core/scheduling.ts | 155 +++++ packages/contrail/src/core/sources.ts | 129 +++- packages/contrail/src/core/verification.ts | 133 ++++ packages/contrail/src/index.ts | 3 + .../contrail/tests/bootstrap-sources.test.ts | 98 ++- .../tests/database-bootstrap-target.test.ts | 128 +++- .../tests/jetstream-change-source.test.ts | 189 ++++++ .../tests/pds-snapshot-source.test.ts | 261 ++++++++ 13 files changed, 2175 insertions(+), 197 deletions(-) create mode 100644 packages/contrail/src/core/jetstream-source.ts create mode 100644 packages/contrail/src/core/pds-snapshot.ts create mode 100644 packages/contrail/src/core/scheduling.ts create mode 100644 packages/contrail/src/core/verification.ts create mode 100644 packages/contrail/tests/jetstream-change-source.test.ts create mode 100644 packages/contrail/tests/pds-snapshot-source.test.ts diff --git a/.changeset/bootstrap-sources.md b/.changeset/bootstrap-sources.md index 65c7198..7952067 100644 --- a/.changeset/bootstrap-sources.md +++ b/.changeset/bootstrap-sources.md @@ -2,4 +2,4 @@ "@atmo-dev/contrail": minor --- -Add source-neutral snapshot and ordered-change contracts, capture-first bootstrap orchestration, and a database-backed target that commits projection progress atomically for fresh generations. Persist source continuity epochs, and anchor legacy PDS discovery before its first relay request so newly created repositories remain recoverable through Jetstream replay. +Add source-neutral snapshot and ordered-change contracts, capture-first bootstrap orchestration, and a database-backed target that commits projection progress atomically for fresh generations. Persist source continuity epochs and the capture mark before snapshot preparation. Add resumable, host-aware relay/PDS snapshots plus source-confirmed Jetstream marks, bounded ordered replay, retention-expiry detection, required source-semantics gates, durable bounded failure categories, and aggregate candidate verification before completion. diff --git a/packages/contrail/src/core/backfill.ts b/packages/contrail/src/core/backfill.ts index 85383bb..9873c7b 100644 --- a/packages/contrail/src/core/backfill.ts +++ b/packages/contrail/src/core/backfill.ts @@ -23,6 +23,7 @@ import { heartbeatBackfillRun, tryStartBackfillRun, } from "./status"; +import { createStreamingHostScheduler, drainQueue } from "./scheduling"; const PAGE_SIZE = 100; const DEFAULT_MAX_ATTEMPTS = 1; @@ -121,11 +122,16 @@ async function withRetry( fn: (signal: AbortSignal) => Promise, label: string, maxRetries = 3, - timeoutMs = REQUEST_TIMEOUT_MS + timeoutMs = REQUEST_TIMEOUT_MS, + parentSignal?: AbortSignal, ): Promise { let lastError: unknown; for (let attempt = 0; attempt <= maxRetries; attempt++) { + if (parentSignal?.aborted) throw parentSignal.reason; const controller = new AbortController(); + const abort = () => controller.abort(parentSignal?.reason); + parentSignal?.addEventListener("abort", abort, { once: true }); + if (parentSignal?.aborted) abort(); const timeout = setTimeout( () => controller.abort(new Error(`Timeout: ${label}`)), timeoutMs @@ -133,13 +139,28 @@ async function withRetry( try { return await fn(controller.signal); } catch (err) { + if (parentSignal?.aborted) throw parentSignal.reason; lastError = err; } finally { clearTimeout(timeout); + parentSignal?.removeEventListener("abort", abort); } if (attempt < maxRetries) { - const delay = Math.min(1000 * 2 ** attempt, 10000); - await new Promise((r) => setTimeout(r, delay)); + const milliseconds = Math.min(1000 * 2 ** attempt, 10000); + await new Promise((resolve, reject) => { + const finish = () => { + parentSignal?.removeEventListener("abort", cancel); + resolve(); + }; + const timer = setTimeout(finish, milliseconds); + const cancel = () => { + clearTimeout(timer); + parentSignal?.removeEventListener("abort", cancel); + reject(parentSignal?.reason ?? new Error("Retry cancelled")); + }; + parentSignal?.addEventListener("abort", cancel, { once: true }); + if (parentSignal?.aborted) cancel(); + }); } } throw lastError; @@ -492,158 +513,6 @@ export interface BackfillAllOptions { onMetrics?: (metrics: BackfillRunMetrics) => void; } -/** Keep a fixed number of jobs active without batch barriers. Returning true - * from consume requeues only that item, allowing paginated repositories to - * yield fairly while completed slots refill immediately. */ -async function drainQueue( - items: TItem[], - concurrency: number, - run: (item: TItem) => Promise, - consume: (result: TResult) => boolean | void -): Promise { - if (items.length === 0) return; - const queue = [...items]; - let nextIndex = 0; - let active = 0; - let settled = false; - - await new Promise((resolve, reject) => { - const pump = () => { - if (settled) return; - while (active < concurrency && nextIndex < queue.length) { - const item = queue[nextIndex++]; - active++; - run(item).then( - (result) => { - active--; - if (consume(result) === true) queue.push(item); - if (active === 0 && nextIndex >= queue.length) { - settled = true; - resolve(); - } else { - pump(); - } - }, - (error) => { - settled = true; - reject(error); - } - ); - } - }; - pump(); - }); -} - -function createStreamingHostScheduler( - hostConcurrency: number, - didsPerHost: number, - run: (pds: string, did: string) => Promise, - consume: (result: TResult) => boolean | void -): { - add(pds: string, did: string): void; - finish(): Promise; -} { - type HostState = { - pending: string[]; - active: number; - queued: boolean; - running: boolean; - }; - const hosts = new Map(); - const waiting: string[] = []; - let activeHosts = 0; - let producerDone = false; - let settled = false; - let resolveFinished!: () => void; - let rejectFinished!: (error: unknown) => void; - const finished = new Promise((resolve, reject) => { - resolveFinished = resolve; - rejectFinished = reject; - }); - - const fail = (error: unknown) => { - if (settled) return; - settled = true; - rejectFinished(error); - }; - - const maybeFinish = () => { - if (settled || !producerDone) return; - if (activeHosts === 0 && waiting.length === 0) { - settled = true; - resolveFinished(); - } - }; - - const pumpHosts = () => { - if (settled) return; - while (activeHosts < hostConcurrency && waiting.length > 0) { - const pds = waiting.shift()!; - const state = hosts.get(pds)!; - state.queued = false; - if (state.running || state.pending.length === 0) continue; - state.running = true; - activeHosts++; - pumpDids(pds, state); - } - maybeFinish(); - }; - - const finishHost = (state: HostState) => { - if (!state.running) return; - state.running = false; - activeHosts--; - pumpHosts(); - }; - - function pumpDids(pds: string, state: HostState): void { - if (settled) return; - while (state.active < didsPerHost && state.pending.length > 0) { - const did = state.pending.shift()!; - state.active++; - run(pds, did).then( - (result) => { - state.active--; - if (consume(result) === true) state.pending.push(did); - pumpDids(pds, state); - }, - fail - ); - } - if (state.active === 0 && state.pending.length === 0) { - finishHost(state); - } - } - - return { - add(pds, did) { - if (producerDone) throw new Error("Cannot add work after scheduler finish"); - let state = hosts.get(pds); - if (!state) { - state = { pending: [], active: 0, queued: false, running: false }; - hosts.set(pds, state); - } - state.pending.push(did); - if (state.running) { - pumpDids(pds, state); - } else if (!state.queued) { - state.queued = true; - waiting.push(pds); - pumpHosts(); - } - }, - finish() { - producerDone = true; - for (const [pds, state] of hosts) { - if (state.running) pumpDids(pds, state); - } - pumpHosts(); - return finished; - }, - }; -} - async function flushIngestDiagnostics( db: Database, counts: IngestDiagnosticCounts, @@ -1035,7 +904,8 @@ interface DiscoveryPage { async function fetchPage( relay: string, collection: string, - cursor?: string + cursor?: string, + parentSignal?: AbortSignal, ): Promise { const url = new URL( `/xrpc/com.atproto.sync.listReposByCollection`, @@ -1067,7 +937,10 @@ async function fetchPage( } return body as DiscoveryPage; }, - `fetchPage(${relay}, ${collection})` + `fetchPage(${relay}, ${collection})`, + 3, + REQUEST_TIMEOUT_MS, + parentSignal, ); } @@ -1160,10 +1033,18 @@ async function ensureDiscoveryRows( } } +export interface DiscoverDIDsOptions { + /** Legacy direct callers anchor the live cursor before discovery. A bootstrap + * coordinator already owns a separate durable capture position. */ + captureReplayBoundary?: boolean; + signal?: AbortSignal; +} + export async function discoverDIDs( db: Database, config: ContrailConfig, - deadline: number + deadline: number, + options: DiscoverDIDsOptions = {}, ): Promise { const collections = getDiscoverableNsids(config); const relays = config.relays ?? DEFAULT_RELAYS; @@ -1172,12 +1053,15 @@ export async function discoverDIDs( // Capture before relay discovery as well as before PDS crawling. Otherwise a // repository created after its relay page was scanned but before the later // PDS phase could fall before the live cursor and disappear from both paths. - await ensureInitialReplayBoundary(db); + if (options.captureReplayBoundary !== false) { + await ensureInitialReplayBoundary(db); + } const discovered: string[] = []; await ensureDiscoveryRows(db, collections, relays); for (const collection of collections) { + if (options.signal?.aborted) throw options.signal.reason; if (Date.now() >= deadline) break; let data: DiscoveryPage | null = null; @@ -1200,10 +1084,16 @@ export async function discoverDIDs( if (row?.next_retry_at && row.next_retry_at > Date.now()) continue; try { - data = await fetchPage(r, collection, row?.cursor ?? undefined); + data = await fetchPage( + r, + collection, + row?.cursor ?? undefined, + options.signal, + ); relay = r; break; } catch (error) { + if (options.signal?.aborted) throw options.signal.reason; await markDiscoveryFailed( db, collection, diff --git a/packages/contrail/src/core/bootstrap.ts b/packages/contrail/src/core/bootstrap.ts index 5a8f388..df39997 100644 --- a/packages/contrail/src/core/bootstrap.ts +++ b/packages/contrail/src/core/bootstrap.ts @@ -2,7 +2,13 @@ import type { ContrailConfig, Database, IngestEvent, Statement } from "./types"; import { recordTimeUs, createIngestEvent, ingestRecords } from "./ingest"; import { getDependentNsids, recordsTableName } from "./types"; import { rebuildDerivedProjections } from "./db/records"; +import { + BOOTSTRAP_VERIFICATION_META_KEY, + BootstrapVerificationError, + verifyBootstrapCandidate, +} from "./verification"; import type { + BootstrapFailureCategory, BootstrapRunState, BootstrapTarget, MutationBatch, @@ -13,6 +19,23 @@ import type { } from "./sources"; const BOOTSTRAP_STATE_ID = 1; +/** Schema 9 predates the capture-before-prepare split. Keep its non-null + * snapshot column and phase constraint compatible by using a private sentinel. */ +const PREPARING_SNAPSHOT_JSON = "null"; +const BOOTSTRAP_FAILURE_META_KEY = "bootstrap_last_failure"; +const BOOTSTRAP_FAILURE_CATEGORIES = new Set([ + "snapshot-incomplete", + "catchup-incomplete", + "source-history-expired", + "verification-failed", + "bootstrap-failed", +]); + +export interface BootstrapFailureReport { + category: BootstrapFailureCategory; + failedAt: number; + attempts: number; +} interface BootstrapStateRow { phase: BootstrapRunState["phase"]; @@ -80,6 +103,38 @@ function parsePreparedSnapshot(serialized: string): PreparedSnapshot { return value as unknown as PreparedSnapshot; } +export async function getBootstrapFailure( + db: Database, +): Promise { + const row = await db + .prepare("SELECT value FROM _contrail_meta WHERE key = ?") + .bind(BOOTSTRAP_FAILURE_META_KEY) + .first<{ value: string }>(); + if (!row) return null; + let value: unknown; + try { + value = JSON.parse(row.value); + } catch { + throw new Error("Durable bootstrap failure is not valid JSON"); + } + if ( + !isObject(value) || + typeof value.category !== "string" || + !BOOTSTRAP_FAILURE_CATEGORIES.has( + value.category as BootstrapFailureCategory, + ) || + typeof value.failedAt !== "number" || + !Number.isSafeInteger(value.failedAt) || + value.failedAt < 0 || + typeof value.attempts !== "number" || + !Number.isSafeInteger(value.attempts) || + value.attempts < 1 + ) { + throw new Error("Durable bootstrap failure is malformed"); + } + return value as unknown as BootstrapFailureReport; +} + function sourceTimeUs(value: number, label: string): number { if (!Number.isSafeInteger(value) || value < 0) { throw new TypeError(`${label} must be a non-negative safe microsecond value`); @@ -126,7 +181,8 @@ export class DatabaseBootstrapTarget implements BootstrapTarget { if (!(["snapshot", "catchup", "complete"] as string[]).includes(row.phase)) { throw new Error(`Invalid durable bootstrap phase: ${row.phase}`); } - const snapshot = parsePreparedSnapshot(row.snapshot_json); + const preparing = row.snapshot_json === PREPARING_SNAPSHOT_JSON; + const snapshot = preparing ? null : parsePreparedSnapshot(row.snapshot_json); const progressRows = await this.db .prepare( "SELECT partition, cursor, completed FROM bootstrap_snapshot_progress WHERE bootstrap_id = ? ORDER BY partition", @@ -141,7 +197,7 @@ export class DatabaseBootstrapTarget implements BootstrapTarget { ); if (!captureFrom) throw new Error("Durable bootstrap capture is missing"); return { - phase: row.phase, + phase: preparing ? "preparing" : row.phase, snapshot, captureFrom, snapshotProgress: (progressRows.results ?? []).map((item) => ({ @@ -165,10 +221,7 @@ export class DatabaseBootstrapTarget implements BootstrapTarget { }; } - async begin( - snapshot: PreparedSnapshot, - captureFrom: SourcePosition, - ): Promise { + async beginCapture(captureFrom: SourcePosition): Promise { const now = Date.now(); await this.db .prepare( @@ -179,7 +232,7 @@ export class DatabaseBootstrapTarget implements BootstrapTarget { ) .bind( BOOTSTRAP_STATE_ID, - JSON.stringify(snapshot), + PREPARING_SNAPSHOT_JSON, captureFrom.source, captureFrom.epoch, captureFrom.cursor, @@ -189,6 +242,33 @@ export class DatabaseBootstrapTarget implements BootstrapTarget { .run(); } + async setSnapshot( + snapshot: PreparedSnapshot, + captureFrom: SourcePosition, + ): Promise { + await this.db + .prepare( + `UPDATE bootstrap_state + SET phase = 'snapshot', snapshot_json = ?, capture_source = ?, + capture_epoch = ?, capture_cursor = ?, updated_at = ? + WHERE id = ? AND phase = 'snapshot' AND snapshot_json = ?`, + ) + .bind( + JSON.stringify(snapshot), + captureFrom.source, + captureFrom.epoch, + captureFrom.cursor, + Date.now(), + BOOTSTRAP_STATE_ID, + PREPARING_SNAPSHOT_JSON, + ) + .run(); + const state = await this.load(); + if (state?.phase !== "snapshot" || !state.snapshot) { + throw new Error("Could not pin the prepared bootstrap snapshot"); + } + } + async applySnapshotBatch( snapshot: PreparedSnapshot, batch: SnapshotBatch, @@ -327,10 +407,37 @@ export class DatabaseBootstrapTarget implements BootstrapTarget { } async complete(): Promise { + const before = await this.load(); + if ( + before?.phase !== "catchup" || + !before.catchupThrough || + !before.changeCheckpoint || + before.catchupThrough.source !== before.changeCheckpoint.source || + before.catchupThrough.epoch !== before.changeCheckpoint.epoch || + before.catchupThrough.cursor !== before.changeCheckpoint.cursor + ) { + throw new Error("Cannot complete bootstrap before catch-up reaches its target"); + } if (this.options.deferDerivedProjections) { await rebuildDerivedProjections(this.db, this.config); } - await this.db + const report = await verifyBootstrapCandidate(this.db, this.config); + const verification = this.db + .prepare( + `INSERT INTO _contrail_meta (key, value) VALUES (?, ?) + ON CONFLICT(key) DO UPDATE SET value = excluded.value`, + ) + .bind(BOOTSTRAP_VERIFICATION_META_KEY, JSON.stringify(report)); + if (!report.ok) { + await verification.run(); + throw new BootstrapVerificationError(report); + } + + const now = Date.now(); + const clearFailure = this.db + .prepare("DELETE FROM _contrail_meta WHERE key = ?") + .bind(BOOTSTRAP_FAILURE_META_KEY); + const completion = this.db .prepare( `UPDATE bootstrap_state SET phase = 'complete', finished_at = ?, updated_at = ? @@ -339,14 +446,30 @@ export class DatabaseBootstrapTarget implements BootstrapTarget { AND change_epoch = catchup_epoch AND change_cursor = catchup_cursor`, ) - .bind(Date.now(), Date.now(), BOOTSTRAP_STATE_ID) - .run(); + .bind(now, now, BOOTSTRAP_STATE_ID); + await this.db.batch([verification, clearFailure, completion]); const state = await this.load(); if (state?.phase !== "complete") { throw new Error("Cannot complete bootstrap before catch-up reaches its target"); } } + async recordFailure(category: BootstrapFailureCategory): Promise { + const prior = await getBootstrapFailure(this.db); + const report: BootstrapFailureReport = { + category, + failedAt: Date.now(), + attempts: (prior?.attempts ?? 0) + 1, + }; + await this.db + .prepare( + `INSERT INTO _contrail_meta (key, value) VALUES (?, ?) + ON CONFLICT(key) DO UPDATE SET value = excluded.value`, + ) + .bind(BOOTSTRAP_FAILURE_META_KEY, JSON.stringify(report)) + .run(); + } + private async apply( events: IngestEvent[], checkpoints: Statement[], @@ -372,6 +495,13 @@ export class DatabaseBootstrapTarget implements BootstrapTarget { .prepare("SELECT did FROM identities") .all<{ did: string }>(); for (const row of identityRows.results ?? []) known.add(row.did); + // Relay discovery defines acquisition scope before concurrent PDS workers + // necessarily resolve every identity. Load it directly so a dependent + // collection arriving first is not mistaken for an unknown actor. + const backfillRows = await this.db + .prepare("SELECT DISTINCT did FROM backfills") + .all<{ did: string }>(); + for (const row of backfillRows.results ?? []) known.add(row.did); for (const [shortName, collection] of Object.entries( this.config.collections, )) { diff --git a/packages/contrail/src/core/jetstream-source.ts b/packages/contrail/src/core/jetstream-source.ts new file mode 100644 index 0000000..c7cdaad --- /dev/null +++ b/packages/contrail/src/core/jetstream-source.ts @@ -0,0 +1,318 @@ +import { + JetstreamSubscription, + type JetstreamEvent, + type JetstreamSubscriptionOptions, +} from "@atcute/jetstream"; +import type { ChangeSource, MutationBatch, SourceMutation, SourcePosition } from "./sources"; +import type { ContrailConfig } from "./types"; +import { DEFAULT_JETSTREAMS, jetstreamUrlOption } from "./types"; + +const BATCH_EVENTS = 50; +const DEFAULT_MARK_TIMEOUT_MS = 15_000; +const DEFAULT_READ_IDLE_TIMEOUT_MS = 30_000; +const DEFAULT_REPLAY_OVERLAP_US = 10_000_000; + +interface JetstreamReader extends AsyncIterable {} + +type SubscriptionFactory = ( + options: JetstreamSubscriptionOptions, +) => JetstreamReader; + +export interface JetstreamChangeSourceOptions { + /** Operator-owned continuity epoch for this endpoint set. Change it whenever + * history continuity or cursor meaning may have changed. */ + epoch: string; + /** Additional busy collections used only as ordered watermarks. Their records + * are never projected unless they were also requested by the bootstrap. */ + watermarkCollections?: string[]; + /** Guaranteed source retention. Replay older than this fails before connect. */ + retentionUs: number; + markTimeoutMs?: number; + readIdleTimeoutMs?: number; + replayOverlapUs?: number; + /** @internal Deterministic transport seam for conformance tests. */ + subscriptionFactory?: SubscriptionFactory; +} + +export class SourceHistoryExpiredError extends Error { + constructor(message: string, options?: ErrorOptions) { + super(message, options); + this.name = "SourceHistoryExpiredError"; + } +} + +export class SourceCatchupIncompleteError extends Error { + constructor(message: string, options?: ErrorOptions) { + super(message, options); + this.name = "SourceCatchupIncompleteError"; + } +} + +function cursorNumber(position: SourcePosition, label: string): number { + const value = Number(position.cursor); + if (!Number.isSafeInteger(value) || value < 0) { + throw new Error(`${label} is not a valid Jetstream cursor`); + } + return value; +} + +function closeIterator(iterator: AsyncIterator): void { + Promise.resolve(iterator.return?.()).catch(() => {}); +} + +async function nextWithTimeout( + iterator: AsyncIterator, + timeoutMs: number, + signal?: AbortSignal, +): Promise> { + if (signal?.aborted) throw signal.reason; + const next = iterator.next(); + next.catch(() => {}); + return await new Promise>((resolve, reject) => { + let settled = false; + const cleanup = () => { + clearTimeout(timer); + signal?.removeEventListener("abort", abort); + }; + const succeed = (value: IteratorResult) => { + if (settled) return; + settled = true; + cleanup(); + resolve(value); + }; + const fail = (error: unknown) => { + if (settled) return; + settled = true; + cleanup(); + reject(error); + }; + const abort = () => fail(signal?.reason ?? new Error("Replay cancelled")); + const timer = setTimeout( + () => + fail(new SourceCatchupIncompleteError("Jetstream watermark timed out")), + timeoutMs, + ); + signal?.addEventListener("abort", abort, { once: true }); + if (signal?.aborted) abort(); + next.then(succeed, fail); + }); +} + +function wantedCollections( + requested: string[], + watermarks: string[] | undefined, +): string[] { + return [...new Set([...requested, ...(watermarks ?? [])])].sort(); +} + +function position(source: string, epoch: string, cursor: number): SourcePosition { + return { source, epoch, cursor: String(cursor) }; +} + +function mutationFromEvent( + event: Extract, + source: string, + epoch: string, +): SourceMutation { + const commit = event.commit; + const base = { + uri: `at://${event.did}/${commit.collection}/${commit.rkey}`, + did: event.did, + collection: commit.collection, + rkey: commit.rkey, + revision: commit.rev, + sourceTimeUs: event.time_us, + position: position(source, epoch, event.time_us), + }; + return commit.operation === "delete" + ? { ...base, operation: "delete" } + : { + ...base, + operation: "put", + cid: commit.cid, + value: commit.record, + }; +} + +/** Ordered replay adapter for Jetstream's microsecond cursor protocol. + * + * Jetstream has no separate head endpoint. Marks therefore wait for a real + * event from the same filtered stream, and replay waits for an event strictly + * beyond that mark. Quiet application collections can add a busy watermark + * collection without projecting its records. This avoids claiming catch-up + * from wall clock or a short idle period. */ +export class JetstreamChangeSource implements ChangeSource { + readonly id = "jetstream"; + readonly semantics = { + ordinaryRecords: true, + ordinaryDeletes: true, + accountLifecycle: false, + repositoryReplacement: false, + verifiedCommits: false, + explicitHead: true, + } as const; + private readonly createSubscription: SubscriptionFactory; + + constructor( + private readonly config: ContrailConfig, + private readonly options: JetstreamChangeSourceOptions, + ) { + if (!options.epoch) throw new TypeError("Jetstream source epoch is required"); + if (!Number.isSafeInteger(options.retentionUs) || options.retentionUs <= 0) { + throw new TypeError("Jetstream retentionUs must be a positive safe integer"); + } + this.createSubscription = + options.subscriptionFactory ?? + ((subscriptionOptions) => new JetstreamSubscription(subscriptionOptions)); + } + + async mark(options: { + collections: string[]; + signal?: AbortSignal; + }): Promise { + const subscription = this.createSubscription({ + url: jetstreamUrlOption( + this.config.jetstreams ?? DEFAULT_JETSTREAMS, + ), + wantedCollections: wantedCollections( + options.collections, + this.options.watermarkCollections, + ), + }); + const iterator = subscription[Symbol.asyncIterator](); + try { + const step = await nextWithTimeout( + iterator, + this.options.markTimeoutMs ?? DEFAULT_MARK_TIMEOUT_MS, + options.signal, + ); + if (step.done) { + throw new SourceCatchupIncompleteError( + "Jetstream ended before producing a capture watermark", + ); + } + return position(this.id, this.options.epoch, step.value.time_us); + } catch (error) { + if (options.signal?.aborted) throw options.signal.reason; + if (error instanceof SourceCatchupIncompleteError) throw error; + throw new SourceCatchupIncompleteError( + `Jetstream could not produce a capture watermark: ${String(error)}`, + { cause: error }, + ); + } finally { + closeIterator(iterator); + } + } + + async *read(options: { + collections: string[]; + after: SourcePosition; + through: SourcePosition; + signal?: AbortSignal; + }): AsyncIterable { + this.assertPosition(options.after, "Replay start"); + this.assertPosition(options.through, "Replay target"); + const after = cursorNumber(options.after, "Replay start"); + const through = cursorNumber(options.through, "Replay target"); + if (through < after) { + throw new Error("Jetstream replay target precedes its start"); + } + if (Date.now() * 1000 - after > this.options.retentionUs) { + throw new SourceHistoryExpiredError( + "Jetstream replay start is older than the configured retention guarantee", + ); + } + if (through === after) return; + + const requested = new Set(options.collections); + const overlap = Math.max( + 0, + Math.floor(this.options.replayOverlapUs ?? DEFAULT_REPLAY_OVERLAP_US), + ); + const subscription = this.createSubscription({ + url: jetstreamUrlOption( + this.config.jetstreams ?? DEFAULT_JETSTREAMS, + ), + wantedCollections: wantedCollections( + options.collections, + this.options.watermarkCollections, + ), + cursor: Math.max(0, after - overlap), + }); + const iterator = subscription[Symbol.asyncIterator](); + let mutations: SourceMutation[] = []; + let observed = 0; + let checkpoint = after; + + try { + for (;;) { + const step = await nextWithTimeout( + iterator, + this.options.readIdleTimeoutMs ?? DEFAULT_READ_IDLE_TIMEOUT_MS, + options.signal, + ); + if (step.done) { + throw new SourceCatchupIncompleteError( + "Jetstream ended before the catch-up watermark", + ); + } + const event = step.value; + if (event.time_us <= after) continue; + if (event.time_us < checkpoint) { + throw new SourceCatchupIncompleteError( + "Jetstream replay moved backwards within one continuity epoch", + ); + } + if (event.time_us > through) { + yield { + mutations, + checkpoint: options.through, + caughtUp: true, + }; + return; + } + + checkpoint = event.time_us; + observed++; + if ( + event.kind === "commit" && + requested.has(event.commit.collection) + ) { + mutations.push(mutationFromEvent(event, this.id, this.options.epoch)); + } + if (observed >= BATCH_EVENTS) { + yield { + mutations, + checkpoint: position(this.id, this.options.epoch, checkpoint), + caughtUp: false, + }; + mutations = []; + observed = 0; + } + } + } catch (error) { + if (options.signal?.aborted) throw options.signal.reason; + if ( + error instanceof SourceCatchupIncompleteError || + error instanceof SourceHistoryExpiredError + ) { + throw error; + } + throw new SourceCatchupIncompleteError( + `Jetstream replay failed before the target watermark: ${String(error)}`, + { cause: error }, + ); + } finally { + closeIterator(iterator); + } + } + + private assertPosition(value: SourcePosition, label: string): void { + if (value.source !== this.id || value.epoch !== this.options.epoch) { + throw new Error( + `${label} belongs to ${value.source}/${value.epoch}, expected ` + + `${this.id}/${this.options.epoch}`, + ); + } + } +} diff --git a/packages/contrail/src/core/pds-snapshot.ts b/packages/contrail/src/core/pds-snapshot.ts new file mode 100644 index 0000000..e3e3b29 --- /dev/null +++ b/packages/contrail/src/core/pds-snapshot.ts @@ -0,0 +1,596 @@ +import type {} from "@atcute/atproto"; +import type { Client } from "@atcute/client"; +import { type Did, type Nsid } from "@atcute/lexicons"; +import { + isDid, + isNsid, + parseCanonicalResourceUri, +} from "@atcute/lexicons/syntax"; +import type { ContrailConfig, Database } from "./types"; +import { + DEFAULT_RELAYS, + getCollectionNsids, + getDiscoverableNsids, +} from "./types"; +import { discoverDIDs } from "./backfill"; +import { createStreamingHostScheduler, drainQueue } from "./scheduling"; +import { createPdsClient, getPDS } from "./client"; +import type { + CollectionCoverage, + PreparedSnapshot, + SnapshotBatch, + SnapshotProgress, + SnapshotSource, +} from "./sources"; + +const PAGE_SIZE = 100; +const DEFAULT_REQUEST_TIMEOUT_MS = 10_000; +const DEFAULT_RESOLUTION_CONCURRENCY = 100; +const DEFAULT_PDS_CONCURRENCY = 20; +const DEFAULT_DIDS_PER_PDS = 3; + +export interface PdsSnapshotSourceOptions { + /** Concurrent DID-to-PDS resolutions. Default: 100. */ + concurrency?: number; + /** PDS hosts allowed to fetch concurrently. Default: 20. */ + pdsConcurrency?: number; + /** Repositories allowed to fetch concurrently from one PDS. Default: 3. */ + didsPerPds?: number; + requestTimeoutMs?: number; + maxRetries?: number; +} + +/** A PDS snapshot stays resumable but not ready while any partition fails. */ +export class PdsSnapshotIncompleteError extends Error { + constructor(message: string, options?: ErrorOptions) { + super(message, options); + this.name = "PdsSnapshotIncompleteError"; + } +} + +function incomplete(label: string, error: unknown): PdsSnapshotIncompleteError { + return error instanceof PdsSnapshotIncompleteError + ? error + : new PdsSnapshotIncompleteError(`${label}: ${String(error)}`, { + cause: error, + }); +} + +interface BackfillPartition { + did: string; + collection: string; +} + +interface PartitionWork extends BackfillPartition { + cursor: string | undefined; + complete: boolean; +} + +interface QueuedBatch { + batch: SnapshotBatch; + acknowledge(): void; +} + +function positiveInteger(value: number | undefined, fallback: number): number { + return typeof value === "number" && Number.isFinite(value) && value > 0 + ? Math.floor(value) + : fallback; +} + +class BoundedAsyncQueue implements AsyncIterable { + private readonly items: T[] = []; + private readonly readers: Array<{ + resolve: (result: IteratorResult) => void; + reject: (error: unknown) => void; + }> = []; + private readonly writers: Array<{ + item: T; + resolve: () => void; + reject: (error: unknown) => void; + }> = []; + private closed = false; + private failure: unknown = null; + + constructor(private readonly capacity: number) {} + + push(item: T): Promise { + if (this.failure !== null) return Promise.reject(this.failure); + if (this.closed) return Promise.reject(new Error("Snapshot queue is closed")); + const reader = this.readers.shift(); + if (reader) { + reader.resolve({ value: item, done: false }); + return Promise.resolve(); + } + if (this.items.length < this.capacity) { + this.items.push(item); + return Promise.resolve(); + } + return new Promise((resolve, reject) => { + this.writers.push({ item, resolve, reject }); + }); + } + + close(): void { + if (this.closed || this.failure !== null) return; + this.closed = true; + while (this.readers.length > 0) { + this.readers.shift()!.resolve({ value: undefined, done: true }); + } + } + + fail(error: unknown): void { + if (this.failure !== null || this.closed) return; + this.failure = error; + while (this.readers.length > 0) this.readers.shift()!.reject(error); + while (this.writers.length > 0) this.writers.shift()!.reject(error); + } + + private next(): Promise> { + const item = this.items.shift(); + if (item !== undefined) { + const writer = this.writers.shift(); + if (writer) { + this.items.push(writer.item); + writer.resolve(); + } + return Promise.resolve({ value: item, done: false }); + } + if (this.failure !== null) return Promise.reject(this.failure); + if (this.closed) return Promise.resolve({ value: undefined, done: true }); + return new Promise>((resolve, reject) => { + this.readers.push({ resolve, reject }); + }); + } + + [Symbol.asyncIterator](): AsyncIterator { + return { next: () => this.next() }; + } +} + +function partitionKey(did: string, collection: string): string { + return JSON.stringify([did, collection]); +} + +function snapshotRecord( + item: { uri: string; cid: string; value: unknown }, + expected: BackfillPartition, +) { + let parsed: ReturnType; + try { + parsed = parseCanonicalResourceUri(item.uri); + } catch { + throw new PdsSnapshotIncompleteError( + `PDS returned an invalid record URI for ${expected.did}/${expected.collection}`, + ); + } + if (parsed.repo !== expected.did || parsed.collection !== expected.collection) { + throw new PdsSnapshotIncompleteError( + `PDS returned a record outside ${expected.did}/${expected.collection}`, + ); + } + return { + uri: item.uri, + did: parsed.repo, + collection: parsed.collection, + rkey: parsed.rkey, + cid: item.cid, + value: item.value, + }; +} + +function progressMap( + progress: SnapshotProgress[] | undefined, +): Map { + return new Map((progress ?? []).map((item) => [item.partition, item])); +} + +function configuredCollectionSet(config: ContrailConfig): Set { + return new Set(getCollectionNsids(config)); +} + +function discoverableCollections( + config: ContrailConfig, + requested: ReadonlySet, +): string[] { + return getDiscoverableNsids(config).filter((collection) => + requested.has(collection), + ); +} + +function placeholders(values: unknown[]): string { + return values.map(() => "?").join(","); +} + +async function pendingDiscoveryCount( + db: Database, + collections: string[], +): Promise { + if (collections.length === 0) return 0; + const row = await db + .prepare( + `SELECT COUNT(*) AS count FROM discovery + WHERE collection IN (${placeholders(collections)}) AND completed = 0`, + ) + .bind(...collections) + .first<{ count: number | string }>(); + return Number(row?.count ?? 0); +} + +async function dueDiscoveryCount( + db: Database, + collections: string[], +): Promise { + if (collections.length === 0) return 0; + const row = await db + .prepare( + `SELECT COUNT(*) AS count FROM discovery + WHERE collection IN (${placeholders(collections)}) AND completed = 0 + AND (next_retry_at IS NULL OR next_retry_at <= ?)`, + ) + .bind(...collections, Date.now()) + .first<{ count: number | string }>(); + return Number(row?.count ?? 0); +} + +async function delay(milliseconds: number, signal?: AbortSignal): Promise { + if (signal?.aborted) throw signal.reason; + await new Promise((resolve, reject) => { + const finish = () => { + signal?.removeEventListener("abort", abort); + resolve(); + }; + const timer = setTimeout(finish, milliseconds); + const abort = () => { + clearTimeout(timer); + signal?.removeEventListener("abort", abort); + reject(signal?.reason ?? new Error("Snapshot cancelled")); + }; + signal?.addEventListener("abort", abort, { once: true }); + if (signal?.aborted) abort(); + }); +} + +async function withRetry( + operation: (signal: AbortSignal) => Promise, + options: { + label: string; + parent?: AbortSignal; + timeoutMs: number; + maxRetries: number; + }, +): Promise { + let lastError: unknown; + for (let attempt = 0; attempt <= options.maxRetries; attempt++) { + if (options.parent?.aborted) throw options.parent.reason; + const controller = new AbortController(); + const abort = () => controller.abort(options.parent?.reason); + options.parent?.addEventListener("abort", abort, { once: true }); + if (options.parent?.aborted) abort(); + const timer = setTimeout( + () => controller.abort(new Error(`Timeout: ${options.label}`)), + options.timeoutMs, + ); + try { + return await operation(controller.signal); + } catch (error) { + lastError = error; + } finally { + clearTimeout(timer); + options.parent?.removeEventListener("abort", abort); + } + if (attempt < options.maxRetries) { + await delay(Math.min(1000 * 2 ** attempt, 10_000), options.parent); + } + } + throw lastError; +} + +/** Current-state snapshot provider backed by relay discovery and PDS listRecords. */ +export class PdsSnapshotSource implements SnapshotSource { + readonly id = "pds"; + + constructor( + private readonly db: Database, + private readonly config: ContrailConfig, + private readonly options: PdsSnapshotSourceOptions = {}, + ) {} + + async prepare(options: { + collections: string[]; + signal?: AbortSignal; + }): Promise { + const configured = configuredCollectionSet(this.config); + for (const collection of options.collections) { + if (!configured.has(collection)) { + throw new Error(`PDS snapshot requested unknown collection ${collection}`); + } + if (!isNsid(collection)) { + throw new Error(`PDS snapshot requested invalid collection ${collection}`); + } + } + + const requested = new Set(options.collections); + const discoverable = discoverableCollections(this.config, requested); + const relays = this.config.relays ?? DEFAULT_RELAYS; + if (discoverable.length > 0 && relays.length === 0) { + throw new PdsSnapshotIncompleteError( + "PDS snapshot cannot discover repositories without a relay", + ); + } + + // The bootstrap coordinator persisted its change-source mark before calling + // prepare. Discovery may therefore resume repeatedly without opening a gap. + for (;;) { + if (options.signal?.aborted) throw options.signal.reason; + try { + await discoverDIDs(this.db, this.config, Infinity, { + captureReplayBoundary: false, + signal: options.signal, + }); + } catch (error) { + if (options.signal?.aborted) throw options.signal.reason; + throw incomplete("Relay discovery failed", error); + } + const pending = await pendingDiscoveryCount(this.db, discoverable); + if (pending === 0) break; + if ((await dueDiscoveryCount(this.db, discoverable)) === 0) { + throw new PdsSnapshotIncompleteError( + `${pending} relay discovery partition(s) remain retryable`, + ); + } + } + + const coverage = Object.fromEntries( + options.collections.map((collection) => [ + collection, + { state: "complete" } satisfies CollectionCoverage, + ]), + ); + return { + id: `pds-${Date.now().toString(36)}-${crypto.randomUUID()}`, + provider: this.id, + consistency: "sampled-current-state", + collections: coverage, + semantics: { + ordinaryRecords: true, + ordinaryDeletes: false, + accountLifecycle: false, + repositoryReplacement: false, + verifiedCommits: false, + explicitHead: false, + }, + }; + } + + async *read(options: { + snapshot: PreparedSnapshot; + progress?: SnapshotProgress[]; + signal?: AbortSignal; + }): AsyncIterable { + if (options.snapshot.provider !== this.id) { + throw new Error( + `PDS source cannot read snapshot from ${options.snapshot.provider}`, + ); + } + const collections = Object.keys(options.snapshot.collections); + const rows = await this.partitions(collections); + const progress = progressMap(options.progress); + const byDid = new Map(); + for (const row of rows) { + const prior = progress.get(partitionKey(row.did, row.collection)); + if (prior?.complete) continue; + const work = byDid.get(row.did) ?? []; + work.push({ + ...row, + cursor: prior?.cursor ?? undefined, + complete: false, + }); + byDid.set(row.did, work); + } + + const pdsConcurrency = positiveInteger( + this.options.pdsConcurrency, + DEFAULT_PDS_CONCURRENCY, + ); + const didsPerPds = positiveInteger( + this.options.didsPerPds, + DEFAULT_DIDS_PER_PDS, + ); + const resolutionConcurrency = positiveInteger( + this.options.concurrency, + DEFAULT_RESOLUTION_CONCURRENCY, + ); + const requestTimeoutMs = positiveInteger( + this.options.requestTimeoutMs, + DEFAULT_REQUEST_TIMEOUT_MS, + ); + const maxRetries = + typeof this.options.maxRetries === "number" && + Number.isFinite(this.options.maxRetries) && + this.options.maxRetries >= 0 + ? Math.floor(this.options.maxRetries) + : 3; + const queue = new BoundedAsyncQueue( + Math.max(2, pdsConcurrency * didsPerPds * 2), + ); + const controller = new AbortController(); + const abort = () => controller.abort(options.signal?.reason); + options.signal?.addEventListener("abort", abort, { once: true }); + if (options.signal?.aborted) abort(); + + const producer = (async () => { + const clients = new Map(); + const scheduler = createStreamingHostScheduler( + pdsConcurrency, + didsPerPds, + async (pds, did) => { + if (controller.signal.aborted) throw controller.signal.reason; + let client = clients.get(pds); + if (!client) { + client = createPdsClient(pds); + clients.set(pds, client); + } + const batches: SnapshotBatch[] = []; + const works = byDid.get(did)!; + for (const work of works) { + if (work.complete) continue; + let response; + try { + response = await withRetry( + (signal) => + client!.get("com.atproto.repo.listRecords", { + params: { + repo: did as Did, + collection: work.collection as Nsid, + limit: PAGE_SIZE, + cursor: work.cursor, + }, + signal, + }), + { + label: `listRecords(${did}/${work.collection})`, + parent: controller.signal, + timeoutMs: requestTimeoutMs, + maxRetries, + }, + ); + } catch (error) { + if (controller.signal.aborted) throw controller.signal.reason; + throw incomplete( + `listRecords failed for ${did}/${work.collection}`, + error, + ); + } + if (!response.ok) { + const detail = response.data.message + ? `${response.data.error}: ${response.data.message}` + : response.data.error; + throw new PdsSnapshotIncompleteError( + `listRecords failed for ${did}/${work.collection}: ` + + `${response.status} ${detail}`, + ); + } + + const nextCursor = response.data.cursor ?? null; + if (nextCursor !== null && nextCursor === work.cursor) { + throw new PdsSnapshotIncompleteError( + `listRecords repeated cursor for ${did}/${work.collection}`, + ); + } + work.cursor = nextCursor ?? undefined; + work.complete = nextCursor === null; + batches.push({ + records: response.data.records.map((record) => + snapshotRecord(record, work), + ), + sourceTimeUs: Date.now() * 1000, + progress: { + partition: partitionKey(did, work.collection), + cursor: nextCursor, + complete: work.complete, + }, + done: false, + }); + } + return { + batches, + complete: works.every((work) => work.complete), + }; + }, + async (result) => { + for (const batch of result.batches) { + let acknowledge!: () => void; + let reject!: (error: unknown) => void; + const acknowledged = new Promise((resolve, rejectPromise) => { + acknowledge = resolve; + reject = rejectPromise; + }); + const cancel = () => reject(controller.signal.reason); + controller.signal.addEventListener("abort", cancel, { once: true }); + await queue.push({ batch, acknowledge }); + try { + await acknowledged; + } finally { + controller.signal.removeEventListener("abort", cancel); + } + } + return !result.complete; + }, + ); + + await drainQueue( + [...byDid.keys()], + resolutionConcurrency, + async (did) => { + let pds: string | null | undefined; + try { + pds = await withRetry( + (signal) => getPDS(did as Did, this.db, this.config, signal), + { + label: `getPDS(${did})`, + parent: controller.signal, + timeoutMs: requestTimeoutMs, + maxRetries: Math.min(maxRetries, 1), + }, + ); + } catch (error) { + if (controller.signal.aborted) throw controller.signal.reason; + throw incomplete(`PDS resolution failed for ${did}`, error); + } + if (!pds) { + throw new PdsSnapshotIncompleteError(`PDS not found for ${did}`); + } + return { did, pds: pds.replace(/\/+$/, "") }; + }, + ({ did, pds }) => scheduler.add(pds, did), + ); + await scheduler.finish(); + await queue.push({ + batch: { + records: [], + sourceTimeUs: Date.now() * 1000, + progress: { partition: "snapshot", cursor: null, complete: true }, + done: true, + }, + acknowledge() {}, + }); + queue.close(); + })().catch((error) => queue.fail(error)); + + try { + for await (const queued of queue) { + if (controller.signal.aborted) throw controller.signal.reason; + yield queued.batch; + queued.acknowledge(); + } + await producer; + } finally { + controller.abort(new Error("PDS snapshot reader closed")); + queue.fail(controller.signal.reason); + options.signal?.removeEventListener("abort", abort); + await producer; + } + } + + private async partitions(collections: string[]): Promise { + if (collections.length === 0) return []; + const rows = await this.db + .prepare( + `SELECT did, collection FROM backfills + WHERE collection IN (${placeholders(collections)}) + ORDER BY did, collection`, + ) + .bind(...collections) + .all(); + const valid: BackfillPartition[] = []; + for (const row of rows.results ?? []) { + if (!isDid(row.did) || !isNsid(row.collection)) { + throw new PdsSnapshotIncompleteError( + `Invalid discovered PDS partition ${row.did}/${row.collection}`, + ); + } + valid.push(row); + } + return valid; + } +} diff --git a/packages/contrail/src/core/scheduling.ts b/packages/contrail/src/core/scheduling.ts new file mode 100644 index 0000000..96ebf32 --- /dev/null +++ b/packages/contrail/src/core/scheduling.ts @@ -0,0 +1,155 @@ +/** Internal fair-work scheduling shared by legacy and generation PDS crawlers. */ + +/** Keep a fixed number of jobs active without batch barriers. Returning true + * from consume requeues only that item, allowing paginated work to yield fairly. */ +export async function drainQueue( + items: TItem[], + concurrency: number, + run: (item: TItem) => Promise, + consume: (result: TResult) => boolean | void, +): Promise { + if (items.length === 0) return; + const queue = [...items]; + let nextIndex = 0; + let active = 0; + let settled = false; + + await new Promise((resolve, reject) => { + const pump = () => { + if (settled) return; + while (active < concurrency && nextIndex < queue.length) { + const item = queue[nextIndex++]; + active++; + run(item).then( + (result) => { + active--; + if (consume(result) === true) queue.push(item); + if (active === 0 && nextIndex >= queue.length) { + settled = true; + resolve(); + } else { + pump(); + } + }, + (error) => { + settled = true; + reject(error); + }, + ); + } + }; + pump(); + }); +} + +export function createStreamingHostScheduler( + hostConcurrency: number, + didsPerHost: number, + run: (pds: string, did: string) => Promise, + consume: (result: TResult) => + | boolean + | void + | Promise, +): { + add(pds: string, did: string): void; + finish(): Promise; +} { + type HostState = { + pending: string[]; + active: number; + queued: boolean; + running: boolean; + }; + const hosts = new Map(); + const waiting: string[] = []; + let activeHosts = 0; + let producerDone = false; + let settled = false; + let resolveFinished!: () => void; + let rejectFinished!: (error: unknown) => void; + const finished = new Promise((resolve, reject) => { + resolveFinished = resolve; + rejectFinished = reject; + }); + + const fail = (error: unknown) => { + if (settled) return; + settled = true; + rejectFinished(error); + }; + + const maybeFinish = () => { + if (settled || !producerDone) return; + if (activeHosts === 0 && waiting.length === 0) { + settled = true; + resolveFinished(); + } + }; + + const pumpHosts = () => { + if (settled) return; + while (activeHosts < hostConcurrency && waiting.length > 0) { + const pds = waiting.shift()!; + const state = hosts.get(pds)!; + state.queued = false; + if (state.running || state.pending.length === 0) continue; + state.running = true; + activeHosts++; + pumpDids(pds, state); + } + maybeFinish(); + }; + + const finishHost = (state: HostState) => { + if (!state.running) return; + state.running = false; + activeHosts--; + pumpHosts(); + }; + + function pumpDids(pds: string, state: HostState): void { + if (settled) return; + while (state.active < didsPerHost && state.pending.length > 0) { + const did = state.pending.shift()!; + state.active++; + run(pds, did) + .then(async (result) => { + const requeue = await consume(result); + state.active--; + if (requeue === true) state.pending.push(did); + pumpDids(pds, state); + }) + .catch(fail); + } + if (state.active === 0 && state.pending.length === 0) { + finishHost(state); + } + } + + return { + add(pds, did) { + if (producerDone) throw new Error("Cannot add work after scheduler finish"); + let state = hosts.get(pds); + if (!state) { + state = { pending: [], active: 0, queued: false, running: false }; + hosts.set(pds, state); + } + state.pending.push(did); + if (state.running) { + pumpDids(pds, state); + } else if (!state.queued) { + state.queued = true; + waiting.push(pds); + pumpHosts(); + } + }, + finish() { + producerDone = true; + for (const [pds, state] of hosts) { + if (state.running) pumpDids(pds, state); + } + pumpHosts(); + return finished; + }, + }; +} diff --git a/packages/contrail/src/core/sources.ts b/packages/contrail/src/core/sources.ts index ca86817..69a821c 100644 --- a/packages/contrail/src/core/sources.ts +++ b/packages/contrail/src/core/sources.ts @@ -108,6 +108,7 @@ export interface MutationBatch { export interface ChangeSource { readonly id: string; + readonly semantics: SourceSemantics; /** Return a durable replay coordinate near the current source head. */ mark(options: { collections: string[]; @@ -121,13 +122,17 @@ export interface ChangeSource { }): AsyncIterable; } -export type BootstrapPhase = "snapshot" | "catchup" | "complete"; +export type BootstrapPhase = + | "preparing" + | "snapshot" + | "catchup" + | "complete"; /** Durable coordinator state. Snapshot progress and mutation checkpoints are * separate because they belong to different cursor namespaces. */ export interface BootstrapRunState { phase: BootstrapPhase; - snapshot: PreparedSnapshot; + snapshot: PreparedSnapshot | null; captureFrom: SourcePosition; snapshotProgress: SnapshotProgress[]; snapshotComplete: boolean; @@ -137,9 +142,19 @@ export interface BootstrapRunState { /** Projection-owned persistence seam. Implementations commit records and the * accompanying progress/checkpoint atomically in the destination database. */ +export type BootstrapFailureCategory = + | "snapshot-incomplete" + | "catchup-incomplete" + | "source-history-expired" + | "verification-failed" + | "bootstrap-failed"; + export interface BootstrapTarget { load(): Promise; - begin(snapshot: PreparedSnapshot, captureFrom: SourcePosition): Promise; + /** Persist the capture boundary before snapshot preparation performs network work. */ + beginCapture(captureFrom: SourcePosition): Promise; + /** Pin the prepared snapshot, optionally replacing capture with its own boundary. */ + setSnapshot(snapshot: PreparedSnapshot, captureFrom: SourcePosition): Promise; applySnapshotBatch( snapshot: PreparedSnapshot, batch: SnapshotBatch, @@ -147,6 +162,8 @@ export interface BootstrapTarget { beginCatchup(through: SourcePosition): Promise; applyMutationBatch(batch: MutationBatch): Promise; complete(): Promise; + /** Persist only a bounded category; raw upstream errors stay private. */ + recordFailure?(category: BootstrapFailureCategory): Promise; } export interface BootstrapResult { @@ -179,6 +196,40 @@ function assertCompatiblePosition( } } +function assertSemantics( + snapshot: PreparedSnapshot, + changes: ChangeSource, + required: Partial | undefined, +): void { + if (!snapshot.semantics.ordinaryRecords) { + throw new Error(`Snapshot ${snapshot.id} does not guarantee ordinary records`); + } + for (const capability of [ + "ordinaryRecords", + "ordinaryDeletes", + "explicitHead", + ] as const) { + if (!changes.semantics[capability]) { + throw new Error( + `Change source ${changes.id} does not guarantee ${capability}`, + ); + } + } + for (const capability of Object.keys(required ?? {}) as Array< + keyof SourceSemantics + >) { + if (required?.[capability] !== true) continue; + if ( + !snapshot.semantics[capability] || + !changes.semantics[capability] + ) { + throw new Error( + `Bootstrap sources do not jointly guarantee required ${capability}`, + ); + } + } +} + function assertCoverage( snapshot: PreparedSnapshot, collections: string[], @@ -210,12 +261,13 @@ function assertCoverage( * A prepared point-in-time snapshot may supply its own upstream boundary; * sampled scans use the position marked before preparation begins. */ -export async function bootstrapFreshProjection(options: { +async function runBootstrapFreshProjection(options: { collections: string[]; snapshotSource: SnapshotSource; changeSource: ChangeSource; target: BootstrapTarget; allowPartial?: boolean; + requiredSemantics?: Partial; signal?: AbortSignal; }): Promise { const { @@ -228,40 +280,49 @@ export async function bootstrapFreshProjection(options: { let state = await target.load(); if (!state) { - // Mark before snapshot preparation so even a provider that performs relay - // discovery while preparing cannot open a capture gap. + // Persist the mark before snapshot preparation so discovery failure cannot + // move a later retry past repositories or records observed in the meantime. const marked = await changeSource.mark({ collections, signal }); - const snapshot = await snapshotSource.prepare({ collections, signal }); - assertCoverage(snapshot, collections, options.allowPartial === true); - const captureFrom = snapshot.through ?? marked; - assertCompatiblePosition(captureFrom, marked, "Snapshot boundary"); - await target.begin(snapshot, captureFrom); + await target.beginCapture(marked); state = { - phase: "snapshot", - snapshot, - captureFrom, + phase: "preparing", + snapshot: null, + captureFrom: marked, snapshotProgress: [], snapshotComplete: false, catchupThrough: null, changeCheckpoint: null, }; + } + + if (!state.snapshot) { + const snapshot = await snapshotSource.prepare({ collections, signal }); + assertCoverage(snapshot, collections, options.allowPartial === true); + const captureFrom = snapshot.through ?? state.captureFrom; + assertCompatiblePosition(captureFrom, state.captureFrom, "Snapshot boundary"); + await target.setSnapshot(snapshot, captureFrom); + state.snapshot = snapshot; + state.captureFrom = captureFrom; + state.phase = "snapshot"; } else { assertCoverage(state.snapshot, collections, options.allowPartial === true); } + const snapshot = state.snapshot; + assertSemantics(snapshot, changeSource, options.requiredSemantics); if (!state.snapshotComplete) { let sawDone = false; for await (const batch of snapshotSource.read({ - snapshot: state.snapshot, + snapshot, ...(state.snapshotProgress.length === 0 ? {} : { progress: state.snapshotProgress }), signal, })) { if (sawDone) { - throw new Error(`Snapshot ${state.snapshot.id} emitted data after done`); + throw new Error(`Snapshot ${snapshot.id} emitted data after done`); } - await target.applySnapshotBatch(state.snapshot, batch); + await target.applySnapshotBatch(snapshot, batch); const progressIndex = state.snapshotProgress.findIndex( (item) => item.partition === batch.progress.partition, ); @@ -271,7 +332,7 @@ export async function bootstrapFreshProjection(options: { sawDone = batch.done; } if (!state.snapshotComplete) { - throw new Error(`Snapshot ${state.snapshot.id} ended before done`); + throw new Error(`Snapshot ${snapshot.id} ended before done`); } } @@ -318,8 +379,38 @@ export async function bootstrapFreshProjection(options: { } return { - snapshot: state.snapshot, + snapshot, captureFrom: state.captureFrom, through: state.catchupThrough, }; } + +function failureCategory(error: unknown): BootstrapFailureCategory { + const name = error instanceof Error ? error.name : ""; + if (name === "PdsSnapshotIncompleteError") return "snapshot-incomplete"; + if (name === "SourceCatchupIncompleteError") return "catchup-incomplete"; + if (name === "SourceHistoryExpiredError") return "source-history-expired"; + if (name === "BootstrapVerificationError") return "verification-failed"; + return "bootstrap-failed"; +} + +export async function bootstrapFreshProjection(options: { + collections: string[]; + snapshotSource: SnapshotSource; + changeSource: ChangeSource; + target: BootstrapTarget; + allowPartial?: boolean; + requiredSemantics?: Partial; + signal?: AbortSignal; +}): Promise { + try { + return await runBootstrapFreshProjection(options); + } catch (error) { + try { + await options.target.recordFailure?.(failureCategory(error)); + } catch { + // Failure telemetry must not replace the canonical source/projection error. + } + throw error; + } +} diff --git a/packages/contrail/src/core/verification.ts b/packages/contrail/src/core/verification.ts new file mode 100644 index 0000000..8dfaa01 --- /dev/null +++ b/packages/contrail/src/core/verification.ts @@ -0,0 +1,133 @@ +import type { ContrailConfig, Database } from "./types"; +import { recordsTableName } from "./types"; + +export const BOOTSTRAP_VERIFICATION_META_KEY = "bootstrap_verification"; + +export interface BootstrapVerificationCheck { + name: string; + ok: boolean; + failures: number; +} + +export interface BootstrapVerificationReport { + ok: boolean; + verifiedAt: number; + checks: BootstrapVerificationCheck[]; +} + +export class BootstrapVerificationError extends Error { + constructor(readonly report: BootstrapVerificationReport) { + const failed = report.checks + .filter((check) => !check.ok) + .map((check) => `${check.name}=${check.failures}`) + .join(", "); + super(`Bootstrap candidate verification failed: ${failed}`); + this.name = "BootstrapVerificationError"; + } +} + +async function count(db: Database, sql: string, bindings: unknown[] = []) { + const row = await db + .prepare(sql) + .bind(...bindings) + .first<{ count: number | string }>(); + return Number(row?.count ?? 0); +} + +function check(name: string, failures: number): BootstrapVerificationCheck { + return { name, ok: failures === 0, failures }; +} + +/** Aggregate-only integrity checks for an unpublished candidate database. */ +export async function verifyBootstrapCandidate( + db: Database, + config: ContrailConfig, +): Promise { + const checks: BootstrapVerificationCheck[] = []; + checks.push( + check( + "snapshot-partitions", + await count( + db, + "SELECT COUNT(*) AS count FROM bootstrap_snapshot_progress WHERE completed = 0", + ), + ), + ); + + for (const [shortName, collectionConfig] of Object.entries( + config.collections, + )) { + const collection = collectionConfig.collection ?? shortName; + const table = recordsTableName(shortName); + checks.push( + check( + `visible-version:${shortName}`, + await count( + db, + `SELECT COUNT(*) AS count FROM ${table} AS record + LEFT JOIN record_versions AS version ON version.uri = record.uri + WHERE version.uri IS NULL OR version.operation = 'delete' + OR version.collection <> ? OR version.did <> record.did + OR version.rkey <> record.rkey`, + [collection], + ), + ), + ); + checks.push( + check( + `version-visible:${shortName}`, + await count( + db, + `SELECT COUNT(*) AS count FROM record_versions AS version + LEFT JOIN ${table} AS record ON record.uri = version.uri + WHERE version.collection = ? AND version.operation <> 'delete' + AND record.uri IS NULL`, + [collection], + ), + ), + ); + } + + return { + ok: checks.every((item) => item.ok), + verifiedAt: Date.now(), + checks, + }; +} + +export async function getBootstrapVerification( + db: Database, +): Promise { + const row = await db + .prepare("SELECT value FROM _contrail_meta WHERE key = ?") + .bind(BOOTSTRAP_VERIFICATION_META_KEY) + .first<{ value: string }>(); + if (!row) return null; + let value: unknown; + try { + value = JSON.parse(row.value); + } catch { + throw new Error("Durable bootstrap verification is not valid JSON"); + } + const report = value as BootstrapVerificationReport; + if ( + !value || + typeof value !== "object" || + typeof report.ok !== "boolean" || + !Number.isSafeInteger(report.verifiedAt) || + report.verifiedAt < 0 || + !Array.isArray(report.checks) || + !report.checks.every( + (item) => + item && + typeof item === "object" && + typeof item.name === "string" && + typeof item.ok === "boolean" && + Number.isSafeInteger(item.failures) && + item.failures >= 0, + ) + ) { + throw new Error("Durable bootstrap verification is malformed"); + } + return report; +} diff --git a/packages/contrail/src/index.ts b/packages/contrail/src/index.ts index 038290f..c1d828e 100644 --- a/packages/contrail/src/index.ts +++ b/packages/contrail/src/index.ts @@ -19,6 +19,9 @@ export type { ResolvedIdentity } from "./core/client"; export * from "./core/ingest"; export * from "./core/sources"; export * from "./core/bootstrap"; +export * from "./core/verification"; +export * from "./core/pds-snapshot"; +export * from "./core/jetstream-source"; export * from "./core/jetstream"; export * from "./core/persistent"; export * from "./core/backfill"; diff --git a/packages/contrail/tests/bootstrap-sources.test.ts b/packages/contrail/tests/bootstrap-sources.test.ts index 1f32ea8..7abef53 100644 --- a/packages/contrail/tests/bootstrap-sources.test.ts +++ b/packages/contrail/tests/bootstrap-sources.test.ts @@ -72,10 +72,10 @@ class MemoryTarget implements BootstrapTarget { return this.state ? structuredClone(this.state) : null; } - async begin(snapshot: PreparedSnapshot, captureFrom: SourcePosition) { + async beginCapture(captureFrom: SourcePosition) { this.state = { - phase: "snapshot", - snapshot: structuredClone(snapshot), + phase: "preparing", + snapshot: null, captureFrom: structuredClone(captureFrom), snapshotProgress: [], snapshotComplete: false, @@ -84,6 +84,12 @@ class MemoryTarget implements BootstrapTarget { }; } + async setSnapshot(snapshot: PreparedSnapshot, captureFrom: SourcePosition) { + this.state!.phase = "snapshot"; + this.state!.snapshot = structuredClone(snapshot); + this.state!.captureFrom = structuredClone(captureFrom); + } + async applySnapshotBatch(_snapshot: PreparedSnapshot, batch: SnapshotBatch) { for (const item of batch.records) this.records.set(item.uri, item); this.snapshotBatches++; @@ -151,6 +157,7 @@ function changeSource(options: { let markIndex = 0; return { id: "changes", + semantics, async mark() { const cursor = options.marks[markIndex++]; if (cursor === undefined) throw new Error("Unexpected mark"); @@ -281,6 +288,58 @@ describe("bootstrap source orchestration", () => { ]); }); + it("retains the original capture mark when snapshot preparation is retried", async () => { + const calls: string[] = []; + let attempts = 0; + const source: SnapshotSource = { + id: "retrying-snapshot", + async prepare() { + calls.push(`prepare:${++attempts}`); + if (attempts === 1) throw new Error("discovery unavailable"); + return prepared(); + }, + async *read() { + yield { + records: [], + sourceTimeUs: 1, + progress: { partition: "main", cursor: null, complete: true }, + done: true, + }; + }, + }; + const changes = changeSource({ marks: [1, 3], mutations: [], calls }); + const target = new MemoryTarget(); + + await expect( + bootstrapFreshProjection({ + collections: [COLLECTION], + snapshotSource: source, + changeSource: changes, + target, + }), + ).rejects.toThrow("discovery unavailable"); + expect(target.state).toMatchObject({ + phase: "preparing", + captureFrom: position(1), + }); + + await bootstrapFreshProjection({ + collections: [COLLECTION], + snapshotSource: source, + changeSource: changes, + target, + }); + + expect(calls).toEqual([ + "mark:1", + "prepare:1", + "prepare:2", + "mark:3", + "changes:1-3", + ]); + expect(target.state?.phase).toBe("complete"); + }); + it("resumes a pinned snapshot from its last committed progress", async () => { const reads: Array = []; let failSecondBatch = true; @@ -398,7 +457,33 @@ describe("bootstrap source orchestration", () => { target, }), ).rejects.toThrow("Snapshot boundary belongs to test-stream/old"); - expect(target.state).toBeNull(); + expect(target.state).toMatchObject({ + phase: "preparing", + captureFrom: position(10), + }); + }); + + it("blocks readiness when required source semantics are not guaranteed", async () => { + const target = new MemoryTarget(); + const changes = changeSource({ marks: [1], mutations: [] }); + + await expect( + bootstrapFreshProjection({ + collections: [COLLECTION], + snapshotSource: snapshotSource({ + snapshot: prepared(), + batches: () => [], + }), + changeSource: changes, + target, + requiredSemantics: { accountLifecycle: true }, + }), + ).rejects.toThrow("required accountLifecycle"); + + expect(target.state).toMatchObject({ + phase: "snapshot", + snapshotComplete: false, + }); }); it("refuses partial and gapped coverage by default", async () => { @@ -420,7 +505,10 @@ describe("bootstrap source orchestration", () => { target, }), ).rejects.toThrow(coverage.reason); - expect(target.state).toBeNull(); + expect(target.state).toMatchObject({ + phase: "preparing", + captureFrom: position(1), + }); } }); }); diff --git a/packages/contrail/tests/database-bootstrap-target.test.ts b/packages/contrail/tests/database-bootstrap-target.test.ts index c269a44..6caac93 100644 --- a/packages/contrail/tests/database-bootstrap-target.test.ts +++ b/packages/contrail/tests/database-bootstrap-target.test.ts @@ -1,7 +1,11 @@ import { describe, expect, it } from "vitest"; import { + BootstrapVerificationError, DatabaseBootstrapTarget, + PdsSnapshotIncompleteError, bootstrapFreshProjection, + getBootstrapFailure, + getBootstrapVerification, initSchema, queryRecords, resolveConfig, @@ -61,6 +65,89 @@ function config() { } describe("database bootstrap target", () => { + it("durably represents capture before a snapshot descriptor exists", async () => { + const resolved = config(); + const db = createSqliteDatabase(":memory:"); + await initSchema(db, resolved); + const target = new DatabaseBootstrapTarget(db, resolved); + + await target.beginCapture(sourcePosition(1)); + + expect(await target.load()).toMatchObject({ + phase: "preparing", + snapshot: null, + captureFrom: sourcePosition(1), + snapshotProgress: [], + snapshotComplete: false, + }); + }); + + it("persists a bounded failure category without raw upstream details", async () => { + const resolved = config(); + const db = createSqliteDatabase(":memory:"); + await initSchema(db, resolved); + const target = new DatabaseBootstrapTarget(db, resolved); + const changes: ChangeSource = { + id: "jetstream", + semantics: snapshot().semantics, + async mark() { + return sourcePosition(1); + }, + async *read() {}, + }; + const failingSnapshot: SnapshotSource = { + id: "pds", + async prepare() { + throw new PdsSnapshotIncompleteError( + `private upstream failed for ${DID}`, + ); + }, + async *read() {}, + }; + + await expect( + bootstrapFreshProjection({ + collections: [COLLECTION], + snapshotSource: failingSnapshot, + changeSource: changes, + target, + }), + ).rejects.toThrow("private upstream failed"); + + expect(await getBootstrapFailure(db)).toMatchObject({ + category: "snapshot-incomplete", + attempts: 1, + }); + const raw = await db + .prepare("SELECT value FROM _contrail_meta WHERE key = ?") + .bind("bootstrap_last_failure") + .first<{ value: string }>(); + expect(raw?.value).not.toContain(DID); + expect((await target.load())?.phase).toBe("preparing"); + + const recoveredSnapshot: SnapshotSource = { + id: "pds", + async prepare() { + return snapshot(); + }, + async *read() { + yield { + records: [], + sourceTimeUs: 1, + progress: { partition: "pds", cursor: null, complete: true }, + done: true, + }; + }, + }; + await bootstrapFreshProjection({ + collections: [COLLECTION], + snapshotSource: recoveredSnapshot, + changeSource: changes, + target, + }); + expect(await getBootstrapFailure(db)).toBeNull(); + }); + it("projects a sampled snapshot and ordered tail with durable epochs", async () => { const resolved = config(); const db = createSqliteDatabase(":memory:"); @@ -83,6 +170,7 @@ describe("database bootstrap target", () => { let mark = 0; const changeSource: ChangeSource = { id: "jetstream", + semantics: prepared.semantics, async mark() { mark++; return sourcePosition(mark === 1 ? 1 : 3); @@ -147,6 +235,40 @@ describe("database bootstrap target", () => { source_epoch: "epoch-one", source_cursor: "2", }); + expect(await getBootstrapVerification(db)).toMatchObject({ ok: true }); + }); + + it("blocks completion and persists aggregate verification failures", async () => { + const resolved = config(); + const db = createSqliteDatabase(":memory:"); + await initSchema(db, resolved); + const target = new DatabaseBootstrapTarget(db, resolved); + const prepared = snapshot(); + await target.beginCapture(sourcePosition(1)); + await target.setSnapshot(prepared, sourcePosition(1)); + await target.applySnapshotBatch(prepared, { + records: [record("a", "orphan")], + sourceTimeUs: 1, + progress: { partition: "pds", cursor: null, complete: true }, + done: true, + }); + await target.beginCatchup(sourcePosition(1)); + await db + .prepare("DELETE FROM record_versions WHERE uri = ?") + .bind(`at://${DID}/${COLLECTION}/a`) + .run(); + + await expect(target.complete()).rejects.toBeInstanceOf( + BootstrapVerificationError, + ); + + expect((await target.load())?.phase).toBe("catchup"); + expect(await getBootstrapVerification(db)).toMatchObject({ + ok: false, + checks: expect.arrayContaining([ + { name: "visible-version:event", ok: false, failures: 1 }, + ]), + }); }); it("rolls projection back when snapshot progress cannot commit", async () => { @@ -155,7 +277,8 @@ describe("database bootstrap target", () => { await initSchema(db, resolved); const target = new DatabaseBootstrapTarget(db, resolved); const prepared = snapshot(); - await target.begin(prepared, sourcePosition(1)); + await target.beginCapture(sourcePosition(1)); + await target.setSnapshot(prepared, sourcePosition(1)); await db .prepare( `CREATE TRIGGER fail_bootstrap_progress @@ -200,7 +323,8 @@ describe("database bootstrap target", () => { await initSchema(db, resolved); const target = new DatabaseBootstrapTarget(db, resolved); const prepared = snapshot(); - await target.begin(prepared, sourcePosition(1)); + await target.beginCapture(sourcePosition(1)); + await target.setSnapshot(prepared, sourcePosition(1)); await target.applySnapshotBatch(prepared, { records: [], sourceTimeUs: 1, diff --git a/packages/contrail/tests/jetstream-change-source.test.ts b/packages/contrail/tests/jetstream-change-source.test.ts new file mode 100644 index 0000000..09e32dc --- /dev/null +++ b/packages/contrail/tests/jetstream-change-source.test.ts @@ -0,0 +1,189 @@ +import type { + JetstreamEvent, + JetstreamSubscriptionOptions, +} from "@atcute/jetstream"; +import { describe, expect, it } from "vitest"; +import { + JetstreamChangeSource, + SourceHistoryExpiredError, + resolveConfig, + type SourcePosition, +} from "../src/index"; + +const COLLECTION = "com.example.event"; +const WATERMARK = "app.bsky.feed.post"; +const DID = "did:plc:jetstream-source"; + +function commit( + timeUs: number, + collection: string, + rkey: string, + operation: "create" | "delete" = "create", +): JetstreamEvent { + return { + kind: "commit", + did: DID, + time_us: timeUs, + commit: + operation === "delete" + ? { + operation, + rev: "3kzfcijpj2z2a", + collection, + rkey, + } + : { + operation, + rev: "3kzfcijpj2z2a", + collection, + rkey, + cid: "bafyreiclp443lav4udnztz7msp6j2wkxwp4m5mns2njm6zq7so4au6druq", + record: { $type: collection, name: rkey }, + }, + } as JetstreamEvent; +} + +function reader(events: JetstreamEvent[]): AsyncIterable { + return { + async *[Symbol.asyncIterator]() { + yield* events; + }, + }; +} + +function config() { + return resolveConfig({ + namespace: "com.example", + profiles: [], + constellation: false, + jetstreams: ["wss://jetstream.test"], + collections: { event: { collection: COLLECTION } }, + }); +} + +describe("Jetstream change source", () => { + it("uses real stream events as marks and replays through the exact watermark", async () => { + const nowUs = Date.now() * 1000; + const start = nowUs - 20_000; + const through = start + 10_000; + const subscriptions: JetstreamSubscriptionOptions[] = []; + const streams = [ + [commit(start, WATERMARK, "capture")], + [commit(through, WATERMARK, "target")], + [ + commit(start - 1, COLLECTION, "overlap"), + commit(start + 1, COLLECTION, "keep"), + commit(start + 2, WATERMARK, "ignore"), + commit(through, COLLECTION, "remove", "delete"), + commit(through + 1, WATERMARK, "past-target"), + ], + ]; + const source = new JetstreamChangeSource(config(), { + epoch: "test-epoch", + retentionUs: 60_000_000, + replayOverlapUs: 100, + watermarkCollections: [WATERMARK], + subscriptionFactory(options) { + subscriptions.push(options); + return reader(streams.shift() ?? []); + }, + }); + + const capture = await source.mark({ collections: [COLLECTION] }); + const target = await source.mark({ collections: [COLLECTION] }); + const batches = []; + for await (const batch of source.read({ + collections: [COLLECTION], + after: capture, + through: target, + })) { + batches.push(batch); + } + + expect(capture).toEqual({ + source: "jetstream", + epoch: "test-epoch", + cursor: String(start), + }); + expect(target.cursor).toBe(String(through)); + expect(subscriptions).toHaveLength(3); + expect(subscriptions[0].wantedCollections).toEqual([WATERMARK, COLLECTION]); + expect(subscriptions[2].cursor).toBe(start - 100); + expect(batches).toHaveLength(1); + expect(batches[0].checkpoint).toEqual(target); + expect(batches[0].caughtUp).toBe(true); + expect(batches[0].mutations.map((mutation) => mutation.rkey)).toEqual([ + "keep", + "remove", + ]); + expect(batches[0].mutations.map((mutation) => mutation.operation)).toEqual([ + "put", + "delete", + ]); + }); + + it("emits empty durable checkpoints for accounted-for watermark traffic", async () => { + const nowUs = Date.now() * 1000; + const after: SourcePosition = { + source: "jetstream", + epoch: "test-epoch", + cursor: String(nowUs - 100), + }; + const through: SourcePosition = { + ...after, + cursor: String(nowUs), + }; + const events = Array.from({ length: 50 }, (_, index) => + commit(nowUs - 99 + index, WATERMARK, `watermark-${index}`), + ); + events.push(commit(nowUs + 1, WATERMARK, "past-target")); + const source = new JetstreamChangeSource(config(), { + epoch: "test-epoch", + retentionUs: 60_000_000, + watermarkCollections: [WATERMARK], + subscriptionFactory: () => reader(events), + }); + + const batches = []; + for await (const batch of source.read({ + collections: [COLLECTION], + after, + through, + })) { + batches.push(batch); + } + + expect(batches).toHaveLength(2); + expect(batches[0]).toMatchObject({ mutations: [], caughtUp: false }); + expect(batches[1]).toEqual({ + mutations: [], + checkpoint: through, + caughtUp: true, + }); + }); + + it("blocks replay once the configured source-history guarantee expires", async () => { + const source = new JetstreamChangeSource(config(), { + epoch: "test-epoch", + retentionUs: 1_000, + subscriptionFactory: () => reader([]), + }); + const after: SourcePosition = { + source: "jetstream", + epoch: "test-epoch", + cursor: "1", + }; + const through = { ...after, cursor: "2" }; + const drain = async () => { + for await (const _batch of source.read({ + collections: [COLLECTION], + after, + through, + })) { + // Drain the source. + } + }; + + await expect(drain()).rejects.toBeInstanceOf(SourceHistoryExpiredError); + }); +}); diff --git a/packages/contrail/tests/pds-snapshot-source.test.ts b/packages/contrail/tests/pds-snapshot-source.test.ts new file mode 100644 index 0000000..eb774f2 --- /dev/null +++ b/packages/contrail/tests/pds-snapshot-source.test.ts @@ -0,0 +1,261 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { + DatabaseBootstrapTarget, + PdsSnapshotIncompleteError, + PdsSnapshotSource, + bootstrapFreshProjection, + initSchema, + queryRecords, + resolveConfig, + type ChangeSource, +} from "../src/index"; +import { createSqliteDatabase } from "../src/adapters/sqlite"; +import { __resetPdsCachesForTests } from "../src/core/client"; + +const DID = "did:plc:pds-snapshot"; +const COLLECTION = "com.example.event"; + +function config() { + return resolveConfig({ + namespace: "com.example", + profiles: [], + constellation: false, + relays: ["https://relay.test"], + networkOverrides: { additionalAllowedHosts: ["pds.allowed.test"] }, + collections: { event: { collection: COLLECTION } }, + }); +} + +describe("PDS snapshot source", () => { + let fetchSpy: ReturnType; + + beforeEach(() => { + __resetPdsCachesForTests(); + fetchSpy = vi.spyOn(globalThis, "fetch"); + }); + + afterEach(() => { + fetchSpy.mockRestore(); + __resetPdsCachesForTests(); + }); + + it("discovers partitions and resumes listRecords from committed partition progress", async () => { + const resolved = config(); + const db = createSqliteDatabase(":memory:"); + await initSchema(db, resolved); + await db + .prepare( + "INSERT INTO identities (did, handle, pds, resolved_at) VALUES (?, ?, ?, ?)", + ) + .bind(DID, "snapshot.test", "https://pds.allowed.test", Date.now()) + .run(); + + const requestedCursors: Array = []; + fetchSpy.mockImplementation(async (input) => { + const url = new URL(input instanceof Request ? input.url : String(input)); + if (url.hostname === "relay.test") { + return new Response(JSON.stringify({ repos: [{ did: DID }] }), { + status: 200, + headers: { "content-type": "application/json" }, + }); + } + if (url.hostname === "pds.allowed.test") { + const cursor = url.searchParams.get("cursor"); + requestedCursors.push(cursor); + const suffix = cursor === null ? "one" : "two"; + return new Response( + JSON.stringify({ + records: [ + { + uri: `at://${DID}/${COLLECTION}/${suffix}`, + cid: `cid-${suffix}`, + value: { $type: COLLECTION, name: suffix }, + }, + ], + ...(cursor === null ? { cursor: "next-page" } : {}), + }), + { status: 200, headers: { "content-type": "application/json" } }, + ); + } + throw new Error(`Unexpected request: ${url}`); + }); + + const source = new PdsSnapshotSource(db, resolved, { + maxRetries: 0, + }); + const prepared = await source.prepare({ collections: [COLLECTION] }); + const firstRead = source.read({ snapshot: prepared })[Symbol.asyncIterator](); + const first = await firstRead.next(); + expect(first.done).toBe(false); + expect(first.value.records.map((record) => record.rkey)).toEqual(["one"]); + expect(first.value.progress).toEqual({ + partition: JSON.stringify([DID, COLLECTION]), + cursor: "next-page", + complete: false, + }); + expect(first.value.done).toBe(false); + await firstRead.return?.(); + + const resumed = []; + for await (const batch of source.read({ + snapshot: prepared, + progress: [first.value.progress], + })) { + resumed.push(batch); + } + + expect(requestedCursors).toEqual([null, "next-page"]); + expect(resumed).toHaveLength(2); + expect(resumed[0].records.map((record) => record.rkey)).toEqual(["two"]); + expect(resumed[0].progress).toEqual({ + partition: JSON.stringify([DID, COLLECTION]), + cursor: null, + complete: true, + }); + expect(resumed[0].done).toBe(false); + expect(resumed[1]).toMatchObject({ + records: [], + progress: { partition: "snapshot", cursor: null, complete: true }, + done: true, + }); + }); + + it("builds a database candidate and replays changes after capture-first discovery", async () => { + const resolved = config(); + const db = createSqliteDatabase(":memory:"); + await initSchema(db, resolved); + await db + .prepare( + "INSERT INTO identities (did, handle, pds, resolved_at) VALUES (?, ?, ?, ?)", + ) + .bind(DID, "snapshot.test", "https://pds.allowed.test", Date.now()) + .run(); + + const calls: string[] = []; + fetchSpy.mockImplementation(async (input) => { + const url = new URL(input instanceof Request ? input.url : String(input)); + if (url.hostname === "relay.test") { + calls.push("relay"); + return new Response(JSON.stringify({ repos: [{ did: DID }] }), { + status: 200, + headers: { "content-type": "application/json" }, + }); + } + calls.push("pds"); + return new Response( + JSON.stringify({ + records: [ + { + uri: `at://${DID}/${COLLECTION}/candidate`, + cid: "bafyreiclp443lav4udnztz7msp6j2wkxwp4m5mns2njm6zq7so4au6druq", + value: { $type: COLLECTION, name: "snapshot" }, + }, + ], + }), + { status: 200, headers: { "content-type": "application/json" } }, + ); + }); + + const base = Date.now() * 1000; + let marks = 0; + const changes: ChangeSource = { + id: "test-changes", + semantics: { + ordinaryRecords: true, + ordinaryDeletes: true, + accountLifecycle: false, + repositoryReplacement: false, + verifiedCommits: false, + explicitHead: true, + }, + async mark() { + marks++; + calls.push(`mark:${marks}`); + return { + source: "test-changes", + epoch: "one", + cursor: String(base + (marks === 1 ? 0 : 2_000_000)), + }; + }, + async *read({ through }) { + calls.push("changes"); + yield { + mutations: [ + { + operation: "put", + uri: `at://${DID}/${COLLECTION}/candidate`, + did: DID, + collection: COLLECTION, + rkey: "candidate", + cid: "bafyreig7jv2h5c3xw3dkf5m7zqxf2spwzvug3k4j5qvbcjjr5w6m2wr6li", + value: { $type: COLLECTION, name: "tail" }, + sourceTimeUs: base + 1_000_000, + position: { + source: "test-changes", + epoch: "one", + cursor: String(base + 1_000_000), + }, + }, + ], + checkpoint: through, + caughtUp: true, + }; + }, + }; + + await bootstrapFreshProjection({ + collections: [COLLECTION], + snapshotSource: new PdsSnapshotSource(db, resolved, { maxRetries: 0 }), + changeSource: changes, + target: new DatabaseBootstrapTarget(db, resolved), + }); + + expect(calls).toEqual(["mark:1", "relay", "pds", "mark:2", "changes"]); + const records = await queryRecords(db, resolved, { collection: "event" }); + expect(records.records).toHaveLength(1); + expect(JSON.parse(records.records[0].record).name).toBe("tail"); + }); + + it("rejects a PDS record outside its requested repository partition", async () => { + const resolved = config(); + const db = createSqliteDatabase(":memory:"); + await initSchema(db, resolved); + await db + .prepare( + "INSERT INTO identities (did, handle, pds, resolved_at) VALUES (?, ?, ?, ?)", + ) + .bind(DID, "snapshot.test", "https://pds.allowed.test", Date.now()) + .run(); + + fetchSpy.mockImplementation(async (input) => { + const url = new URL(input instanceof Request ? input.url : String(input)); + if (url.hostname === "relay.test") { + return new Response(JSON.stringify({ repos: [{ did: DID }] }), { + status: 200, + headers: { "content-type": "application/json" }, + }); + } + return new Response( + JSON.stringify({ + records: [ + { + uri: `at://did:plc:someone-else/${COLLECTION}/wrong`, + cid: "cid-wrong", + value: { $type: COLLECTION, name: "wrong" }, + }, + ], + }), + { status: 200, headers: { "content-type": "application/json" } }, + ); + }); + + const source = new PdsSnapshotSource(db, resolved, { maxRetries: 0 }); + const prepared = await source.prepare({ collections: [COLLECTION] }); + const read = async () => { + for await (const _batch of source.read({ snapshot: prepared })) { + // Drain the source. + } + }; + await expect(read()).rejects.toBeInstanceOf(PdsSnapshotIncompleteError); + }); +}); -- 2.51.2