Fix code block highlighting resulting in shared promises (#195)

* Fix code block highlighting resulting in shared promises

* Lint

* Fix tests

src/components/DocumentView/CodeBlock/highlight.ts

@@ -3,6 +3,7 @@ import { getHighlighter, loadWasm, bundledLanguages, Highlighter, ThemedToken }
 // @ts-ignore - onigWasm is a Wasm module
 import onigWasm from 'shikiji/onig.wasm?module';
 
+import { singleton } from '@/lib/async';
 import { getNodeText } from '@/lib/document';
 
 export type HighlightLine = {
@@ -297,37 +298,19 @@ function isEmptyPositionedToken(token: PositionedToken): boolean {
     return token.start === token.end;
 }
 
-let highlighter: Highlighter | null = null;
-let highlighterPromise: Promise<Highlighter> | null = null;
-
 /**
  * Load the highlighter, only once, and reuse it.
  * It makes sure to handle concurrent calls.
  */
-async function loadHighlighter() {
-    if (highlighter) {
-        return highlighter;
+const loadHighlighter = singleton(async () => {
+    if (typeof onigWasm !== 'string') {
+        // When running bun test, the import is a string, we ignore it and let the module
+        // loads it on its own.
+        //
+        // Otherwise for Vercel/Cloudflare, we need to load it ourselves.
+        await loadWasm((obj) => WebAssembly.instantiate(onigWasm, obj));
     }
-
-    if (highlighterPromise) {
-        return highlighterPromise;
-    }
-
-    highlighterPromise = (async () => {
-        if (typeof onigWasm !== 'string') {
-            // When running bun test, the import is a string, we ignore it and let the module
-            // loads it on its own.
-            //
-            // Otherwise for Vercel/Cloudflare, we need to load it ourselves.
-            await loadWasm((obj) => WebAssembly.instantiate(onigWasm, obj));
-        }
-        const instance = await getHighlighter();
-        await instance.loadTheme('css-variables');
-        return instance;
-    })();
-
-    highlighter = await highlighterPromise;
-    highlighterPromise = null;
-
-    return highlighter;
-}
+    const instance = await getHighlighter();
+    await instance.loadTheme('css-variables');
+    return instance;
+});

src/lib/async.ts

@@ -1,4 +1,4 @@
-import { waitUntil } from './waitUntil';
+import { waitUntil, getGlobalContext } from './waitUntil';
 
 /**
  * Execute a function for each input in parallel and return the first result.
@@ -48,3 +48,35 @@ export async function race<I, R>(
 
     return result;
 }
+
+const UndefinedSymbol = Symbol('Undefined');
+
+/**
+ * Wrap a singleton operation in a safe way for Cloudflare worker
+ * where I/O cannot be performed on behalf of a different request.
+ */
+export function singleton<R>(execute: () => Promise<R>): () => Promise<R> {
+    let cachedResult: R | typeof UndefinedSymbol = UndefinedSymbol;
+    const states = new WeakMap<object, Promise<R>>();
+
+    return async () => {
+        if (cachedResult !== UndefinedSymbol) {
+            // Result is actually shared between requests
+            return cachedResult;
+        }
+
+        // Promises are not shared between requests in Cloudflare Workers
+        const ctx = await getGlobalContext();
+        const current = states.get(ctx);
+        if (current) {
+            return current;
+        }
+
+        const promise = execute();
+        states.set(ctx, promise);
+
+        const result = await promise;
+        cachedResult = result;
+        return result;
+    };
+}

src/lib/waitUntil.ts

@@ -3,6 +3,11 @@
  * This object can be used as a key to store request-specific data in a WeakMap.
  */
 export async function getGlobalContext(): Promise<object> {
+    if (process.env.NODE_ENV === 'test') {
+        // Do not try loading the next-on-pages package in tests as it'll fail
+        return globalThis;
+    }
+
     // We lazy-load the next-on-pages package to avoid errors when running tests because of 'server-only'.
     const { getOptionalRequestContext } = await import('@cloudflare/next-on-pages');
     return getOptionalRequestContext()?.ctx ?? globalThis;