refactor: use jotai-effect (#4641)

apps/core/src/adapters/local/index.tsx

@@ -19,10 +19,11 @@ import {
   globalBlockSuiteSchema,
 } from '@affine/workspace/manager';
 import { createIndexedDBDownloadProvider } from '@affine/workspace/providers';
-import { useStaticBlockSuiteWorkspace } from '@toeverything/infra/__internal__/react';
+import { getBlockSuiteWorkspaceAtom } from '@toeverything/infra/__internal__/workspace';
 import { getCurrentStore } from '@toeverything/infra/atom';
 import { initEmptyPage } from '@toeverything/infra/blocksuite';
 import { buildShowcaseWorkspace } from '@toeverything/infra/blocksuite';
+import { useAtomValue } from 'jotai';
 import { nanoid } from 'nanoid';
 import { useCallback } from 'react';
 
@@ -87,7 +88,8 @@ export const LocalAdapter: WorkspaceAdapter<WorkspaceFlavour.LOCAL> = {
     Header: WorkspaceHeader,
     Provider,
     PageDetail: ({ currentWorkspaceId, currentPageId, onLoadEditor }) => {
-      const workspace = useStaticBlockSuiteWorkspace(currentWorkspaceId);
+      const [workspaceAtom] = getBlockSuiteWorkspaceAtom(currentWorkspaceId);
+      const workspace = useAtomValue(workspaceAtom);
       const page = workspace.getPage(currentPageId);
       if (!page) {
         throw new PageNotFoundError(workspace, currentPageId);

apps/core/src/components/affine/setting-modal/setting-sidebar/index.tsx

@@ -11,7 +11,7 @@ import { Avatar } from '@toeverything/components/avatar';
 import { Tooltip } from '@toeverything/components/tooltip';
 import { useBlockSuiteWorkspaceAvatarUrl } from '@toeverything/hooks/use-block-suite-workspace-avatar-url';
 import { useBlockSuiteWorkspaceName } from '@toeverything/hooks/use-block-suite-workspace-name';
-import { useStaticBlockSuiteWorkspace } from '@toeverything/infra/__internal__/react';
+import { getBlockSuiteWorkspaceAtom } from '@toeverything/infra/__internal__/workspace';
 import clsx from 'clsx';
 import { useAtom, useAtomValue } from 'jotai/react';
 import { type ReactElement, Suspense, useCallback, useMemo } from 'react';
@@ -209,7 +209,8 @@ const WorkspaceListItem = ({
   isCurrent: boolean;
   isActive: boolean;
 }) => {
-  const workspace = useStaticBlockSuiteWorkspace(meta.id);
+  const [workspaceAtom] = getBlockSuiteWorkspaceAtom(meta.id);
+  const workspace = useAtomValue(workspaceAtom);
   const [workspaceAvatar] = useBlockSuiteWorkspaceAvatarUrl(workspace);
   const [workspaceName] = useBlockSuiteWorkspaceName(workspace);
 

apps/core/src/components/affine/setting-modal/workspace-setting/index.tsx

@@ -3,7 +3,6 @@ import { WorkspaceSubPath } from '@affine/env/workspace';
 import { useAFFiNEI18N } from '@affine/i18n/hooks';
 import { rootWorkspacesMetadataAtom } from '@affine/workspace/atom';
 import { useBlockSuiteWorkspaceName } from '@toeverything/hooks/use-block-suite-workspace-name';
-import { usePassiveWorkspaceEffect } from '@toeverything/infra/__internal__/react';
 import { useSetAtom } from 'jotai';
 import { useAtomValue } from 'jotai';
 import { useCallback } from 'react';
@@ -34,7 +33,6 @@ export const WorkspaceSetting = ({ workspaceId }: { workspaceId: string }) => {
   const pushNotification = useSetAtom(pushNotificationAtom);
 
   const leaveWorkspace = useLeaveWorkspace();
-  usePassiveWorkspaceEffect(workspace.blockSuiteWorkspace);
   const setSettingModal = useSetAtom(openSettingModalAtom);
   const { deleteWorkspace } = useAppHelper();
 

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

@@ -7,7 +7,7 @@ import { useAtom, useSetAtom } from 'jotai';
 import { useCallback, useEffect } from 'react';
 
 import type { AllWorkspace } from '../../shared';
-import { useWorkspace } from '../use-workspace';
+import { useWorkspace, useWorkspaceEffect } from '../use-workspace';
 
 declare global {
   /**
@@ -27,6 +27,8 @@ export function useCurrentWorkspace(): [
   const [id, setId] = useAtom(currentWorkspaceIdAtom);
   assertExists(id);
   const currentWorkspace = useWorkspace(id);
+  // when you call current workspace, effect is always called
+  useWorkspaceEffect(currentWorkspace.id);
   useEffect(() => {
     globalThis.currentWorkspace = currentWorkspace;
     globalThis.dispatchEvent(

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

@@ -2,7 +2,7 @@ import type { AffineOfficialWorkspace } from '@affine/env/workspace';
 import { rootWorkspacesMetadataAtom } from '@affine/workspace/atom';
 import { assertExists } from '@blocksuite/global/utils';
 import type { Workspace } from '@blocksuite/store';
-import { useStaticBlockSuiteWorkspace } from '@toeverything/infra/__internal__/react';
+import { getBlockSuiteWorkspaceAtom } from '@toeverything/infra/__internal__/workspace';
 import type { Atom } from 'jotai';
 import { atom, useAtomValue } from 'jotai';
 
@@ -11,25 +11,32 @@ const workspaceWeakMap = new WeakMap<
   Atom<Promise<AffineOfficialWorkspace>>
 >();
 
+// workspace effect is the side effect like connect to the server/indexeddb,
+//  this will save the workspace updates permanently.
+export function useWorkspaceEffect(workspaceId: string): void {
+  const [, effectAtom] = getBlockSuiteWorkspaceAtom(workspaceId);
+  useAtomValue(effectAtom);
+}
+
+// todo(himself65): remove this hook
 export function useWorkspace(workspaceId: string): AffineOfficialWorkspace {
-  const blockSuiteWorkspace = useStaticBlockSuiteWorkspace(workspaceId);
-  if (!workspaceWeakMap.has(blockSuiteWorkspace)) {
+  const [workspaceAtom] = getBlockSuiteWorkspaceAtom(workspaceId);
+  const workspace = useAtomValue(workspaceAtom);
+  if (!workspaceWeakMap.has(workspace)) {
     const baseAtom = atom(async get => {
       const metadata = await get(rootWorkspacesMetadataAtom);
       const flavour = metadata.find(({ id }) => id === workspaceId)?.flavour;
       assertExists(flavour, 'workspace flavour not found');
       return {
         id: workspaceId,
         flavour,
-        blockSuiteWorkspace,
+        blockSuiteWorkspace: workspace,
       };
     });
-    workspaceWeakMap.set(blockSuiteWorkspace, baseAtom);
+    workspaceWeakMap.set(workspace, baseAtom);
   }
 
   return useAtomValue(
-    workspaceWeakMap.get(blockSuiteWorkspace) as Atom<
-      Promise<AffineOfficialWorkspace>
-    >
+    workspaceWeakMap.get(workspace) as Atom<Promise<AffineOfficialWorkspace>>
   );
 }

apps/core/src/layouts/workspace-layout.tsx

@@ -28,7 +28,6 @@ import {
   useSensors,
 } from '@dnd-kit/core';
 import { useBlockSuitePageMeta } from '@toeverything/hooks/use-block-suite-page-meta';
-import { usePassiveWorkspaceEffect } from '@toeverything/infra/__internal__/react';
 import { currentWorkspaceIdAtom } from '@toeverything/infra/atom';
 import { useAtom, useAtomValue, useSetAtom } from 'jotai';
 import { nanoid } from 'nanoid';
@@ -172,8 +171,6 @@ export const WorkspaceLayoutInner = ({
     }
   }, [currentWorkspace.blockSuiteWorkspace.doc]);
 
-  usePassiveWorkspaceEffect(currentWorkspace.blockSuiteWorkspace);
-
   const handleCreatePage = useCallback(() => {
     const id = nanoid();
     pageHelper.createPage(id);

apps/core/src/pages/workspace/all-page.tsx

@@ -1,7 +1,7 @@
 import { useCollectionManager } from '@affine/component/page-list';
 import { WorkspaceSubPath } from '@affine/env/workspace';
 import { assertExists } from '@blocksuite/global/utils';
-import { getActiveBlockSuiteWorkspaceAtom } from '@toeverything/infra/__internal__/workspace';
+import { getBlockSuiteWorkspaceAtom } from '@toeverything/infra/__internal__/workspace';
 import { getCurrentStore } from '@toeverything/infra/atom';
 import { useCallback } from 'react';
 import type { LoaderFunction } from 'react-router-dom';
@@ -16,7 +16,7 @@ export const loader: LoaderFunction = async args => {
   const rootStore = getCurrentStore();
   const workspaceId = args.params.workspaceId;
   assertExists(workspaceId);
-  const workspaceAtom = getActiveBlockSuiteWorkspaceAtom(workspaceId);
+  const [workspaceAtom] = getBlockSuiteWorkspaceAtom(workspaceId);
   const workspace = await rootStore.get(workspaceAtom);
   for (const pageId of workspace.pages.keys()) {
     const page = workspace.getPage(pageId);

apps/core/src/pages/workspace/index.tsx

@@ -1,6 +1,6 @@
 import { WorkspaceFlavour } from '@affine/env/workspace';
 import { rootWorkspacesMetadataAtom } from '@affine/workspace/atom';
-import { getActiveBlockSuiteWorkspaceAtom } from '@toeverything/infra/__internal__/workspace';
+import { getBlockSuiteWorkspaceAtom } from '@toeverything/infra/__internal__/workspace';
 import {
   currentPageIdAtom,
   currentWorkspaceIdAtom,
@@ -31,7 +31,7 @@ export const loader: LoaderFunction = async args => {
     rootStore.set(currentPageIdAtom, null);
   }
   if (currentMetadata.flavour === WorkspaceFlavour.AFFINE_CLOUD) {
-    const workspaceAtom = getActiveBlockSuiteWorkspaceAtom(currentMetadata.id);
+    const [workspaceAtom] = getBlockSuiteWorkspaceAtom(currentMetadata.id);
     const workspace = await rootStore.get(workspaceAtom);
     return (() => {
       const blockVersions = workspace.meta.blockVersions;

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

@@ -8,7 +8,8 @@ import { Divider } from '@toeverything/components/divider';
 import { Tooltip } from '@toeverything/components/tooltip';
 import { useBlockSuiteWorkspaceAvatarUrl } from '@toeverything/hooks/use-block-suite-workspace-avatar-url';
 import { useBlockSuiteWorkspaceName } from '@toeverything/hooks/use-block-suite-workspace-name';
-import { useStaticBlockSuiteWorkspace } from '@toeverything/infra/__internal__/react';
+import { getBlockSuiteWorkspaceAtom } from '@toeverything/infra/__internal__/workspace';
+import { useAtomValue } from 'jotai/react';
 import { useCallback } from 'react';
 
 import {
@@ -97,7 +98,8 @@ export const WorkspaceCard = ({
   meta,
   isOwner = true,
 }: WorkspaceCardProps) => {
-  const workspace = useStaticBlockSuiteWorkspace(meta.id);
+  const [workspaceAtom] = getBlockSuiteWorkspaceAtom(meta.id);
+  const workspace = useAtomValue(workspaceAtom);
   const [name] = useBlockSuiteWorkspaceName(workspace);
   const [workspaceAvatar] = useBlockSuiteWorkspaceAvatarUrl(workspace);
   return (

packages/infra/package.json

@@ -59,6 +59,7 @@
     "@blocksuite/global": "0.0.0-20230926212737-6d4b1569-nightly",
     "@blocksuite/store": "0.0.0-20230926212737-6d4b1569-nightly",
     "jotai": "^2.4.3",
+    "jotai-effect": "^0.1.0",
     "tinykeys": "^2.1.0",
     "zod": "^3.22.4"
   },

packages/infra/src/__internal__/react.ts

@@ -1,22 +0,0 @@
-import type { Workspace } from '@blocksuite/store';
-import { useAtomValue } from 'jotai/react';
-import { useEffect } from 'react';
-
-import {
-  disablePassiveProviders,
-  enablePassiveProviders,
-  getActiveBlockSuiteWorkspaceAtom,
-} from './workspace.js';
-
-export function useStaticBlockSuiteWorkspace(id: string): Workspace {
-  return useAtomValue(getActiveBlockSuiteWorkspaceAtom(id));
-}
-
-export function usePassiveWorkspaceEffect(workspace: Workspace) {
-  useEffect(() => {
-    enablePassiveProviders(workspace);
-    return () => {
-      disablePassiveProviders(workspace);
-    };
-  }, [workspace]);
-}

packages/infra/src/__internal__/workspace.ts

@@ -2,9 +2,9 @@ import type { ActiveDocProvider, Workspace } from '@blocksuite/store';
 import type { PassiveDocProvider } from '@blocksuite/store';
 import type { Atom } from 'jotai/vanilla';
 import { atom } from 'jotai/vanilla';
+import { atomEffect } from 'jotai-effect';
 
 /**
- * DO NOT ACCESS THIS MAP IN PRODUCTION, OR YOU WILL BE FIRED
  * Map: guid -> Workspace
  */
 export const INTERNAL_BLOCKSUITE_HASH_MAP = new Map<string, Workspace>([]);
@@ -14,49 +14,8 @@ const workspaceActiveAtomWeakMap = new WeakMap<
   Atom<Promise<Workspace>>
 >();
 
-// Whether the workspace is active to use
 const workspaceActiveWeakMap = new WeakMap<Workspace, boolean>();
-
-/**
- * Whether the workspace has been enabled the passive effect (background)
- *
- * @internal
- */
-export const workspacePassiveEffectWeakMap = new WeakMap<Workspace, number>();
-
-export function enablePassiveProviders(workspace: Workspace) {
-  const value = workspacePassiveEffectWeakMap.get(workspace);
-  if (value !== undefined && value !== 0) {
-    workspacePassiveEffectWeakMap.set(workspace, value + 1);
-    return;
-  }
-  const providers = workspace.providers.filter(
-    (provider): provider is PassiveDocProvider =>
-      'passive' in provider && provider.passive === true
-  );
-  providers.forEach(provider => {
-    provider.connect();
-  });
-  workspacePassiveEffectWeakMap.set(workspace, 1);
-}
-
-export function disablePassiveProviders(workspace: Workspace) {
-  const value = workspacePassiveEffectWeakMap.get(workspace);
-  if (value && value > 0) {
-    workspacePassiveEffectWeakMap.set(workspace, value - 1);
-    if (value - 1 === 0) {
-      const providers = workspace.providers.filter(
-        (provider): provider is PassiveDocProvider =>
-          'passive' in provider && provider.passive === true
-      );
-      providers.forEach(provider => {
-        provider.disconnect();
-      });
-      workspacePassiveEffectWeakMap.delete(workspace);
-    }
-    return;
-  }
-}
+const workspaceEffectAtomWeakMap = new WeakMap<Workspace, Atom<void>>();
 
 export async function waitForWorkspace(workspace: Workspace) {
   if (workspaceActiveWeakMap.get(workspace) !== true) {
@@ -69,6 +28,7 @@ export async function waitForWorkspace(workspace: Workspace) {
       // we will wait for the necessary providers to be ready
       await provider.whenReady;
     }
+    // timeout is INFINITE
     workspaceActiveWeakMap.set(workspace, true);
   }
 }
@@ -80,9 +40,9 @@ export function getWorkspace(id: string) {
   return INTERNAL_BLOCKSUITE_HASH_MAP.get(id) as Workspace;
 }
 
-export function getActiveBlockSuiteWorkspaceAtom(
+export function getBlockSuiteWorkspaceAtom(
   id: string
-): Atom<Promise<Workspace>> {
+): [workspaceAtom: Atom<Promise<Workspace>>, workspaceEffectAtom: Atom<void>] {
   if (!INTERNAL_BLOCKSUITE_HASH_MAP.has(id)) {
     throw new Error('Workspace not found');
   }
@@ -94,5 +54,26 @@ export function getActiveBlockSuiteWorkspaceAtom(
     });
     workspaceActiveAtomWeakMap.set(workspace, baseAtom);
   }
-  return workspaceActiveAtomWeakMap.get(workspace) as Atom<Promise<Workspace>>;
+  if (!workspaceEffectAtomWeakMap.has(workspace)) {
+    const effectAtom = atomEffect(() => {
+      const providers = workspace.providers.filter(
+        (provider): provider is PassiveDocProvider =>
+          'passive' in provider && provider.passive === true
+      );
+      providers.forEach(provider => {
+        provider.connect();
+      });
+      return () => {
+        providers.forEach(provider => {
+          provider.disconnect();
+        });
+      };
+    });
+    workspaceEffectAtomWeakMap.set(workspace, effectAtom);
+  }
+
+  return [
+    workspaceActiveAtomWeakMap.get(workspace) as Atom<Promise<Workspace>>,
+    workspaceEffectAtomWeakMap.get(workspace) as Atom<void>,
+  ];
 }

packages/infra/src/__tests__/use-block-suite-workspace.spec.ts → packages/infra/src/__tests__/blocksuite-atom.ts

@@ -2,20 +2,16 @@
  * @vitest-environment happy-dom
  */
 import { Schema, Workspace } from '@blocksuite/store';
-import { renderHook } from '@testing-library/react';
+import { waitFor } from '@testing-library/react';
 import { getDefaultStore } from 'jotai/vanilla';
 import { expect, test, vi } from 'vitest';
 
 import {
-  usePassiveWorkspaceEffect,
-  useStaticBlockSuiteWorkspace,
-} from '../__internal__/react.js';
-import {
-  getActiveBlockSuiteWorkspaceAtom,
+  getBlockSuiteWorkspaceAtom,
   INTERNAL_BLOCKSUITE_HASH_MAP,
 } from '../__internal__/workspace.js';
 
-test('useStaticBlockSuiteWorkspace', async () => {
+test('blocksuite atom', async () => {
   const sync = vi.fn();
   let connected = false;
   const connect = vi.fn(() => (connected = true));
@@ -45,27 +41,15 @@ test('useStaticBlockSuiteWorkspace', async () => {
   INTERNAL_BLOCKSUITE_HASH_MAP.set('1', workspace);
 
   {
-    const workspaceHook = renderHook(() => useStaticBlockSuiteWorkspace('1'));
-    // wait for suspense to resolve
-    await new Promise(resolve => setTimeout(resolve, 100));
-    expect(workspaceHook.result.current).toBe(workspace);
-    expect(sync).toBeCalledTimes(1);
-    expect(connect).not.toHaveBeenCalled();
-  }
-
-  {
-    const atom = getActiveBlockSuiteWorkspaceAtom('1');
+    const [atom, effectAtom] = getBlockSuiteWorkspaceAtom('1');
     const store = getDefaultStore();
     const result = await store.get(atom);
     expect(result).toBe(workspace);
     expect(sync).toBeCalledTimes(1);
     expect(connect).not.toHaveBeenCalled();
-  }
 
-  {
-    renderHook(() => usePassiveWorkspaceEffect(workspace));
-    expect(sync).toBeCalledTimes(1);
-    expect(connect).toBeCalledTimes(1);
+    store.sub(effectAtom, vi.fn());
+    await waitFor(() => expect(connect).toBeCalledTimes(1));
     expect(connected).toBe(true);
   }
 });