feat: improve pdf rendering (#14171)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Bitmap caching for PDF pages to speed up rendering and reduce repeated
work.
* Automatic prefetching of adjacent pages and expanded viewport overscan
for smoother scrolling.

* **Performance**
* LRU-style in-memory cache with eviction to manage memory and improve
responsiveness.
* Reusable-bitmap lookup and error-tolerant fallbacks for more reliable,
faster page display.

<sub>✏️ Tip: You can customize this high-level summary in your review
settings.</sub>
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

packages/frontend/core/src/blocksuite/attachment-viewer/pdf/pdf-viewer.tsx

@@ -1,6 +1,7 @@
 import { IconButton, Menu, observeResize } from '@affine/component';
 import type { PDF, PDFMeta, PDFRendererState } from '@affine/core/modules/pdf';
 import { PDFService, PDFStatus } from '@affine/core/modules/pdf';
+import { cacheBitmap } from '@affine/core/modules/pdf/cache/bitmap-cache';
 import {
   Item,
   List,
@@ -36,6 +37,7 @@ import {
   Virtuoso,
   type VirtuosoHandle,
 } from 'react-virtuoso';
+import type { Subscription } from 'rxjs';
 
 import type { AttachmentViewerProps } from '../types';
 import * as styles from './styles.css';
@@ -67,6 +69,7 @@ export const PDFViewerInner = ({ pdf, meta }: PDFViewerInnerProps) => {
   const pagesScrollerRef = useRef<HTMLElement | null>(null);
   const pagesScrollerHandleRef = useRef<VirtuosoHandle>(null);
   const thumbnailsScrollerHandleRef = useRef<VirtuosoHandle>(null);
+  const prefetching = useRef<Map<string, () => void>>(new Map());
 
   const updateScrollerRef = useCallback(
     (scroller: HTMLElement | Window | null) => {
@@ -173,6 +176,12 @@ export const PDFViewerInner = ({ pdf, meta }: PDFViewerInnerProps) => {
     };
   }, [meta, viewportInfo, onPageSelect]);
 
+  const overscan = useMemo(() => {
+    const h = meta.maxSize.height || 0;
+    // Keep roughly one page above and below rendered to reduce flicker on re-entry.
+    return Math.max(Math.ceil(h * 1.2), 800);
+  }, [meta.maxSize.height]);
+
   // 1. works fine if they are the same size
   // 2. uses the `observeIntersection` when targeting different sizes
   const scrollSeekConfig = useMemo<ScrollSeekConfiguration>(() => {
@@ -190,6 +199,77 @@ export const PDFViewerInner = ({ pdf, meta }: PDFViewerInnerProps) => {
     );
   }, []);
 
+  const prefetchPage = useCallback(
+    (index: number) => {
+      if (index < 0 || index >= meta.pageCount) return;
+      if (!viewportInfo.width || !viewportInfo.height) return;
+
+      const actualSize = meta.pageSizes[index];
+      const renderSize = fitToPage(
+        {
+          width: viewportInfo.width - 40,
+          height: viewportInfo.height - 40,
+        },
+        actualSize,
+        meta.maxSize
+      );
+      const scale = window.devicePixelRatio;
+      const key = `${index}:${renderSize.width}:${renderSize.height}:${scale}`;
+      if (prefetching.current.has(key)) return;
+
+      const { page, release } = pdf.page(
+        index,
+        `${renderSize.width}:${renderSize.height}:${scale}`
+      );
+      let stopped = false;
+      let subscription: Subscription | undefined;
+
+      const stop = () => {
+        if (stopped) return;
+        stopped = true;
+        prefetching.current.delete(key);
+        page.render.unsubscribe();
+        release();
+        subscription?.unsubscribe();
+      };
+
+      subscription = page.bitmap$.subscribe(bitmap => {
+        if (!bitmap) return;
+        cacheBitmap(
+          {
+            blobId: pdf.id,
+            pageNum: index,
+            width: renderSize.width,
+            height: renderSize.height,
+            scale,
+          },
+          bitmap
+        )
+          .catch(e => console.error('Failed to cache bitmap', e))
+          .finally(stop);
+      });
+
+      prefetching.current.set(key, stop);
+      page.render({
+        width: renderSize.width,
+        height: renderSize.height,
+        scale,
+      });
+    },
+    [meta, pdf, viewportInfo]
+  );
+
+  useEffect(() => {
+    prefetchPage(cursor - 1);
+    prefetchPage(cursor + 1);
+
+    const map = prefetching.current;
+    return () => {
+      map.forEach(stop => stop());
+      map.clear();
+    };
+  }, [cursor, prefetchPage]);
+
   return (
     <div
       ref={viewerRef}
@@ -212,6 +292,7 @@ export const PDFViewerInner = ({ pdf, meta }: PDFViewerInnerProps) => {
           Footer: ListPadding,
           ScrollSeekPlaceholder,
         }}
+        increaseViewportBy={{ top: overscan, bottom: overscan }}
         context={{
           viewportInfo: {
             width: viewportInfo.width - 40,

packages/frontend/core/src/modules/pdf/cache/bitmap-cache.ts

@@ -0,0 +1,165 @@
+type CacheKey = string;
+
+type CacheEntry = {
+  key: CacheKey;
+  blobId: string;
+  pageNum: number;
+  width: number;
+  height: number;
+  scale: number;
+  blob: Blob;
+};
+
+type CacheParams = {
+  blobId: string;
+  pageNum: number;
+  width: number;
+  height: number;
+  scale: number;
+};
+
+class BitmapLRU {
+  private readonly map = new Map<CacheKey, CacheEntry>();
+
+  constructor(private readonly maxEntries: number) {}
+
+  has(key: CacheKey) {
+    return this.map.has(key);
+  }
+
+  get(key: CacheKey): CacheEntry | null {
+    const entry = this.map.get(key);
+    if (!entry) return null;
+    this.map.delete(key);
+    this.map.set(key, entry);
+    return entry;
+  }
+
+  findReusable(
+    params: CacheParams,
+    upscaleThreshold: number
+  ): CacheEntry | null {
+    let best: CacheEntry | null = null;
+
+    for (const entry of this.map.values()) {
+      if (
+        entry.blobId !== params.blobId ||
+        entry.pageNum !== params.pageNum ||
+        entry.scale !== params.scale
+      ) {
+        continue;
+      }
+
+      const ratio = Math.max(
+        params.width / entry.width,
+        params.height / entry.height
+      );
+
+      if (ratio > upscaleThreshold) continue;
+
+      if (!best || entry.width * entry.height > best.width * best.height) {
+        best = entry;
+      }
+    }
+
+    if (!best) return null;
+
+    this.map.delete(best.key);
+    this.map.set(best.key, best);
+    return best;
+  }
+
+  set(entry: CacheEntry) {
+    if (this.map.has(entry.key)) {
+      this.map.delete(entry.key);
+    }
+
+    this.map.set(entry.key, entry);
+
+    while (this.map.size > this.maxEntries) {
+      const oldest = this.map.keys().next().value as CacheKey | undefined;
+      if (oldest === undefined) break;
+      this.map.delete(oldest);
+    }
+  }
+}
+
+const MAX_ENTRIES = 64;
+const UPSCALE_THRESHOLD = 1.3;
+const QUALITY = 0.72;
+
+const cache = new BitmapLRU(MAX_ENTRIES);
+
+const normalize = (value: number) => Math.round(value);
+const toKey = ({
+  blobId,
+  pageNum,
+  width,
+  height,
+  scale,
+}: CacheParams): CacheKey =>
+  `${blobId}:${pageNum}:${normalize(width)}:${normalize(height)}:${scale}`;
+
+export async function getReusableBitmap(
+  params: CacheParams
+): Promise<ImageBitmap | null> {
+  const exact = cache.get(toKey(params));
+  if (exact) {
+    try {
+      return await createImageBitmap(exact.blob);
+    } catch {
+      return null;
+    }
+  }
+
+  const reusable = cache.findReusable(params, UPSCALE_THRESHOLD);
+  if (!reusable) return null;
+
+  try {
+    return await createImageBitmap(reusable.blob);
+  } catch {
+    return null;
+  }
+}
+
+export async function cacheBitmap(params: CacheParams, bitmap: ImageBitmap) {
+  const key = toKey(params);
+  if (cache.has(key)) return;
+
+  try {
+    const blob = await bitmapToWebp(bitmap);
+    if (!blob) return;
+    cache.set({ key, ...params, blob });
+  } catch (e) {
+    console.error('Failed to convert bitmap', e);
+  }
+}
+
+async function bitmapToWebp(bitmap: ImageBitmap): Promise<Blob | null> {
+  const width = bitmap.width;
+  const height = bitmap.height;
+
+  if (typeof OffscreenCanvas !== 'undefined') {
+    const canvas = new OffscreenCanvas(width, height);
+    const ctx = canvas.getContext('2d');
+    if (!ctx) return null;
+    ctx.drawImage(bitmap, 0, 0, width, height);
+    return canvas.convertToBlob({
+      type: 'image/webp',
+      quality: QUALITY,
+    });
+  }
+
+  if (typeof document === 'undefined') return null;
+
+  const canvas = document.createElement('canvas');
+  canvas.width = width;
+  canvas.height = height;
+  const ctx = canvas.getContext('2d');
+  if (!ctx) return null;
+  ctx.drawImage(bitmap, 0, 0, width, height);
+
+  return new Promise(resolve => {
+    canvas.toBlob(blob => resolve(blob), 'image/webp', QUALITY);
+  });
+}

packages/frontend/core/src/modules/pdf/views/page-renderer.tsx

@@ -4,6 +4,7 @@ import { useLiveData } from '@toeverything/infra';
 import { debounce } from 'lodash-es';
 import { forwardRef, useEffect, useMemo, useRef, useState } from 'react';
 
+import { cacheBitmap, getReusableBitmap } from '../cache/bitmap-cache';
 import type { PDF } from '../entities/pdf';
 import type { PDFPage } from '../entities/pdf-page';
 import type { PageSize } from '../renderer/types';
@@ -28,6 +29,89 @@ interface PDFPageProps {
   isThumbnail?: boolean;
 }
 
+function usePDFPage({
+  pdf,
+  pageNum,
+  width,
+  height,
+  scale,
+  visibility,
+}: {
+  pdf: PDF;
+  pageNum: number;
+  width: number;
+  height: number;
+  scale: number;
+  visibility: boolean;
+}) {
+  const [page, setPage] = useState<PDFPage | null>(null);
+  const [cachedBitmap, setCachedBitmap] = useState<ImageBitmap | null>(null);
+  const img = useLiveData(useMemo(() => (page ? page.bitmap$ : null), [page]));
+  const error = useLiveData(page?.error$ ?? null);
+
+  // Consolidated effect to handle loading strategy (Cache vs Render)
+  useEffect(() => {
+    if (!visibility || !width || !height) {
+      setPage(null);
+      return;
+    }
+
+    let active = true;
+    let releasePage: (() => void) | undefined;
+
+    const load = async () => {
+      try {
+        // 1. Try cache
+        const compressed = await getReusableBitmap({
+          blobId: pdf.id,
+          pageNum,
+          width,
+          height,
+          scale,
+        });
+
+        if (!active) return;
+
+        if (compressed) {
+          setCachedBitmap(compressed);
+          setPage(null);
+        } else {
+          // 2. Load Page
+          setCachedBitmap(null); // Clear stale cache
+
+          const key = `${width}:${height}:${scale}`;
+          const { page: newPage, release } = pdf.page(pageNum, key);
+          releasePage = release;
+
+          setPage(newPage);
+          newPage.render({ width, height, scale });
+        }
+      } catch (e) {
+        console.error('Failed to load PDF page', e);
+      }
+    };
+
+    load().catch(console.error);
+
+    return () => {
+      active = false;
+      releasePage?.();
+      setPage(null);
+    };
+  }, [visibility, pdf, pageNum, width, height, scale]);
+
+  // Cache new bitmap when generated
+  useEffect(() => {
+    if (!img || !page) return;
+
+    cacheBitmap({ blobId: pdf.id, pageNum, width, height, scale }, img).catch(
+      e => console.error('Failed to cache bitmap', e)
+    );
+  }, [img, page, pdf.id, pageNum, width, height, scale]);
+
+  return { displayImg: cachedBitmap ?? img, error };
+}
+
 export const PDFPageRenderer = ({
   pdf,
   pageNum,
@@ -43,62 +127,32 @@ export const PDFPageRenderer = ({
   const t = useI18n();
   const pageViewRef = useRef<HTMLDivElement>(null);
   const canvasRef = useRef<HTMLCanvasElement>(null);
-  const [page, setPage] = useState<PDFPage | null>(null);
-  const img = useLiveData(useMemo(() => (page ? page.bitmap$ : null), [page]));
-  const error = useLiveData(page?.error$ ?? null);
   const size = useMemo(
     () => resize(viewportInfo, actualSize, maxSize, isThumbnail),
     [resize, viewportInfo, actualSize, maxSize, isThumbnail]
   );
   const [visibility, setVisibility] = useState(false);
 
+  const { displayImg, error } = usePDFPage({
+    pdf,
+    pageNum,
+    width: size.width,
+    height: size.height,
+    scale,
+    visibility,
+  });
+
   useEffect(() => {
     const canvas = canvasRef.current;
     if (!canvas) return;
-    if (!img) return;
+    if (!displayImg) return;
     const ctx = canvas.getContext('2d');
     if (!ctx) return;
 
-    canvas.width = img.width;
-    canvas.height = img.height;
-    ctx.drawImage(img, 0, 0);
-  }, [img]);
-
-  useEffect(() => {
-    if (!visibility) return;
-    if (!page) return;
-
-    const width = size.width;
-    const height = size.height;
-    if (width * height === 0) return;
-
-    page.render({
-      width,
-      height,
-      scale,
-    });
-
-    return () => {
-      page.render.unsubscribe();
-    };
-  }, [visibility, page, size, scale]);
-
-  useEffect(() => {
-    if (!visibility) return;
-    if (!pdf) return;
-
-    const width = size.width;
-    const height = size.height;
-    const key = `${width}:${height}:${scale}`;
-    const { page, release } = pdf.page(pageNum, key);
-
-    setPage(page);
-
-    return () => {
-      release();
-      setPage(null);
-    };
-  }, [visibility, pdf, pageNum, size, scale]);
+    canvas.width = displayImg.width;
+    canvas.height = displayImg.height;
+    ctx.drawImage(displayImg, 0, 0);
+  }, [displayImg]);
 
   useEffect(() => {
     const pageView = pageViewRef.current;
@@ -126,7 +180,7 @@ export const PDFPageRenderer = ({
       onClick={() => onSelect?.(pageNum)}
     >
       <PageRendererInner
-        img={img}
+        img={displayImg}
         ref={canvasRef}
         err={error ? t['com.affine.pdf.page.render.error']() : null}
         scale={scale}