feat: bulk-load documents from the mirror bundle with prefetch

README.md

@@ -78,13 +78,21 @@ file is always refetched. Store errors and corrupt entries are ignored; eviction
 
 A base URL tried before GitHub by `listAll` and `getFile`, such as
 `https://docs.nbtca.space/docs-api`. It serves `index.json` in the shape of GitHub's recursive tree
-response and each file at `raw/<path>`. Any mirror failure, invalid or truncated index, or wait past
+response, each file at `raw/<path>`, and optionally `bundle.json` as
+`{ files: [{ path, sha, content }] }`. Any mirror failure, invalid or truncated index, or wait past
 5 seconds falls back to GitHub. The GitHub token is never sent to the mirror.
 
+### `docs.prefetch()`
+
+Loads every document from the mirror's `bundle.json` in one request and resolves to how many were
+loaded. Only files whose content hashes to the blob id in the current tree are kept. Resolves to `0`
+without a mirror or when the bundle is unavailable.
+
 ### `docs.search(query, options?)`
 
 Searches paths, titles, summaries, Markdown text, and semantic component attributes. Results are
-ranked and include excerpts. Use `pathPrefix` to scope a search and `limit` to cap results.
+ranked and include excerpts. With a mirror, a search that would fetch many uncached documents calls
+`prefetch()` first. Use `pathPrefix` to scope a search and `limit` to cap results.
 
 ### `docs.clear()`
 

src/__tests__/client.test.ts

@@ -1,3 +1,4 @@
+import { createHash } from 'node:crypto';
 import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
 import { createDocsClient } from '../client.js';
 import { DocsFetchError } from '../types.js';
@@ -1207,6 +1208,124 @@ describe('mirror', () => {
     }
   });
 
