refactor: lazy load workspaces (#3091)

apps/web/src/hooks/__tests__/index.spec.tsx

@@ -3,14 +3,7 @@
  */
 import 'fake-indexeddb/auto';
 
-import assert from 'node:assert';
-
-import type { WorkspaceAdapter } from '@affine/env/workspace';
-import { WorkspaceFlavour, WorkspaceSubPath } from '@affine/env/workspace';
-import {
-  rootCurrentWorkspaceIdAtom,
-  workspaceAdaptersAtom,
-} from '@affine/workspace/atom';
+import { WorkspaceSubPath } from '@affine/env/workspace';
 import type { PageBlockModel } from '@blocksuite/blocks';
 import { __unstableSchemas, AffineSchemas } from '@blocksuite/blocks/models';
 import { assertExists } from '@blocksuite/global/utils';
@@ -20,18 +13,12 @@ import {
   useBlockSuitePageMeta,
   usePageMetaHelper,
 } from '@toeverything/hooks/use-block-suite-page-meta';
-import { createStore, Provider } from 'jotai';
 import routerMock from 'next-router-mock';
 import { createDynamicRouteParser } from 'next-router-mock/dynamic-routes';
 import type React from 'react';
 import { beforeAll, beforeEach, describe, expect, test, vi } from 'vitest';
 
-import { WorkspaceAdapters } from '../../adapters/workspace';
-import { workspacesAtom } from '../../atoms';
-import { rootCurrentWorkspaceAtom } from '../../atoms/root';
 import { BlockSuiteWorkspace } from '../../shared';
-import { useCurrentWorkspace } from '../current/use-current-workspace';
-import { useAppHelper, useWorkspaces } from '../use-workspaces';
 
 vi.mock(
   '../../components/blocksuite/header/editor-mode-switch/CustomLottie',
@@ -56,28 +43,6 @@ beforeEach(() => {
   localStorage.clear();
 });
 
-async function getJotaiContext() {
-  const store = createStore();
-  store.set(
-    workspaceAdaptersAtom,
-    WorkspaceAdapters as Record<
-      WorkspaceFlavour,
-      WorkspaceAdapter<WorkspaceFlavour>
-    >
-  );
-  const ProviderWrapper: React.FC<React.PropsWithChildren> =
-    function ProviderWrapper({ children }) {
-      return <Provider store={store}>{children}</Provider>;
-    };
-  const workspaces = await store.get(workspacesAtom);
-  expect(workspaces.length).toBe(1);
-  return {
-    store,
-    ProviderWrapper,
-    initialWorkspaces: workspaces,
-  } as const;
-}
-
 beforeEach(async () => {
   blockSuiteWorkspace = new BlockSuiteWorkspace({ id: 'test' })
     .register(AffineSchemas)
@@ -163,67 +128,3 @@ describe('usePageMetas', async () => {
     expect(result.current[0].title).toBe('test');
   });
 });
-
-describe('useWorkspacesHelper', () => {
-  test('basic', async () => {
-    const { ProviderWrapper, store } = await getJotaiContext();
-    const workspaceHelperHook = renderHook(() => useAppHelper(), {
-      wrapper: ProviderWrapper,
-    });
-    const id = await workspaceHelperHook.result.current.createLocalWorkspace(
-      'test'
-    );
-    const workspaces = await store.get(workspacesAtom);
-    expect(workspaces.length).toBe(2);
-    expect(workspaces[1].id).toBe(id);
-    const workspacesHook = renderHook(() => useWorkspaces(), {
-      wrapper: ProviderWrapper,
-    });
-    store.set(rootCurrentWorkspaceIdAtom, workspacesHook.result.current[1].id);
-    await store.get(rootCurrentWorkspaceAtom);
-    const currentWorkspaceHook = renderHook(() => useCurrentWorkspace(), {
-      wrapper: ProviderWrapper,
-    });
-    currentWorkspaceHook.result.current[1](workspacesHook.result.current[1].id);
-  });
-});
-
-describe('useWorkspaces', () => {
-  test('basic', async () => {
-    const { ProviderWrapper } = await getJotaiContext();
-    const { result } = renderHook(() => useWorkspaces(), {
-      wrapper: ProviderWrapper,
-    });
-    expect(result.current).toEqual([
-      {
-        id: expect.stringContaining(''),
-        flavour: WorkspaceFlavour.LOCAL,
-        blockSuiteWorkspace: expect.anything(),
-      },
-    ]);
-  });
-
-  test('mutation', async () => {
-    const { ProviderWrapper, store } = await getJotaiContext();
-    const { result } = renderHook(() => useAppHelper(), {
-      wrapper: ProviderWrapper,
-    });
-    {
-      const workspaces = await store.get(workspacesAtom);
-      expect(workspaces.length).toEqual(1);
-    }
-    await result.current.createLocalWorkspace('test');
-    {
-      const workspaces = await store.get(workspacesAtom);
-      expect(workspaces.length).toEqual(2);
-    }
-    const { result: result2 } = renderHook(() => useWorkspaces(), {
-      wrapper: ProviderWrapper,
-    });
-    expect(result2.current.length).toEqual(2);
-    const secondWorkspace = result2.current[1];
-    expect(secondWorkspace.flavour).toBe('local');
-    assert(secondWorkspace.flavour === WorkspaceFlavour.LOCAL);
-    expect(secondWorkspace.blockSuiteWorkspace.meta.name).toBe('test');
-  });
-});

apps/web/src/hooks/current/use-current-workspace.ts

@@ -1,26 +1,47 @@
-import { isBrowser } from '@affine/env/constant';
 import {
   rootCurrentPageIdAtom,
   rootCurrentWorkspaceIdAtom,
 } from '@affine/workspace/atom';
-import { useAtom, useAtomValue } from 'jotai';
-import { useCallback } from 'react';
+import { assertExists } from '@blocksuite/global/utils';
+import type { PassiveDocProvider, Workspace } from '@blocksuite/store';
+import { useAtom, useSetAtom } from 'jotai';
+import { useCallback, useEffect } from 'react';
 
-import { rootCurrentWorkspaceAtom } from '../../atoms/root';
 import type { AllWorkspace } from '../../shared';
+import { useWorkspace } from '../use-workspace';
+
+declare global {
+  /**
+   * @internal debug only
+   */
+  // eslint-disable-next-line no-var
+  var currentWorkspace: AllWorkspace | undefined;
+  interface WindowEventMap {
+    'affine:workspace:change': CustomEvent<{ id: string }>;
+  }
+}
 
 export function useCurrentWorkspace(): [
   AllWorkspace,
   (id: string | null) => void,
 ] {
-  const currentWorkspace = useAtomValue(rootCurrentWorkspaceAtom);
-  const [, setId] = useAtom(rootCurrentWorkspaceIdAtom);
-  const [, setPageId] = useAtom(rootCurrentPageIdAtom);
+  const [id, setId] = useAtom(rootCurrentWorkspaceIdAtom);
+  assertExists(id);
+  const currentWorkspace = useWorkspace(id);
+  useEffect(() => {
+    globalThis.currentWorkspace = currentWorkspace;
+    globalThis.dispatchEvent(
+      new CustomEvent('affine:workspace:change', {
+        detail: { id: currentWorkspace.id },
+      })
+    );
+  }, [currentWorkspace]);
+  const setPageId = useSetAtom(rootCurrentPageIdAtom);
   return [
     currentWorkspace,
     useCallback(
       (id: string | null) => {
-        if (isBrowser && id) {
+        if (environment.isBrowser && id) {
           localStorage.setItem('last_workspace_id', id);
         }
         setPageId(null);
@@ -30,3 +51,27 @@ export function useCurrentWorkspace(): [
     ),
   ];
 }
+
+const activeWorkspaceWeakMap = new WeakMap<Workspace, boolean>();
+
+export function usePassiveWorkspaceEffect(workspace: Workspace) {
+  useEffect(() => {
+    if (activeWorkspaceWeakMap.get(workspace) === true) {
+      return;
+    }
+    const providers = workspace.providers.filter(
+      (provider): provider is PassiveDocProvider =>
+        'passive' in provider && provider.passive === true
+    );
+    providers.forEach(provider => {
+      provider.connect();
+    });
+    activeWorkspaceWeakMap.set(workspace, true);
+    return () => {
+      providers.forEach(provider => {
+        provider.disconnect();
+      });
+      activeWorkspaceWeakMap.delete(workspace);
+    };
+  }, [workspace]);
+}

apps/web/src/hooks/use-get-page-info.ts

@@ -4,10 +4,10 @@ import { useAtomValue } from 'jotai';
 import { useMemo } from 'react';
 
 import { pageSettingsAtom } from '../atoms';
-import { rootCurrentWorkspaceAtom } from '../atoms/root';
+import { useCurrentWorkspace } from './current/use-current-workspace';
 
 export const useGetPageInfoById = (): GetPageInfoById => {
-  const currentWorkspace = useAtomValue(rootCurrentWorkspaceAtom);
+  const [currentWorkspace] = useCurrentWorkspace();
   const pageMetas = useBlockSuitePageMeta(currentWorkspace.blockSuiteWorkspace);
   const pageMap = useMemo(
     () => Object.fromEntries(pageMetas.map(page => [page.id, page])),

apps/web/src/hooks/use-workspace.ts

@@ -0,0 +1,36 @@
+import { rootWorkspacesMetadataAtom } from '@affine/workspace/atom';
+import { useStaticBlockSuiteWorkspace } from '@affine/workspace/utils';
+import { assertExists } from '@blocksuite/global/utils';
+import type { Workspace } from '@blocksuite/store';
+import type { Atom } from 'jotai';
+import { atom, useAtomValue } from 'jotai';
+
+import type { AffineOfficialWorkspace } from '../shared';
+
+const workspaceWeakMap = new WeakMap<
+  Workspace,
+  Atom<Promise<AffineOfficialWorkspace>>
+>();
+
+export function useWorkspace(workspaceId: string): AffineOfficialWorkspace {
+  const blockSuiteWorkspace = useStaticBlockSuiteWorkspace(workspaceId);
+  if (!workspaceWeakMap.has(blockSuiteWorkspace)) {
+    const baseAtom = atom(async get => {
+      const metadata = await get(rootWorkspacesMetadataAtom);
+      const flavour = metadata.find(({ id }) => id === workspaceId)?.flavour;
+      assertExists(flavour);
+      return {
+        id: workspaceId,
+        flavour,
+        blockSuiteWorkspace,
+      };
+    });
+    workspaceWeakMap.set(blockSuiteWorkspace, baseAtom);
+  }
+
+  return useAtomValue(
+    workspaceWeakMap.get(blockSuiteWorkspace) as Atom<
+      Promise<AffineOfficialWorkspace>
+    >
+  );
+}

apps/web/src/hooks/use-workspaces.ts

@@ -2,19 +2,16 @@ import { DebugLogger } from '@affine/debug';
 import { WorkspaceFlavour, WorkspaceVersion } from '@affine/env/workspace';
 import { rootWorkspacesMetadataAtom } from '@affine/workspace/atom';
 import { saveWorkspaceToLocalStorage } from '@affine/workspace/local/crud';
-import { createEmptyBlockSuiteWorkspace } from '@affine/workspace/utils';
+import {
+  createEmptyBlockSuiteWorkspace,
+  getWorkspace,
+} from '@affine/workspace/utils';
 import { nanoid } from '@blocksuite/store';
 import { useAtomValue, useSetAtom } from 'jotai';
 import { useCallback } from 'react';
 
 import { LocalAdapter } from '../adapters/local';
 import { WorkspaceAdapters } from '../adapters/workspace';
-import { workspacesAtom } from '../atoms';
-import type { AllWorkspace } from '../shared';
-
-export function useWorkspaces(): AllWorkspace[] {
-  return useAtomValue(workspacesAtom);
-}
 
 const logger = new DebugLogger('use-workspaces');
 
@@ -22,12 +19,12 @@ const logger = new DebugLogger('use-workspaces');
  * This hook has the permission to all workspaces. Be careful when using it.
  */
 export function useAppHelper() {
-  const workspaces = useWorkspaces();
   const jotaiWorkspaces = useAtomValue(rootWorkspacesMetadataAtom);
   const set = useSetAtom(rootWorkspacesMetadataAtom);
   return {
     addLocalWorkspace: useCallback(
       async (workspaceId: string): Promise<string> => {
+        createEmptyBlockSuiteWorkspace(workspaceId, WorkspaceFlavour.LOCAL);
         saveWorkspaceToLocalStorage(workspaceId);
         await set(workspaces => [
           ...workspaces,
@@ -68,20 +65,20 @@ export function useAppHelper() {
         const targetJotaiWorkspace = jotaiWorkspaces.find(
           ws => ws.id === workspaceId
         );
-        const targetWorkspace = workspaces.find(ws => ws.id === workspaceId);
-        if (!targetJotaiWorkspace || !targetWorkspace) {
+        if (!targetJotaiWorkspace) {
           throw new Error('page cannot be found');
         }
 
+        const targetWorkspace = getWorkspace(targetJotaiWorkspace.id);
+
         // delete workspace from plugin
-        await WorkspaceAdapters[targetWorkspace.flavour].CRUD.delete(
-          // fixme: type casting
-          targetWorkspace as any
+        await WorkspaceAdapters[targetJotaiWorkspace.flavour].CRUD.delete(
+          targetWorkspace
         );
         // delete workspace from jotai storage
         await set(workspaces => workspaces.filter(ws => ws.id !== workspaceId));
       },
-      [jotaiWorkspaces, set, workspaces]
+      [jotaiWorkspaces, set]
     ),
   };
 }