Protect against many renders

Author: emmerich

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 (
         <PrintPage id={pagePDFContainerId(page)}>
@@ -238,6 +239,7 @@ async function PDFPageDocument(props: {
                         contentRefContext: refContext,
                         resolveContentRef: (ref) => resolveContentRef(ref, refContext),
                         getId: (id) => pagePDFContainerId(page, id),
+                        shouldHighlightCode,
                     }}
                 />
             ) : null}

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<DocumentBlockCode>) {
     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!;
 

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={[]}

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;
+    };
+}

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<HighlightLine
         // Language not found, fallback to plain highlighting
         return plainHighlighting(block);
     }
-
-    blockCount++;
-    if (blockCount > 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);

src/components/DocumentView/CodeBlock/index.ts

@@ -1,2 +1,3 @@
 export * from './CodeBlock';
 export * from './PlainCodeBlock';
+export * from './createHighlightingContext';

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 {

src/components/DocumentView/index.ts

@@ -1 +1,2 @@
 export * from './DocumentView';
+export { createHighlightingContext } from './CodeBlock';

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,
                         }}
                     />
                 ) : (

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']}
             />