refactor: rootWorkspacesMetadataAtom loading logic (#2882)

packages/env/src/workspace.ts

@@ -206,3 +206,14 @@ export interface AppEvents {
   // request to revoke access to workspace plugin
   'workspace:revoke': () => Promise<void>;
 }
+
+export interface WorkspaceAdapter<Flavour extends WorkspaceFlavour> {
+  releaseType: ReleaseType;
+  flavour: Flavour;
+  // The Adapter will be loaded according to the priority
+  loadPriority: LoadPriority;
+  Events: Partial<AppEvents>;
+  // Fetch necessary data for the first render
+  CRUD: WorkspaceCRUD<Flavour>;
+  UI: WorkspaceUISchema<Flavour>;
+}

packages/workspace/src/affine/sync.ts

@@ -52,7 +52,7 @@ export function createAffineGlobalChannel(
 
       // If the workspace is not in the current workspace list, remove it
       if (workspaceIndex === -1) {
-        rootStore.set(rootWorkspacesMetadataAtom, workspaces => {
+        await rootStore.set(rootWorkspacesMetadataAtom, workspaces => {
           const idx = workspaces.findIndex(workspace => workspace.id === id);
           workspaces.splice(idx, 1);
           return [...workspaces];

packages/workspace/src/atom.ts

@@ -1,28 +1,54 @@
-import { isBrowser } from '@affine/env/constant';
-import type { WorkspaceFlavour } from '@affine/env/workspace';
-import type { WorkspaceVersion } from '@affine/env/workspace';
+import type { WorkspaceAdapter } from '@affine/env/workspace';
+import { WorkspaceFlavour, WorkspaceVersion } from '@affine/env/workspace';
 import type { EditorContainer } from '@blocksuite/editor';
+import { assertExists } from '@blocksuite/global/utils';
 import { atom } from 'jotai';
-import { atomWithStorage } from 'jotai/utils';
 import Router from 'next/router';
+import { z } from 'zod';
 
-export type RootWorkspaceMetadataV2 = {
-  id: string;
-  flavour: WorkspaceFlavour;
-  version: WorkspaceVersion;
-};
+const rootWorkspaceMetadataV1Schema = z.object({
+  id: z.string(),
+  flavour: z.nativeEnum(WorkspaceFlavour),
+});
 
-export type RootWorkspaceMetadataV1 = {
-  id: string;
-  flavour: WorkspaceFlavour;
-  // force type check
-  version: undefined;
-};
+const rootWorkspaceMetadataV2Schema = rootWorkspaceMetadataV1Schema.extend({
+  version: z.nativeEnum(WorkspaceVersion),
+});
+
+const rootWorkspaceMetadataArraySchema = z.array(
+  z.union([rootWorkspaceMetadataV1Schema, rootWorkspaceMetadataV2Schema])
+);
+
+export type RootWorkspaceMetadataV2 = z.infer<
+  typeof rootWorkspaceMetadataV2Schema
+>;
+
+export type RootWorkspaceMetadataV1 = z.infer<
+  typeof rootWorkspaceMetadataV1Schema
+>;
 
 export type RootWorkspaceMetadata =
   | RootWorkspaceMetadataV1
   | RootWorkspaceMetadataV2;
 
+export const workspaceAdaptersAtom = atom<
+  Record<
+    WorkspaceFlavour,
+    Pick<
+      WorkspaceAdapter<WorkspaceFlavour>,
+      'CRUD' | 'Events' | 'flavour' | 'loadPriority'
+    >
+  >
+>(
+  null as unknown as Record<
+    WorkspaceFlavour,
+    Pick<
+      WorkspaceAdapter<WorkspaceFlavour>,
+      'CRUD' | 'Events' | 'flavour' | 'loadPriority'
+    >
+  >
+);
+
 // #region root atoms
 // root primitive atom that stores the necessary data for the whole app
 // be careful when you use this atom,
@@ -32,20 +58,170 @@ export type RootWorkspaceMetadata =
  * this atom stores the metadata of all workspaces,
  * which is `id` and `flavor`, that is enough to load the real workspace data
  */
-export const rootWorkspacesMetadataAtom = atomWithStorage<
-  RootWorkspaceMetadata[]
+const METADATA_STORAGE_KEY = 'jotai-workspaces';
+const rootWorkspacesMetadataPrimitiveAtom = atom<
+  RootWorkspaceMetadata[] | null
+>(null);
+const rootWorkspacesMetadataPromiseAtom = atom<
+  Promise<RootWorkspaceMetadata[]>
+>(async (get, { signal }) => {
+  const WorkspaceAdapters = get(workspaceAdaptersAtom);
+  assertExists(WorkspaceAdapters, 'workspace adapter should be defined');
+  const maybeMetadata = get(rootWorkspacesMetadataPrimitiveAtom);
+  if (maybeMetadata !== null) {
+    return maybeMetadata;
+  }
+  const createFirst = (): RootWorkspaceMetadataV2[] => {
+    if (signal.aborted) {
+      return [];
+    }
+
+    const Plugins = Object.values(WorkspaceAdapters).sort(
+      (a, b) => a.loadPriority - b.loadPriority
+    );
+
+    return Plugins.flatMap(Plugin => {
+      return Plugin.Events['app:init']?.().map(
+        id =>
+          ({
+            id,
+            flavour: Plugin.flavour,
+            // new workspace should all support sub-doc feature
+            version: WorkspaceVersion.SubDoc,
+          } satisfies RootWorkspaceMetadataV2)
+      );
+    }).filter((ids): ids is RootWorkspaceMetadataV2 => !!ids);
+  };
+
+  if (environment.isServer) {
+    // return a promise in SSR to avoid the hydration mismatch
+    return Promise.resolve([]);
+  } else {
+    const metadata: RootWorkspaceMetadata[] = [];
+
+    // fixme(himself65): we might not need step 1
+    // step 1: try load metadata from localStorage
+    {
+      // don't change this key,
+      // otherwise it will cause the data loss in the production
+      const primitiveMetadata = localStorage.getItem(METADATA_STORAGE_KEY);
+      if (primitiveMetadata) {
+        try {
+          const items = JSON.parse(primitiveMetadata) as z.infer<
+            typeof rootWorkspaceMetadataArraySchema
+          >;
+          rootWorkspaceMetadataArraySchema.parse(items);
+          metadata.push(...items);
+        } catch (e) {
+          console.error('cannot parse worksapce', e);
+        }
+      }
+
+      // migration step, only data in `METADATA_STORAGE_KEY` will be migrated
+      if (metadata.some(meta => !('version' in meta))) {
+        await new Promise<void>((resolve, reject) => {
+          signal.addEventListener('abort', () => reject(), { once: true });
+          window.addEventListener('migration-done', () => resolve(), {
+            once: true,
+          });
+        });
+      }
+    }
+    // step 2: fetch from adapters
+    {
+      const lists = Object.values(WorkspaceAdapters)
+        .sort((a, b) => a.loadPriority - b.loadPriority)
+        .map(({ CRUD }) => CRUD.list);
+
+      for (const list of lists) {
+        try {
+          const item = await list();
+          if (metadata.length) {
+            item.sort((a, b) => {
+              return (
+                metadata.findIndex(x => x.id === a.id) -
+                metadata.findIndex(x => x.id === b.id)
+              );
+            });
+          }
+          metadata.push(
+            ...item.map(x => ({
+              id: x.id,
+              flavour: x.flavour,
+              version: WorkspaceVersion.SubDoc,
+            }))
+          );
+        } catch (e) {
+          console.error('list data error:', e);
+        }
+      }
+    }
+    // step 3: create initial workspaces
+    {
+      if (
+        metadata.length === 0 &&
+        localStorage.getItem('is-first-open') === null
+      ) {
+        metadata.push(...createFirst());
+        console.info('create first workspace', metadata);
+        localStorage.setItem('is-first-open', 'false');
+      }
+    }
+    const metadataMap = new Map(metadata.map(x => [x.id, x]));
+    return Array.from(metadataMap.values());
+  }
+});
+
+type SetStateAction<Value> = Value | ((prev: Value) => Value);
+
+export const rootWorkspacesMetadataAtom = atom<
+  Promise<RootWorkspaceMetadata[]>,
+  [SetStateAction<RootWorkspaceMetadata[]>],
+  Promise<RootWorkspaceMetadata[]>
 >(
-  // don't change this key,
-  // otherwise it will cause the data loss in the production
-  'jotai-workspaces',
-  []
+  async get => {
+    if (environment.isServer) {
+      return Promise.resolve([]);
+    }
+    const maybeMetadata = get(rootWorkspacesMetadataPrimitiveAtom);
+    if (maybeMetadata !== null) {
+      return maybeMetadata;
+    }
+    return get(rootWorkspacesMetadataPromiseAtom);
+  },
+  async (get, set, action) => {
+    // get metadata
+    let metadata: RootWorkspaceMetadata[];
+    const maybeMetadata = get(rootWorkspacesMetadataPrimitiveAtom);
+    if (maybeMetadata !== null) {
+      metadata = maybeMetadata;
+    } else {
+      metadata = await get(rootWorkspacesMetadataPromiseAtom);
+    }
+
+    // update metadata
+    if (typeof action === 'function') {
+      metadata = action(metadata);
+    } else {
+      metadata = action;
+    }
+
+    const metadataMap = new Map(metadata.map(x => [x.id, x]));
+    metadata = Array.from(metadataMap.values());
+
+    // write back to localStorage
+    rootWorkspaceMetadataArraySchema.parse(metadata);
+    localStorage.setItem(METADATA_STORAGE_KEY, JSON.stringify(metadata));
+    set(rootWorkspacesMetadataPrimitiveAtom, metadata);
+    return metadata;
+  }
 );
 
 // two more atoms to store the current workspace and page
 export const rootCurrentWorkspaceIdAtom = atom<string | null>(null);
 
 rootCurrentWorkspaceIdAtom.onMount = set => {
-  if (isBrowser) {
+  if (environment.isBrowser) {
     const callback = (url: string) => {
       const value = url.split('/')[2];
       if (value) {
@@ -67,7 +243,7 @@ rootCurrentWorkspaceIdAtom.onMount = set => {
 export const rootCurrentPageIdAtom = atom<string | null>(null);
 
 rootCurrentPageIdAtom.onMount = set => {
-  if (isBrowser) {
+  if (environment.isBrowser) {
     const callback = (url: string) => {
       const value = url.split('/')[3];
       if (value) {

packages/workspace/src/migration/index.ts

@@ -26,7 +26,7 @@ export function upgradeV1ToV2(oldWorkspace: LocalWorkspace): LocalWorkspace {
         }
       });
     });
-    console.log(newBlockSuiteWorkspace.doc.toJSON());
+    console.log('migration result', newBlockSuiteWorkspace.doc.toJSON());
 
     return {
       blockSuiteWorkspace: newBlockSuiteWorkspace,

packages/workspace/src/utils.ts

@@ -20,9 +20,11 @@ import { createAffineBlobStorage } from './blob';
 import { createSQLiteStorage } from './blob/sqlite-blob-storage';
 
 export function cleanupWorkspace(flavour: WorkspaceFlavour) {
-  rootStore.set(rootWorkspacesMetadataAtom, metas =>
-    metas.filter(meta => meta.flavour !== flavour)
-  );
+  rootStore
+    .set(rootWorkspacesMetadataAtom, metas =>
+      metas.filter(meta => meta.flavour !== flavour)
+    )
+    .catch(console.error);
 }
 
 function setEditorFlags(workspace: Workspace) {