Try custom cache (#9)

* Try custom cache

* Fix redis cache

* Comment

.env.example

@@ -16,3 +16,7 @@ GITBOOK_MODE=multi
 
 # Use a different GitBook environment by changing the API URL
 # GITBOOK_API_URL=https://api.gitbook.com
+
+# Use upstash for caching
+# UPSTASH_REDIS_REST_URL
+# UPSTASH_REDIS_REST_TOKEN

package.json

@@ -16,6 +16,7 @@
         "@geist-ui/icons": "^1.0.2",
         "@gitbook/api": "^0.15.0",
         "@readme/openapi-parser": "^2.5.0",
+        "@upstash/redis": "^1.25.1",
         "ajv": "^8.12.0",
         "assert-never": "^1.2.1",
         "bun-types": "^1.0.7",

src/lib/api.ts

@@ -1,9 +1,10 @@
 import 'server-only';
 
 import { ContentVisibility, GitBookAPI, JSONDocument } from '@gitbook/api';
-import { unstable_cache } from 'next/cache';
 import { headers } from 'next/headers';
 
+import { cache } from './cache';
+
 /**
  * Create an API client for the current request.
  */
@@ -28,85 +29,74 @@ export function api(): GitBookAPI {
 /**
  * Get a space by its ID.
  */
-export const getSpace = unstable_cache(
-    async (spaceId: string) => {
-        const { data } = await api().spaces.getSpaceById(spaceId);
-        return data;
-    },
-    ['api', 'spaces'],
-    {
-        tags: ['api', 'spaces'],
-    },
-);
+export const getSpace = cache('api.getSpace', async (spaceId: string) => {
+    const { data } = await api().spaces.getSpaceById(spaceId, {
+        cache: 'no-store',
+    });
+    return data;
+});
 
 /**
  * Get the current revision of a space
  */
-export const getCurrentRevision = unstable_cache(
-    async (spaceId: string) => {
-        const { data } = await api().spaces.getCurrentRevision(spaceId);
-        return data;
-    },
-    ['api', 'revisions'],
-    {
-        tags: ['api', 'revisions'],
-    },
-);
+export const getCurrentRevision = cache('api.getCurrentRevision', async (spaceId: string) => {
+    const { data } = await api().spaces.getCurrentRevision(spaceId, {
+        cache: 'no-store',
+    });
+    return data;
+});
 
 /**
  * Get the document for a page.
  */
-export const getPageDocument = unstable_cache(
+export const getPageDocument = cache(
+    'api.getPageDocument',
     async (spaceId: string, revisionId: string, pageId: string) => {
-        const { data } = await api().spaces.getPageInRevisionById(spaceId, revisionId, pageId);
+        const { data } = await api().spaces.getPageInRevisionById(
+            spaceId,
+            revisionId,
+            pageId,
+            {},
+            {
+                cache: 'no-store',
+            },
+        );
         // @ts-ignore
         return data.document as JSONDocument;
     },
-    ['api', 'documents'],
-    {
-        tags: ['api', 'documents'],
-    },
 );
 
 /**
  * Get the customization settings for a space.
  */
-export const getSpaceCustomization = unstable_cache(
-    async (spaceId: string) => {
-        const { data } = await api().spaces.getSpacePublishingCustomizationById(spaceId);
-        return data;
-    },
-    ['api', 'customization'],
-    {
-        tags: ['api', 'customization'],
-    },
-);
+export const getSpaceCustomization = cache('api.getSpaceCustomization', async (spaceId: string) => {
+    const { data } = await api().spaces.getSpacePublishingCustomizationById(spaceId, {
+        cache: 'no-store',
+    });
+    return data;
+});
 
 /**
  * Get the infos about a collection by its ID.
  */
-export const getCollection = unstable_cache(
-    async (spaceId: string) => {
-        const { data } = await api().collections.getCollectionById(spaceId);
-        return data;
-    },
-    ['api', 'collections'],
-    {
-        tags: ['api', 'collections'],
-    },
-);
+export const getCollection = cache('api.getCollection', async (spaceId: string) => {
+    const { data } = await api().collections.getCollectionById(spaceId, {
+        cache: 'no-store',
+    });
+    return data;
+});
 
 /**
  * List all the spaces variants published in a collection.
  */
-export const getCollectionSpaces = unstable_cache(
-    async (spaceId: string) => {
-        const { data } = await api().collections.listSpacesInCollectionById(spaceId);
-        // TODO: do this filtering on the API side
-        return data.items.filter((space) => space.visibility === ContentVisibility.InCollection);
-    },
-    ['api', 'collections', 'spaces'],
-    {
-        tags: ['api', 'collections'],
-    },
-);
+export const getCollectionSpaces = cache('api.getCollectionSpaces', async (spaceId: string) => {
+    const { data } = await api().collections.listSpacesInCollectionById(
+        spaceId,
+        {},
+        {
+            cache: 'no-store',
+        },
+    );
+    // TODO: do this filtering on the API side
+    return data.items.filter((space) => space.visibility === ContentVisibility.InCollection);
+});

src/lib/cache.ts

@@ -0,0 +1,72 @@
+import { Redis } from '@upstash/redis';
+
+const redis =
+    process.env.UPSTASH_REDIS_REST_URL && process.env.UPSTASH_REDIS_REST_TOKEN
+        ? new Redis({
+              url: process.env.UPSTASH_REDIS_REST_URL,
+              token: process.env.UPSTASH_REDIS_REST_TOKEN,
+          })
+        : null;
+
+const ttl = 60 * 60 * 24;
+const memoryCache = new Map<string, any>();
+
+/**
+ * Cache data from an async function.
+ * We don't use the next.js cache because it has a 2MB limit.
+ */
+export function cache<Args extends string[], Result>(
+    fnName: string,
+    fn: (...args: Args) => Promise<Result>,
+): (...args: Args) => Promise<Result> {
+    return async (...args: Args) => {
+        const key = getCacheKey(fnName, args);
+
+        const cachedValue = await getCacheValue(key);
+        if (cachedValue !== null) {
+            return cachedValue;
+        }
+
+        const result = await fn(...args);
+
+        await setCacheValue(key, result);
+
+        return result;
+    };
+}
+
+/**
+ * Create a cache key from a function name and its arguments.
+ */
+function getCacheKey(fnName: string, args: string[]) {
+    return `${fnName}(${args.join(',')})`;
+}
+
+/**
+ * Get a value from the cache.
+ */
+async function getCacheValue(key: string) {
+    if (memoryCache.has(key)) {
+        return memoryCache.get(key);
+    }
+
+    if (redis) {
+        const value = await redis.get(key);
+        return value;
+    }
+
+    return null;
+}
+
+/**
+ * Set a value in the cache.
+ */
+async function setCacheValue(key: string, value: any) {
+    memoryCache.set(key, value);
+
+    if (redis) {
+        await redis.set(key, value, {
+            ex: ttl,
+        });
+    }
+}

src/middleware.ts

@@ -163,7 +163,10 @@ async function lookupSpaceInMultiPathMode(
 
             return {
                 target: 'content',
-                redirect: new URL(`/` + redirect.hostname + redirect.pathname + redirect.search, url).toString(),
+                redirect: new URL(
+                    `/` + redirect.hostname + redirect.pathname + redirect.search,
+                    url,
+                ).toString(),
             };
         }