feat: improve workspace hook (#3099)

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

@@ -17,11 +17,9 @@ import {
   saveWorkspaceToLocalStorage,
 } from '@affine/workspace/local/crud';
 import { createIndexedDBDownloadProvider } from '@affine/workspace/providers';
-import {
-  createEmptyBlockSuiteWorkspace,
-  useStaticBlockSuiteWorkspace,
-} from '@affine/workspace/utils';
+import { createEmptyBlockSuiteWorkspace } from '@affine/workspace/utils';
 import { nanoid } from '@blocksuite/store';
+import { useStaticBlockSuiteWorkspace } from '@toeverything/hooks/use-block-suite-workspace';
 
 import {
   BlockSuitePageList,

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

@@ -2,7 +2,7 @@ import { UserAvatar } from '@affine/component/user-avatar';
 import { WorkspaceAvatar } from '@affine/component/workspace-avatar';
 import { useAFFiNEI18N } from '@affine/i18n/hooks';
 import type { RootWorkspaceMetadata } from '@affine/workspace/atom';
-import { useStaticBlockSuiteWorkspace } from '@affine/workspace/utils';
+import { useStaticBlockSuiteWorkspace } from '@toeverything/hooks/use-block-suite-workspace';
 import { useBlockSuiteWorkspaceName } from '@toeverything/hooks/use-block-suite-workspace-name';
 import clsx from 'clsx';
 

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

@@ -1,7 +1,7 @@
+import { usePassiveWorkspaceEffect } from '@toeverything/hooks/use-block-suite-workspace';
 import { Suspense, useCallback } from 'react';
 
 import { getUIAdapter } from '../../../../adapters/workspace';
-import { usePassiveWorkspaceEffect } from '../../../../hooks/current/use-current-workspace';
 import { useOnTransformWorkspace } from '../../../../hooks/root/use-on-transform-workspace';
 import { useWorkspace } from '../../../../hooks/use-workspace';
 import { useAppHelper } from '../../../../hooks/use-workspaces';

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

@@ -3,7 +3,6 @@ import {
   rootCurrentWorkspaceIdAtom,
 } from '@affine/workspace/atom';
 import { assertExists } from '@blocksuite/global/utils';
-import type { PassiveDocProvider, Workspace } from '@blocksuite/store';
 import { useAtom, useSetAtom } from 'jotai';
 import { useCallback, useEffect } from 'react';
 
@@ -51,27 +50,3 @@ 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-workspace.ts

@@ -1,7 +1,7 @@
 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 { useStaticBlockSuiteWorkspace } from '@toeverything/hooks/use-block-suite-workspace';
 import type { Atom } from 'jotai';
 import { atom, useAtomValue } from 'jotai';
 

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

@@ -2,11 +2,9 @@ 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,
-  getWorkspace,
-} from '@affine/workspace/utils';
+import { createEmptyBlockSuiteWorkspace } from '@affine/workspace/utils';
 import { nanoid } from '@blocksuite/store';
+import { getWorkspace } from '@toeverything/hooks/use-block-suite-workspace';
 import { useAtomValue, useSetAtom } from 'jotai';
 import { useCallback } from 'react';
 

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

@@ -31,6 +31,7 @@ import {
   useSensor,
   useSensors,
 } from '@dnd-kit/core';
+import { usePassiveWorkspaceEffect } from '@toeverything/hooks/use-block-suite-workspace';
 import { useBlockSuiteWorkspaceHelper } from '@toeverything/hooks/use-block-suite-workspace-helper';
 import { useAtom, useAtomValue, useSetAtom } from 'jotai';
 import Head from 'next/head';
@@ -54,10 +55,7 @@ import {
   RootAppSidebar,
 } from '../components/root-app-sidebar';
 import { useBlockSuiteMetaHelper } from '../hooks/affine/use-block-suite-meta-helper';
-import {
-  useCurrentWorkspace,
-  usePassiveWorkspaceEffect,
-} from '../hooks/current/use-current-workspace';
+import { useCurrentWorkspace } from '../hooks/current/use-current-workspace';
 import { useRouterHelper } from '../hooks/use-router-helper';
 import { useRouterTitle } from '../hooks/use-router-title';
 import {

apps/web/src/pages/index.tsx

@@ -2,7 +2,7 @@ import { WorkspaceFallback } from '@affine/component/workspace';
 import { DebugLogger } from '@affine/debug';
 import { WorkspaceSubPath } from '@affine/env/workspace';
 import { rootWorkspacesMetadataAtom } from '@affine/workspace/atom';
-import { getWorkspace } from '@affine/workspace/utils';
+import { getWorkspace } from '@toeverything/hooks/use-block-suite-workspace';
 import { useAtomValue } from 'jotai';
 import type { NextPage } from 'next';
 import { useRouter } from 'next/router';

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

@@ -1,7 +1,6 @@
 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,
@@ -9,6 +8,7 @@ import {
   LocalDataIcon as DefaultLocalDataIcon,
   LocalWorkspaceIcon as DefaultLocalWorkspaceIcon,
 } from '@blocksuite/icons';
+import { useStaticBlockSuiteWorkspace } from '@toeverything/hooks/use-block-suite-workspace';
 import { useBlockSuiteWorkspaceName } from '@toeverything/hooks/use-block-suite-workspace-name';
 import type { FC } from 'react';
 import { useCallback } from 'react';

packages/hooks/src/__tests__/use-block-suite-workspace.spec.ts

@@ -0,0 +1,67 @@
+/**
+ * @vitest-environment happy-dom
+ */
+import { Workspace } from '@blocksuite/store';
+import { renderHook } from '@testing-library/react';
+import {
+  getActiveBlockSuiteWorkspaceAtom,
+  INTERNAL_BLOCKSUITE_HASH_MAP,
+  usePassiveWorkspaceEffect,
+  useStaticBlockSuiteWorkspace,
+} from '@toeverything/hooks/use-block-suite-workspace';
+import { getDefaultStore } from 'jotai/vanilla';
+import { expect, test, vi } from 'vitest';
+
+test('useStaticBlockSuiteWorkspace', async () => {
+  const sync = vi.fn();
+  let connected = false;
+  const connect = vi.fn(() => (connected = true));
+  const workspace = new Workspace({
+    id: '1',
+    providerCreators: [
+      () => ({
+        flavour: 'fake',
+        active: true,
+        sync,
+        get whenReady(): Promise<void> {
+          return Promise.resolve();
+        },
+      }),
+      () => ({
+        flavour: 'fake-2',
+        passive: true,
+        get connected() {
+          return connected;
+        },
+        connect,
+        disconnect: vi.fn(),
+      }),
+    ],
+  });
+  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 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);
+    expect(connected).toBe(true);
+  }
+});

