Adapt GBO cache logic to support users sessions info (#2606)

This commit is contained in:
spastorelli
2024-12-10 14:11:59 +01:00
committed by GitHub
parent 664debc0bc
commit 31396d8109
9 changed files with 192 additions and 55 deletions
Executable → Regular
View File
@@ -42,7 +42,7 @@ export default async function ContentLayout(props: { children: React.ReactNode }
sections,
} = await fetchContentData();
ReactDOM.preconnect(api().endpoint);
ReactDOM.preconnect(api().client.endpoint);
if (assetsDomain) {
ReactDOM.preconnect(assetsDomain);
}
@@ -148,7 +148,7 @@ export function PageBody(props: {
sitePointer={pointer}
spaceId={space.id}
pageId={page.id}
apiHost={api().endpoint}
apiHost={api().client.endpoint}
/>
) : null}
</>
@@ -18,7 +18,7 @@ export async function postPageFeedback(args: {
`No siteSpaceId in pointer. organizationId: ${organizationId}, siteId: ${siteId}, pageId: ${args.pageId}`,
);
await api().orgs.createSitesPageFeedback(
await api().client.orgs.createSitesPageFeedback(
organizationId,
siteId,
siteSpaceId,
@@ -158,7 +158,7 @@ export const streamAskQuestion = streamResponse(async function* (
siteSpaceId: string | null,
question: string,
) {
const stream = api.api().orgs.streamAskInSite(
const stream = api.api().client.orgs.streamAskInSite(
organizationId,
siteId,
{
+57 -29
View File
@@ -87,14 +87,30 @@ const cacheTtl_1day = {
ttl: 24 * 60 * 60,
};
const apiSyncStorage = new AsyncLocalStorage<GitBookAPI>();
export type GitBookAPIContext = {
/**
* Instance of the GitBook API client.
*/
client: GitBookAPI;
/**
* Context ID representing a hash of the visitor's attributes/assertions that are
* included in the claims property of the content API JWT token.
*
* It serves as a suffix for the cache key to ensure that the content cache is invalidated
* when these attributees/assertions change.
*/
contextId: string | undefined;
};
const apiSyncStorage = new AsyncLocalStorage<GitBookAPIContext>();
export const DEFAULT_API_ENDPOINT = process.env.GITBOOK_API_URL ?? 'https://api.gitbook.com';
/**
* Create a new API client with a token.
*/
export function apiWithToken(apiToken: string): GitBookAPI {
export function apiWithToken(apiToken: string, contextId: string | undefined): GitBookAPIContext {
const headersList = headers();
const apiEndpoint = headersList.get('x-gitbook-api') ?? DEFAULT_API_ENDPOINT;
@@ -104,13 +120,13 @@ export function apiWithToken(apiToken: string): GitBookAPI {
userAgent: userAgent(),
});
return gitbook;
return { client: gitbook, contextId };
}
/**
* Create an API client for the current request.
*/
export function api(): GitBookAPI {
export function api(): GitBookAPIContext {
const existing = apiSyncStorage.getStore();
if (existing) {
return existing;
@@ -118,6 +134,7 @@ export function api(): GitBookAPI {
const headersList = headers();
const apiToken = headersList.get('x-gitbook-token');
const contextId = headersList.get('x-gitbook-token-context') ?? undefined;
if (!apiToken) {
throw new Error(
@@ -125,13 +142,13 @@ export function api(): GitBookAPI {
);
}
return apiWithToken(apiToken);
return apiWithToken(apiToken, contextId);
}
/**
* Use an API client for an async function.
*/
export function withAPI<T>(client: GitBookAPI, fn: () => Promise<T>): Promise<T> {
export function withAPI<T>(client: GitBookAPIContext, fn: () => Promise<T>): Promise<T> {
return apiSyncStorage.run(client, fn);
}
@@ -164,7 +181,7 @@ export const getUserById = cache({
}),
get: async (userId: string, options: CacheFunctionOptions) => {
try {
const response = await api().users.getUserById(userId, {
const response = await api().client.users.getUserById(userId, {
signal: options.signal,
...noCacheFetchOptions,
});
@@ -200,7 +217,7 @@ export const getPublishedContentByUrl = cache({
options: CacheFunctionOptions,
) => {
try {
const response = await api().urls.getPublishedContentByUrl(
const response = await api().client.urls.getPublishedContentByUrl(
{
url,
visitorAuthToken,
@@ -249,7 +266,7 @@ export const getSpace = cache({
name: 'api.getSpace',
tag: (spaceId) => getAPICacheTag({ tag: 'space', space: spaceId }),
get: async (spaceId: string, shareKey: string | undefined, options: CacheFunctionOptions) => {
const response = await api().spaces.getSpaceById(
const response = await api().client.spaces.getSpaceById(
spaceId,
{
shareKey,
@@ -273,7 +290,7 @@ export const getChangeRequest = cache({
tag: (spaceId, changeRequestId) =>
getAPICacheTag({ tag: 'change-request', space: spaceId, changeRequest: changeRequestId }),
get: async (spaceId: string, changeRequestId: string, options: CacheFunctionOptions) => {
const response = await api().spaces.getChangeRequestById(spaceId, changeRequestId, {
const response = await api().client.spaces.getChangeRequestById(spaceId, changeRequestId, {
...noCacheFetchOptions,
signal: options.signal,
});
@@ -301,13 +318,14 @@ export const getRevision = cache({
name: 'api.getRevision.v2',
tag: (spaceId, revisionId) =>
getAPICacheTag({ tag: 'revision', space: spaceId, revision: revisionId }),
getKeySuffix: () => api().contextId,
get: async (
spaceId: string,
revisionId: string,
fetchOptions: GetRevisionOptions,
options: CacheFunctionOptions,
) => {
const response = await api().spaces.getRevisionById(
const response = await api().client.spaces.getRevisionById(
spaceId,
revisionId,
{
@@ -331,13 +349,14 @@ export const getRevisionPages = cache({
name: 'api.getRevisionPages.v4',
tag: (spaceId, revisionId) =>
getAPICacheTag({ tag: 'revision', space: spaceId, revision: revisionId }),
getKeySuffix: () => api().contextId,
get: async (
spaceId: string,
revisionId: string,
fetchOptions: GetRevisionOptions,
options: CacheFunctionOptions,
) => {
const response = await api().spaces.listPagesInRevisionById(
const response = await api().client.spaces.listPagesInRevisionById(
spaceId,
revisionId,
{
@@ -364,6 +383,7 @@ export const getRevisionPageByPath = cache({
name: 'api.getRevisionPageByPath.v3',
tag: (spaceId, revisionId) =>
getAPICacheTag({ tag: 'revision', space: spaceId, revision: revisionId }),
getKeySuffix: () => api().contextId,
get: async (
spaceId: string,
revisionId: string,
@@ -373,7 +393,7 @@ export const getRevisionPageByPath = cache({
const encodedPath = encodeURIComponent(pagePath);
try {
const response = await api().spaces.getPageInRevisionByPath(
const response = await api().client.spaces.getPageInRevisionByPath(
spaceId,
revisionId,
encodedPath,
@@ -416,7 +436,7 @@ const getRevisionFileById = cache({
) => {
try {
const response = await (async () => {
return api().spaces.getFileInRevisionById(
return api().client.spaces.getFileInRevisionById(
spaceId,
revisionId,
fileId,
@@ -445,6 +465,7 @@ const getRevisionReusableContentById = cache({
name: 'api.getRevisionReusableContentById.v1',
tag: (spaceId, revisionId) =>
getAPICacheTag({ tag: 'revision', space: spaceId, revision: revisionId }),
getKeySuffix: () => api().contextId,
get: async (
spaceId: string,
revisionId: string,
@@ -453,7 +474,7 @@ const getRevisionReusableContentById = cache({
) => {
try {
const response = await (async () => {
return api().spaces.getReusableContentInRevisionById(
return api().client.spaces.getReusableContentInRevisionById(
spaceId,
revisionId,
reusableContentId,
@@ -489,7 +510,7 @@ const getRevisionAllFiles = cache({
get: async (spaceId: string, revisionId: string, options: CacheFunctionOptions) => {
const response = await getAll(
(params) =>
api().spaces.listFilesInRevisionById(
api().client.spaces.listFilesInRevisionById(
spaceId,
revisionId,
{
@@ -600,8 +621,9 @@ export const getDocument = cache({
name: 'api.getDocument.v2',
tag: (spaceId, documentId) =>
getAPICacheTag({ tag: 'document', space: spaceId, document: documentId }),
getKeySuffix: () => api().contextId,
get: async (spaceId: string, documentId: string, options: CacheFunctionOptions) => {
const response = await api().spaces.getDocumentById(
const response = await api().client.spaces.getDocumentById(
spaceId,
documentId,
{
@@ -626,6 +648,7 @@ export const getDocument = cache({
export const getSiteRedirectBySource = cache({
name: 'api.getSiteRedirectBySource',
tag: ({ siteId }) => getAPICacheTag({ tag: 'site', site: siteId }),
getKeySuffix: () => api().contextId,
get: async (
args: {
organizationId: string;
@@ -637,7 +660,7 @@ export const getSiteRedirectBySource = cache({
options: CacheFunctionOptions,
) => {
try {
const response = await api().orgs.getSiteRedirectBySource(
const response = await api().client.orgs.getSiteRedirectBySource(
args.organizationId,
args.siteId,
{
@@ -669,8 +692,9 @@ export const getSiteRedirectBySource = cache({
export const getSite = cache({
name: 'api.getSite',
tag: (organizationId, siteId) => getAPICacheTag({ tag: 'site', site: siteId }),
getKeySuffix: () => api().contextId,
get: async (organizationId: string, siteId: string, options: CacheFunctionOptions) => {
const response = await api().orgs.getSiteById(organizationId, siteId, {
const response = await api().client.orgs.getSiteById(organizationId, siteId, {
...noCacheFetchOptions,
signal: options.signal,
});
@@ -686,6 +710,7 @@ export const getSite = cache({
export const getPublishedContentSite = cache({
name: 'api.getPublishedContentSite',
tag: ({ siteId }) => getAPICacheTag({ tag: 'site', site: siteId }),
getKeySuffix: () => api().contextId,
get: async (
args: {
organizationId: string;
@@ -694,7 +719,7 @@ export const getPublishedContentSite = cache({
},
options: CacheFunctionOptions,
) => {
const response = await api().orgs.getPublishedContentSite(
const response = await api().client.orgs.getPublishedContentSite(
args.organizationId,
args.siteId,
{
@@ -826,7 +851,7 @@ export const getCollection = cache({
name: 'api.getCollection',
tag: (collectionId) => getAPICacheTag({ tag: 'collection', collection: collectionId }),
get: async (collectionId: string, options: CacheFunctionOptions) => {
const response = await api().collections.getCollectionById(collectionId, {
const response = await api().client.collections.getCollectionById(collectionId, {
...noCacheFetchOptions,
signal: options.signal,
});
@@ -844,7 +869,7 @@ export const getCollectionSpaces = cache({
tag: (collectionId) => getAPICacheTag({ tag: 'collection', collection: collectionId }),
get: async (collectionId: string, options: CacheFunctionOptions) => {
const response = await getAll((params) =>
api().collections.listSpacesInCollectionById(collectionId, params, {
api().client.collections.listSpacesInCollectionById(collectionId, params, {
...noCacheFetchOptions,
signal: options.signal,
}),
@@ -896,6 +921,7 @@ export async function getSpaceContentData(
export const searchSpaceContent = cache({
name: 'api.searchSpaceContent',
tag: (spaceId) => getAPICacheTag({ tag: 'space', space: spaceId }),
getKeySuffix: () => api().contextId,
get: async (
spaceId: string,
/** The revision ID is used as a cache bust key, to avoid revalidating lot of cache entries by tags */
@@ -903,7 +929,7 @@ export const searchSpaceContent = cache({
query: string,
options: CacheFunctionOptions,
) => {
const response = await api().spaces.searchSpaceContent(
const response = await api().client.spaces.searchSpaceContent(
spaceId,
{ query },
{
@@ -921,8 +947,9 @@ export const searchSpaceContent = cache({
export const searchParentContent = cache({
name: 'api.searchParentContent',
tag: (spaceId) => getAPICacheTag({ tag: 'space', space: spaceId }),
getKeySuffix: () => api().contextId,
get: async (parentId: string, query: string, options: CacheFunctionOptions) => {
const response = await api().search.searchContent(
const response = await api().client.search.searchContent(
{ query },
{
...noCacheFetchOptions,
@@ -941,6 +968,7 @@ export const searchParentContent = cache({
export const searchSiteContent = cache({
name: 'api.searchSiteContent',
tag: (organizationId, siteId) => getAPICacheTag({ tag: 'site', site: siteId }),
getKeySuffix: () => api().contextId,
get: async (
organizationId: string,
siteId: string,
@@ -953,7 +981,7 @@ export const searchSiteContent = cache({
cacheBust?: string,
options?: CacheFunctionOptions,
) => {
const response = await api().orgs.searchSiteContent(
const response = await api().client.orgs.searchSiteContent(
organizationId,
siteId,
{
@@ -980,7 +1008,7 @@ export const getRecommendedQuestionsInSpace = cache({
name: 'api.getRecommendedQuestionsInSpace',
tag: (spaceId) => getAPICacheTag({ tag: 'space', space: spaceId }),
get: async (spaceId: string, options: CacheFunctionOptions) => {
const response = await api().spaces.getRecommendedQuestionsInSpace(spaceId, {
const response = await api().client.spaces.getRecommendedQuestionsInSpace(spaceId, {
...noCacheFetchOptions,
signal: options.signal,
});
@@ -999,7 +1027,7 @@ export const renderIntegrationUi = cache({
request: RequestRenderIntegrationUI,
options: CacheFunctionOptions,
) => {
const response = await api().integrations.renderIntegrationUiWithPost(
const response = await api().client.integrations.renderIntegrationUiWithPost(
integrationName,
request,
{
@@ -1018,7 +1046,7 @@ export const renderIntegrationUi = cache({
export const getEmbedByUrl = cache({
name: 'api.getEmbedByUrl',
get: async (url: string, options: CacheFunctionOptions) => {
const response = await api().urls.getEmbedByUrl(
const response = await api().client.urls.getEmbedByUrl(
{ url },
{
...noCacheFetchOptions,
@@ -1036,7 +1064,7 @@ export const getEmbedByUrlInSpace = cache({
name: 'api.getEmbedByUrlInSpace',
tag: (spaceId) => getAPICacheTag({ tag: 'space', space: spaceId }),
get: async (spaceId: string, url: string, options: CacheFunctionOptions) => {
const response = await api().spaces.getEmbedByUrlInSpace(
const response = await api().client.spaces.getEmbedByUrlInSpace(
spaceId,
{ url },
{
+79 -2
View File
@@ -1,6 +1,7 @@
import { Mock, beforeEach, describe, expect, it, mock } from 'bun:test';
import { beforeEach, describe, expect, it, Mock, mock } from 'bun:test';
import hash from 'object-hash';
import { CacheFunction, CacheFunctionOptions, cache } from './cache';
import { CacheDefinition, CacheFunction, CacheFunctionOptions, cache } from './cache';
describe('cache', () => {
const impl = mock((arg: string) => 'test-' + arg);
@@ -69,3 +70,79 @@ describe('cache', () => {
expect(impl).toHaveBeenCalledTimes(2);
});
});
describe('cache with suffix key', () => {
const impl = mock((arg: string) => 'test-' + arg);
const getKeySuffixImpl: Mock<NonNullable<CacheDefinition<[string], string>['getKeySuffix']>> =
mock(() => hash({ test: 1 }));
let fn: CacheFunction<[string], string>;
let testId = 0;
let getTtl: () => number;
beforeEach(() => {
impl.mockClear();
testId += 1;
getTtl = () => 1000;
fn = cache({
name: `cache-with-suffix-${testId}`,
getKeySuffix: getKeySuffixImpl,
tag: (arg) => 'test',
get: async (arg: string, options: CacheFunctionOptions) => {
await new Promise((resolve) => setTimeout(resolve, 20));
return {
data: impl(arg),
ttl: getTtl(),
};
},
});
});
it('should execute when the key suffix changes even if the args do not change', async () => {
const result = await Promise.all([fn('a'), fn('b')]);
expect(result).toEqual(['test-a', 'test-b']);
expect(impl).toHaveBeenCalled();
expect(impl).toHaveBeenCalledTimes(2);
expect(await fn('a')).toEqual('test-a');
expect(impl).toHaveBeenCalledTimes(2);
getKeySuffixImpl.mockImplementation(() => hash({ test: 2 }));
expect(await fn('a')).toEqual('test-a');
expect(impl).toHaveBeenCalledTimes(3);
});
it('should preserve behaviour even when the returned key suffix is undefined', async () => {
// Start with the returned suffix being undefined
getKeySuffixImpl.mockImplementation(() => undefined);
const result = await fn('a');
expect(result).toEqual('test-a');
expect(impl).toHaveBeenCalled();
expect(impl).toHaveBeenCalledTimes(1);
// The returned suffix changes so we should get the value computed by the function
getKeySuffixImpl.mockImplementation(() => hash({ test: 1 }));
expect(await fn('a')).toEqual('test-a');
expect(impl).toHaveBeenCalledTimes(2);
// The returned suffix is undefined again so we should get the value from a previous cache entry
getKeySuffixImpl.mockImplementation(() => undefined);
expect(await fn('a')).toEqual('test-a');
expect(impl).toHaveBeenCalledTimes(2);
// The returned suffix goes back to a previous hash so we should the value from a previous cache entry
getKeySuffixImpl.mockImplementation(() => hash({ test: 1 }));
expect(await fn('a')).toEqual('test-a');
expect(impl).toHaveBeenCalledTimes(2);
});
});
+16 -5
View File
@@ -55,6 +55,9 @@ export interface CacheDefinition<Args extends any[], Result> {
/** Filter the arguments that should be taken into consideration for the cache key */
getKeyArgs?: (args: Args) => any[];
/** Returns a precomputed hash that is used alongside arguments to generate the cache key */
getKeySuffix?: () => string | undefined;
/** Default ttl (in seconds) */
defaultTtl?: number;
@@ -222,7 +225,8 @@ export function cache<Args extends any[], Result>(
const [args, { signal }] = extractCacheFunctionOptions<Args>(rawArgs);
const cacheArgs = cacheDef.getKeyArgs ? cacheDef.getKeyArgs(args) : args;
const key = getCacheKey(cacheDef.name, cacheArgs);
const cacheKeySuffix = cacheDef.getKeySuffix ? cacheDef.getKeySuffix() : undefined;
const key = getCacheKey(cacheDef.name, cacheArgs, cacheKeySuffix);
return await trace(
{
@@ -239,7 +243,8 @@ export function cache<Args extends any[], Result>(
cacheFn.revalidate = async (...rawArgs: Args | [...Args, CacheFunctionOptions]) => {
const [args, { signal }] = extractCacheFunctionOptions<Args>(rawArgs);
const cacheArgs = cacheDef.getKeyArgs ? cacheDef.getKeyArgs(args) : args;
const key = getCacheKey(cacheDef.name, cacheArgs);
const cacheKeySuffix = cacheDef.getKeySuffix ? cacheDef.getKeySuffix() : undefined;
const key = getCacheKey(cacheDef.name, cacheArgs, cacheKeySuffix);
const result = await revalidate(key, signal, ...args);
return result.data;
@@ -247,7 +252,8 @@ export function cache<Args extends any[], Result>(
cacheFn.hasInMemory = async (...args: Args) => {
const cacheArgs = cacheDef.getKeyArgs ? cacheDef.getKeyArgs(args) : args;
const key = getCacheKey(cacheDef.name, cacheArgs);
const cacheKeySuffix = cacheDef.getKeySuffix ? cacheDef.getKeySuffix() : undefined;
const key = getCacheKey(cacheDef.name, cacheArgs, cacheKeySuffix);
const tag = cacheDef.tag?.(...args);
const memoryEntry = await memoryCache.get({ key, tag });
@@ -275,8 +281,13 @@ export function getCache(name: string): CacheFunction<any[], any> | null {
/**
* Get a cache key for a function and its arguments.
*/
export function getCacheKey(fnName: string, args: any[]) {
let innerKey = args.map((arg) => hashValue(arg)).join(',');
export function getCacheKey(fnName: string, args: any[], suffix: string | undefined) {
const hashedArgs = args.map((arg) => hashValue(arg));
if (suffix) {
hashedArgs.push(suffix);
}
let innerKey = hashedArgs.join(',');
// Avoid crazy long keys, by fallbacking to a hash
if (innerKey.length > 400) {
+36 -15
View File
@@ -1,9 +1,10 @@
import { GitBookAPI, ContentAPITokenPayload } from '@gitbook/api';
import { ContentAPITokenPayload, GitBookAPI } from '@gitbook/api';
import { setTag, setContext } from '@sentry/nextjs';
import assertNever from 'assert-never';
import jwt from 'jsonwebtoken';
import type { ResponseCookie } from 'next/dist/compiled/@edge-runtime/cookies';
import { NextResponse, NextRequest } from 'next/server';
import hash from 'object-hash';
import {
PublishedContentWithCache,
@@ -108,11 +109,14 @@ export async function middleware(request: NextRequest) {
const inputURL = stripURLBasePath(url, originBasePath);
const resolved = await withAPI(
new GitBookAPI({
endpoint: apiEndpoint,
authToken: getDefaultAPIToken(apiEndpoint),
userAgent: userAgent(),
}),
{
client: new GitBookAPI({
endpoint: apiEndpoint,
authToken: getDefaultAPIToken(apiEndpoint),
userAgent: userAgent(),
}),
contextId: undefined,
},
() => lookupSpaceForURL(mode, request, inputURL),
);
if ('error' in resolved) {
@@ -153,13 +157,17 @@ export async function middleware(request: NextRequest) {
// Resolution might have changed the API endpoint
apiEndpoint = resolved.apiEndpoint ?? apiEndpoint;
const contextId = 'site' in resolved ? resolved.contextId : undefined;
const nonce = createContentSecurityPolicyNonce();
const csp = await withAPI(
new GitBookAPI({
endpoint: apiEndpoint,
authToken: resolved.apiToken,
userAgent: userAgent(),
}),
{
client: new GitBookAPI({
endpoint: apiEndpoint,
authToken: resolved.apiToken,
userAgent: userAgent(),
}),
contextId,
},
async () => {
const [siteData] = await Promise.all([
'site' in resolved
@@ -198,6 +206,9 @@ export async function middleware(request: NextRequest) {
headers.set('x-forwarded-host', inputURL.host);
headers.set('origin', inputURL.origin);
headers.set('x-gitbook-token', resolved.apiToken);
if (contextId) {
headers.set('x-gitbook-token-context', contextId);
}
headers.set('x-gitbook-mode', mode);
headers.set('x-gitbook-origin-basepath', originBasePath);
headers.set('x-gitbook-basepath', joinPath(originBasePath, resolved.basePath));
@@ -359,7 +370,7 @@ async function lookupSpaceInSingleMode(url: URL): Promise<LookupResult> {
);
}
const apiToken = getDefaultAPIToken(api().endpoint);
const apiToken = getDefaultAPIToken(api().client.endpoint);
if (!apiToken) {
throw new Error(
`Missing GITBOOK_TOKEN environment variable. It should be passed when using GITBOOK_MODE=single.`,
@@ -470,8 +481,16 @@ async function lookupSiteOrSpaceInMultiIdMode(
throw new Error('Collection is not supported in multi-id mode');
}
// The claims property in the content API token is included when
// visitor attributes/assertions are passed to the site preview URL.
//
// When it's present, we generate a hash using the same method as
// getPublishedContentByURL to get the context ID so the cache can be
// invalidated when trying to preview the site with different visitor
// attributes.
const contextId = decoded.claims ? hash(decoded.claims) : undefined;
const gitbookAPI = new GitBookAPI({
endpoint: apiEndpoint ?? api().endpoint,
endpoint: apiEndpoint ?? api().client.endpoint,
authToken: apiToken,
userAgent: userAgent(),
});
@@ -479,13 +498,15 @@ async function lookupSiteOrSpaceInMultiIdMode(
// Verify access to the space to avoid leaking cached data in this mode
// (the cache is not dependend on the auth token, so it could leak data)
if (source.kind === 'space') {
await withAPI(gitbookAPI, () => getSpace.revalidate(source.id, undefined));
await withAPI({ client: gitbookAPI, contextId }, () =>
getSpace.revalidate(source.id, undefined),
);
}
// Verify access to the site to avoid leaking cached data in this mode
// (the cache is not dependend on the auth token, so it could leak data)
if (source.kind === 'site') {
await withAPI(gitbookAPI, () =>
await withAPI({ client: gitbookAPI, contextId }, () =>
getPublishedContentSite.revalidate({
organizationId: decoded.organization,
siteId: source.id,