refactor: workspace provider (#2218)

apps/web/src/atoms/__tests__/atom.spec.ts

@@ -0,0 +1,71 @@
+/**
+ * @vitest-environment happy-dom
+ */
+import 'fake-indexeddb/auto';
+
+import { initPage } from '@affine/env/blocksuite';
+import {
+  rootCurrentWorkspaceIdAtom,
+  rootWorkspacesMetadataAtom,
+} from '@affine/workspace/atom';
+import { createIndexedDBDownloadProvider } from '@affine/workspace/providers';
+import { WorkspaceFlavour } from '@affine/workspace/type';
+import {
+  _cleanupBlockSuiteWorkspaceCache,
+  createEmptyBlockSuiteWorkspace,
+} from '@affine/workspace/utils';
+import type { ParagraphBlockModel } from '@blocksuite/blocks/models';
+import type { Page } from '@blocksuite/store';
+import { createStore } from 'jotai';
+import { describe, expect, test } from 'vitest';
+
+import { WorkspacePlugins } from '../../plugins';
+import { rootCurrentWorkspaceAtom } from '../root';
+
+describe('currentWorkspace atom', () => {
+  test('should be defined', async () => {
+    const store = createStore();
+    let id: string;
+    {
+      const workspace = createEmptyBlockSuiteWorkspace(
+        'test',
+        WorkspaceFlavour.LOCAL
+      );
+      const page = workspace.createPage('page0');
+      initPage(page);
+      const frameId = page.getBlockByFlavour('affine:frame').at(0)
+        ?.id as string;
+      id = page.addBlock(
+        'affine:paragraph',
+        {
+          text: new page.Text('test 1'),
+        },
+        frameId
+      );
+      const provider = createIndexedDBDownloadProvider(workspace);
+      provider.sync();
+      await provider.whenReady;
+      const workspaceId = await WorkspacePlugins[
+        WorkspaceFlavour.LOCAL
+      ].CRUD.create(workspace);
+      store.set(rootWorkspacesMetadataAtom, [
+        {
+          id: workspaceId,
+          flavour: WorkspaceFlavour.LOCAL,
+        },
+      ]);
+      _cleanupBlockSuiteWorkspaceCache();
+    }
+    store.set(
+      rootCurrentWorkspaceIdAtom,
+      store.get(rootWorkspacesMetadataAtom)[0].id
+    );
+    const workspace = await store.get(rootCurrentWorkspaceAtom);
+    expect(workspace).toBeDefined();
+    const page = workspace.blockSuiteWorkspace.getPage('page0') as Page;
+    expect(page).not.toBeNull();
+    const paragraphBlock = page.getBlockById(id) as ParagraphBlockModel;
+    expect(paragraphBlock).not.toBeNull();
+    expect(paragraphBlock.text.toString()).toBe('test 1');
+  });
+});

apps/web/src/atoms/root.ts

@@ -5,6 +5,10 @@ import {
   rootCurrentWorkspaceIdAtom,
   rootWorkspacesMetadataAtom,
 } from '@affine/workspace/atom';
+import type {
+  NecessaryProvider,
+  WorkspaceRegistry,
+} from '@affine/workspace/type';
 import { WorkspaceFlavour } from '@affine/workspace/type';
 import { assertExists } from '@blocksuite/store';
 import { atom } from 'jotai';
@@ -37,18 +41,38 @@ export const workspacesAtom = atom<Promise<AllWorkspace[]>>(async get => {
         WorkspacePlugins[workspace.flavour as keyof typeof WorkspacePlugins];
       assertExists(plugin);
       const { CRUD } = plugin;
-      return CRUD.get(workspace.id);
+      return CRUD.get(workspace.id).then(workspace => {
+        if (workspace === null) {
+          console.warn(
+            'workspace is null. this should not happen. If you see this error, please report it to the developer.'
+          );
+        }
+        return workspace;
+      });
     })
+  ).then(workspaces =>
+    workspaces.filter(
+      (workspace): workspace is WorkspaceRegistry['affine' | 'local'] =>
+        workspace !== null
+    )
   );
