From 61d2cec9131ab5a724e4af82e83332e757aa5fe7 Mon Sep 17 00:00:00 2001
From: pengx17 <pengxiao@outlook.com>
Date: Fri, 7 Feb 2025 02:30:22 +0000
Subject: [PATCH] fix(core): remove suspense from image preview (#9981)

---
 .../peek-view/view/image-preview/index.tsx    | 184 ++++++++++--------
 1 file changed, 102 insertions(+), 82 deletions(-)

diff --git a/packages/frontend/core/src/modules/peek-view/view/image-preview/index.tsx b/packages/frontend/core/src/modules/peek-view/view/image-preview/index.tsx
index f7511c4516..f38e5fe0db 100644
--- a/packages/frontend/core/src/modules/peek-view/view/image-preview/index.tsx
+++ b/packages/frontend/core/src/modules/peek-view/view/image-preview/index.tsx
@@ -1,9 +1,9 @@
-import { Divider, toast } from '@affine/component';
+import { Divider, Loading, toast } from '@affine/component';
 import { Button, IconButton } from '@affine/component/ui/button';
 import { useAsyncCallback } from '@affine/core/components/hooks/affine-async-hooks';
+import { useI18n } from '@affine/i18n';
 import type { ImageBlockModel } from '@blocksuite/affine/blocks';
-import { assertExists } from '@blocksuite/affine/global/utils';
-import type { BlockModel } from '@blocksuite/affine/store';
+import type { BlockModel, Workspace } from '@blocksuite/affine/store';
 import {
   ArrowLeftSmallIcon,
   ArrowRightSmallIcon,
@@ -17,18 +17,15 @@ import {
 } from '@blocksuite/icons/rc';
 import { useService } from '@toeverything/infra';
 import clsx from 'clsx';
-import { useErrorBoundary } from 'foxact/use-error-boundary';
-import type { PropsWithChildren, ReactElement } from 'react';
+import type { ImgHTMLAttributes, ReactElement } from 'react';
 import {
-  Suspense,
+  forwardRef,
   useCallback,
   useEffect,
   useMemo,
   useRef,
   useState,
 } from 'react';
-import type { FallbackProps } from 'react-error-boundary';
-import { ErrorBoundary } from 'react-error-boundary';
 import useSWR from 'swr';
 
 import {
@@ -63,6 +60,80 @@ export type ImagePreviewModalProps = {
   blockId: string;
 };
 
+function useImageBlob(
+  docCollection: Workspace,
+  docId: string,
+  blockId: string
+) {
+  const { data, error, isLoading } = useSWR(
+    ['workspace', 'image', docId, blockId],
+    {
+      fetcher: async ([_, __, pageId, blockId]) => {
+        const page = docCollection.getDoc(pageId);
+        const block = page?.getBlock(blockId);
+        if (!block) {
+          return null;
+        }
+        const blockModel = block.model as ImageBlockModel;
+        return await docCollection.blobSync.get(blockModel.sourceId as string);
+      },
+      suspense: false,
+    }
+  );
+
+  return { data, error, isLoading };
+}
+
+const ImagePreview = forwardRef<
+  HTMLImageElement,
+  {
+    docCollection: Workspace;
+    docId: string;
+    blockId: string;
+  } & ImgHTMLAttributes<HTMLImageElement>
+>(function ImagePreview({ docCollection, docId, blockId, ...props }, ref) {
+  const { data, error, isLoading } = useImageBlob(
+    docCollection,
+    docId,
+    blockId
+  );
+
+  const [blobUrl, setBlobUrl] = useState<string | null>(null);
+
+  const t = useI18n();
+
+  useEffect(() => {
+    let blobUrl = null;
+    if (data) {
+      blobUrl = URL.createObjectURL(data);
+      setBlobUrl(blobUrl);
+    }
+    return () => {
+      if (blobUrl) {
+        URL.revokeObjectURL(blobUrl);
+      }
+    };
+  }, [data]);
+
+  if (error) {
+    return <div>{t['error.NOT_FOUND']()}</div>;
+  }
+
+  if (!blobUrl || isLoading) {
+    return <Loading size={24} />;
+  }
+
+  return (
+    <img
+      data-blob-id={blockId}
+      data-testid="image-content"
+      src={blobUrl}
+      ref={ref}
+      {...props}
+    />
+  );
+});
+
 const ImagePreviewModalImpl = ({
   docId,
   blockId,
@@ -103,9 +174,6 @@ const ImagePreviewModalImpl = ({
 
   const goto = useCallback(
     (index: number) => {
-      const page = docCollection.getDoc(docId);
-      assertExists(page);
-
       const block = blocks[index];
 
       if (!block) return;
@@ -114,7 +182,7 @@ const ImagePreviewModalImpl = ({
       onBlockIdChange(block.id);
       resetZoom();
     },
-    [docCollection, docId, blocks, onBlockIdChange, resetZoom]
+    [blocks, onBlockIdChange, resetZoom]
   );
 
   const deleteHandler = useCallback(
@@ -179,21 +247,6 @@ const ImagePreviewModalImpl = ({
     setCursor(blocks.length ? prevs.length : 0);
   }, [setBlocks, blockModel, blocksuiteDoc]);
 
-  const { data, error } = useSWR(['workspace', 'image', docId, blockId], {
-    fetcher: ([_, __, pageId, blockId]) => {
-      const page = docCollection.getDoc(pageId);
-      assertExists(page);
-
-      const block = page.getBlock(blockId);
-      if (!block) {
-        return null;
-      }
-      const blockModel = block.model as ImageBlockModel;
-      return docCollection.blobSync.get(blockModel.sourceId as string);
-    },
-    suspense: true,
-  });
-
   useEffect(() => {
     const handleKeyUp = (event: KeyboardEvent) => {
       if (!blocksuiteDoc || !blockModel) {
@@ -242,26 +295,6 @@ const ImagePreviewModalImpl = ({
     };
   }, [blockModel, blocksuiteDoc, copyHandler, onBlockIdChange]);
 
-  useErrorBoundary(error);
-
-  const [prevData, setPrevData] = useState<string | null>(() => data);
-  const [url, setUrl] = useState<string | null>(null);
-
-  if (data === null) {
-    return null;
-  } else if (prevData !== data) {
-    if (url) {
-      URL.revokeObjectURL(url);
-    }
-    setUrl(URL.createObjectURL(data));
-
-    setPrevData(data);
-  } else if (!url) {
-    setUrl(URL.createObjectURL(data));
-  }
-  if (!url) {
-    return null;
-  }
   return (
     <div
       data-testid="image-preview-modal"
@@ -274,10 +307,12 @@ const ImagePreviewModalImpl = ({
           ref={zoomRef}
         >
           <div className={styles.imagePreviewModalCenterStyle}>
-            <img
+            <ImagePreview
               data-blob-id={blockId}
               data-testid="image-content"
-              src={url}
+              docCollection={docCollection}
+              docId={docId}
+              blockId={blockId}
               alt={caption}
               tabIndex={0}
               ref={imageRef}
@@ -385,19 +420,6 @@ const ImagePreviewModalImpl = ({
   );
 };
 
-const ErrorLogger = (props: FallbackProps) => {
-  console.error('image preview modal error', props.error);
-  return null;
-};
-
-export const ImagePreviewErrorBoundary = (
-  props: PropsWithChildren
-): ReactElement => {
-  return (
-    <ErrorBoundary fallbackRender={ErrorLogger}>{props.children}</ErrorBoundary>
-  );
-};
-
 export const ImagePreviewPeekView = (
   props: ImagePreviewModalProps
 ): ReactElement | null => {
@@ -411,25 +433,23 @@ export const ImagePreviewPeekView = (
   }, [props.blockId]);
 
   return (
-    <ImagePreviewErrorBoundary>
-      <Suspense>
-        {blockId ? (
-          <ImagePreviewModalImpl
-            {...props}
-            onClose={onClose}
-            blockId={blockId}
-            onBlockIdChange={setBlockId}
-          />
-        ) : null}
-        <button
-          ref={buttonRef}
-          data-testid="image-preview-close-button"
-          onClick={onClose}
-          className={styles.imagePreviewModalCloseButtonStyle}
-        >
-          <CloseIcon />
-        </button>
-      </Suspense>
-    </ImagePreviewErrorBoundary>
+    <>
+      {blockId ? (
+        <ImagePreviewModalImpl
+          {...props}
+          onClose={onClose}
+          blockId={blockId}
+          onBlockIdChange={setBlockId}
+        />
+      ) : null}
+      <button
+        ref={buttonRef}
+        data-testid="image-preview-close-button"
+        onClick={onClose}
+        className={styles.imagePreviewModalCloseButtonStyle}
+      >
+        <CloseIcon />
+      </button>
+    </>
   );
 };