diff --git a/src/server/management/usage-aggregate-cache.ts b/src/server/management/usage-aggregate-cache.ts index b54ed7d9685..17e530fa866 100644 --- a/src/server/management/usage-aggregate-cache.ts +++ b/src/server/management/usage-aggregate-cache.ts @@ -15,6 +15,7 @@ import { readUsageFilterCandidatesPassive, requestUsageFilterIndexRebuild, resetUsageFilterIndexForTests, + UsageFilterIndexPageError, type UsageFilterCandidateRead, type UsageFilterIndexQuery, type UsageFilterIndexSnapshot, @@ -110,7 +111,6 @@ const MAX_REBUILD_ATTEMPTS = 2; const MAX_FILTER_INDEX_ATTEMPTS = 3; // Selected rows are folded in bounded batches, so a filter that matches most of a large ledger // costs the same peak memory as the streaming scan it replaces. -const FILTER_ROW_READ_BATCH = 2_000; const MAX_RETAINED_FILTERED_AGGREGATES = 4; let retainedAggregate: RetainedUsageAggregate | null = null; @@ -396,7 +396,7 @@ function filterStats( ): UsageFilterQueryStats { return { totalRows: read.snapshot.indexedRows, - candidateRows: read.offsets.length, + candidateRows: read.candidateRows, scope: update, ingestedRows: read.snapshot.ingestedRows, ingestPasses: read.snapshot.ingestPasses, @@ -465,33 +465,45 @@ async function readCandidates( : readUsageFilterCandidates(query); } +/** + * Folds every candidate page into the accumulator, yielding between pages and re-reading rows + * against the snapshot identity; closes the reader when done. + */ async function accumulateCandidates( accumulator: UsageSummaryAccumulator, read: UsageFilterCandidateRead, filter: NormalizedUsageFilter, ): Promise { - // Always read at least once, even with no candidates: the reader validates the snapshot - // against the ledger, and an empty selection must not skip that check. - const batchCount = Math.max(1, Math.ceil(read.offsets.length / FILTER_ROW_READ_BATCH)); - for (let batch = 0; batch < batchCount; batch += 1) { - const start = batch * FILTER_ROW_READ_BATCH; - const rows = await readUsageLedgerRowsAtOffsets({ - offsets: read.offsets.slice(start, start + FILTER_ROW_READ_BATCH), - expectedIdentityKey: read.snapshot.identityKey, - // The offsets belong to this index snapshot. If the ledger moved since the snapshot, - // the reader re-validates the boundary before trusting any of them, and a rewrite of - // the indexed tail becomes a rebuild instead of a silently different row set. - expectedRevisionKey: usageLogRevisionKey(read.snapshot.revision), - expectedProcessedThroughBytes: read.snapshot.indexedOffset, - expectedProcessedThroughDigest: read.snapshot.processedThroughDigest, - // The physical size at publish time; a torn suffix makes it larger than the boundary. - expectedFileSize: read.snapshot.revision?.size, - }); - if (rows.unreadableOffsets.length > 0) { - // Serving only the rows that happened to parse would silently under-report the filter. - throw new UsageFilterIndexCorruptError(); + try { + let first = true; + for (;;) { + // Candidate selection and row decoding are synchronous bounded chunks. A timer yield + // between them prevents a broad authenticated query from monopolizing the event loop. + if (!first) await new Promise(resolve => setTimeout(resolve, 0)); + const offsets = read.nextPage(); + if (!first && offsets.length === 0) break; + first = false; + const rows = await readUsageLedgerRowsAtOffsets({ + offsets, + expectedIdentityKey: read.snapshot.identityKey, + // The offsets belong to this index snapshot. If the ledger moved since the snapshot, + // the reader re-validates the boundary before trusting any of them, and a rewrite of + // the indexed tail becomes a rebuild instead of a silently different row set. + expectedRevisionKey: usageLogRevisionKey(read.snapshot.revision), + expectedProcessedThroughBytes: read.snapshot.indexedOffset, + expectedProcessedThroughDigest: read.snapshot.processedThroughDigest, + // The physical size at publish time; a torn suffix makes it larger than the boundary. + expectedFileSize: read.snapshot.revision?.size, + }); + if (rows.unreadableOffsets.length > 0) { + // Serving only the rows that happened to parse would silently under-report the filter. + throw new UsageFilterIndexCorruptError(); + } + for (const entry of rows.entries) accumulator.add(entry); + if (offsets.length === 0) break; } - for (const entry of rows.entries) accumulator.add(entry); + } finally { + read.close(); } // Without a source filter the reported span covers the whole ledger, including rows this // filter never selected. A source filter derives its span from the rows that match it, so @@ -626,57 +638,80 @@ async function appendFilteredAggregate( } } +/** + * Serves the filtered report from the retained aggregate, appending new candidates or rebuilding + * as needed; a candidate-page failure clears the retained state and retries the streaming scan. + */ async function refreshFilteredAggregate( key: string, filter: NormalizedUsageFilter, window?: UsageTimeWindow, ): Promise { - return withFilterIndexRetry(async () => { - const overlayVersion = userCostOverlayVersion(); - const timeZone = currentTimeZone(); - const state = retainedFilteredAggregates.get(key); - if (!state - || state.overlayVersion !== overlayVersion - || state.timeZone !== timeZone - || state.indexGeneration === null) { - return rebuildFilteredAggregate(key, filter, window, overlayVersion, timeZone); - } - const read = await readCandidates({ - ...filterQuery(filter, window), - afterOffset: state.processedThroughBytes, + try { + return await withFilterIndexRetry(async () => { + const overlayVersion = userCostOverlayVersion(); + const timeZone = currentTimeZone(); + const state = retainedFilteredAggregates.get(key); + if (!state + || state.overlayVersion !== overlayVersion + || state.timeZone !== timeZone + || state.indexGeneration === null) { + return rebuildFilteredAggregate(key, filter, window, overlayVersion, timeZone); + } + const read = await readCandidates({ + ...filterQuery(filter, window), + afterOffset: state.processedThroughBytes, + }); + if (!read) { + retainedFilteredAggregates.delete(key); + return rebuildFilteredAggregate(key, filter, window, overlayVersion, timeZone); + } + if (state.indexGeneration !== read.snapshot.generation + || state.identityKey !== read.snapshot.identityKey + || state.processedThroughBytes > read.snapshot.indexedOffset) { + // This branch never pages, so it owns closing the reader before the rebuild discards it. + read.close(); + retainedFilteredAggregates.delete(key); + return rebuildFilteredAggregate(key, filter, window, overlayVersion, timeZone); + } + if (state.processedThroughBytes === read.snapshot.indexedOffset) { + // Even an unchanged retained projection validates the ledger against the snapshot, so a + // rewrite in the window between the index refresh and this return is not served stale. + try { + await readUsageLedgerRowsAtOffsets({ + offsets: [], + expectedIdentityKey: read.snapshot.identityKey, + expectedRevisionKey: usageLogRevisionKey(read.snapshot.revision), + expectedProcessedThroughBytes: read.snapshot.indexedOffset, + expectedProcessedThroughDigest: read.snapshot.processedThroughDigest, + expectedFileSize: read.snapshot.revision?.size, + }); + } finally { + // This branch never pages, so it owns closing the reader before publishing its stats. + read.close(); + } + state.retainedAt = Date.now(); + // Whole-ledger diagnostics can move without the boundary moving: a torn oversized or + // malformed suffix is observed by every append scan and never advances the offset. + state.usageIncomplete = read.snapshot.oversizedRows > 0; + state.usageInvalidRows = read.snapshot.invalidRows > 0; + state.revision = read.snapshot.revision; + state.revisionKey = usageLogRevisionKey(read.snapshot.revision); + return publishFilteredAggregate(key, state, "unchanged", read); + } + return appendFilteredAggregate(key, state, filter, read, overlayVersion, timeZone); }); - if (!read) { - retainedFilteredAggregates.delete(key); - return rebuildFilteredAggregate(key, filter, window, overlayVersion, timeZone); - } - if (state.indexGeneration !== read.snapshot.generation - || state.identityKey !== read.snapshot.identityKey - || state.processedThroughBytes > read.snapshot.indexedOffset) { + } catch (error) { + if (error instanceof UsageFilterIndexPageError) { + // A page failed after its reader was handed out: the index cannot serve this query, + // so answer from the authoritative ledger scan instead of failing the request. retainedFilteredAggregates.delete(key); - return rebuildFilteredAggregate(key, filter, window, overlayVersion, timeZone); - } - if (state.processedThroughBytes === read.snapshot.indexedOffset) { - // Even an unchanged retained projection validates the ledger against the snapshot, so a - // rewrite in the window between the index refresh and this return is not served stale. - await readUsageLedgerRowsAtOffsets({ - offsets: [], - expectedIdentityKey: read.snapshot.identityKey, - expectedRevisionKey: usageLogRevisionKey(read.snapshot.revision), - expectedProcessedThroughBytes: read.snapshot.indexedOffset, - expectedProcessedThroughDigest: read.snapshot.processedThroughDigest, - expectedFileSize: read.snapshot.revision?.size, - }); - state.retainedAt = Date.now(); - // Whole-ledger diagnostics can move without the boundary moving: a torn oversized or - // malformed suffix is observed by every append scan and never advances the offset. - state.usageIncomplete = read.snapshot.oversizedRows > 0; - state.usageInvalidRows = read.snapshot.invalidRows > 0; - state.revision = read.snapshot.revision; - state.revisionKey = usageLogRevisionKey(read.snapshot.revision); - return publishFilteredAggregate(key, state, "unchanged", read); + return withFilterIndexRetry(() => + rebuildFilteredAggregateStreaming(key, filter, window, userCostOverlayVersion(), currentTimeZone()), + ); } - return appendFilteredAggregate(key, state, filter, read, overlayVersion, timeZone); - }); + throw error; + } } function jevStatsResultFrom( diff --git a/src/usage/filter-index.ts b/src/usage/filter-index.ts index 68eb1a0257b..52e162f1e09 100644 --- a/src/usage/filter-index.ts +++ b/src/usage/filter-index.ts @@ -91,7 +91,10 @@ const SQLITE_BUSY_TIMEOUT_MS = 250; const RETAINED_WAL_LIMIT_BYTES = 4 * 1024 * 1024; const BATCH_BUSY_RETRIES = 3; const BATCH_BUSY_RETRY_MS = 5; -const BUSY_RETRY_DELAYS_MS = [10, 25, 50] as const; +const BUSY_RETRY_DELAYS_MS = [10, 25, 50, 100, 200, 400, 800, 1_600, 3_200] as const; +// A rebuild-required error resolves as soon as the competing publish settles, so it retries on +// the original short schedule instead of parking behind the long busy backoff. +const REBUILD_RETRY_DELAYS_MS = [10, 25, 50] as const; const SCHEMA_DDL = ` CREATE TABLE IF NOT EXISTS ${META_TABLE} ( @@ -184,6 +187,17 @@ export class UsageFilterIndexHomeChangedError extends Error { } } +/** + * A candidate page could not be read from the index after the reader was handed out. The + * projection cannot serve this query, so the caller streams the ledger instead of failing. + */ +export class UsageFilterIndexPageError extends Error { + constructor(readonly pageError: unknown) { + super("usage filter index candidate page could not be read"); + this.name = "UsageFilterIndexPageError"; + } +} + /** * A retained filtered projection read this generation. A full rebuild replaces every offset, * so a projection from an older generation must be discarded even when the file identity and @@ -220,8 +234,15 @@ export interface UsageFilterIndexQuery { export interface UsageFilterCandidateRead { snapshot: UsageFilterIndexSnapshot; - /** Ascending row-start offsets selected by the index. */ - offsets: number[]; + /** Number of candidates returned by pages consumed so far. */ + candidateRows: number; + /** Returns the next bounded page of ascending row-start offsets. */ + nextPage(): number[]; + /** + * Closes the database handle. Pages are not read from one long snapshot: each `nextPage()` + * re-checks schema and generation inside its own short transaction. + */ + close(): void; } /** @@ -1145,10 +1166,21 @@ async function refreshUsageFilterIndex( } } -function buildCandidateSelect( +interface CandidateBranch { + sql: string; + values: Array; +} + +/** + * One UNION branch of the candidate select. When `paged` is set each branch ends with + * `byte_offset > ?` so the pager can seek from its cursor instead of re-reading the candidates + * before it; the caller appends the cursor value at the end of each branch's values. + */ +function candidateBranches( query: UsageFilterIndexQuery, snapshot: UsageFilterIndexSnapshot, -): { sql: string; values: Array } { + paged: boolean, +): CandidateBranch[] { const attrClauses: string[] = []; const attrValues: Array = []; const rowClauses: string[] = []; @@ -1223,18 +1255,15 @@ function buildCandidateSelect( // estimate prefers the row primary key, which examines every row in the ledger even when // one API key or client matches a handful. const baseIndex = rowAccessIndexWithoutClient(query); - const branches = clientAlternatives.map(alternative => { + const cursor = paged ? " AND r.byte_offset > ?" : ""; + return clientAlternatives.map(alternative => { const index = alternative.clauses.length > 0 ? "idx_usage_rows_client" : baseIndex; const clauses = [...rowClauses, ...alternative.clauses]; return { - sql: `SELECT r.byte_offset AS byteOffset FROM usage_rows r${index ? ` INDEXED BY ${index}` : ""} WHERE r.generation = ? AND ${clauses.join(" AND ")}`, + sql: `SELECT r.byte_offset AS byteOffset FROM usage_rows r${index ? ` INDEXED BY ${index}` : ""} WHERE r.generation = ? AND ${clauses.join(" AND ")}${cursor}`, values: [snapshot.generation, ...rowValues, ...alternative.values], }; }); - return { - sql: `${branches.map(branch => branch.sql).join(" UNION ")} ORDER BY byteOffset`, - values: branches.flatMap(branch => branch.values), - }; } // Drive from the attribution index rather than scanning the row range: SQLite's default // estimate otherwise prefers the row primary key, which examines every row in the ledger. @@ -1242,18 +1271,36 @@ function buildCandidateSelect( // Rows with no usable attribution are selected too: the accumulator is what decides, and a // hand-edited row with a missing model must fail a provider/model filter exactly as the // streaming read does rather than disappear from the result. + return [ + { + sql: `SELECT DISTINCT a.byte_offset AS byteOffset FROM usage_attributions a INDEXED BY ${attrIndex} WHERE a.generation = ? AND ${attrClauses.join(" AND ")} AND a.byte_offset <= ? AND EXISTS (SELECT 1 FROM usage_rows r WHERE r.generation = a.generation AND r.byte_offset = a.byte_offset AND ${rowSql})${paged ? " AND a.byte_offset > ?" : ""}`, + values: [ + snapshot.generation, + ...attrValues, + snapshot.indexedOffset, + ...rowValues, + ...clientAlternatives.flatMap(alternative => alternative.values), + ], + }, + { + sql: `SELECT r.byte_offset AS byteOffset FROM usage_rows r INDEXED BY idx_usage_rows_unattributable WHERE r.generation = ? AND r.unattributable = 1 AND ${rowSql}${paged ? " AND r.byte_offset > ?" : ""}`, + values: [ + snapshot.generation, + ...rowValues, + ...clientAlternatives.flatMap(alternative => alternative.values), + ], + }, + ]; +} + +function buildCandidateSelect( + query: UsageFilterIndexQuery, + snapshot: UsageFilterIndexSnapshot, +): { sql: string; values: Array } { + const branches = candidateBranches(query, snapshot, false); return { - sql: `SELECT DISTINCT a.byte_offset AS byteOffset FROM usage_attributions a INDEXED BY ${attrIndex} WHERE a.generation = ? AND ${attrClauses.join(" AND ")} AND a.byte_offset <= ? AND EXISTS (SELECT 1 FROM usage_rows r WHERE r.generation = a.generation AND r.byte_offset = a.byte_offset AND ${rowSql}) UNION SELECT r.byte_offset AS byteOffset FROM usage_rows r INDEXED BY idx_usage_rows_unattributable WHERE r.generation = ? AND r.unattributable = 1 AND ${rowSql} ORDER BY byteOffset`, - values: [ - snapshot.generation, - ...attrValues, - snapshot.indexedOffset, - ...rowValues, - ...clientAlternatives.flatMap(alternative => alternative.values), - snapshot.generation, - ...rowValues, - ...clientAlternatives.flatMap(alternative => alternative.values), - ], + sql: `${branches.map(branch => branch.sql).join(" UNION ")} ORDER BY byteOffset`, + values: branches.flatMap(branch => branch.values), }; } @@ -1267,14 +1314,74 @@ function rowAccessIndexWithoutClient(query: UsageFilterIndexQuery): string | nul return null; } -function selectCandidates( +const CANDIDATE_PAGE_SIZE = 2_000; + +/** + * The page statement behind the candidate pager: each UNION branch seeks from the page cursor + * and stops at the page size, so a late page reads at most `CANDIDATE_PAGE_SIZE` rows per branch + * rather than re-materializing every candidate before the cursor. The cursor value is appended + * at the end of each branch's values when the statement runs. + */ +function candidatePageSql( + query: UsageFilterIndexQuery, + snapshot: UsageFilterIndexSnapshot, +): { sql: string; branchValues: Array> } { + // A compound operand cannot carry its own ORDER BY/LIMIT, so each branch is wrapped before + // the UNION deduplicates and the outer ORDER BY/LIMIT merges the bounded branch pages. + const branches = candidateBranches(query, snapshot, true); + const paged = branches.map( + branch => `SELECT * FROM (${branch.sql} ORDER BY byteOffset LIMIT ${CANDIDATE_PAGE_SIZE})`, + ); + return { + sql: `SELECT byteOffset FROM (${paged.join(" UNION ")}) ORDER BY byteOffset LIMIT ?`, + branchValues: branches.map(branch => branch.values), + }; +} + +/** Keyset pager over the candidate select: bounded pages, each validated in its own transaction. */ +function candidatePager( db: Database, snapshot: UsageFilterIndexSnapshot, query: UsageFilterIndexQuery, -): number[] { - const { sql, values } = buildCandidateSelect(query, snapshot); - const rows = db.query(sql).all(...values) as Array<{ byteOffset: number }>; - return rows.map(row => row.byteOffset); +): Pick { + const { sql, branchValues } = candidatePageSql(query, snapshot); + const statement = db.query(sql); + let cursor = -1; + const pager = { + candidateRows: 0, + /** + * Returns the next bounded page of ascending offsets, re-validating schema and generation + * inside its own transaction; a changed generation throws `UsageLedgerRebuildRequiredError`. + */ + nextPage(): number[] { + try { + db.exec("BEGIN"); + const meta = readMeta(db); + if (meta.schemaVersion !== USAGE_FILTER_INDEX_SCHEMA_VERSION + || meta.generation !== snapshot.generation) { + throw new UsageLedgerRebuildRequiredError("content_changed"); + } + const rows = statement.all( + ...branchValues.flatMap(values => [...values, cursor]), + CANDIDATE_PAGE_SIZE, + ) as Array<{ byteOffset: number }>; + const offsets = rows.map(row => row.byteOffset); + db.exec("COMMIT"); + if (offsets.length > 0) cursor = offsets[offsets.length - 1]!; + pager.candidateRows += offsets.length; + return offsets; + } catch (error) { + try { db.exec("ROLLBACK"); } catch { /* nothing to roll back */ } + // A generation change keeps the rebuild retry; every other index-side failure means + // this projection cannot serve the page, so the caller streams the ledger instead. + if (error instanceof UsageLedgerRebuildRequiredError) throw error; + // Paging cannot repair a damaged projection; ask the next active refresh to rebuild it. + if (isCorruptError(error)) requestUsageFilterIndexRebuild(); + throw new UsageFilterIndexPageError(error); + } + }, + }; + return pager; } /** A published generation with a nonzero row count must actually hold rows. */ @@ -1283,6 +1390,10 @@ function hasPublishedRows(db: Database, generation: string): boolean { .get(generation) !== null; } +/** + * Refreshes the projection, then returns a paged reader for the published generation. Throws + * `UsageLedgerRebuildRequiredError` when the published meta moved under the open. + */ async function readActive( dir: string, path: string, @@ -1304,22 +1415,28 @@ async function readActive( ingestedRows: refreshed.ingestedRows, ingestPasses: refreshed.ingestPasses, }; - const offsets = selectCandidates(db, snapshot, query); - if (offsets.length === 0 && meta.indexedRows > 0 && !hasPublishedRows(db, meta.generation)) { + if (meta.indexedRows > 0 && !hasPublishedRows(db, meta.generation)) { // The published generation claims rows it no longer has. Never serve that as a // complete zero; drop the projection and rebuild. The refresh alone cannot see this // (the cursor and identity still match), so the rebuild must be requested explicitly. requestUsageFilterIndexRebuild(); throw new UsageLedgerRebuildRequiredError("content_changed"); } + const pager = candidatePager(db, snapshot, query); db.exec("COMMIT"); - return { snapshot, offsets }; + return { + snapshot, + get candidateRows() { return pager.candidateRows; }, + nextPage: () => pager.nextPage(), + close: () => db.close(), + }; } catch (error) { try { db.exec("ROLLBACK"); } catch { /* nothing to roll back */ } throw error; } - } finally { + } catch (error) { db.close(); + throw error; } } @@ -1343,22 +1460,26 @@ export async function readUsageFilterCandidates( return await readActive(dir, path, query); } catch (error) { if (error instanceof UsageFilterIndexSourceError) throw error; - if (error instanceof UsageFilterIndexForeignError - || error instanceof UnsafeSqlitePathError - || error instanceof UsageFilterIndexHomeChangedError) { + // A changed home never heals on a re-check; every later open targets the new directory. + if (error instanceof UsageFilterIndexHomeChangedError) { return null; } - if (error instanceof UsageLedgerRebuildRequiredError) { - const delay = BUSY_RETRY_DELAYS_MS[attempt]; + // A foreign classification or unsafe-path flag can still be a transient view: the file at + // this path can be replaced mid-flight, or a sidecar appears or vanishes mid-stat while + // another writer transacts. Give them the short re-check schedule before degrading, exactly + // like a moved publish. + if (error instanceof UsageLedgerRebuildRequiredError + || error instanceof UsageFilterIndexForeignError + || error instanceof UnsafeSqlitePathError) { + const delay = REBUILD_RETRY_DELAYS_MS[attempt]; if (delay === undefined) return null; await new Promise(resolve => setTimeout(resolve, delay)); continue; } if (isBusyError(error)) { // Contention, including a peer still creating the projection, is waited out rather than - // read as a foreign file. The bound is four busy timeouts (the fourth attempt has no - // back-off) plus these delays, about 1.1 s, before the request falls back to the - // streaming read. + // read as a foreign file. The bound is the BUSY_RETRY_DELAYS_MS sum (about 6.4 s) plus up + // to ten busy timeouts before the request falls back to the streaming read. const delay = BUSY_RETRY_DELAYS_MS[attempt]; if (delay === undefined) return null; testHooks?.beforeBusyRetry?.({ dbPath: path, attempt }); @@ -1396,16 +1517,21 @@ export async function readUsageFilterCandidatesPassive( || !tables.has("usage_rows") || !tables.has("usage_attributions") || !tables.has("usage_filter_index_writers")) { - return null; + throw new UsageFilterIndexForeignError(); } db.exec("BEGIN"); try { const meta = readMeta(db); const revision = currentUsageLogRevision(); - if (meta.schemaVersion !== USAGE_FILTER_INDEX_SCHEMA_VERSION || !meta.generation) return null; - if (usageLogIdentityKey(meta.revision) !== usageLogIdentityKey(revision)) return null; - if (!revision || meta.indexedOffset !== revision.size) return null; - if (usageLogRevisionKey(meta.revision) !== usageLogRevisionKey(revision)) return null; + if (meta.schemaVersion !== USAGE_FILTER_INDEX_SCHEMA_VERSION || !meta.generation) { + throw new UsageLedgerRebuildRequiredError("content_changed"); + } + if (usageLogIdentityKey(meta.revision) !== usageLogIdentityKey(revision) + || !revision + || meta.indexedOffset !== revision.size + || usageLogRevisionKey(meta.revision) !== usageLogRevisionKey(revision)) { + throw new UsageLedgerRebuildRequiredError("content_changed"); + } const snapshot = snapshotFrom(meta, 0, 0); // The exact revision match above already proves the ledger is the one indexed; the // empty read keeps the passive path on the same validation as the active one. @@ -1417,30 +1543,48 @@ export async function readUsageFilterCandidatesPassive( expectedProcessedThroughDigest: snapshot.processedThroughDigest, expectedFileSize: snapshot.revision?.size, }); - const offsets = selectCandidates(db, snapshot, query); - if (offsets.length === 0 && meta.indexedRows > 0 && !hasPublishedRows(db, meta.generation)) { + if (meta.indexedRows > 0 && !hasPublishedRows(db, meta.generation)) { // Same guard as the active read: a published generation that lost its rows must not // be reported as a complete zero. - return null; + throw new UsageLedgerRebuildRequiredError("content_changed"); } - return { snapshot, offsets }; - } finally { - try { db.exec("COMMIT"); } catch { /* read-only transaction */ } + const pager = candidatePager(db, snapshot, query); + db.exec("COMMIT"); + return { + snapshot, + get candidateRows() { return pager.candidateRows; }, + nextPage: () => pager.nextPage(), + close: () => db.close(), + }; + } catch (error) { + try { db.exec("ROLLBACK"); } catch { /* read-only transaction */ } + throw error; } } catch { - return null; - } finally { + try { db.exec("ROLLBACK"); } catch { /* transaction may not have started */ } db.close(); + return null; } } -/** Test seam: the SQLite plan for one filter's candidate select. */ -export function usageFilterCandidatePlan(query: UsageFilterIndexQuery): string[] { +/** Test seam: the SQLite plan for one filter's candidate select; `paged` explains the pager statement. */ +export function usageFilterCandidatePlan( + query: UsageFilterIndexQuery, + options: { paged?: boolean } = {}, +): string[] { const path = indexPathForCurrentHome(); const db = new Database(path, { readonly: true }); try { const meta = readMeta(db); const snapshot = snapshotFrom(meta, 0, 0); + if (options.paged) { + const { sql, branchValues } = candidatePageSql(query, snapshot); + const rows = db.query(`EXPLAIN QUERY PLAN ${sql}`).all( + ...branchValues.flatMap(values => [...values, 0]), + CANDIDATE_PAGE_SIZE, + ) as Array<{ detail: string }>; + return rows.map(row => row.detail); + } const { sql, values } = buildCandidateSelect(query, snapshot); const rows = db.query(`EXPLAIN QUERY PLAN ${sql}`).all(...values) as Array<{ detail: string }>; return rows.map(row => row.detail); diff --git a/structure/usage-client-attribution.md b/structure/usage-client-attribution.md index 51a2fd97ffd..2be55ce16a3 100644 --- a/structure/usage-client-attribution.md +++ b/structure/usage-client-attribution.md @@ -39,7 +39,10 @@ match. `src/usage/ledger-scanner.ts` supplies both halves: `scanUsageLedgerCoope row offset during ingest, and `readUsageLedgerRowsAtOffsets` re-reads only the selected rows under the same identity, row-ceiling and torn-suffix rules. The first filtered query after the index is missing or invalid pays one cooperative full ingest; appends ingest only new bytes, and every later filter -reuses the projection. A replaced, truncated or same-size rewritten ledger rebuilds it, a write +reuses the projection. Candidate offsets are selected and consumed in bounded keyset pages, with a +timer yield between pages; each page rechecks the published generation before it is read, so broad +filters neither retain the complete offset set nor monopolize the event loop. A replaced, truncated +or same-size rewritten ledger rebuilds it, a write observed during a row read re-verifies the selected bytes before publishing them, and a recorded offset that no longer frames a complete row drops the projection instead of serving a partial answer. `usage.jsonl` stays authoritative and is never truncated or rewritten. @@ -82,8 +85,9 @@ writer's own publish fails the fence and retries, and a reader only ever sees on generation. A crash before publish leaves the previously published generation complete. An append writes into the published generation and publishes its cursor, digest, counts and span in the same transaction, -fenced on the cursor so a stale append cannot regress it. Readers take meta and the candidate -select in one read transaction and reject a generation that changed under them. A published +fenced on the cursor so a stale append cannot regress it. Each candidate page re-reads meta and +selects inside its own short read transaction, and a generation or schema change aborts the read -- +the per-page recheck described with the paging contract above. A published generation that claims rows it no longer holds is rebuilt rather than served as a complete zero, in both the active and the passive read. @@ -113,10 +117,20 @@ existing file that classified as this projection, once they appear; a file at th not this projection is never registered or modified. A passive management read opens an existing projection read-only, requires this projection's full table set rather than a coincidental meta table, and only uses it when it is current for the exact ledger revision; otherwise it streams, -and it never creates, writes or registers anything. Any other index-side failure -- unwritable home, -full disk, contention, a corrupt or foreign file -- falls back to the authoritative streaming scan +and it never creates, writes or registers anything. The passive read makes a single attempt and +degrades to the streaming scan on any error; it does not retry. Only the active index read +(`readUsageFilterCandidates`) retries contention on nine delays totaling about 6.4 s and runs a +foreign or unsafe open and a rebuild-required read through three short delays before degrading. +Any other +index-side failure -- unwritable home, +full disk, a corrupt file -- falls back to the authoritative streaming scan instead of failing a request the scan can answer; a failure to read `usage.jsonl` itself is -propagated. A nonselective filter (`client=all`, `device=all`, no provider/model/key) still reads +propagated. A `UsageFilterIndexPageError` from a `src/usage/filter-index.ts` page read clears the +retained filtered aggregate, then retries the authoritative streaming scan under +`withFilterIndexRetry` (`MAX_FILTER_INDEX_ATTEMPTS` attempts) in +`src/server/management/usage-aggregate-cache.ts`: only a concurrent ledger replacement or an +aggregation-input change is retried, the unusable index is never re-entered, and unrelated ledger +I/O failures still propagate. A nonselective filter (`client=all`, `device=all`, no provider/model/key) still reads every row in its window; `candidateRows` counts the rows selected and normalized for that call and `scope` says whether that is the whole candidate set or only the appended suffix, not how many SQL index entries were examined. `tests/usage/usage-filter-index.test.ts` asserts candidate counts against the whole diff --git a/tests/usage/usage-filter-index.test.ts b/tests/usage/usage-filter-index.test.ts index 0804d4e2fd1..900d313f94c 100644 --- a/tests/usage/usage-filter-index.test.ts +++ b/tests/usage/usage-filter-index.test.ts @@ -4,6 +4,7 @@ import { appendFileSync, chmodSync, existsSync, + mkdirSync, mkdtempSync, readFileSync, readdirSync, @@ -40,7 +41,9 @@ import { USAGE_FILTER_INDEX_INSERT_BATCH, usageFilterCandidatePlan, usageFilterIndexPath, + UsageFilterIndexPageError, } from "../../src/usage/filter-index"; +import * as filterIndexModule from "../../src/usage/filter-index"; import { currentUsageLogRevision, resetUsageReadCacheForTests, @@ -213,6 +216,47 @@ describe("filtered usage candidate index", () => { expect(summary(result).summary.requests).toBe(inside); }); + test("a broad filter consumes bounded candidate pages with timer yields", async () => { + writeRows(Array.from({ length: SYNTHETIC_ROWS }, (_, index) => syntheticRow(index))); + await getFilteredUsageAggregate({ provider: "kimi" }); + resetUsageAggregateCacheForTests(); + + const originalRead = ledgerScannerModule.readUsageLedgerRowsAtOffsets; + const pageSizes: number[] = []; + let yieldedBetweenPages = false; + let timerFired = false; + const readSpy = spyOn(ledgerScannerModule, "readUsageLedgerRowsAtOffsets") + .mockImplementation(async options => { + if (pageSizes.length > 0) yieldedBetweenPages ||= timerFired; + pageSizes.push(options.offsets.length); + const result = await originalRead(options); + if (pageSizes.length === 1) setTimeout(() => { timerFired = true; }, 0); + return result; + }); + try { + const result = await getFilteredUsageAggregate({}); + expect(result.filterIndex?.candidateRows).toBe(SYNTHETIC_ROWS); + expect(pageSizes).toEqual([2_000, 2_000]); + expect(yieldedBetweenPages).toBe(true); + } finally { + readSpy.mockRestore(); + } + }); + + test("a broad provider filter answers complete totals across bounded pages", async () => { + // More matching candidates than one page (2,000) forces late pages; each must still seek + // from the page cursor rather than re-reading the candidates before it. + const rows = Array.from({ length: 4_100 }, (_, index) => ({ + ...syntheticRow(index), + provider: "openai", + })); + writeRows(rows); + const result = await getFilteredUsageAggregate({ provider: "openai" }); + expect(result.filterIndex?.candidateRows).toBe(rows.length); + expect(summary(result)).toEqual((await fullScan({ provider: "openai" })).summarize("all", NOW)); + expect(summary(result).summary.requests).toBe(rows.length); + }); + test("an append ingests only new bytes and a retained filter folds only new matches", async () => { const rows = Array.from({ length: SYNTHETIC_ROWS }, (_, index) => syntheticRow(index)); writeRows(rows); @@ -626,7 +670,9 @@ console.log(JSON.stringify({ requests: summary.summary.requests, candidates: res } expect(indexGenerations()).toHaveLength(1); expect(indexRowCount(publishedGeneration())).toBe(rows.length); - }); + // Two processes contend on one 30k-row rebuild; the loser's busy-retry backoff plus the + // double scan legitimately exceeds the default 5s timeout. + }, 30_000); test("a projection another process is still creating is waited for, not taken as foreign", async () => { const rows = Array.from({ length: 20 }, (_, index) => syntheticRow(index)); @@ -768,7 +814,7 @@ console.log(JSON.stringify({ requests: summary.summary.requests, candidates: res expect(read.entries).toHaveLength(1); expect(read.unreadableOffsets).toHaveLength(49); }); -+ + test("a rewritten indexed row makes the next append rebuild instead of streaming forever", async () => { const rows = Array.from({ length: 200 }, (_, index) => syntheticRow(index)); writeRows(rows); @@ -815,7 +861,7 @@ console.log(JSON.stringify({ requests: summary.summary.requests, candidates: res setUsageFilterIndexTestHooksForTests(null); } }); -+ + test("a same-size rewrite before the trailing boundary forces a rebuild", async () => { // Large enough that row zero sits far outside the trailing 64 KiB digest window. const rows = Array.from({ length: 4000 }, (_, index) => syntheticRow(index)); @@ -901,7 +947,7 @@ console.log(JSON.stringify({ requests: summary.summary.requests, candidates: res for (const seen of observed) expect(seen).toBe("none/none"); expect(totalIndexRows()).toBe(ledgerRowCount()); }); -+ + test("a replaced ledger between the index refresh and the read is caught", async () => { const rows = Array.from({ length: 20 }, (_, index) => syntheticRow(index)); writeRows(rows); @@ -988,7 +1034,7 @@ console.log(JSON.stringify({ requests: summary.summary.requests, candidates: res probe.close(); } }); -+ + test("an unchanged retained projection still validates the ledger", async () => { const rows = Array.from({ length: 40 }, (_, index) => syntheticRow(index)); writeRows(rows); @@ -1012,7 +1058,7 @@ console.log(JSON.stringify({ requests: summary.summary.requests, candidates: res setUsageFilterIndexTestHooksForTests(null); } }); -+ + test("a torn suffix does not let a same-size rewrite masquerade as an append", async () => { const rows = Array.from({ length: 1000 }, (_, index) => ( index === 0 ? { ...syntheticRow(index), provider: "kimi" } : syntheticRow(index) @@ -1110,6 +1156,64 @@ describe("filtered usage index availability", () => { expect(existsSync(`${indexPath}-wal`)).toBe(false); }); + test("a transient foreign classification recovers on the open retry", async () => { + const rows = Array.from({ length: 5 }, (_, index) => syntheticRow(index)); + writeRows(rows); + const expected = rows.filter(row => row.provider === "openai").length; + const indexPath = usageFilterIndexPath(testDir); + const foreign = "this is not a sqlite database\n"; + writeFileSync(indexPath, foreign); + + const request = getFilteredUsageAggregate({ provider: "openai" }); + // Early attempts still see the foreign file; clear it inside the short re-check window so + // a later attempt builds the projection instead of degrading to the streaming read. + const clear = setTimeout(() => rmSync(indexPath, { force: true }), 60); + try { + const result = await request; + expect(summary(result).summary.requests).toBe(expected); + expect(result.filterIndex?.totalRows).toBe(rows.length); + expect(result.filterIndex?.candidateRows).toBe(expected); + } finally { + clearTimeout(clear); + } + }); + + test("a transient unsafe sidecar recovers on the open retry", async () => { + const rows = Array.from({ length: 5 }, (_, index) => syntheticRow(index)); + writeRows(rows); + const expected = rows.filter(row => row.provider === "openai").length; + const indexPath = usageFilterIndexPath(testDir); + const sidecar = `${indexPath}-wal`; + // A directory at a sidecar path is not a private regular file, so the open refuses it. + mkdirSync(sidecar); + + const request = getFilteredUsageAggregate({ provider: "openai" }); + const clear = setTimeout(() => rmSync(sidecar, { recursive: true, force: true }), 60); + try { + const result = await request; + expect(summary(result).summary.requests).toBe(expected); + expect(result.filterIndex?.totalRows).toBe(rows.length); + expect(result.filterIndex?.candidateRows).toBe(expected); + } finally { + clearTimeout(clear); + } + }); + + test("a permanently unsafe sidecar still falls back without being modified", async () => { + const rows = Array.from({ length: 5 }, (_, index) => syntheticRow(index)); + writeRows(rows); + const expected = rows.filter(row => row.provider === "openai").length; + const indexPath = usageFilterIndexPath(testDir); + const sidecar = `${indexPath}-wal`; + mkdirSync(sidecar); + + const result = await getFilteredUsageAggregate({ provider: "openai" }); + expect(summary(result).summary.requests).toBe(expected); + expect(result.filterIndex).toBeUndefined(); + expect(statSync(sidecar).isDirectory()).toBe(true); + expect(existsSync(indexPath)).toBe(false); + }); + // chmod cannot make a Windows directory unwritable, so there this would build the index and // assert nothing about permissions. The degrade itself is the platform-neutral catch-all in // readUsageFilterCandidates, which the vanishing-after-refresh case below reaches on every @@ -1233,6 +1337,259 @@ describe("filtered usage index availability", () => { lock.exec("ROLLBACK"); lock.close(); } + // Sustained contention exhausts the extended busy-retry window before degrading to the + // streaming read, which takes several seconds by design; the default 5s timeout is too tight. + }, 20_000); + + test("an unchanged filtered read releases its index reader", async () => { + const rows = Array.from({ length: 5 }, (_, index) => syntheticRow(index)); + writeRows(rows); + await getFilteredUsageAggregate({ provider: "openai" }); + + const realRead = filterIndexModule.readUsageFilterCandidates; + let closeCalls = 0; + const spy = spyOn(filterIndexModule, "readUsageFilterCandidates").mockImplementation(async (query) => { + const read = await realRead(query); + if (read) { + const originalClose = read.close; + read.close = () => { + closeCalls += 1; + originalClose(); + }; + } + return read; + }); + try { + // The retained projection is unchanged, so this exercises the branch that validates the + // ledger without paging and must still release the reader it opened. + const result = await getFilteredUsageAggregate({ provider: "openai" }); + expect(result.update).toBe("unchanged"); + expect(closeCalls).toBe(1); + } finally { + spy.mockRestore(); + } + }); + + test("a candidate page failure falls back to the streaming read", async () => { + const rows = Array.from({ length: 5 }, (_, index) => syntheticRow(index)); + writeRows(rows); + await getFilteredUsageAggregate({ provider: "openai" }); + resetUsageAggregateCacheForTests(); + + const expected = rows.filter(row => row.provider === "openai").length; + const spy = spyOn(filterIndexModule, "readUsageFilterCandidates").mockImplementation(async () => ({ + snapshot: { + revision: null, + identityKey: "", + generation: "", + indexedOffset: 0, + processedThroughDigest: "", + indexedRows: 0, + invalidRows: 0, + oversizedRows: 0, + minTimestamp: null, + maxTimestamp: null, + ingestPasses: 0, + ingestedRows: 0, + }, + candidateRows: 0, + nextPage() { + throw new UsageFilterIndexPageError(new Error("SQLITE_CORRUPT")); + }, + close() {}, + })); + try { + const result = await getFilteredUsageAggregate({ provider: "openai" }); + expect(summary(result).summary.requests).toBe(expected); + expect(result.filterIndex).toBeUndefined(); + } finally { + spy.mockRestore(); + } + }); + + test("page fallback retries a concurrent ledger replacement", async () => { + const rows = Array.from({ length: 5 }, (_, index) => syntheticRow(index)); + writeRows(rows); + await getFilteredUsageAggregate({ provider: "openai" }); + resetUsageAggregateCacheForTests(); + + const expected = rows.filter(row => row.provider === "openai").length; + const spy = spyOn(filterIndexModule, "readUsageFilterCandidates").mockImplementation(async () => ({ + snapshot: { + revision: null, + identityKey: "", + generation: "", + indexedOffset: 0, + processedThroughDigest: "", + indexedRows: 0, + invalidRows: 0, + oversizedRows: 0, + minTimestamp: null, + maxTimestamp: null, + ingestPasses: 0, + ingestedRows: 0, + }, + candidateRows: 0, + nextPage() { + throw new UsageFilterIndexPageError(new Error("SQLITE_CORRUPT")); + }, + close() {}, + })); + const realScan = ledgerScannerModule.scanUsageLedgerCooperatively; + let scans = 0; + const scanSpy = spyOn(ledgerScannerModule, "scanUsageLedgerCooperatively").mockImplementation(async options => { + if (++scans === 1) throw new ledgerScannerModule.UsageLedgerRebuildRequiredError("content_changed"); + return realScan(options); + }); + try { + const result = await getFilteredUsageAggregate({ provider: "openai" }); + expect(summary(result).summary.requests).toBe(expected); + expect(result.filterIndex).toBeUndefined(); + expect(scans).toBe(2); + } finally { + spy.mockRestore(); + scanSpy.mockRestore(); + } + }); + + test("page fallback retries an overlay change", async () => { + const rows = Array.from({ length: 5 }, (_, index) => syntheticRow(index)); + writeRows(rows); + await getFilteredUsageAggregate({ provider: "openai" }); + resetUsageAggregateCacheForTests(); + + const expected = rows.filter(row => row.provider === "openai").length; + const spy = spyOn(filterIndexModule, "readUsageFilterCandidates").mockImplementation(async () => ({ + snapshot: { + revision: null, + identityKey: "", + generation: "", + indexedOffset: 0, + processedThroughDigest: "", + indexedRows: 0, + invalidRows: 0, + oversizedRows: 0, + minTimestamp: null, + maxTimestamp: null, + ingestPasses: 0, + ingestedRows: 0, + }, + candidateRows: 0, + nextPage() { + throw new UsageFilterIndexPageError(new Error("SQLITE_CORRUPT")); + }, + close() {}, + })); + const realScan = ledgerScannerModule.scanUsageLedgerCooperatively; + let scans = 0; + const scanSpy = spyOn(ledgerScannerModule, "scanUsageLedgerCooperatively").mockImplementation(async options => { + const result = await realScan(options); + if (++scans === 1) refreshUserCostOverlays({ providers: { openai: {} } } as unknown as OcxConfig); + return result; + }); + try { + const result = await getFilteredUsageAggregate({ provider: "openai" }); + expect(summary(result).summary.requests).toBe(expected); + expect(result.filterIndex).toBeUndefined(); + expect(scans).toBe(2); + } finally { + spy.mockRestore(); + scanSpy.mockRestore(); + } + }); + + for (const passive of [false, true]) { + test(`real SQLite page failure rolls back and wraps the error (passive=${passive})`, async () => { + writeRows(Array.from({ length: 5 }, (_, i) => syntheticRow(i))); + await getFilteredUsageAggregate({ provider: "openai" }); + // Explicit nulls match the normalized query shape; a bare `{}` is not a valid caller query. + const query = { provider: null, model: null, apiKeyId: null }; + const read = await (passive + ? filterIndexModule.readUsageFilterCandidatesPassive(query) + : filterIndexModule.readUsageFilterCandidates(query)); + expect(read).not.toBeNull(); + const writer = new Database(usageFilterIndexPath(testDir)); + writer.exec("DROP TABLE usage_rows"); + writer.close(); + const realExec = Database.prototype.exec; + let rolledBack = false; + const execSpy = spyOn(Database.prototype, "exec").mockImplementation(function (this: Database, sql: string) { + const result = realExec.call(this, sql); + if (sql === "ROLLBACK") rolledBack = !this.inTransaction; + return result; + }); + try { + let caught: unknown; + try { read!.nextPage(); } catch (error) { caught = error; } + expect(caught).toBeInstanceOf(UsageFilterIndexPageError); + expect(String((caught as UsageFilterIndexPageError).pageError)).toContain("no such table"); + expect(rolledBack).toBe(true); + } finally { + execSpy.mockRestore(); + read!.close(); + } + }); + } + + test("a generation change between pages keeps the rebuild retry and correct totals", async () => { + const rows = Array.from({ length: 2_500 }, (_, index) => syntheticRow(index)); + writeRows(rows); + await getFilteredUsageAggregate({}); + + // Explicit nulls match the normalized query shape; a bare `{}` is not a valid caller query. + const read = await filterIndexModule.readUsageFilterCandidates({ + provider: null, model: null, apiKeyId: null, + }); + expect(read).not.toBeNull(); + // A full first page proves this reader has another page left to validate. + expect(read!.nextPage()).toHaveLength(2_000); + + try { + // Publish a new generation under the open reader through the normal rebuild flow. + filterIndexModule.requestUsageFilterIndexRebuild(); + await getFilteredUsageAggregate({}); + let caught: unknown; + try { read!.nextPage(); } catch (error) { caught = error; } + expect(caught).toBeInstanceOf(ledgerScannerModule.UsageLedgerRebuildRequiredError); + expect(caught).not.toBeInstanceOf(UsageFilterIndexPageError); + } finally { + read!.close(); + } + + // The retried aggregate still answers from the fresh generation, matching the oracle. + resetUsageAggregateCacheForTests(); + await expectMatchesFullScan({}); + }); + + test("a rebuild mid-aggregate restarts paging with fresh totals", async () => { + const rows = Array.from({ length: 2_500 }, (_, index) => syntheticRow(index)); + writeRows(rows); + await getFilteredUsageAggregate({}); + resetUsageAggregateCacheForTests(); + + const realReadRows = ledgerScannerModule.readUsageLedgerRowsAtOffsets; + let rebuilt = false; + const rowsSpy = spyOn(ledgerScannerModule, "readUsageLedgerRowsAtOffsets").mockImplementation(async options => { + if (!rebuilt) { + rebuilt = true; + // Publish a new generation between the first and second candidate pages. + filterIndexModule.requestUsageFilterIndexRebuild(); + const fresh = await filterIndexModule.readUsageFilterCandidates({ + provider: null, model: null, apiKeyId: null, + }); + fresh?.close(); + } + return realReadRows(options); + }); + try { + const result = await getFilteredUsageAggregate({}); + const oracle = await fullScan({}); + expect(summary(result)).toEqual(oracle.summarize("all", NOW)); + // The retry consumed the rebuilt generation end to end, not partial old-generation rows. + expect(result.filterIndex?.candidateRows).toBe(rows.length); + } finally { + rowsSpy.mockRestore(); + } }); /** Delete the projection at one point of a filtered query and expect a streaming answer. */ @@ -1306,7 +1663,7 @@ describe("filtered usage index availability", () => { // The original home's projection was not advanced with another home's ledger. expect(publishedOffset()).toBe(offsetBefore); }); -+ + test("a passive read does not report a published generation that lost its rows", async () => { const rows = Array.from({ length: 5 }, (_, index) => syntheticRow(index)); writeRows(rows); @@ -1323,7 +1680,7 @@ describe("filtered usage index availability", () => { expect(summary(passive).summary.requests) .toBe(rows.filter(row => row.provider === "openai").length); }); -+ + test("a foreign SQLite file with a coincidental meta table is not adopted", async () => { const rows = Array.from({ length: 5 }, (_, index) => syntheticRow(index)); writeRows(rows); @@ -1415,6 +1772,22 @@ describe("filtered usage index selection evidence", () => { .toContain("idx_usage_rows_ts"); }); + test("the paged candidate select seeks each branch from the page cursor", async () => { + writeRows(Array.from({ length: 50 }, (_, index) => syntheticRow(index))); + await getFilteredUsageAggregate({ provider: "openai" }); + + const paged = (query: Parameters[0]) => + usageFilterCandidatePlan(query, { paged: true }).join("\n"); + // Every UNION branch carries the `byte_offset > cursor` bound and its own LIMIT, so a late + // page seeks from the cursor instead of re-materializing the candidates before it. + const providerOnly = paged({ provider: "openai", model: null, apiKeyId: null }); + expect(providerOnly).toContain("idx_usage_attr_provider"); + expect(providerOnly).toContain("idx_usage_rows_unattributable"); + expect(providerOnly).toContain("byte_offset>"); + expect(paged({ provider: null, model: null, apiKeyId: null })) + .toContain("PRIMARY KEY (generation=? AND byte_offset>"); + }); + test("every source filter selects exactly the rows matchesUsageSource keeps", async () => { const rows: PersistedUsageEntry[] = []; let ordinal = 0;