refactor(workspace): split workspace interface and implementation (#5463)

@affine/workspace -> (@affine/workspace, @affine/workspace-impl)

packages/frontend/workspace-impl/src/cloud/awareness.ts

@@ -0,0 +1,108 @@
+import { DebugLogger } from '@affine/debug';
+import type { AwarenessProvider } from '@affine/workspace';
+import {
+  applyAwarenessUpdate,
+  type Awareness,
+  encodeAwarenessUpdate,
+  removeAwarenessStates,
+} from 'y-protocols/awareness';
+
+import { getIoManager } from '../utils/affine-io';
+import { base64ToUint8Array, uint8ArrayToBase64 } from '../utils/base64';
+
+const logger = new DebugLogger('affine:awareness:socketio');
+
+type AwarenessChanges = Record<'added' | 'updated' | 'removed', number[]>;
+
+export function createCloudAwarenessProvider(
+  workspaceId: string,
+  awareness: Awareness
+): AwarenessProvider {
+  const socket = getIoManager().socket('/');
+
+  const awarenessBroadcast = ({
+    workspaceId: wsId,
+    awarenessUpdate,
+  }: {
+    workspaceId: string;
+    awarenessUpdate: string;
+  }) => {
+    if (wsId !== workspaceId) {
+      return;
+    }
+    applyAwarenessUpdate(
+      awareness,
+      base64ToUint8Array(awarenessUpdate),
+      'remote'
+    );
+  };
+
+  const awarenessUpdate = (changes: AwarenessChanges, origin: unknown) => {
+    if (origin === 'remote') {
+      return;
+    }
+
+    const changedClients = Object.values(changes).reduce((res, cur) =>
+      res.concat(cur)
+    );
+
+    const update = encodeAwarenessUpdate(awareness, changedClients);
+    uint8ArrayToBase64(update)
+      .then(encodedUpdate => {
+        socket.emit('awareness-update', {
+          workspaceId: workspaceId,
+          awarenessUpdate: encodedUpdate,
+        });
+      })
+      .catch(err => logger.error(err));
+  };
+
+  const newClientAwarenessInitHandler = () => {
+    const awarenessUpdate = encodeAwarenessUpdate(awareness, [
+      awareness.clientID,
+    ]);
+    uint8ArrayToBase64(awarenessUpdate)
+      .then(encodedAwarenessUpdate => {
+        socket.emit('awareness-update', {
+          guid: workspaceId,
+          awarenessUpdate: encodedAwarenessUpdate,
+        });
+      })
+      .catch(err => logger.error(err));
+  };
+
+  const windowBeforeUnloadHandler = () => {
+    removeAwarenessStates(awareness, [awareness.clientID], 'window unload');
+  };
+
+  function handleConnect() {
+    socket.emit('client-handshake-awareness', workspaceId);
+    socket.emit('awareness-init', workspaceId);
+  }
+
+  return {
+    connect: () => {
+      socket.on('server-awareness-broadcast', awarenessBroadcast);
+      socket.on('new-client-awareness-init', newClientAwarenessInitHandler);
+      awareness.on('update', awarenessUpdate);
+
+      window.addEventListener('beforeunload', windowBeforeUnloadHandler);
+
+      socket.connect();
+
+      socket.on('connect', handleConnect);
+
+      socket.emit('client-handshake-awareness', workspaceId);
+      socket.emit('awareness-init', workspaceId);
+    },
+    disconnect: () => {
+      removeAwarenessStates(awareness, [awareness.clientID], 'disconnect');
+      awareness.off('update', awarenessUpdate);
+      socket.emit('client-leave-awareness', workspaceId);
+      socket.off('server-awareness-broadcast', awarenessBroadcast);
+      socket.off('new-client-awareness-init', newClientAwarenessInitHandler);
+      socket.off('connect', handleConnect);
+      window.removeEventListener('unload', windowBeforeUnloadHandler);
+    },
+  };
+}

packages/frontend/workspace-impl/src/cloud/blob.ts

@@ -0,0 +1,77 @@
+import {
+  checkBlobSizesQuery,
+  deleteBlobMutation,
+  fetchWithTraceReport,
+  getBaseUrl,
+  listBlobsQuery,
+  setBlobMutation,
+} from '@affine/graphql';
+import { fetcher } from '@affine/graphql';
+import type { BlobStorage } from '@affine/workspace';
+
+import { bufferToBlob } from '../utils/buffer-to-blob';
+
+export const createAffineCloudBlobStorage = (
+  workspaceId: string
+): BlobStorage => {
+  return {
+    name: 'affine-cloud',
+    readonly: false,
+    get: async key => {
+      const suffix = key.startsWith('/')
+        ? key
+        : `/api/workspaces/${workspaceId}/blobs/${key}`;
+
+      return fetchWithTraceReport(getBaseUrl() + suffix).then(async res => {
+        if (!res.ok) {
+          // status not in the range 200-299
+          return null;
+        }
+        return bufferToBlob(await res.arrayBuffer());
+      });
+    },
+    set: async (key, value) => {
+      const {
+        checkBlobSize: { size },
+      } = await fetcher({
+        query: checkBlobSizesQuery,
+        variables: {
+          workspaceId,
+          size: value.size,
+        },
+      });
+
+      if (size <= 0) {
+        throw new Error('Blob size limit exceeded');
+      }
+
+      const result = await fetcher({
+        query: setBlobMutation,
+        variables: {
+          workspaceId,
+          blob: new File([value], key),
+        },
+      });
+      console.assert(result.setBlob === key, 'Blob hash mismatch');
+      return result.setBlob;
+    },
+    list: async () => {
+      const result = await fetcher({
+        query: listBlobsQuery,
+        variables: {
+          workspaceId,
+        },
+      });
+      return result.listBlobs;
+    },
+    delete: async (key: string) => {
+      await fetcher({
+        query: deleteBlobMutation,
+        variables: {
+          workspaceId,
+          hash: key,
+        },
+      });
+    },
+  };
+};

packages/frontend/workspace-impl/src/cloud/consts.ts

@@ -0,0 +1,2 @@
+export const CLOUD_WORKSPACE_CHANGED_BROADCAST_CHANNEL_KEY =
+  'affine-cloud-workspace-changed';

packages/frontend/workspace-impl/src/cloud/index.ts

@@ -0,0 +1,6 @@
+export * from './awareness';
+export * from './blob';
+export * from './consts';
+export * from './list';
+export * from './sync';
+export * from './workspace-factory';

packages/frontend/workspace-impl/src/cloud/list.ts

@@ -0,0 +1,155 @@
+import { WorkspaceFlavour } from '@affine/env/workspace';
+import {
+  createWorkspaceMutation,
+  deleteWorkspaceMutation,
+  getWorkspacesQuery,
+} from '@affine/graphql';
+import { fetcher } from '@affine/graphql';
+import type { WorkspaceListProvider } from '@affine/workspace';
+import { globalBlockSuiteSchema } from '@affine/workspace';
+import { Workspace as BlockSuiteWorkspace } from '@blocksuite/store';
+import { difference } from 'lodash-es';
+import { nanoid } from 'nanoid';
+import { applyUpdate, encodeStateAsUpdate } from 'yjs';
+
+import { createLocalBlobStorage } from '../local/blob';
+import { createLocalStorage } from '../local/sync';
+import { CLOUD_WORKSPACE_CHANGED_BROADCAST_CHANNEL_KEY } from './consts';
+import { createAffineStaticStorage } from './sync';
+
+async function getCloudWorkspaceList() {
+  try {
+    const { workspaces } = await fetcher({
+      query: getWorkspacesQuery,
+    });
+    const ids = workspaces.map(({ id }) => id);
+    return ids.map(id => ({
+      id,
+      flavour: WorkspaceFlavour.AFFINE_CLOUD,
+    }));
+  } catch (err) {
+    if (err instanceof Array && err[0]?.message === 'Forbidden resource') {
+      // user not logged in
+      return [];
+    }
+    throw err;
+  }
+}
+
+export function createCloudWorkspaceListProvider(): WorkspaceListProvider {
+  const notifyChannel = new BroadcastChannel(
+    CLOUD_WORKSPACE_CHANGED_BROADCAST_CHANNEL_KEY
+  );
+
+  return {
+    name: WorkspaceFlavour.AFFINE_CLOUD,
+    async getList() {
+      return getCloudWorkspaceList();
+    },
+    async create(initial) {
+      const tempId = nanoid();
+
+      const workspace = new BlockSuiteWorkspace({
+        id: tempId,
+        idGenerator: () => nanoid(),
+        schema: globalBlockSuiteSchema,
+      });
+
+      // create workspace on cloud, get workspace id
+      const {
+        createWorkspace: { id: workspaceId },
+      } = await fetcher({
+        query: createWorkspaceMutation,
+      });
+
+      // save the initial state to local storage, then sync to cloud
+      const blobStorage = createLocalBlobStorage(workspaceId);
+      const syncStorage = createLocalStorage(workspaceId);
+
+      // apply initial state
+      await initial(workspace, blobStorage);
+
+      // save workspace to local storage, should be vary fast
+      await syncStorage.push(workspaceId, encodeStateAsUpdate(workspace.doc));
+      for (const subdocs of workspace.doc.getSubdocs()) {
+        await syncStorage.push(subdocs.guid, encodeStateAsUpdate(subdocs));
+      }
+
+      // notify all browser tabs, so they can update their workspace list
+      notifyChannel.postMessage(null);
+
+      return workspaceId;
+    },
+    async delete(id) {
+      await fetcher({
+        query: deleteWorkspaceMutation,
+        variables: {
+          id,
+        },
+      });
+      // notify all browser tabs, so they can update their workspace list
+      notifyChannel.postMessage(null);
+    },
+    subscribe(callback) {
+      let lastWorkspaceIDs: string[] = [];
+
+      function scan() {
+        (async () => {
+          const allWorkspaceIDs = (await getCloudWorkspaceList()).map(
+            workspace => workspace.id
+          );
+          const added = difference(allWorkspaceIDs, lastWorkspaceIDs);
+          const deleted = difference(lastWorkspaceIDs, allWorkspaceIDs);
+          lastWorkspaceIDs = allWorkspaceIDs;
+          callback({
+            added: added.map(id => ({
+              id,
+              flavour: WorkspaceFlavour.AFFINE_CLOUD,
+            })),
+            deleted: deleted.map(id => ({
+              id,
+              flavour: WorkspaceFlavour.AFFINE_CLOUD,
+            })),
+          });
+        })().catch(err => {
+          console.error(err);
+        });
+      }
+
+      scan();
+
+      // rescan if other tabs notify us
+      notifyChannel.addEventListener('message', scan);
+      return () => {
+        notifyChannel.removeEventListener('message', scan);
+      };
+    },
+    async getInformation(id) {
+      // get information from both cloud and local storage
+
+      // we use affine 'static' storage here, which use http protocol, no need to websocket.
+      const cloudStorage = createAffineStaticStorage(id);
+      const localStorage = createLocalStorage(id);
+      // download root doc
+      const localData = await localStorage.pull(id, new Uint8Array([]));
+      const cloudData = await cloudStorage.pull(id, new Uint8Array([]));
+
+      if (!cloudData && !localData) {
+        return;
+      }
+
+      const bs = new BlockSuiteWorkspace({
+        id,
+        schema: globalBlockSuiteSchema,
+      });
+
+      if (localData) applyUpdate(bs.doc, localData.data);
+      if (cloudData) applyUpdate(bs.doc, cloudData.data);
+
+      return {
+        name: bs.meta.name,
+        avatar: bs.meta.avatar,
+      };
+    },
+  };
+}

packages/frontend/workspace-impl/src/cloud/sync/batch-sync-sender.ts

@@ -0,0 +1,107 @@
+interface SyncUpdateSender {
+  (
+    guid: string,
+    updates: Uint8Array[]
+  ): Promise<{
+    accepted: boolean;
+    retry: boolean;
+  }>;
+}
+
+/**
+ * BatchSyncSender is simple wrapper with vanilla update sync with several advanced features:
+ *   - ACK mechanism, send updates sequentially with previous sync request correctly responds with ACK
+ *   - batching updates, when waiting for previous ACK, new updates will be buffered and sent in single sync request
+ *   - retryable, allow retry when previous sync request failed but with retry flag been set to true
+ */
+export class BatchSyncSender {
+  private readonly buffered: Uint8Array[] = [];
+  private job: Promise<void> | null = null;
+  private started = true;
+
+  constructor(
+    private readonly guid: string,
+    private readonly rawSender: SyncUpdateSender
+  ) {}
+
+  send(update: Uint8Array) {
+    this.buffered.push(update);
+    this.next();
+    return Promise.resolve();
+  }
+
+  stop() {
+    this.started = false;
+  }
+
+  start() {
+    this.started = true;
+    this.next();
+  }
+
+  private next() {
+    if (!this.started || this.job || !this.buffered.length) {
+      return;
+    }
+
+    const lastIndex = Math.min(
+      this.buffered.length - 1,
+      99 /* max batch updates size */
+    );
+    const updates = this.buffered.slice(0, lastIndex + 1);
+
+    if (updates.length) {
+      this.job = this.rawSender(this.guid, updates)
+        .then(({ accepted, retry }) => {
+          // remove pending updates if updates are accepted
+          if (accepted) {
+            this.buffered.splice(0, lastIndex + 1);
+          }
+
+          // stop when previous sending failed and non-recoverable
+          if (accepted || retry) {
+            // avoid call stack overflow
+            setTimeout(() => {
+              this.next();
+            }, 0);
+          } else {
+            this.stop();
+          }
+        })
+        .catch(() => {
+          this.stop();
+        })
+        .finally(() => {
+          this.job = null;
+        });
+    }
+  }
+}
+
+export class MultipleBatchSyncSender {
+  private senders: Record<string, BatchSyncSender> = {};
+
+  constructor(private readonly rawSender: SyncUpdateSender) {}
+
+  async send(guid: string, update: Uint8Array) {
+    return this.getSender(guid).send(update);
+  }
+
+  private getSender(guid: string) {
+    let sender = this.senders[guid];
+    if (!sender) {
+      sender = new BatchSyncSender(guid, this.rawSender);
+      this.senders[guid] = sender;
+    }
+
+    return sender;
+  }
+
+  start() {
+    Object.values(this.senders).forEach(sender => sender.start());
+  }
+
+  stop() {
+    Object.values(this.senders).forEach(sender => sender.stop());
+  }
+}

packages/frontend/workspace-impl/src/cloud/sync/index.ts

@@ -0,0 +1,196 @@
+import { DebugLogger } from '@affine/debug';
+import { fetchWithTraceReport } from '@affine/graphql';
+import type { SyncStorage } from '@affine/workspace';
+
+import { getIoManager } from '../../utils/affine-io';
+import { base64ToUint8Array, uint8ArrayToBase64 } from '../../utils/base64';
+import { MultipleBatchSyncSender } from './batch-sync-sender';
+
+const logger = new DebugLogger('affine:storage:socketio');
+
+export function createAffineStorage(
+  workspaceId: string
+): SyncStorage & { disconnect: () => void } {
+  logger.debug('createAffineStorage', workspaceId);
+  const socket = getIoManager().socket('/');
+
+  const syncSender = new MultipleBatchSyncSender(async (guid, updates) => {
+    const payload = await Promise.all(
+      updates.map(update => uint8ArrayToBase64(update))
+    );
+
+    return new Promise(resolve => {
+      socket.emit(
+        'client-update-v2',
+        {
+          workspaceId,
+          guid,
+          updates: payload,
+        },
+        (response: {
+          // TODO: reuse `EventError` with server
+          error?: any;
+          data: any;
+        }) => {
+          // TODO: raise error with different code to users
+          if (response.error) {
+            logger.error('client-update-v2 error', {
+              workspaceId,
+              guid,
+              response,
+            });
+          }
+
+          resolve({
+            accepted: !response.error,
+            // TODO: reuse `EventError` with server
+            retry: response.error?.code === 'INTERNAL',
+          });
+        }
+      );
+    });
+  });
+
+  function handleConnect() {
+    socket.emit(
+      'client-handshake-sync',
+      workspaceId,
+      (response: { error?: any }) => {
+        if (!response.error) {
+          syncSender.start();
+        }
+      }
+    );
+  }
+
+  socket.on('connect', handleConnect);
+
+  socket.connect();
+
+  socket.emit(
+    'client-handshake-sync',
+    workspaceId,
+    (response: { error?: any }) => {
+      if (!response.error) {
+        syncSender.start();
+      }
+    }
+  );
+
+  return {
+    name: 'affine-cloud',
+    async pull(docId, state) {
+      const stateVector = state ? await uint8ArrayToBase64(state) : undefined;
+
+      return new Promise((resolve, reject) => {
+        logger.debug('doc-load-v2', {
+          workspaceId: workspaceId,
+          guid: docId,
+          stateVector,
+        });
+        socket.emit(
+          'doc-load-v2',
+          {
+            workspaceId: workspaceId,
+            guid: docId,
+            stateVector,
+          },
+          (
+            response: // TODO: reuse `EventError` with server
+            { error: any } | { data: { missing: string; state: string } }
+          ) => {
+            logger.debug('doc-load callback', {
+              workspaceId: workspaceId,
+              guid: docId,
+              stateVector,
+              response,
+            });
+
+            if ('error' in response) {
+              // TODO: result `EventError` with server
+              if (response.error.code === 'DOC_NOT_FOUND') {
+                resolve(null);
+              } else {
+                reject(new Error(response.error.message));
+              }
+            } else {
+              resolve({
+                data: base64ToUint8Array(response.data.missing),
+                state: response.data.state
+                  ? base64ToUint8Array(response.data.state)
+                  : undefined,
+              });
+            }
+          }
+        );
+      });
+    },
+    async push(docId, update) {
+      logger.debug('client-update-v2', {
+        workspaceId,
+        guid: docId,
+        update,
+      });
+
+      await syncSender.send(docId, update);
+    },
+    async subscribe(cb, disconnect) {
+      const handleUpdate = async (message: {
+        workspaceId: string;
+        guid: string;
+        updates: string[];
+      }) => {
+        if (message.workspaceId === workspaceId) {
+          message.updates.forEach(update => {
+            cb(message.guid, base64ToUint8Array(update));
+          });
+        }
+      };
+      socket.on('server-updates', handleUpdate);
+
+      socket.on('disconnect', reason => {
+        socket.off('server-updates', handleUpdate);
+        disconnect(reason);
+      });
+
+      return () => {
+        socket.off('server-updates', handleUpdate);
+      };
+    },
+    disconnect() {
+      syncSender.stop();
+      socket.emit('client-leave-sync', workspaceId);
+      socket.off('connect', handleConnect);
+    },
+  };
+}
+
+export function createAffineStaticStorage(workspaceId: string): SyncStorage {
+  logger.debug('createAffineStaticStorage', workspaceId);
+
+  return {
+    name: 'affine-cloud-static',
+    async pull(docId) {
+      const response = await fetchWithTraceReport(
+        runtimeConfig.serverUrlPrefix +
+          `/api/workspaces/${workspaceId}/docs/${docId}`,
+        {
+          priority: 'high',
+        }
+      );
+      if (response.ok) {
+        const arrayBuffer = await response.arrayBuffer();
+
+        return { data: new Uint8Array(arrayBuffer) };
+      }
+
+      return null;
+    },
+    async push() {
+      throw new Error('Not implemented');
+    },
+    async subscribe() {
+      throw new Error('Not implemented');
+    },
+  };
+}

packages/frontend/workspace-impl/src/cloud/workspace-factory.ts

@@ -0,0 +1,77 @@
+import { setupEditorFlags } from '@affine/env/global';
+import type { WorkspaceFactory } from '@affine/workspace';
+import { BlobEngine, SyncEngine, WorkspaceEngine } from '@affine/workspace';
+import { globalBlockSuiteSchema } from '@affine/workspace';
+import { Workspace } from '@affine/workspace';
+import { Workspace as BlockSuiteWorkspace } from '@blocksuite/store';
+import { nanoid } from 'nanoid';
+
+import { createBroadcastChannelAwarenessProvider } from '../local/awareness';
+import { createLocalBlobStorage } from '../local/blob';
+import { createStaticBlobStorage } from '../local/blob-static';
+import { createLocalStorage } from '../local/sync';
+import { createCloudAwarenessProvider } from './awareness';
+import { createAffineCloudBlobStorage } from './blob';
+import { createAffineStorage } from './sync';
+
+export const cloudWorkspaceFactory: WorkspaceFactory = {
+  name: 'affine-cloud',
+  openWorkspace(metadata) {
+    const blobEngine = new BlobEngine(createLocalBlobStorage(metadata.id), [
+      createAffineCloudBlobStorage(metadata.id),
+      createStaticBlobStorage(),
+    ]);
+
+    // create blocksuite workspace
+    const bs = new BlockSuiteWorkspace({
+      id: metadata.id,
+      blobStorages: [
+        () => ({
+          crud: blobEngine,
+        }),
+      ],
+      idGenerator: () => nanoid(),
+      schema: globalBlockSuiteSchema,
+    });
+
+    const affineStorage = createAffineStorage(metadata.id);
+    const syncEngine = new SyncEngine(bs.doc, createLocalStorage(metadata.id), [
+      affineStorage,
+    ]);
+
+    const awarenessProviders = [
+      createBroadcastChannelAwarenessProvider(
+        metadata.id,
+        bs.awarenessStore.awareness
+      ),
+      createCloudAwarenessProvider(metadata.id, bs.awarenessStore.awareness),
+    ];
+    const engine = new WorkspaceEngine(
+      blobEngine,
+      syncEngine,
+      awarenessProviders
+    );
+
+    setupEditorFlags(bs);
+
+    const workspace = new Workspace(metadata, engine, bs);
+
+    workspace.onStop.once(() => {
+      // affine sync storage need manually disconnect
+      affineStorage.disconnect();
+    });
+
+    return workspace;
+  },
+  async getWorkspaceBlob(id: string, blobKey: string): Promise<Blob | null> {
+    // try to get blob from local storage first
+    const localBlobStorage = createLocalBlobStorage(id);
+    const localBlob = await localBlobStorage.get(blobKey);
+    if (localBlob) {
+      return localBlob;
+    }
+
+    const blobStorage = createAffineCloudBlobStorage(id);
+    return await blobStorage.get(blobKey);
+  },
+};