Batch request for getRevisionFile to avoid too many sub-requests (#259)

* Try approach with batching

* Improve it to support getRevision

* Fix

* Real fix

* Implement skip

* Lint

src/lib/api.ts

@@ -10,11 +10,13 @@ import {
     List,
     PublishedContentLookup,
     RequestRenderIntegrationUI,
+    RevisionFile,
 } from '@gitbook/api';
 import assertNever from 'assert-never';
 import { headers } from 'next/headers';
 import rison from 'rison';
 
+import { batch } from './async';
 import { buildVersion } from './build';
 import {
     CacheFunctionOptions,
@@ -415,8 +417,9 @@ export const getRevisionPageByPath = cache(
 
 /**
  * Resolve a file by its ID.
+ * It should not be used directly, use `getRevisionFile` instead.
  */
-export const getRevisionFile = cache(
+const getRevisionFileById = cache(
     'api.getRevisionFile.v2',
     async (spaceId: string, revisionId: string, fileId: string, options: CacheFunctionOptions) => {
         try {
@@ -446,6 +449,94 @@ export const getRevisionFile = cache(
     },
 );
 
+/**
+ * Get all the files in a revision of a space.
+ * It should not be used directly, use `getRevisionFile` instead.
+ */
+const getRevisionAllFiles = cache(
+    'api.getRevisionAllFiles',
+    async (spaceId: string, revisionId: string, options: CacheFunctionOptions) => {
+        const response = await getAll(
+            (params) =>
+                api().spaces.listFilesInRevisionById(
+                    spaceId,
+                    revisionId,
+                    {
+                        ...params,
+                        metadata: false,
+                    },
+                    {
+                        ...noCacheFetchOptions,
+                        signal: options.signal,
+                    },
+                ),
+            {
+                limit: 1000,
+            },
+        );
+
+        const files: { [fileId: string]: RevisionFile } = {};
+        response.data.items.forEach((file) => {
+            files[file.id] = file;
+        });
+
+        return cacheResponse(response, { ...immutableCacheTtl_7days, data: files });
+    },
+);
+
+/**
+ * Resolve a file by its ID.
+ * The approach is optimized to use the entire list of files in the revision if it has been fetched
+ * or to use a per-file approach if not.
+ */
+export const getRevisionFile = batch<[string, string, string], RevisionFile | null>(
+    async (executions) => {
+        const [spaceId, revisionId] = executions[0];
+
+        const hasRevisionInMemory = await getRevision.hasInMemory(spaceId, revisionId, {
+            metadata: false,
+        });
+        const hasRevisionFilesInMemory = await getRevisionAllFiles.hasInMemory(spaceId, revisionId);
+
+        // When fetching more than 5 files, we should bundle them all into one call to get the entire revision
+        if (executions.length > 5 || hasRevisionFilesInMemory || hasRevisionInMemory) {
+            let files: Record<string, RevisionFile> = {};
+
+            if (hasRevisionInMemory) {
+                const revision = await getRevision(spaceId, revisionId, { metadata: false });
+                files = {};
+                revision.files.forEach((file) => {
+                    files[file.id] = file;
+                });
+            } else {
+                files = await getRevisionAllFiles(spaceId, revisionId);
+            }
+
+            return executions.map(([spaceId, revisionId, fileId]) => files[fileId] ?? null);
+        } else {
+            // Fetch file individually
+            return Promise.all(
+                executions.map(([spaceId, revisionId, fileId]) =>
+                    getRevisionFileById(spaceId, revisionId, fileId),
+                ),
+            );
+        }
+    },
+    {
+        delay: 20,
+        groupBy: (spaceId, revisionId) => spaceId + '/' + revisionId,
+        skip: async (spaceId, revisionId, fileId) => {
+            return (
+                (await getRevision.hasInMemory(spaceId, revisionId, {
+                    metadata: false,
+                })) ||
+                (await getRevisionAllFiles.hasInMemory(spaceId, revisionId)) ||
+                (await getRevisionFileById.hasInMemory(spaceId, revisionId, fileId))
+            );
+        },
+    },
+);
+
 /**
  * Get a document by its ID.
  */
@@ -774,6 +865,9 @@ async function getAll<T, E>(
             E
         >
     >,
+    options: {
+        limit?: number;
+    } = {},
 ): Promise<
     HttpResponse<
         List & {
@@ -782,7 +876,7 @@ async function getAll<T, E>(
         E
     >
 > {
-    const limit = 100;
+    const { limit = 100 } = options;
 
     let page: string | undefined = undefined;
     const result: T[] = [];

src/lib/async.ts

@@ -1,3 +1,5 @@
+import { MaybePromise } from 'p-map';
+
 import { waitUntil, getGlobalContext } from './waitUntil';
 
 /**
@@ -251,15 +253,22 @@ export function singleton<R>(execute: () => Promise<R>): () => Promise<R> {
     };
 }
 
+type SingletonFunction<Args extends any[], Result> = ((
+    key: string,
+    ...args: Args
+) => Promise<Result>) & {
+    isRunning(key: string): Promise<boolean>;
+};
+
 /**
  * Create a map of singleton operations in a safe way for Cloudflare worker
  */
 export function singletonMap<Args extends any[], Result>(
     execute: (key: string, ...args: Args) => Promise<Result>,
-): (key: string, ...args: Args) => Promise<Result> {
+): SingletonFunction<Args, Result> {
     const states = new WeakMap<object, Map<string, Promise<Result>>>();
 
-    return async (key, ...args) => {
+    const fn: SingletonFunction<Args, Result> = async (key, ...args) => {
         const ctx = await getGlobalContext();
         let current = states.get(ctx);
         if (current) {
@@ -281,4 +290,74 @@ export function singletonMap<Args extends any[], Result>(
 
         return promise;
     };
+
+    fn.isRunning = async (key: string) => {
+        const ctx = await getGlobalContext();
+        const current = states.get(ctx);
+        return current?.has(key) ?? false;
+    };
+
+    return fn;
+}
+
+/**
+ * Batch the calls to a function and resolve them all ar once
+ */
+export function batch<Args extends any[], R>(
+    fn: (executions: Args[]) => Promise<R[]>,
+    options: {
+        /**
+         * Maximum delay in milliseconds before the batch is resolved.
+         */
+        delay: number;
+
+        /**
+         * Group the calls by a key.
+         * @default () => 'default'
+         */
+        groupBy?: (...args: Args) => string;
+
+        /**
+         * Skip the batching for a single call.
+         */
+        skip?: (...args: Args) => MaybePromise<boolean>;
+    },
+): (...args: Args) => Promise<R> {
+    const { delay, groupBy = () => 'default', skip = () => false } = options;
+
+    const groups = new Map<string, Array<[Args, (r: R) => void, (error: Error) => void]>>();
+    let timeoutId: NodeJS.Timeout | null = null;
+
+    return async (...args) => {
+        if (await skip(...args)) {
+            const results = await fn([args]);
+            return results[0];
+        }
+
+        return new Promise<R>((resolve, reject) => {
+            const groupId = groupBy(...args);
+            const executions = groups.get(groupId) ?? [];
+            groups.set(groupId, executions);
+
+            executions.push([args, resolve, reject]);
+            if (executions.length === 1) {
+                timeoutId = setTimeout(() => {
+                    const currentExecutions = executions.splice(0, executions.length);
+                    const batchArgs = currentExecutions.map(([args]) => args);
+                    fn(batchArgs).then(
+                        (results) => {
+                            currentExecutions.forEach(([, resolve], index) => {
+                                resolve(results[index]);
+                            });
+                        },
+                        (error) => {
+                            currentExecutions.forEach(([, , reject]) => {
+                                reject(error);
+                            });
+                        },
+                    );
+                }, delay);
+            }
+        });
+    };
 }

src/lib/cache/cache.ts

@@ -1,6 +1,7 @@
 import hash from 'object-hash';
 
 import { cacheBackends } from './backends';
+import { memoryCache } from './memory';
 import { CacheEntry } from './types';
 import { race, singletonMap } from '../async';
 import { TraceSpan, trace } from '../tracing';
@@ -17,6 +18,11 @@ export type CacheFunction<Args extends any[], Result> = ((
      * Refetch the data and update the cache.
      */
     revalidate: (...args: Args | [...Args, CacheFunctionOptions]) => Promise<void>;
+
+    /**
+     * Check if a value is in the memory cache.
+     */
+    hasInMemory: (...args: Args) => Promise<boolean>;
 };
 
 /**
@@ -112,9 +118,7 @@ export function cache<Args extends any[], Result>(
             let result: readonly [CacheEntry, string] | null = null;
 
             // Try the memory backend, independently of the other backends as it doesn't have a network cost
-            const memoryEntry = await cacheBackends
-                .find((backend) => backend.name === 'memory')
-                ?.get(key);
+            const memoryEntry = await memoryCache.get(key);
             if (memoryEntry) {
                 span.setAttribute('memory', true);
                 result = [memoryEntry, 'memory'] as const;
@@ -227,6 +231,18 @@ export function cache<Args extends any[], Result>(
         await revalidate(key, signal, ...args);
     };
 
+    cacheFn.hasInMemory = async (...args: Args) => {
+        const cacheArgs = options.extractArgs ? options.extractArgs(args) : args;
+        const key = getCacheKey(cacheName, cacheArgs);
+
+        const memoryEntry = await memoryCache.get(key);
+        if (memoryEntry) {
+            return true;
+        }
+
+        return fetchValue.isRunning(key);
+    };
+
     // @ts-ignore
     registeredCaches.set(cacheName, cacheFn);