-  logger.info('workspaces', workspaces);
-  workspaces.forEach(workspace => {
-    if (workspace === null) {
-      console.warn(
-        'workspace is null. this should not happen. If you see this error, please report it to the developer.'
-      );
+  const workspaceProviders = workspaces.map(workspace =>
+    workspace.providers.filter(
+      (provider): provider is NecessaryProvider =>
+        'necessary' in provider && provider.necessary
+    )
+  );
+  const promises: Promise<void>[] = [];
+  for (const providers of workspaceProviders) {
+    for (const provider of providers) {
+      provider.sync();
+      promises.push(provider.whenReady);
     }
-  });
-  return workspaces.filter(workspace => workspace !== null) as AllWorkspace[];
+  }
+  // we will wait for all the necessary providers to be ready
+  await Promise.all(promises);
+  logger.info('workspaces', workspaces);
+  return workspaces;
 });
 
 /**
@@ -77,6 +101,15 @@ export const rootCurrentWorkspaceAtom = atom<Promise<AllWorkspace>>(
         `cannot find the workspace with id ${targetId} in the plugin ${targetWorkspace.flavour}.`
       );
     }
+    const providers = workspace.providers.filter(
+      (provider): provider is NecessaryProvider =>
+        'necessary' in provider && provider.necessary === true
+    );
+    for (const provider of providers) {
+      provider.sync();
+      // we will wait for the necessary providers to be ready
+      await provider.whenReady;
+    }
     return workspace;
   }
 );

apps/web/src/blocksuite/__tests__/index.spec.ts

@@ -1,30 +0,0 @@
-/**
- * @vitest-environment happy-dom
- */
-import 'fake-indexeddb/auto';
-
-import { beforeEach, describe, expect, test } from 'vitest';
-
-import { BlockSuiteWorkspace } from '../../shared';
-import { createAffineProviders, createLocalProviders } from '..';
-
-let blockSuiteWorkspace: BlockSuiteWorkspace;
-
-beforeEach(() => {
-  blockSuiteWorkspace = new BlockSuiteWorkspace({ id: 'test' });
-});
-
-describe('blocksuite providers', () => {
-  test('should be valid provider', () => {
-    [createLocalProviders, createAffineProviders].forEach(createProviders => {
-      createProviders(blockSuiteWorkspace).forEach(provider => {
-        expect(provider).toBeTypeOf('object');
-        expect(provider).toHaveProperty('flavour');
-        expect(provider).toHaveProperty('connect');
-        expect(provider.connect).toBeTypeOf('function');
-        expect(provider).toHaveProperty('disconnect');
-        expect(provider.disconnect).toBeTypeOf('function');
-      });
-    });
-  });
-});

apps/web/src/blocksuite/index.ts

@@ -1,13 +1,15 @@
 import { config } from '@affine/env';
 import {
-  createIndexedDBProvider,
+  createIndexedDBDownloadProvider,
   createLocalProviders,
 } from '@affine/workspace/providers';
-import { createBroadCastChannelProvider } from '@affine/workspace/providers';
+import {
+  createAffineWebSocketProvider,
+  createBroadCastChannelProvider,
+} from '@affine/workspace/providers';
 import type { Provider } from '@affine/workspace/type';
 
 import type { BlockSuiteWorkspace } from '../shared';
-import { createAffineWebSocketProvider } from './providers';
 import { createAffineDownloadProvider } from './providers/affine';
 
 export const createAffineProviders = (
@@ -19,7 +21,7 @@ export const createAffineProviders = (
       createAffineWebSocketProvider(blockSuiteWorkspace),
       config.enableBroadCastChannelProvider &&
         createBroadCastChannelProvider(blockSuiteWorkspace),
-      createIndexedDBProvider(blockSuiteWorkspace),
+      createIndexedDBDownloadProvider(blockSuiteWorkspace),
     ] as any[]
   ).filter(v => Boolean(v));
 };

apps/web/src/blocksuite/providers/affine/index.ts

