From d0eebceacf19286ee6248301feaffdd5aa9c4b4e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Samy=20Pess=C3=A9?= Date: Wed, 13 Mar 2024 17:10:07 +0000 Subject: [PATCH] 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 | 98 +++++++++++++++++++++++++++++++++++++++++- src/lib/async.ts | 83 ++++++++++++++++++++++++++++++++++- src/lib/cache/cache.ts | 22 ++++++++-- 3 files changed, 196 insertions(+), 7 deletions(-) diff --git a/src/lib/api.ts b/src/lib/api.ts index db5c644ad..27491e797 100644 --- a/src/lib/api.ts +++ b/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 = {}; + + 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( E > >, + options: { + limit?: number; + } = {}, ): Promise< HttpResponse< List & { @@ -782,7 +876,7 @@ async function getAll( E > > { - const limit = 100; + const { limit = 100 } = options; let page: string | undefined = undefined; const result: T[] = []; diff --git a/src/lib/async.ts b/src/lib/async.ts index 16276c145..33788de94 100644 --- a/src/lib/async.ts +++ b/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(execute: () => Promise): () => Promise { }; } +type SingletonFunction = (( + key: string, + ...args: Args +) => Promise) & { + isRunning(key: string): Promise; +}; + /** * Create a map of singleton operations in a safe way for Cloudflare worker */ export function singletonMap( execute: (key: string, ...args: Args) => Promise, -): (key: string, ...args: Args) => Promise { +): SingletonFunction { const states = new WeakMap>>(); - return async (key, ...args) => { + const fn: SingletonFunction = async (key, ...args) => { const ctx = await getGlobalContext(); let current = states.get(ctx); if (current) { @@ -281,4 +290,74 @@ export function singletonMap( 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( + fn: (executions: Args[]) => Promise, + 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; + }, +): (...args: Args) => Promise { + const { delay, groupBy = () => 'default', skip = () => false } = options; + + const groups = new Map 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((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); + } + }); + }; } diff --git a/src/lib/cache/cache.ts b/src/lib/cache/cache.ts index f15df5ef7..7f020314f 100644 --- a/src/lib/cache/cache.ts +++ b/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 = (( * Refetch the data and update the cache. */ revalidate: (...args: Args | [...Args, CacheFunctionOptions]) => Promise; + + /** + * Check if a value is in the memory cache. + */ + hasInMemory: (...args: Args) => Promise; }; /** @@ -112,9 +118,7 @@ export function cache( 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( 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);