refactor: lazy load workspaces (#3091)

packages/component/src/components/app-sidebar/index.tsx

@@ -1,4 +1,4 @@
-import { Skeleton } from '@mui/material';
+import { NoSsr, Skeleton } from '@mui/material';
 import { assignInlineVars } from '@vanilla-extract/dynamic';
 import clsx from 'clsx';
 import { useAtom, useAtomValue } from 'jotai';
@@ -107,7 +107,9 @@ export function AppSidebar(props: AppSidebarProps): ReactElement {
         data-enable-animation={enableAnimation && !isResizing}
       >
         <nav className={navStyle} ref={navRef} data-testid="app-sidebar">
-          <SidebarHeader router={props.router} />
+          <NoSsr>
+            <SidebarHeader router={props.router} />
+          </NoSsr>
           <div className={navBodyStyle} data-testid="sliderBar-inner">
             {props.children}
           </div>

packages/component/src/components/card/workspace-card/index.tsx

@@ -1,10 +1,14 @@
-import type {
-  AffineCloudWorkspace,
-  LocalWorkspace,
-} from '@affine/env/workspace';
 import { WorkspaceFlavour } from '@affine/env/workspace';
 import { useAFFiNEI18N } from '@affine/i18n/hooks';
+import type { RootWorkspaceMetadata } from '@affine/workspace/atom';
+import { useStaticBlockSuiteWorkspace } from '@affine/workspace/utils';
 import { SettingsIcon } from '@blocksuite/icons';
+import {
+  CloudWorkspaceIcon as DefaultCloudWorkspaceIcon,
+  CollaborationIcon as DefaultJoinedWorkspaceIcon,
+  LocalDataIcon as DefaultLocalDataIcon,
+  LocalWorkspaceIcon as DefaultLocalWorkspaceIcon,
+} from '@blocksuite/icons';
 import { useBlockSuiteWorkspaceName } from '@toeverything/hooks/use-block-suite-workspace-name';
 import type { FC } from 'react';
 import { useCallback } from 'react';
@@ -18,16 +22,8 @@ import {
 } from './styles';
 
 export type WorkspaceTypeProps = {
-  workspace: AffineCloudWorkspace | LocalWorkspace;
+  flavour: WorkspaceFlavour;
 };
-
-import {
-  CloudWorkspaceIcon as DefaultCloudWorkspaceIcon,
-  CollaborationIcon as DefaultJoinedWorkspaceIcon,
-  LocalDataIcon as DefaultLocalDataIcon,
-  LocalWorkspaceIcon as DefaultLocalWorkspaceIcon,
-} from '@blocksuite/icons';
-
 const JoinedWorkspaceIcon = () => {
   return <DefaultJoinedWorkspaceIcon style={{ color: '#FF646B' }} />;
 };
@@ -43,12 +39,12 @@ const LocalDataIcon = () => {
   return <DefaultLocalDataIcon style={{ color: '#62CD80' }} />;
 };
 
-const WorkspaceType: FC<WorkspaceTypeProps> = ({ workspace }) => {
+const WorkspaceType: FC<WorkspaceTypeProps> = ({ flavour }) => {
   const t = useAFFiNEI18N();
   // fixme: cloud regression
   const isOwner = true;
 
-  if (workspace.flavour === WorkspaceFlavour.LOCAL) {
+  if (flavour === WorkspaceFlavour.LOCAL) {
     return (
       <p title={t['Local Workspace']()}>
         <LocalWorkspaceIcon />
@@ -72,51 +68,46 @@ const WorkspaceType: FC<WorkspaceTypeProps> = ({ workspace }) => {
 
 export type WorkspaceCardProps = {
   currentWorkspaceId: string | null;
-  workspace: AffineCloudWorkspace | LocalWorkspace;
-  onClick: (workspace: AffineCloudWorkspace | LocalWorkspace) => void;
-  onSettingClick: (workspace: AffineCloudWorkspace | LocalWorkspace) => void;
+  meta: RootWorkspaceMetadata;
+  onClick: (workspaceId: string) => void;
+  onSettingClick: (workspaceId: string) => void;
 };
 
 export const WorkspaceCard: FC<WorkspaceCardProps> = ({
-  workspace,
   onClick,
   onSettingClick,
   currentWorkspaceId,
+  meta,
 }) => {
   const t = useAFFiNEI18N();
-  const [name] = useBlockSuiteWorkspaceName(workspace.blockSuiteWorkspace);
+  const workspace = useStaticBlockSuiteWorkspace(meta.id);
+  const [name] = useBlockSuiteWorkspaceName(workspace);
 
   return (
     <StyledCard
       data-testid="workspace-card"
       onClick={useCallback(() => {
-        onClick(workspace);
-      }, [onClick, workspace])}
+        onClick(meta.id);
+      }, [onClick, meta.id])}
       active={workspace.id === currentWorkspaceId}
     >
       <WorkspaceAvatar size={58} workspace={workspace} />
 
       <StyleWorkspaceInfo>
         <StyleWorkspaceTitle>{name}</StyleWorkspaceTitle>
-        <WorkspaceType workspace={workspace} />
-        {workspace.flavour === WorkspaceFlavour.LOCAL && (
+        <WorkspaceType flavour={meta.flavour} />
+        {meta.flavour === WorkspaceFlavour.LOCAL && (
           <p title={t['Available Offline']()}>
             <LocalDataIcon />
             <span>{t['Available Offline']()}</span>
           </p>
         )}
-        {/* {workspace.flavour === WorkspaceFlavour.AFFINE && workspace.public && (
-          <p title={t['Published to Web']()}>
-            <PublishIcon />
-            <span>{t['Published to Web']()}</span>
-          </p>
-        )} */}
       </StyleWorkspaceInfo>
       <StyledSettingLink
         className="setting-entry"
         onClick={e => {
           e.stopPropagation();
-          onSettingClick(workspace);
+          onSettingClick(meta.id);
         }}
       >
         <SettingsIcon />

packages/component/src/components/workspace-avatar/index.tsx

@@ -1,8 +1,3 @@
-import type {
-  AffineCloudWorkspace,
-  AffinePublicWorkspace,
-  LocalWorkspace,
-} from '@affine/env/workspace';
 import type { Workspace } from '@blocksuite/store';
 import * as RadixAvatar from '@radix-ui/react-avatar';
 import { useBlockSuiteWorkspaceAvatarUrl } from '@toeverything/hooks/use-block-suite-workspace-avatar-url';
@@ -15,11 +10,7 @@ import { avatarImageStyle, avatarStyle } from './index.css';
 
 export type WorkspaceAvatarProps = {
   size?: number;
-  workspace:
-    | AffineCloudWorkspace
-    | LocalWorkspace
-    | AffinePublicWorkspace
-    | null;
+  workspace: Workspace | null;
   className?: string;
 };
 
@@ -60,13 +51,9 @@ export const WorkspaceAvatar: React.FC<WorkspaceAvatarProps> = ({
   workspace,
   ...props
 }) => {
-  if (workspace && 'blockSuiteWorkspace' in workspace) {
+  if (workspace) {
     return (
-      <BlockSuiteWorkspaceAvatar
-        {...props}
-        size={size}
-        workspace={workspace.blockSuiteWorkspace}
-      />
+      <BlockSuiteWorkspaceAvatar {...props} size={size} workspace={workspace} />
     );
   }
   return (

packages/component/src/components/workspace-list/index.tsx

@@ -2,6 +2,7 @@ import type {
   AffineCloudWorkspace,
   LocalWorkspace,
 } from '@affine/env/workspace';
+import type { RootWorkspaceMetadata } from '@affine/workspace/atom';
 import type { DragEndEvent } from '@dnd-kit/core';
 import {
   DndContext,
@@ -10,7 +11,8 @@ import {
   useSensors,
 } from '@dnd-kit/core';
 import { SortableContext, useSortable } from '@dnd-kit/sortable';
-import type { FC } from 'react';
+import type { CSSProperties, FC } from 'react';
+import { useMemo } from 'react';
 
 import { WorkspaceCard } from '../../components/card/workspace-card';
 import { workspaceItemStyle } from './index.css';
@@ -19,26 +21,29 @@ export type WorkspaceListProps = {
   disabled?: boolean;
   currentWorkspaceId: string | null;
   items: (AffineCloudWorkspace | LocalWorkspace)[];
-  onClick: (workspace: AffineCloudWorkspace | LocalWorkspace) => void;
-  onSettingClick: (workspace: AffineCloudWorkspace | LocalWorkspace) => void;
+  onClick: (workspaceId: string) => void;
+  onSettingClick: (workspaceId: string) => void;
   onDragEnd: (event: DragEndEvent) => void;
 };
 
 const SortableWorkspaceItem: FC<
   Omit<WorkspaceListProps, 'items'> & {
-    item: AffineCloudWorkspace | LocalWorkspace;
+    item: RootWorkspaceMetadata;
   }
 > = props => {
   const { setNodeRef, attributes, listeners, transform } = useSortable({
     id: props.item.id,
   });
-  const style: React.CSSProperties = {
-    transform: transform
-      ? `translate3d(${transform.x}px, ${transform.y}px, 0)`
-      : undefined,
-    pointerEvents: props.disabled ? 'none' : undefined,
-    opacity: props.disabled ? 0.6 : undefined,
-  };
+  const style: CSSProperties = useMemo(
+    () => ({
+      transform: transform
+        ? `translate3d(${transform.x}px, ${transform.y}px, 0)`
+        : undefined,
+      pointerEvents: props.disabled ? 'none' : undefined,
+      opacity: props.disabled ? 0.6 : undefined,
+    }),
+    [props.disabled, transform]
+  );
   return (
     <div
       className={workspaceItemStyle}
@@ -50,7 +55,7 @@ const SortableWorkspaceItem: FC<
     >
       <WorkspaceCard
         currentWorkspaceId={props.currentWorkspaceId}
-        workspace={props.item}
+        meta={props.item}
         onClick={props.onClick}
         onSettingClick={props.onSettingClick}
       />

packages/component/src/ui/tooltip/tooltip.tsx

@@ -1,8 +1,10 @@
 import type { TooltipProps } from '@mui/material';
+import { NoSsr } from '@mui/material';
 
 import { styled } from '../../styles';
 import { Popper, type PopperProps } from '../popper';
 import StyledPopperContainer from '../shared/container';
+
 const StyledTooltip = styled(StyledPopperContainer)(() => {
   return {
     maxWidth: '320px',
@@ -19,11 +21,13 @@ const StyledTooltip = styled(StyledPopperContainer)(() => {
 export const Tooltip = (props: PopperProps & Omit<TooltipProps, 'title'>) => {
   const { content, placement = 'top-start', children } = props;
   return (
-    <Popper
-      {...props}
-      content={<StyledTooltip placement={placement}>{content}</StyledTooltip>}
-    >
-      {children}
-    </Popper>
+    <NoSsr>
+      <Popper
+        {...props}
+        content={<StyledTooltip placement={placement}>{content}</StyledTooltip>}
+      >
+        {children}
+      </Popper>
+    </NoSsr>
   );
 };

packages/env/src/workspace.ts

@@ -50,18 +50,25 @@ export interface SQLiteDBDownloadProvider extends ActiveDocProvider {
   flavour: 'sqlite-download';
 }
 
-// todo: update type with nest.js
-export type AffineCloudWorkspace = Omit<LocalWorkspace, 'flavour'> & {
-  flavour: WorkspaceFlavour.AFFINE_CLOUD;
+type BaseWorkspace = {
+  flavour: string;
+  id: string;
+  blockSuiteWorkspace: BlockSuiteWorkspace;
 };
 
-export interface LocalWorkspace {
+export interface AffineCloudWorkspace extends BaseWorkspace {
+  flavour: WorkspaceFlavour.AFFINE_CLOUD;
+  id: string;
+  blockSuiteWorkspace: BlockSuiteWorkspace;
+}
+
+export interface LocalWorkspace extends BaseWorkspace {
   flavour: WorkspaceFlavour.LOCAL;
   id: string;
   blockSuiteWorkspace: BlockSuiteWorkspace;
 }
 
-export interface AffinePublicWorkspace {
+export interface AffinePublicWorkspace extends BaseWorkspace {
   flavour: WorkspaceFlavour.PUBLIC;
   id: string;
   blockSuiteWorkspace: BlockSuiteWorkspace;
@@ -107,15 +114,15 @@ export interface WorkspaceRegistry {
 
 export interface WorkspaceCRUD<Flavour extends keyof WorkspaceRegistry> {
   create: (blockSuiteWorkspace: BlockSuiteWorkspace) => Promise<string>;
-  delete: (workspace: WorkspaceRegistry[Flavour]) => Promise<void>;
+  delete: (blockSuiteWorkspace: BlockSuiteWorkspace) => Promise<void>;
   get: (workspaceId: string) => Promise<WorkspaceRegistry[Flavour] | null>;
   // not supported yet
   // update: (workspace: FlavourToWorkspace[Flavour]) => Promise<void>;
   list: () => Promise<WorkspaceRegistry[Flavour][]>;
 }
 
-type UIBaseProps<Flavour extends keyof WorkspaceRegistry> = {
-  currentWorkspace: WorkspaceRegistry[Flavour];
+type UIBaseProps<_Flavour extends keyof WorkspaceRegistry> = {
+  currentWorkspaceId: string;
 };
 
 export type WorkspaceHeaderProps<Flavour extends keyof WorkspaceRegistry> =

packages/workspace/src/local/__tests__/crud.spec.ts

@@ -3,9 +3,10 @@
  */
 import 'fake-indexeddb/auto';
 
-import type { LocalWorkspace, WorkspaceCRUD } from '@affine/env/workspace';
+import type { WorkspaceCRUD } from '@affine/env/workspace';
 import { WorkspaceFlavour } from '@affine/env/workspace';
 import { __unstableSchemas, AffineSchemas } from '@blocksuite/blocks/models';
+import { assertExists } from '@blocksuite/global/utils';
 import { Workspace } from '@blocksuite/store';
 import { afterEach, assertType, describe, expect, test } from 'vitest';
 
@@ -29,11 +30,7 @@ describe('crud', () => {
 
   test('delete not exist', async () => {
     await expect(async () =>
-      CRUD.delete({
-        id: 'not_exist',
-        flavour: WorkspaceFlavour.LOCAL,
-        blockSuiteWorkspace: new Workspace({ id: 'test' }),
-      })
+      CRUD.delete(new Workspace({ id: 'test' }))
     ).rejects.toThrowError();
   });
 
@@ -54,7 +51,8 @@ describe('crud', () => {
     const list = await CRUD.list();
     expect(list.length).toBe(1);
     expect(list[0].id).toBe(id);
-    const localWorkspace = list.at(0) as LocalWorkspace;
+    const localWorkspace = list.at(0);
+    assertExists(localWorkspace);
     expect(localWorkspace.id).toBe(id);
     expect(localWorkspace.flavour).toBe(WorkspaceFlavour.LOCAL);
     expect(localWorkspace.blockSuiteWorkspace.doc.toJSON()).toEqual({
@@ -64,7 +62,7 @@ describe('crud', () => {
       }),
     });
 
-    await CRUD.delete(localWorkspace);
+    await CRUD.delete(localWorkspace.blockSuiteWorkspace);
     expect(await CRUD.get(id)).toBeNull();
     expect(await CRUD.list()).toEqual([]);
   });

packages/workspace/src/utils.ts

@@ -7,25 +7,19 @@ import {
 } from '@affine/workspace/providers';
 import { __unstableSchemas, AffineSchemas } from '@blocksuite/blocks/models';
 import type {
+  ActiveDocProvider,
   DocProviderCreator,
   Generator,
   StoreOptions,
 } from '@blocksuite/store';
 import { createIndexeddbStorage, Workspace } from '@blocksuite/store';
-import { rootStore } from '@toeverything/plugin-infra/manager';
+import { useAtomValue } from 'jotai/react';
+import type { Atom } from 'jotai/vanilla';
+import { atom } from 'jotai/vanilla';
 
-import { rootWorkspacesMetadataAtom } from './atom';
 import { createStaticStorage } from './blob/local-static-storage';
 import { createSQLiteStorage } from './blob/sqlite-blob-storage';
 
-export function cleanupWorkspace(flavour: WorkspaceFlavour) {
-  rootStore
-    .set(rootWorkspacesMetadataAtom, metas =>
-      metas.filter(meta => meta.flavour !== flavour)
-    )
-    .catch(console.error);
-}
-
 function setEditorFlags(workspace: Workspace) {
   Object.entries(runtimeConfig.editorFlags).forEach(([key, value]) => {
     workspace.awarenessStore.setFlag(
@@ -39,12 +33,53 @@ function setEditorFlags(workspace: Workspace) {
   );
 }
 
-const hashMap = new Map<string, Workspace>();
+// guid -> Workspace
+export const workspaceHashMap = new Map<string, Workspace>();
 
-/**
- * @internal test only
- */
-export const _cleanupBlockSuiteWorkspaceCache = () => hashMap.clear();
+const workspacePassiveAtomWeakMap = new WeakMap<
+  Workspace,
+  Atom<Promise<Workspace>>
+>();
+const workspaceActiveWeakMap = new WeakMap<Workspace, boolean>();
+
+export function getWorkspace(id: string) {
+  if (!workspaceHashMap.has(id)) {
+    throw new Error('Workspace not found');
+  }
+  return workspaceHashMap.get(id) as Workspace;
+}
+
+export function getPassiveBlockSuiteWorkspaceAtom(
+  id: string
+): Atom<Promise<Workspace>> {
+  if (!workspaceHashMap.has(id)) {
+    throw new Error('Workspace not found');
+  }
+  const workspace = workspaceHashMap.get(id) as Workspace;
+  if (!workspacePassiveAtomWeakMap.has(workspace)) {
+    const baseAtom = atom(async () => {
+      if (workspaceActiveWeakMap.get(workspace) !== true) {
+        const providers = workspace.providers.filter(
+          (provider): provider is ActiveDocProvider =>
+            'active' in provider && provider.active === true
+        );
+        for (const provider of providers) {
+          provider.sync();
+          // we will wait for the necessary providers to be ready
+          await provider.whenReady;
+        }
+        workspaceActiveWeakMap.set(workspace, true);
+      }
+      return workspace;
+    });
+    workspacePassiveAtomWeakMap.set(workspace, baseAtom);
+  }
+  return workspacePassiveAtomWeakMap.get(workspace) as Atom<Promise<Workspace>>;
+}
+
+export function useStaticBlockSuiteWorkspace(id: string): Workspace {
+  return useAtomValue(getPassiveBlockSuiteWorkspaceAtom(id));
+}
 
 export function createEmptyBlockSuiteWorkspace(
   id: string,
@@ -73,8 +108,8 @@ export function createEmptyBlockSuiteWorkspace(
   const providerCreators: DocProviderCreator[] = [];
   const prefix: string = config?.cachePrefix ?? '';
   const cacheKey = `${prefix}${id}`;
-  if (hashMap.has(cacheKey)) {
-    return hashMap.get(cacheKey) as Workspace;
+  if (workspaceHashMap.has(cacheKey)) {
+    return workspaceHashMap.get(cacheKey) as Workspace;
   }
   const idGenerator = config?.idGenerator;
 
@@ -111,36 +146,6 @@ export function createEmptyBlockSuiteWorkspace(
     .register(AffineSchemas)
     .register(__unstableSchemas);
   setEditorFlags(workspace);
-  hashMap.set(cacheKey, workspace);
+  workspaceHashMap.set(cacheKey, workspace);
   return workspace;
 }
-
-export class CallbackSet extends Set<() => void> {
-  #ready = false;
-
-  get ready(): boolean {
-    return this.#ready;
-  }
-
-  set ready(v: boolean) {
-    this.#ready = v;
-  }
-
-  override add(cb: () => void) {
-    if (this.ready) {
-      cb();
-      return this;
-    }
-    if (this.has(cb)) {
-      return this;
-    }
-    return super.add(cb);
-  }
-
-  override delete(cb: () => void) {
-    if (this.has(cb)) {
-      return super.delete(cb);
-    }
-    return false;
-  }
-}