@@ -12,9 +12,15 @@ export const createAffineDownloadProvider = (
 ): AffineDownloadProvider => {
   assertExists(blockSuiteWorkspace.id);
   const id = blockSuiteWorkspace.id;
+  let connected = false;
+  const callbacks = new Set<() => void>();
   return {
     flavour: 'affine-download',
     background: true,
+    get connected() {
+      return connected;
+    },
+    callbacks,
     connect: () => {
       providerLogger.info('connect download provider', id);
       if (hashMap.has(id)) {
@@ -23,6 +29,7 @@ export const createAffineDownloadProvider = (
           blockSuiteWorkspace.doc,
           new Uint8Array(hashMap.get(id) as ArrayBuffer)
         );
+        connected = true;
         return;
       }
       affineApis
@@ -41,6 +48,7 @@ export const createAffineDownloadProvider = (
     },
     disconnect: () => {
       providerLogger.info('disconnect download provider', id);
+      connected = false;
     },
     cleanup: () => {
       hashMap.delete(id);

apps/web/src/blocksuite/providers/index.ts

@@ -1,48 +0,0 @@
-import { websocketPrefixUrl } from '@affine/env';
-import { KeckProvider } from '@affine/workspace/affine/keck';
-import { getLoginStorage } from '@affine/workspace/affine/login';
-import type { AffineWebSocketProvider } from '@affine/workspace/type';
-import { assertExists } from '@blocksuite/store';
-
-import type { BlockSuiteWorkspace } from '../../shared';
-import { providerLogger } from '../logger';
-
-const createAffineWebSocketProvider = (
-  blockSuiteWorkspace: BlockSuiteWorkspace
-): AffineWebSocketProvider => {
-  let webSocketProvider: KeckProvider | null = null;
-  return {
-    flavour: 'affine-websocket',
-    background: false,
-    cleanup: () => {
-      assertExists(webSocketProvider);
-      webSocketProvider.destroy();
-      webSocketProvider = null;
-    },
-    connect: () => {
-      webSocketProvider = new KeckProvider(
-        websocketPrefixUrl + '/api/sync/',
-        blockSuiteWorkspace.id,
-        blockSuiteWorkspace.doc,
-        {
-          params: { token: getLoginStorage()?.token ?? '' },
-          awareness: blockSuiteWorkspace.awarenessStore.awareness,
-          // we maintain broadcast channel by ourselves
-          // @ts-expect-error
-          disableBc: true,
-          connect: false,
-        }
-      );
-      providerLogger.info('connect', webSocketProvider.url);
-      webSocketProvider.connect();
-    },
-    disconnect: () => {
-      assertExists(webSocketProvider);
-      providerLogger.info('disconnect', webSocketProvider.url);
-      webSocketProvider.destroy();
-      webSocketProvider = null;
-    },
-  };
-};
-
-export { createAffineWebSocketProvider };

apps/web/src/hooks/use-router-with-workspace-id-defense.ts

@@ -22,6 +22,7 @@ export function useRouterWithWorkspaceIdDefense(router: NextRouter) {
     }
     const exist = metadata.find(m => m.id === currentWorkspaceId);
     if (!exist) {
+      console.warn('workspace not exist, redirect to first one');
       // clean up
       setCurrentWorkspaceId(null);
       setCurrentPageId(null);

apps/web/src/hooks/use-sync-router-with-current-workspace-id.ts

@@ -45,8 +45,9 @@ export function useSyncRouterWithCurrentWorkspaceId(router: NextRouter) {
         window.apis?.onWorkspaceChange(targetWorkspace.id);
       }
       void router.push({
-        pathname: '/workspace/[workspaceId]/all',
+        pathname: router.pathname,
         query: {
+          ...router.query,
           workspaceId: targetWorkspace.id,
         },
       });
@@ -56,8 +57,9 @@ export function useSyncRouterWithCurrentWorkspaceId(router: NextRouter) {
         console.log('set workspace id', workspaceId);
         setCurrentWorkspaceId(targetWorkspace.id);
         void router.push({
-          pathname: '/workspace/[workspaceId]/all',
+          pathname: router.pathname,
           query: {
+            ...router.query,
             workspaceId: targetWorkspace.id,
           },
         });

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

@@ -1,6 +1,6 @@
 import { DebugLogger } from '@affine/debug';
 import { DEFAULT_HELLO_WORLD_PAGE_ID } from '@affine/env';
-import { ensureRootPinboard, initPage } from '@affine/env/blocksuite';
+import { initPage } from '@affine/env/blocksuite';
 import { setUpLanguage, useTranslation } from '@affine/i18n';
 import { createAffineGlobalChannel } from '@affine/workspace/affine/sync';
 import {
@@ -9,7 +9,7 @@ import {
   rootStore,
   rootWorkspacesMetadataAtom,
 } from '@affine/workspace/atom';
-import type { LocalIndexedDBProvider } from '@affine/workspace/type';
+import type { BackgroundProvider } from '@affine/workspace/type';
 import { WorkspaceFlavour } from '@affine/workspace/type';
 import { assertEquals, assertExists, nanoid } from '@blocksuite/store';
 import { useBlockSuiteWorkspaceHelper } from '@toeverything/hooks/use-block-suite-workspace-helper';
@@ -17,14 +17,7 @@ import { useAtom, useAtomValue, useSetAtom } from 'jotai';
 import Head from 'next/head';
 import { useRouter } from 'next/router';
 import type { FC, PropsWithChildren, ReactElement } from 'react';
-import {
-  lazy,
-  Suspense,
-  useCallback,
-  useEffect,
-  useMemo,
-  useState,
-} from 'react';
+import { lazy, Suspense, useCallback, useEffect, useMemo } from 'react';
 
 import { openQuickSearchModalAtom, openWorkspacesModalAtom } from '../atoms';
 import {
@@ -127,7 +120,10 @@ export const AllWorkspaceContext = ({
       // ignore current workspace
       .filter(workspace => workspace.id !== currentWorkspaceId)
       .flatMap(workspace =>
-        workspace.providers.filter(provider => provider.background)
+        workspace.providers.filter(
+          (provider): provider is BackgroundProvider =>
+            'background' in provider && provider.background
+        )
       );
     providers.forEach(provider => {
       provider.connect();
@@ -260,69 +256,48 @@ export const WorkspaceLayoutInner: FC<PropsWithChildren> = ({ children }) => {
   const currentPageId = useAtomValue(rootCurrentPageIdAtom);
   const router = useRouter();
   const { jumpToPage } = useRouterHelper(router);
-  const [isLoading, setIsLoading] = useState(true);
   const { t } = useTranslation();
 
   useEffect(() => {
     logger.info('currentWorkspace: ', currentWorkspace);
+    globalThis.currentWorkspace = currentWorkspace;
   }, [currentWorkspace]);
 
-  useEffect(() => {
-    if (currentWorkspace) {
-      globalThis.currentWorkspace = currentWorkspace;
+  //#region init workspace
+  if (currentWorkspace.blockSuiteWorkspace.isEmpty) {
+    // this is a new workspace, so we should redirect to the new page
+    const pageId = nanoid();
+    const page = currentWorkspace.blockSuiteWorkspace.createPage(pageId);
+    assertEquals(page.id, pageId);
+    currentWorkspace.blockSuiteWorkspace.setPageMeta(page.id, {
+      init: true,
+    });
+    initPage(page);
+    if (!router.query.pageId) {
+      setCurrentPageId(pageId);
+      void jumpToPage(currentWorkspace.id, pageId);
     }
-  }, [currentWorkspace]);
+  }
+
+  // fixme: pinboard has been removed,
+  //  the related code should be removed in the future.
+  // no matter the workspace is empty, ensure the root pinboard exists
+  // ensureRootPinboard(currentWorkspace.blockSuiteWorkspace);
+  //#endregion
 
   useEffect(() => {
-    if (currentWorkspace) {
-      currentWorkspace.providers.forEach(provider => {
-        provider.connect();
-      });
-      return () => {
-        currentWorkspace.providers.forEach(provider => {
-          provider.disconnect();
-        });
-      };
-    }
-  }, [currentWorkspace]);
-
-  useEffect(() => {
-    if (!router.isReady) {
-      return;
-    }
-    if (!currentWorkspace) {
-      return;
-    }
-    const localProvider = currentWorkspace.providers.find(
-      provider => provider.flavour === 'local-indexeddb'
+    const backgroundProviders = currentWorkspace.providers.filter(
+      (provider): provider is BackgroundProvider => 'background' in provider
     );
-    if (localProvider && localProvider.flavour === 'local-indexeddb') {
-      const provider = localProvider as LocalIndexedDBProvider;
-      const callback = () => {
-        setIsLoading(false);
-        if (currentWorkspace.blockSuiteWorkspace.isEmpty) {
-          // this is a new workspace, so we should redirect to the new page
-          const pageId = nanoid();
-          const page = currentWorkspace.blockSuiteWorkspace.createPage(pageId);
-          assertEquals(page.id, pageId);
-          currentWorkspace.blockSuiteWorkspace.setPageMeta(page.id, {
-            init: true,
-          });
-          initPage(page);
-          if (!router.query.pageId) {
-            setCurrentPageId(pageId);
-            void jumpToPage(currentWorkspace.id, pageId);
-          }
-        }
-        // no matter the workspace is empty, ensure the root pinboard exists
-        ensureRootPinboard(currentWorkspace.blockSuiteWorkspace);
-      };
-      provider.callbacks.add(callback);
-      return () => {
-        provider.callbacks.delete(callback);
-      };
-    }
-  }, [currentWorkspace, jumpToPage, router, setCurrentPageId]);
+    backgroundProviders.forEach(provider => {
+      provider.connect();
+    });
+    return () => {
+      backgroundProviders.forEach(provider => {
+        provider.disconnect();
+      });
+    };
+  }, [currentWorkspace]);
 
   useEffect(() => {
     if (!currentWorkspace) {
@@ -395,11 +370,7 @@ export const WorkspaceLayoutInner: FC<PropsWithChildren> = ({ children }) => {
         <MainContainerWrapper>
           <MainContainer className="main-container">
             <Suspense fallback={<PageLoading text={t('Page is Loading')} />}>
-              {isLoading ? (
-                <PageLoading text={t('Page is Loading')} />
-              ) : (
-                children
-              )}
+              {children}
             </Suspense>
             <StyledToolWrapper>
               {/* fixme(himself65): remove this */}

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

@@ -3,12 +3,12 @@ import {
   DEFAULT_HELLO_WORLD_PAGE_ID,
   DEFAULT_WORKSPACE_NAME,
 } from '@affine/env';
-import { ensureRootPinboard, initPage } from '@affine/env/blocksuite';
+import { initPage } from '@affine/env/blocksuite';
 import {
   CRUD,
   saveWorkspaceToLocalStorage,
 } from '@affine/workspace/local/crud';
-import { createIndexedDBProvider } from '@affine/workspace/providers';
+import { createIndexedDBBackgroundProvider } from '@affine/workspace/providers';
 import { LoadPriority, WorkspaceFlavour } from '@affine/workspace/type';
 import { createEmptyBlockSuiteWorkspace } from '@affine/workspace/utils';
 import { nanoid } from '@blocksuite/store';
@@ -40,12 +40,11 @@ export const LocalPlugin: WorkspacePlugin<WorkspaceFlavour.LOCAL> = {
       blockSuiteWorkspace.setPageMeta(page.id, {
         jumpOnce: true,
       });
-      const provider = createIndexedDBProvider(blockSuiteWorkspace);
+      const provider = createIndexedDBBackgroundProvider(blockSuiteWorkspace);
       provider.connect();
       provider.callbacks.add(() => {
         provider.disconnect();
       });
-      ensureRootPinboard(blockSuiteWorkspace);
       saveWorkspaceToLocalStorage(blockSuiteWorkspace.id);
       logger.debug('create first workspace');
       return [blockSuiteWorkspace.id];