refactor: use jotai-effect (#4641)

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__/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);
   }
 });

packages/infra/src/__tests__/workspace.spec.ts

@@ -1,90 +0,0 @@
-import { AffineSchemas } from '@blocksuite/blocks/models';
-import type { DocProviderCreator } from '@blocksuite/store';
-import { Schema, Workspace } from '@blocksuite/store';
-import { getDefaultStore } from 'jotai/vanilla';
-import { beforeEach, expect, test } from 'vitest';
-
-import {
-  disablePassiveProviders,
-  enablePassiveProviders,
-  getActiveBlockSuiteWorkspaceAtom,
-  getWorkspace,
-  INTERNAL_BLOCKSUITE_HASH_MAP,
-  workspacePassiveEffectWeakMap,
-} from '../__internal__/workspace';
-
-const schema = new Schema();
-
-schema.register(AffineSchemas);
-
-const activeWorkspaceEnabled = new Set<string>();
-const passiveWorkspaceEnabled = new Set<string>();
-
-beforeEach(() => {
-  activeWorkspaceEnabled.clear();
-});
-
-const createWorkspace = (id: string) => {
-  const activeCreator: DocProviderCreator = () => ({
-    flavour: 'active',
-    active: true,
-    sync() {
-      activeWorkspaceEnabled.add(id);
-    },
-    get whenReady(): Promise<void> {
-      return Promise.resolve();
-    },
-  });
-  const passiveCreator: DocProviderCreator = () => ({
-    flavour: 'passive',
-    passive: true,
-    connect() {
-      passiveWorkspaceEnabled.add(id);
-    },
-    disconnect() {
-      passiveWorkspaceEnabled.delete(id);
-    },
-    get connected() {
-      return false;
-    },
-  });
-
-  return new Workspace({
-    id,
-    schema,
-    providerCreators: [activeCreator, passiveCreator],
-  });
-};
-
-test('workspace passive provider should enable correctly', () => {
-  INTERNAL_BLOCKSUITE_HASH_MAP.set('1', createWorkspace('1'));
-  INTERNAL_BLOCKSUITE_HASH_MAP.set('2', createWorkspace('2'));
-  expect(workspacePassiveEffectWeakMap.get(getWorkspace('1'))).toBe(undefined);
-  enablePassiveProviders(getWorkspace('1'));
-  expect(workspacePassiveEffectWeakMap.get(getWorkspace('1'))).toBe(1);
-  expect(workspacePassiveEffectWeakMap.get(getWorkspace('2'))).toBe(undefined);
-  enablePassiveProviders(getWorkspace('1'));
-  expect(workspacePassiveEffectWeakMap.get(getWorkspace('1'))).toBe(2);
-  disablePassiveProviders(getWorkspace('1'));
-  expect(workspacePassiveEffectWeakMap.get(getWorkspace('1'))).toBe(1);
-  disablePassiveProviders(getWorkspace('1'));
-  expect(workspacePassiveEffectWeakMap.get(getWorkspace('1'))).toBe(undefined);
-});
-
-test('workspace provider should initialize correctly', async () => {
-  INTERNAL_BLOCKSUITE_HASH_MAP.set('1', createWorkspace('1'));
-  {
-    enablePassiveProviders(getWorkspace('1'));
-    expect(activeWorkspaceEnabled.size).toBe(0);
-    expect(passiveWorkspaceEnabled.size).toBe(1);
-    disablePassiveProviders(getWorkspace('1'));
-    expect(activeWorkspaceEnabled.size).toBe(0);
-    expect(passiveWorkspaceEnabled.size).toBe(0);
-  }
-  {
-    const atom = getActiveBlockSuiteWorkspaceAtom('1');
-    await getDefaultStore().get(atom);
-    expect(activeWorkspaceEnabled.size).toBe(1);
-    expect(passiveWorkspaceEnabled.size).toBe(0);
-  }
-});

packages/infra/src/atom.ts

@@ -3,7 +3,7 @@ import { assertExists } from '@blocksuite/global/utils';
 import type { Workspace } from '@blocksuite/store';
 import { atom, createStore } from 'jotai/vanilla';
 
-import { getActiveBlockSuiteWorkspaceAtom } from './__internal__/workspace';
+import { getBlockSuiteWorkspaceAtom } from './__internal__/workspace';
 
 // global store
 let rootStore = createStore();
@@ -26,7 +26,7 @@ export const currentPageIdAtom = atom<string | null>(null);
 export const currentWorkspaceAtom = atom<Promise<Workspace>>(async get => {
   const workspaceId = get(currentWorkspaceIdAtom);
   assertExists(workspaceId);
-  const currentWorkspaceAtom = getActiveBlockSuiteWorkspaceAtom(workspaceId);
+  const [currentWorkspaceAtom] = getBlockSuiteWorkspaceAtom(workspaceId);
   return get(currentWorkspaceAtom);
 });