packages/hooks/src/use-block-suite-workspace.ts

@@ -0,0 +1,84 @@
+// guid -> Workspace
+import type { ActiveDocProvider, Workspace } from '@blocksuite/store';
+import type { PassiveDocProvider } from '@blocksuite/store';
+import { useAtomValue } from 'jotai/react';
+import type { Atom } from 'jotai/vanilla';
+import { atom } from 'jotai/vanilla';
+import { useEffect } from 'react';
+
+/**
+ * DO NOT ACCESS THIS MAP IN PRODUCTION, OR YOU WILL BE FIRED
+ */
+export const INTERNAL_BLOCKSUITE_HASH_MAP = new Map<string, Workspace>([]);
+
+const workspacePassiveAtomWeakMap = new WeakMap<
+  Workspace,
+  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)
+const workspacePassiveEffectWeakMap = new WeakMap<Workspace, boolean>();
+
+export function getWorkspace(id: string) {
+  if (!INTERNAL_BLOCKSUITE_HASH_MAP.has(id)) {
+    throw new Error('Workspace not found');
+  }
+  return INTERNAL_BLOCKSUITE_HASH_MAP.get(id) as Workspace;
+}
+
+export function getActiveBlockSuiteWorkspaceAtom(
+  id: string
+): Atom<Promise<Workspace>> {
+  if (!INTERNAL_BLOCKSUITE_HASH_MAP.has(id)) {
+    throw new Error('Workspace not found');
+  }
+  const workspace = INTERNAL_BLOCKSUITE_HASH_MAP.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(getActiveBlockSuiteWorkspaceAtom(id));
+}
+
+export function usePassiveWorkspaceEffect(workspace: Workspace) {
+  useEffect(() => {
+    if (workspacePassiveEffectWeakMap.get(workspace) === true) {
+      return;
+    }
+    const providers = workspace.providers.filter(
+      (provider): provider is PassiveDocProvider =>
+        'passive' in provider && provider.passive === true
+    );
+    providers.forEach(provider => {
+      provider.connect();
+    });
+    workspacePassiveEffectWeakMap.set(workspace, true);
+    return () => {
+      providers.forEach(provider => {
+        provider.disconnect();
+      });
+      workspacePassiveEffectWeakMap.delete(workspace);
+    };
+  }, [workspace]);
+}

packages/workspace/src/utils.ts

@@ -7,15 +7,12 @@ 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 { useAtomValue } from 'jotai/react';
-import type { Atom } from 'jotai/vanilla';
-import { atom } from 'jotai/vanilla';
+import { INTERNAL_BLOCKSUITE_HASH_MAP } from '@toeverything/hooks/use-block-suite-workspace';
 
 import { createStaticStorage } from './blob/local-static-storage';
 import { createSQLiteStorage } from './blob/sqlite-blob-storage';
@@ -33,54 +30,6 @@ function setEditorFlags(workspace: Workspace) {
   );
 }
 
-// guid -> Workspace
-export const workspaceHashMap = new Map<string, Workspace>();
-
-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,
   flavour: WorkspaceFlavour.AFFINE_CLOUD,
@@ -108,8 +57,8 @@ export function createEmptyBlockSuiteWorkspace(
   const providerCreators: DocProviderCreator[] = [];
   const prefix: string = config?.cachePrefix ?? '';
   const cacheKey = `${prefix}${id}`;
-  if (workspaceHashMap.has(cacheKey)) {
-    return workspaceHashMap.get(cacheKey) as Workspace;
+  if (INTERNAL_BLOCKSUITE_HASH_MAP.has(cacheKey)) {
+    return INTERNAL_BLOCKSUITE_HASH_MAP.get(cacheKey) as Workspace;
   }
   const idGenerator = config?.idGenerator;
 
@@ -122,7 +71,7 @@ export function createEmptyBlockSuiteWorkspace(
         blobStorages.push(createSQLiteStorage);
       }
 
-      // todo: add support for cloud storage
+      // todo(JimmFly): add support for cloud storage
     }
     providerCreators.push(...createAffineProviders());
   } else {
@@ -146,6 +95,6 @@ export function createEmptyBlockSuiteWorkspace(
     .register(AffineSchemas)
     .register(__unstableSchemas);
   setEditorFlags(workspace);
-  workspaceHashMap.set(cacheKey, workspace);
+  INTERNAL_BLOCKSUITE_HASH_MAP.set(cacheKey, workspace);
   return workspace;
 }