+  describe('prefetch', () => {
+    const bundleOf = (files: { path: string; sha: string; content: string }[]) => ({
+      ok: true,
+      json: async () => ({ files }),
+    });
+
+    function mirrorWith(bundle: unknown) {
+      return routes({
+        mirror: (url) =>
+          url.endsWith('/index.json') ? { ok: true, json: async () => mirrorTree } : bundle,
+      });
+    }
+
+    it('loads every verified file from the bundle in one request', async () => {
+      const fetchMock = mirrorWith(
+        bundleOf([{ path: 'repair/guide.md', sha: guideSha, content: '# Guide' }]),
+      );
+      const store = memoryStore();
+      const client = createDocsClient({ mirror, store });
+
+      await expect(client.prefetch()).resolves.toBe(1);
+      await expect(client.getFile('repair/guide.md')).resolves.toBe('# Guide');
+      expect(store.map.get(`blob-${guideSha}`)).toBe('# Guide');
+      expect(fetchMock.mock.calls.map(([url]) => String(url))).toEqual([
+        'https://docs.example.org/docs-api/index.json',
+        'https://docs.example.org/docs-api/bundle.json',
+      ]);
+    });
+
+    it('rejects files whose content or sha does not match the tree', async () => {
+      const staleSha = '0'.repeat(40);
+      mirrorWith(
+        bundleOf([
+          { path: 'repair/guide.md', sha: guideSha, content: '# Tampered' },
+          { path: 'repair/guide.md', sha: staleSha, content: '# Stale' },
+          { path: 'other.md', sha: guideSha, content: '# Guide' },
+        ]),
+      );
+      const store = memoryStore();
+
+      await expect(createDocsClient({ mirror, store }).prefetch()).resolves.toBe(0);
+      expect([...store.map.keys()]).toEqual(['tree']);
+    });
+
+    it.each([
+      ['a malformed bundle', { ok: true, json: async () => ({ files: [{ path: 1 }] }) }],
+      ['a missing bundle', { ok: false, status: 404 }],
+      ['a failing mirror', { ok: false, status: 503 }],
+    ])('loads nothing from %s', async (_label, bundle) => {
+      mirrorWith(bundle);
+      await expect(createDocsClient({ mirror }).prefetch()).resolves.toBe(0);
+    });
+
+    it('keeps reading single files from the mirror when the bundle fails', async () => {
+      const fetchMock = routes({
+        mirror: (url) =>
+          url.endsWith('/index.json')
+            ? { ok: true, json: async () => mirrorTree }
+            : url.endsWith('/bundle.json')
+              ? { ok: false, status: 503 }
+              : { ok: true, text: async () => '# From mirror' },
+      });
+      const client = createDocsClient({ mirror });
+
+      await expect(client.prefetch()).resolves.toBe(0);
+      await expect(client.getFile('repair/guide.md')).resolves.toBe('# From mirror');
+      expect(String(fetchMock.mock.calls.at(-1)?.[0])).toContain('docs.example.org');
+    });
+
+    it('asks the mirror again after clear()', async () => {
+      let mirrorUp = false;
+      const fetchMock = routes({
+        mirror: () =>
+          mirrorUp ? { ok: true, json: async () => mirrorTree } : { ok: false, status: 503 },
+        github: () => ({ ok: true, json: async () => mockTree }),
+      });
+      const client = createDocsClient({ mirror });
+
+      await client.listAll();
+      mirrorUp = true;
+      client.clear();
+      await client.listAll();
+      expect(String(fetchMock.mock.calls.at(-1)?.[0])).toContain('docs.example.org');
+    });
+
+    it('does nothing without a mirror', async () => {
+      const fetchMock = vi.fn();
+      vi.stubGlobal('fetch', fetchMock);
+      await expect(createDocsClient().prefetch()).resolves.toBe(0);
+      expect(fetchMock).not.toHaveBeenCalled();
+    });
+
+    it('lets search read many uncached documents from the bundle', async () => {
+      const files = Array.from({ length: 8 }, (_, index) => {
+        const content = `# Repair ${String(index)}`;
+        const sha = createHash('sha1')
+          .update(`blob ${String(content.length)}\0${content}`)
+          .digest('hex');
+        return { path: `repair/${String(index)}.md`, sha, content };
+      });
+      const fetchMock = routes({
+        mirror: (url) =>
+          url.endsWith('/index.json')
+            ? {
+                ok: true,
+                json: async () => ({
+                  truncated: false,
+                  tree: files.map(({ path, sha }) => ({ path, type: 'blob', sha })),
+                }),
+              }
+            : bundleOf(files),
+      });
+
+      await expect(createDocsClient({ mirror }).search('repair')).resolves.toHaveLength(8);
+      expect(fetchMock).toHaveBeenCalledTimes(2);
+    });
+  });
+
   it('rejects a mirror that is not a plain http(s) URL', () => {
     for (const value of ['docs.example.org', 'ftp://docs.example.org', `${mirror}?v=1`]) {
       expect(() => createDocsClient({ mirror: value })).toThrow(TypeError);

src/client.ts

@@ -43,7 +43,9 @@ const SKIP = new Set([
 ]);
 
 const TREE_KEY = 'tree';
+const BUNDLE_KEY = 'bundle';
 const MIRROR_TIMEOUT_MS = 5_000;
+const BUNDLE_TIMEOUT_MS = 30_000;
 const SHA = /^[0-9a-f]{40}$/;
 const SEARCH_CONCURRENCY = 6;
 const SEARCH_RESULT_LIMIT = 20;
@@ -185,6 +187,12 @@ interface GitHubTreeResponse {
   truncated: boolean;
 }
 
+interface BundleFile {
+  path: string;
+  sha: string;
+  content: string;
+}
+
 function encodePath(path: string): string {
   return path
     .split('/')
@@ -340,6 +348,22 @@ function parseTreeResponse(value: unknown): GitHubTreeResponse {
   return { tree: value.tree, truncated: value.truncated };
 }
 
+function isBundleFile(value: unknown): value is BundleFile {
+  return (
+    isRecord(value) &&
+    typeof value.path === 'string' &&
+    typeof value.sha === 'string' &&
+    typeof value.content === 'string'
+  );
+}
+
+function parseBundle(value: unknown): BundleFile[] {
+  if (!isRecord(value) || !Array.isArray(value.files) || !value.files.every(isBundleFile)) {
+    throw new TypeError('Invalid mirror bundle');
+  }
+  return value.files;
+}
+
 export function createDocsClient(options: DocsClientOptions = {}): DocsClient {
   const owner = options.owner ?? DEFAULTS.owner;
   const repo = options.repo ?? DEFAULTS.repo;
@@ -361,11 +385,12 @@ export function createDocsClient(options: DocsClientOptions = {}): DocsClient {
   const fileTtlMs = cacheTtl(options.cacheTtlMs?.file, DEFAULTS.fileTtlMs, 'cacheTtlMs.file');
 
   const dirCache = new TtlCache<DocItem[]>(dirTtlMs, 30);
-  const fileCache = new TtlCache<string>(fileTtlMs, 200);
+  const fileCache = new TtlCache<string>(fileTtlMs, 1000);
   const treeCache = new TtlCache<DocItem[]>(dirTtlMs, 1);
   const dirRequests = new Map<string, Promise<DocItem[]>>();
   const fileRequests = new Map<string, Promise<string>>();
   const treeRequests = new Map<string, Promise<DocItem[]>>();
+  const bundleRequests = new Map<string, Promise<number>>();
   const store = options.store;
   let storedTree: DocItem[] | undefined;
   let cacheGeneration = 0;
@@ -426,20 +451,22 @@ export function createDocsClient(options: DocsClientOptions = {}): DocsClient {
     base: string,
     path: string,
     read: (response: Response) => Promise<T | undefined>,
+    timeoutMs = MIRROR_TIMEOUT_MS,
+    tripsBreaker = true,
   ): Promise<T | undefined> {
     if (mirrorFailed) return undefined;
     try {
       return await withResponse(
         `${base}/${path}`,
-        MIRROR_TIMEOUT_MS,
+        timeoutMs,
         async (response) => {
-          if (response.status >= 500) mirrorFailed = true;
+          if (response.status >= 500 && tripsBreaker) mirrorFailed = true;
           return response.ok ? await read(response) : undefined;
         },
         {},
       );
     } catch {
-      mirrorFailed = true;
+      if (tripsBreaker) mirrorFailed = true;
       return undefined;
     }
   }
@@ -452,9 +479,14 @@ export function createDocsClient(options: DocsClientOptions = {}): DocsClient {
     return items;
   }
 
-  async function keepFile(path: string, content: string, generation: number): Promise<string> {
+  async function keepFile(
+    path: string,
+    content: string,
+    generation: number,
+    sha?: string,
+  ): Promise<string> {
     if (generation === cacheGeneration) fileCache.set(path, content);
-    if (store) storeWrite(`blob-${await blobSha(content)}`, content);
+    if (store) storeWrite(`blob-${sha ?? (await blobSha(content))}`, content);
     return content;
   }
 
@@ -615,6 +647,37 @@ export function createDocsClient(options: DocsClientOptions = {}): DocsClient {
     return shareRequest(fileRequests, path, () => loadFile(path, cacheGeneration));
   }
 
+  async function loadBundle(base: string, generation: number): Promise<number> {
+    const shas = new Map((await listAll()).map((item) => [item.path, item.sha]));
+    // The bundle is optional: losing it must not stop per-file mirror reads.
+    const files = await fromMirror(
+      base,
+      'bundle.json',
+      async (response) => parseBundle(await response.json()),
+      BUNDLE_TIMEOUT_MS,
+      false,
+    );
+    let loaded = 0;
+    for (const file of files ?? []) {
+      if (file.sha !== shas.get(file.path) || (await blobSha(file.content)) !== file.sha) continue;
+      await keepFile(file.path, file.content, generation, file.sha);
+      loaded += 1;
+    }
+    return loaded;
+  }
+
+  function prefetch(): Promise<number> {
+    if (mirror === undefined) return Promise.resolve(0);
+    return shareRequest(bundleRequests, BUNDLE_KEY, () => loadBundle(mirror, cacheGeneration));
+  }
+
+  function isUncached(item: DocItem): boolean {
+    return (
+      fileCache.get(item.path) === undefined &&
+      (item.sha === undefined || storeRead(`blob-${item.sha}`) === undefined)
+    );
+  }
+
   async function getDocument(path: string): Promise<DocPage> {
     if (!path.toLowerCase().endsWith('.md')) {
       throw new TypeError('path must point to a Markdown document');
@@ -636,6 +699,9 @@ export function createDocsClient(options: DocsClientOptions = {}): DocsClient {
     const candidates = pathPrefix
       ? all.filter((item) => item.path.startsWith(`${pathPrefix}/`))
       : all;
+    if (mirror !== undefined && candidates.filter(isUncached).length > SEARCH_CONCURRENCY) {
+      await prefetch();
+    }
     let loaded = 0;
     let firstFailure: DocsFetchError | undefined;
     const matches = await mapConcurrent(candidates, SEARCH_CONCURRENCY, async (item) => {
@@ -667,7 +733,19 @@ export function createDocsClient(options: DocsClientOptions = {}): DocsClient {
     dirRequests.clear();
     fileRequests.clear();
     treeRequests.clear();
-  }
-
-  return { listDir, listAll, peekAll, listSections, getFile, getDocument, search, clear };
+    bundleRequests.clear();
+    mirrorFailed = false;
+  }
+
+  return {
+    listDir,
+    listAll,
+    peekAll,
+    listSections,
+    getFile,
+    getDocument,
+    prefetch,
+    search,
+    clear,
+  };
 }

src/types.ts

@@ -68,6 +68,7 @@ export interface DocsClient {
   listSections(): Promise<DocSection[]>;
   getFile(path: string): Promise<string>;
   getDocument(path: string): Promise<DocPage>;
+  prefetch(): Promise<number>;
   search(query: string, options?: DocsSearchOptions): Promise<DocsSearchResult[]>;
   clear(): void;
 }