diff --git a/src/app/(space)/(core)/~gitbook/pdf/page.tsx b/src/app/(space)/(core)/~gitbook/pdf/page.tsx index 89efafc439..c034e295b8 100644 --- a/src/app/(space)/(core)/~gitbook/pdf/page.tsx +++ b/src/app/(space)/(core)/~gitbook/pdf/page.tsx @@ -11,7 +11,7 @@ import { notFound } from 'next/navigation'; import * as React from 'react'; import { getContentPointer } from '@/app/(space)/fetch'; -import { DocumentView } from '@/components/DocumentView'; +import { DocumentView, createHighlightingContext } from '@/components/DocumentView'; import { TrademarkLink } from '@/components/TableOfContents/Trademark'; import { PolymorphicComponentProp } from '@/components/utils/types'; import { getSpaceLanguage } from '@/intl/server'; @@ -216,6 +216,7 @@ async function PDFPageDocument(props: { const { space, page, refContext } = props; const document = page.documentId ? await getDocument(space.id, page.documentId) : null; + const shouldHighlightCode = createHighlightingContext(); return ( @@ -238,6 +239,7 @@ async function PDFPageDocument(props: { contentRefContext: refContext, resolveContentRef: (ref) => resolveContentRef(ref, refContext), getId: (id) => pagePDFContainerId(page, id), + shouldHighlightCode, }} /> ) : null} diff --git a/src/components/DocumentView/CodeBlock/CodeBlock.tsx b/src/components/DocumentView/CodeBlock/CodeBlock.tsx index 4d4e67b96e..bb545cb994 100644 --- a/src/components/DocumentView/CodeBlock/CodeBlock.tsx +++ b/src/components/DocumentView/CodeBlock/CodeBlock.tsx @@ -3,7 +3,7 @@ import { DocumentBlockCode, JSONDocument } from '@gitbook/api'; import { tcls } from '@/lib/tailwind'; import { CopyCodeButton } from './CopyCodeButton'; -import { highlight, HighlightLine, HighlightToken } from './highlight'; +import { highlight, HighlightLine, HighlightToken, plainHighlighting } from './highlight'; import { BlockProps } from '../Block'; import { DocumentContext } from '../DocumentView'; import { Inline } from '../Inline'; @@ -15,7 +15,8 @@ import './theme.css'; */ export async function CodeBlock(props: BlockProps) { const { block, document, style, context } = props; - const lines = await highlight(block); + const withHighlighting = context.shouldHighlightCode(); + const lines = withHighlighting ? await highlight(block) : plainHighlighting(block); const id = block.key!; diff --git a/src/components/DocumentView/CodeBlock/PlainCodeBlock.tsx b/src/components/DocumentView/CodeBlock/PlainCodeBlock.tsx index 886bbe6bae..b80ad49c57 100644 --- a/src/components/DocumentView/CodeBlock/PlainCodeBlock.tsx +++ b/src/components/DocumentView/CodeBlock/PlainCodeBlock.tsx @@ -50,6 +50,7 @@ export function PlainCodeBlock(props: { code: string; syntax: string }) { mode: 'default', contentRefContext: null, resolveContentRef: async () => null, + shouldHighlightCode: () => true, }} block={block} ancestorBlocks={[]} diff --git a/src/components/DocumentView/CodeBlock/createHighlightingContext.ts b/src/components/DocumentView/CodeBlock/createHighlightingContext.ts new file mode 100644 index 0000000000..63008f1bfa --- /dev/null +++ b/src/components/DocumentView/CodeBlock/createHighlightingContext.ts @@ -0,0 +1,16 @@ +const CODE_HIGHLIGHT_BLOCK_LIMIT = 50; + +/** + * Protect against memory issues when highlighting a large number of code blocks. + * This context only allows 50 code blocks per render to be highlighted. + * Once highlighting can scale up to a large number of code blocks, it can be removed. + * + * https://linear.app/gitbook-x/issue/RND-3588/gitbook-open-code-syntax-highlighting-runs-out-of-memory-after-a + */ +export function createHighlightingContext() { + let count = 0; + return () => { + count += 1; + return count < CODE_HIGHLIGHT_BLOCK_LIMIT; + }; +} diff --git a/src/components/DocumentView/CodeBlock/highlight.ts b/src/components/DocumentView/CodeBlock/highlight.ts index ad4138cabc..c3c7b36a42 100644 --- a/src/components/DocumentView/CodeBlock/highlight.ts +++ b/src/components/DocumentView/CodeBlock/highlight.ts @@ -15,6 +15,8 @@ import { asyncMutexFunction, singleton } from '@/lib/async'; import { getNodeText } from '@/lib/document'; import { trace } from '@/lib/tracing'; +import { DocumentContext } from '../DocumentView'; + export type HighlightLine = { highlighted: boolean; tokens: HighlightToken[]; @@ -29,14 +31,6 @@ type InlineIndexed = { inline: any; start: number; end: number }; type PositionedToken = ThemedToken & { start: number; end: number }; -/** - * Due to a combination of memory limitations of Cloudflare workers and the memory - * cost of shiki, we need to set a limit on the number of blocks we can highlight - * in a single page. - */ -let blockCount = 0; -const BLOCK_LIMIT = 50; - /** * Highlight a code block while preserving inline elements. */ @@ -47,12 +41,7 @@ export async function highlight(block: DocumentBlockCode): Promise BLOCK_LIMIT) { - // Too many blocks and we risk crashing the worker, fallback to plain highlighting - return plainHighlighting(block); - } + debugger; const inlines: InlineIndexed[] = []; const code = getPlainCodeBlock(block, inlines); diff --git a/src/components/DocumentView/CodeBlock/index.ts b/src/components/DocumentView/CodeBlock/index.ts index f322414c54..bd25d69909 100644 --- a/src/components/DocumentView/CodeBlock/index.ts +++ b/src/components/DocumentView/CodeBlock/index.ts @@ -1,2 +1,3 @@ export * from './CodeBlock'; export * from './PlainCodeBlock'; +export * from './createHighlightingContext'; diff --git a/src/components/DocumentView/DocumentView.tsx b/src/components/DocumentView/DocumentView.tsx index f2ab1f626f..f449b65366 100644 --- a/src/components/DocumentView/DocumentView.tsx +++ b/src/components/DocumentView/DocumentView.tsx @@ -36,6 +36,16 @@ export interface DocumentContext { * Transform an ID to be added to the DOM. */ getId?: (id: string) => string; + + /** + * Returns true if the given code block should be highlighted. + * This function was added to protect against memory issues when highlighting + * a large number of code blocks. + * Once highlighting can scale up to a large number of code blocks, it can be removed. + * + * https://linear.app/gitbook-x/issue/RND-3588/gitbook-open-code-syntax-highlighting-runs-out-of-memory-after-a + */ + shouldHighlightCode: () => boolean; } export interface DocumentContextProps { diff --git a/src/components/DocumentView/index.ts b/src/components/DocumentView/index.ts index 2b742d2d5b..80caab2759 100644 --- a/src/components/DocumentView/index.ts +++ b/src/components/DocumentView/index.ts @@ -1 +1,2 @@ export * from './DocumentView'; +export { createHighlightingContext } from './CodeBlock'; diff --git a/src/components/PageBody/PageBody.tsx b/src/components/PageBody/PageBody.tsx index 0b989cbe15..34085bcbd4 100644 --- a/src/components/PageBody/PageBody.tsx +++ b/src/components/PageBody/PageBody.tsx @@ -13,10 +13,17 @@ import { PageCover } from './PageCover'; import { PageFooterNavigation } from './PageFooterNavigation'; import { PageHeader } from './PageHeader'; import { TrackPageView } from './TrackPageView'; -import { DocumentView } from '../DocumentView'; +import { DocumentView, createHighlightingContext } from '../DocumentView'; import { PageFeedbackForm } from '../PageFeedback'; import { DateRelative } from '../primitives'; +/** + * Due to a combination of memory limitations of Cloudflare workers and the memory + * cost of shiki, we need to set a limit on the number of blocks we can highlight + * in a single page. + */ +const CODE_HIGHLIGHT_BLOCK_LIMIT = 50; + export function PageBody(props: { space: Space; contentTarget: ContentTarget; @@ -32,6 +39,7 @@ export function PageBody(props: { const asFullWidth = document ? hasFullWidthBlock(document) : false; const language = getSpaceLanguage(customization); const updatedAt = page.updatedAt ?? page.createdAt; + const shouldHighlightCode = createHighlightingContext(); return ( <> @@ -73,6 +81,7 @@ export function PageBody(props: { contentRefContext: context, resolveContentRef: (ref, options) => resolveContentRef(ref, context, options), + shouldHighlightCode, }} /> ) : ( diff --git a/src/components/Search/server-actions.tsx b/src/components/Search/server-actions.tsx index af26502648..cbd0395b57 100644 --- a/src/components/Search/server-actions.tsx +++ b/src/components/Search/server-actions.tsx @@ -161,6 +161,7 @@ function transformAnswer( mode: 'default', contentRefContext: null, resolveContentRef: async () => null, + shouldHighlightCode: () => false, }} style={['space-y-5']} />