diff --git a/src/components/DocumentView/CodeBlock/highlight.ts b/src/components/DocumentView/CodeBlock/highlight.ts index 323dfcaee..f8406c604 100644 --- a/src/components/DocumentView/CodeBlock/highlight.ts +++ b/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 | 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; +}); diff --git a/src/lib/async.ts b/src/lib/async.ts index 564f58330..d4a9cb994 100644 --- a/src/lib/async.ts +++ b/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( 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(execute: () => Promise): () => Promise { + let cachedResult: R | typeof UndefinedSymbol = UndefinedSymbol; + const states = new WeakMap>(); + + 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; + }; +} diff --git a/src/lib/waitUntil.ts b/src/lib/waitUntil.ts index 95c384226..676301ab2 100644 --- a/src/lib/waitUntil.ts +++ b/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 { + 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;