refactor: workspace loading logic (#1966)

packages/workspace/src/affine/sync.ts

@@ -6,7 +6,7 @@ import {
 } from '@affine/workspace/affine/api';
 import { WebsocketClient } from '@affine/workspace/affine/channel';
 import { storageChangeSlot } from '@affine/workspace/affine/login';
-import { jotaiStore, jotaiWorkspacesAtom } from '@affine/workspace/atom';
+import { rootStore, rootWorkspacesMetadataAtom } from '@affine/workspace/atom';
 import type { WorkspaceCRUD } from '@affine/workspace/type';
 import type { WorkspaceFlavour } from '@affine/workspace/type';
 import { assertExists } from '@blocksuite/global/utils';
@@ -51,7 +51,7 @@ export function createAffineGlobalChannel(
 
       // If the workspace is not in the current workspace list, remove it
       if (workspaceIndex === -1) {
-        jotaiStore.set(jotaiWorkspacesAtom, workspaces => {
+        rootStore.set(rootWorkspacesMetadataAtom, workspaces => {
           const idx = workspaces.findIndex(workspace => workspace.id === id);
           workspaces.splice(idx, 1);
           return [...workspaces];

packages/workspace/src/atom.ts

@@ -1,18 +1,49 @@
 import { atomWithSyncStorage } from '@affine/jotai';
 import type { WorkspaceFlavour } from '@affine/workspace/type';
-import { createStore } from 'jotai/index';
+import type { EditorContainer } from '@blocksuite/editor';
+import { atom, createStore } from 'jotai';
+import { atomWithStorage, createJSONStorage } from 'jotai/utils';
 
-export type JotaiWorkspace = {
+export type RootWorkspaceMetadata = {
   id: string;
   flavour: WorkspaceFlavour;
 };
+// #region root atoms
+// root primitive atom that stores the necessary data for the whole app
+// be careful when you use this atom,
+// it should be used only in the root component
 
-// root primitive atom that stores the list of workspaces which could be used in the app
-// if a workspace is not in this list, it should not be used in the app
-export const jotaiWorkspacesAtom = atomWithSyncStorage<JotaiWorkspace[]>(
+/**
+ * root workspaces atom
+ * this atom stores the metadata of all workspaces,
+ * which is `id` and `flavour`, that is enough to load the real workspace data
+ */
+export const rootWorkspacesMetadataAtom = atomWithSyncStorage<
+  RootWorkspaceMetadata[]
+>(
+  // don't change this key,
+  // otherwise it will cause the data loss in the production
   'jotai-workspaces',
   []
 );
 
-// global jotai store, which is used to store all the atoms
-export const jotaiStore = createStore();
+// two more atoms to store the current workspace and page
+export const rootCurrentWorkspaceIdAtom = atomWithStorage<string | null>(
+  'root-current-workspace-id',
+  null,
+  createJSONStorage(() => sessionStorage)
+);
+export const rootCurrentPageIdAtom = atomWithStorage<string | null>(
+  'root-current-page-id',
+  null,
+  createJSONStorage(() => sessionStorage)
+);
+
+// current editor atom, each app should have only one editor in the same time
+export const rootCurrentEditorAtom = atom<Readonly<EditorContainer> | null>(
+  null
+);
+//#endregion
+
+// global store
+export const rootStore = createStore();

packages/workspace/src/local/crud.ts

@@ -1,3 +1,4 @@
+import { DebugLogger } from '@affine/debug';
 import { nanoid, Workspace as BlockSuiteWorkspace } from '@blocksuite/store';
 import { createIndexedDBProvider } from '@toeverything/y-indexeddb';
 import { createJSONStorage } from 'jotai/utils';
@@ -13,8 +14,25 @@ const getStorage = () => createJSONStorage(() => localStorage);
 const kStoreKey = 'affine-local-workspace';
 const schema = z.array(z.string());
 
+const logger = new DebugLogger('affine:workspace:local:crud');
+
+/**
+ * @internal
+ */
+export function saveWorkspaceToLocalStorage(workspaceId: string) {
+  const storage = getStorage();
+  !Array.isArray(storage.getItem(kStoreKey)) && storage.setItem(kStoreKey, []);
+  const data = storage.getItem(kStoreKey) as z.infer<typeof schema>;
+  const id = data.find(id => id === workspaceId);
+  if (!id) {
+    logger.debug('saveWorkspaceToLocalStorage', workspaceId);
+    storage.setItem(kStoreKey, [...data, workspaceId]);
+  }
+}
+
 export const CRUD: WorkspaceCRUD<WorkspaceFlavour.LOCAL> = {
   get: async workspaceId => {
+    logger.debug('get', workspaceId);
     const storage = getStorage();
     !Array.isArray(storage.getItem(kStoreKey)) &&
       storage.setItem(kStoreKey, []);
@@ -36,10 +54,10 @@ export const CRUD: WorkspaceCRUD<WorkspaceFlavour.LOCAL> = {
     return workspace;
   },
   create: async ({ doc }) => {
+    logger.debug('create', doc);
     const storage = getStorage();
     !Array.isArray(storage.getItem(kStoreKey)) &&
       storage.setItem(kStoreKey, []);
-    const data = storage.getItem(kStoreKey) as z.infer<typeof schema>;
     const binary = BlockSuiteWorkspace.Y.encodeStateAsUpdateV2(doc);
     const id = nanoid();
     const blockSuiteWorkspace = createEmptyBlockSuiteWorkspace(
@@ -52,11 +70,11 @@ export const CRUD: WorkspaceCRUD<WorkspaceFlavour.LOCAL> = {
     await persistence.whenSynced.then(() => {
       persistence.disconnect();
     });
-    storage.setItem(kStoreKey, [...data, id]);
-    console.log('create', id, storage.getItem(kStoreKey));
+    saveWorkspaceToLocalStorage(id);
     return id;
   },
   delete: async workspace => {
+    logger.debug('delete', workspace);
     const storage = getStorage();
     !Array.isArray(storage.getItem(kStoreKey)) &&
       storage.setItem(kStoreKey, []);
@@ -69,6 +87,7 @@ export const CRUD: WorkspaceCRUD<WorkspaceFlavour.LOCAL> = {
     storage.setItem(kStoreKey, [...data]);
   },
   list: async () => {
+    logger.debug('list');
     const storage = getStorage();
     !Array.isArray(storage.getItem(kStoreKey)) &&
       storage.setItem(kStoreKey, []);

packages/workspace/src/providers/index.ts

@@ -69,6 +69,36 @@ const createAffineWebSocketProvider = (
   return apis;
 };
 
+class CallbackSet extends Set<() => void> {
+  #ready = false;
+
+  get ready(): boolean {
+    return this.#ready;
+  }
+
+  set ready(v: boolean) {
+    this.#ready = v;
+  }
+
+  add(cb: () => void) {
+    if (this.ready) {
+      cb();
+      return this;
+    }
+    if (this.has(cb)) {
+      return this;
+    }
+    return super.add(cb);
+  }
+
+  delete(cb: () => void) {
+    if (this.has(cb)) {
+      return super.delete(cb);
+    }
+    return false;
+  }
+}
+
 const createIndexedDBProvider = (
   blockSuiteWorkspace: BlockSuiteWorkspace
 ): LocalIndexedDBProvider => {
@@ -76,7 +106,7 @@ const createIndexedDBProvider = (
     blockSuiteWorkspace.id,
     blockSuiteWorkspace.doc
   );
-  const callbacks = new Set<() => void>();
+  const callbacks = new CallbackSet();
   return {
     flavour: 'local-indexeddb',
     callbacks,
@@ -93,6 +123,7 @@ const createIndexedDBProvider = (
       indexeddbProvider.connect();
       indexeddbProvider.whenSynced
         .then(() => {
+          callbacks.ready = true;
           callbacks.forEach(cb => cb());
         })
         .catch(error => {
@@ -110,6 +141,7 @@ const createIndexedDBProvider = (
         blockSuiteWorkspace.id
       );
       indexeddbProvider.disconnect();
+      callbacks.ready = false;
     },
   };
 };

packages/workspace/src/type.ts

@@ -2,8 +2,11 @@
 /// <reference path='../../../apps/electron/layers/preload/preload.d.ts' />
 import type { Workspace as RemoteWorkspace } from '@affine/workspace/affine/api';
 import type { Workspace as BlockSuiteWorkspace } from '@blocksuite/store';
+import type { createStore } from 'jotai';
 import type { FC, PropsWithChildren } from 'react';
 
+export type JotaiStore = ReturnType<typeof createStore>;
+
 export type BaseProvider = {
   flavour: string;
   // if this is true, we will connect the provider on the background
@@ -141,9 +144,9 @@ export interface WorkspaceUISchema<Flavour extends keyof WorkspaceRegistry> {
 }
 
 export interface AppEvents {
-  // event when app is first initialized
+  // event there is no workspace
   // usually used to initialize workspace plugin
-  'app:first-init': () => Promise<void>;
+  'app:init': () => string[];
   // request to gain access to workspace plugin
   'workspace:access': () => Promise<void>;
   // request to revoke access to workspace plugin