refactor(plugin-infra): workspace loading (#3222)

packages/plugin-infra/package.json

@@ -18,17 +18,19 @@
       "type": "./dist/type.d.ts",
       "import": "./dist/type.js",
       "require": "./dist/type.cjs"
+    },
+    "./__internal__/workspace": {
+      "type": "./dist/__internal__/workspace.d.ts",
+      "import": "./dist/__internal__/workspace.js",
+      "require": "./dist/__internal__/workspace.cjs"
     }
   },
-  "dependencies": {
+  "devDependencies": {
     "@blocksuite/blocks": "0.0.0-20230713062147-9eb7708c-nightly",
     "@blocksuite/editor": "0.0.0-20230713062147-9eb7708c-nightly",
     "@blocksuite/global": "0.0.0-20230713062147-9eb7708c-nightly",
     "@blocksuite/lit": "0.0.0-20230713062147-9eb7708c-nightly",
     "@blocksuite/store": "0.0.0-20230713062147-9eb7708c-nightly",
-    "jotai": "^2.2.2"
-  },
-  "devDependencies": {
     "jotai": "^2.2.2",
     "vite": "^4.3.9",
     "vite-plugin-dts": "3.0.2"

packages/plugin-infra/src/__internal__/workspace.ts

@@ -0,0 +1,88 @@
+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
+ * Map: guid -> Workspace
+ */
+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
+export const workspaceActiveWeakMap = new WeakMap<Workspace, boolean>();
+
+// Whether the workspace has been enabled the passive effect (background)
+export const workspacePassiveEffectWeakMap = new WeakMap<Workspace, boolean>();
+
+export async function waitForWorkspace(workspace: Workspace) {
+  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);
+  }
+}
+
+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 () => {
+      await waitForWorkspace(workspace);
+      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/plugin-infra/src/__tests__/use-block-suite-workspace.spec.ts

@@ -0,0 +1,68 @@
+/**
+ * @vitest-environment happy-dom
+ */
+import { Workspace } from '@blocksuite/store';
+import { renderHook } from '@testing-library/react';
+import { getDefaultStore } from 'jotai/vanilla';
+import { expect, test, vi } from 'vitest';
+
+import {
+  getActiveBlockSuiteWorkspaceAtom,
+  INTERNAL_BLOCKSUITE_HASH_MAP,
+  usePassiveWorkspaceEffect,
+  useStaticBlockSuiteWorkspace,
+} from '../__internal__/workspace';
+
+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/plugin-infra/src/manager.ts

@@ -1,5 +1,8 @@
+import { assertExists } from '@blocksuite/global/utils';
+import type { Page, Workspace } from '@blocksuite/store';
 import { atom, createStore } from 'jotai/vanilla';
 
+import { getWorkspace, waitForWorkspace } from './__internal__/workspace';
 import type { AffinePlugin, Definition, ServerAdapter } from './type';
 import type { Loader, PluginUIAdapter } from './type';
 import type { PluginBlockSuiteAdapter } from './type';
@@ -12,8 +15,29 @@ export const rootStore = createStore();
 
 // todo: for now every plugin is enabled by default
 export const affinePluginsAtom = atom<Record<string, AffinePlugin<string>>>({});
-export const currentPageIdAtom = atom<string | null>(null);
 export const currentWorkspaceIdAtom = atom<string | null>(null);
+export const currentPageIdAtom = atom<string | null>(null);
+export const currentWorkspaceAtom = atom<Promise<Workspace>>(async get => {
+  const currentWorkspaceId = get(currentWorkspaceIdAtom);
+  assertExists(currentWorkspaceId, 'current workspace id');
+  const workspace = getWorkspace(currentWorkspaceId);
+  await waitForWorkspace(workspace);
+  return workspace;
+});
+export const currentPageAtom = atom<Promise<Page>>(async get => {
+  const currentWorkspaceId = get(currentWorkspaceIdAtom);
+  assertExists(currentWorkspaceId, 'current workspace id');
+  const currentPageId = get(currentPageIdAtom);
+  assertExists(currentPageId, 'current page id');
+  const workspace = getWorkspace(currentWorkspaceId);
+  await waitForWorkspace(workspace);
+  const page = workspace.getPage(currentPageId);
+  assertExists(page);
+  if (!page.loaded) {
+    await page.waitForLoaded();
+  }
+  return page;
+});
 
 export function definePlugin<ID extends string>(
   definition: Definition<ID>,

packages/plugin-infra/vite.config.ts

@@ -13,18 +13,14 @@ export default defineConfig({
       entry: {
         type: resolve(root, 'src/type.ts'),
         manager: resolve(root, 'src/manager.ts'),
+        '__internal__/workspace': resolve(
+          root,
+          'src/__internal__/workspace.ts'
+        ),
       },
     },
     rollupOptions: {
-      external: [
-        'jotai',
-        'jotai/vanilla',
-        '@blocksuite/blocks',
-        '@blocksuite/store',
-        '@blocksuite/global',
-        '@blocksuite/editor',
-        '@blocksuite/lit',
-      ],
+      external: ['react', /^jotai/, /^@blocksuite/],
     },
   },
   plugins: [