feat: reconcile when saving to firebase (#4991)

* naming tweaks

* do not mark local element as duplicate when there's no remote counterpart

* merge instead of overwrite elements when saving to firebase & reconcile local state

* decouple syncing from persistence

* fix ts

* clarify doc

* fix reconciliation not removing duplicates

src/excalidraw-app/app_constants.ts

@@ -11,12 +11,12 @@ export const FILE_UPLOAD_MAX_BYTES = 3 * 1024 * 1024; // 3 MiB
 // 1 year (https://stackoverflow.com/a/25201898/927631)
 export const FILE_CACHE_MAX_AGE_SEC = 31536000;
 
-export const BROADCAST = {
+export const WS_EVENTS = {
   SERVER_VOLATILE: "server-volatile-broadcast",
   SERVER: "server-broadcast",
 };
 
-export enum SCENE {
+export enum WS_SCENE_EVENT_TYPES {
   INIT = "SCENE_INIT",
   UPDATE = "SCENE_UPDATE",
 }

src/excalidraw-app/collab/CollabWrapper.tsx

@@ -22,7 +22,7 @@ import {
   FIREBASE_STORAGE_PREFIXES,
   INITIAL_SCENE_UPDATE_TIMEOUT,
   LOAD_IMAGES_TIMEOUT,
-  SCENE,
+  WS_SCENE_EVENT_TYPES,
   STORAGE_KEYS,
   SYNC_FULL_SCENE_INTERVAL_MS,
 } from "../app_constants";
@@ -88,7 +88,7 @@ export interface CollabAPI {
   onPointerUpdate: CollabInstance["onPointerUpdate"];
   initializeSocketClient: CollabInstance["initializeSocketClient"];
   onCollabButtonClick: CollabInstance["onCollabButtonClick"];
-  broadcastElements: CollabInstance["broadcastElements"];
+  syncElements: CollabInstance["syncElements"];
   fetchImageFilesFromFirebase: CollabInstance["fetchImageFilesFromFirebase"];
   setUsername: (username: string) => void;
 }
@@ -232,12 +232,20 @@ class CollabWrapper extends PureComponent<Props, CollabState> {
   });
 
   saveCollabRoomToFirebase = async (
-    syncableElements: readonly ExcalidrawElement[] = this.getSyncableElements(
-      this.excalidrawAPI.getSceneElementsIncludingDeleted(),
-    ),
+    syncableElements: readonly ExcalidrawElement[],
   ) => {
     try {
-      await saveToFirebase(this.portal, syncableElements);
+      const savedData = await saveToFirebase(
+        this.portal,
+        syncableElements,
+        this.excalidrawAPI.getAppState(),
+      );
+
+      if (this.isCollaborating() && savedData && savedData.reconciledElements) {
+        this.handleRemoteSceneUpdate(
+          this.reconcileElements(savedData.reconciledElements),
+        );
+      }
     } catch (error: any) {
       console.error(error);
     }
@@ -250,9 +258,14 @@ class CollabWrapper extends PureComponent<Props, CollabState> {
 
   closePortal = () => {
     this.queueBroadcastAllElements.cancel();
+    this.queueSaveToFirebase.cancel();
     this.loadImageFiles.cancel();
 
-    this.saveCollabRoomToFirebase();
+    this.saveCollabRoomToFirebase(
+      this.getSyncableElements(
+        this.excalidrawAPI.getSceneElementsIncludingDeleted(),
+      ),
+    );
     if (window.confirm(t("alerts.collabStopOverridePrompt"))) {
       // hack to ensure that we prefer we disregard any new browser state
       // that could have been saved in other tabs while we were collaborating
@@ -400,10 +413,7 @@ class CollabWrapper extends PureComponent<Props, CollabState> {
         commitToHistory: true,
       });
 
-      this.broadcastElements(elements);
-
-      const syncableElements = this.getSyncableElements(elements);
-      this.saveCollabRoomToFirebase(syncableElements);
+      this.saveCollabRoomToFirebase(this.getSyncableElements(elements));
     }
 
     // fallback in case you're not alone in the room but still don't receive
@@ -433,7 +443,7 @@ class CollabWrapper extends PureComponent<Props, CollabState> {
         switch (decryptedData.type) {
           case "INVALID_RESPONSE":
             return;
-          case SCENE.INIT: {
+          case WS_SCENE_EVENT_TYPES.INIT: {
             if (!this.portal.socketInitialized) {
               this.initializeRoom({ fetchScene: false });
               const remoteElements = decryptedData.payload.elements;
@@ -449,7 +459,7 @@ class CollabWrapper extends PureComponent<Props, CollabState> {
             }
             break;
           }
-          case SCENE.UPDATE:
+          case WS_SCENE_EVENT_TYPES.UPDATE:
             this.handleRemoteSceneUpdate(
               this.reconcileElements(decryptedData.payload.elements),
             );
@@ -711,15 +721,20 @@ class CollabWrapper extends PureComponent<Props, CollabState> {
       getSceneVersion(elements) >
       this.getLastBroadcastedOrReceivedSceneVersion()
     ) {
-      this.portal.broadcastScene(SCENE.UPDATE, elements, false);
+      this.portal.broadcastScene(WS_SCENE_EVENT_TYPES.UPDATE, elements, false);
       this.lastBroadcastedOrReceivedSceneVersion = getSceneVersion(elements);
       this.queueBroadcastAllElements();
     }
   };
 
+  syncElements = (elements: readonly ExcalidrawElement[]) => {
+    this.broadcastElements(elements);
+    this.queueSaveToFirebase();
+  };
+
   queueBroadcastAllElements = throttle(() => {
     this.portal.broadcastScene(
-      SCENE.UPDATE,
+      WS_SCENE_EVENT_TYPES.UPDATE,
       this.excalidrawAPI.getSceneElementsIncludingDeleted(),
       true,
     );
@@ -731,6 +746,16 @@ class CollabWrapper extends PureComponent<Props, CollabState> {
     this.setLastBroadcastedOrReceivedSceneVersion(newVersion);
   }, SYNC_FULL_SCENE_INTERVAL_MS);
 
+  queueSaveToFirebase = throttle(() => {
+    if (this.portal.socketInitialized) {
+      this.saveCollabRoomToFirebase(
+        this.getSyncableElements(
+          this.excalidrawAPI.getSceneElementsIncludingDeleted(),
+        ),
+      );
+    }
+  }, SYNC_FULL_SCENE_INTERVAL_MS);
+
   handleClose = () => {
     this.setState({ modalIsShown: false });
   };
@@ -771,7 +796,7 @@ class CollabWrapper extends PureComponent<Props, CollabState> {
     this.contextValue.onPointerUpdate = this.onPointerUpdate;
     this.contextValue.initializeSocketClient = this.initializeSocketClient;
     this.contextValue.onCollabButtonClick = this.onCollabButtonClick;
-    this.contextValue.broadcastElements = this.broadcastElements;
+    this.contextValue.syncElements = this.syncElements;
     this.contextValue.fetchImageFilesFromFirebase =
       this.fetchImageFilesFromFirebase;
     this.contextValue.setUsername = this.setUsername;

src/excalidraw-app/collab/Portal.tsx

@@ -3,7 +3,11 @@ import { SocketUpdateData, SocketUpdateDataSource } from "../data";
 import CollabWrapper from "./CollabWrapper";
 
 import { ExcalidrawElement } from "../../element/types";
-import { BROADCAST, FILE_UPLOAD_TIMEOUT, SCENE } from "../app_constants";
+import {
+  WS_EVENTS,
+  FILE_UPLOAD_TIMEOUT,
+  WS_SCENE_EVENT_TYPES,
+} from "../app_constants";
 import { UserIdleState } from "../../types";
 import { trackEvent } from "../../analytics";
 import { throttle } from "lodash";
@@ -37,7 +41,7 @@ class Portal {
     });
     this.socket.on("new-user", async (_socketId: string) => {
       this.broadcastScene(
-        SCENE.INIT,
+        WS_SCENE_EVENT_TYPES.INIT,
         this.collab.getSceneElementsIncludingDeleted(),
         /* syncAll */ true,
       );
@@ -81,7 +85,7 @@ class Portal {
       const { encryptedBuffer, iv } = await encryptData(this.roomKey!, encoded);
 
       this.socket?.emit(
-        volatile ? BROADCAST.SERVER_VOLATILE : BROADCAST.SERVER,
+        volatile ? WS_EVENTS.SERVER_VOLATILE : WS_EVENTS.SERVER,
         this.roomId,
         encryptedBuffer,
         iv,
@@ -121,11 +125,11 @@ class Portal {
   }, FILE_UPLOAD_TIMEOUT);
 
   broadcastScene = async (
-    sceneType: SCENE.INIT | SCENE.UPDATE,
+    updateType: WS_SCENE_EVENT_TYPES.INIT | WS_SCENE_EVENT_TYPES.UPDATE,
     allElements: readonly ExcalidrawElement[],
     syncAll: boolean,
   ) => {
-    if (sceneType === SCENE.INIT && !syncAll) {
+    if (updateType === WS_SCENE_EVENT_TYPES.INIT && !syncAll) {
       throw new Error("syncAll must be true when sending SCENE.INIT");
     }
 
@@ -152,8 +156,8 @@ class Portal {
       [] as BroadcastedExcalidrawElement[],
     );
 
-    const data: SocketUpdateDataSource[typeof sceneType] = {
-      type: sceneType,
+    const data: SocketUpdateDataSource[typeof updateType] = {
+      type: updateType,
       payload: {
         elements: syncableElements,
       },
@@ -166,20 +170,9 @@ class Portal {
       );
     }
 
-    const broadcastPromise = this._broadcastSocketData(
-      data as SocketUpdateData,
-    );
-
     this.queueFileUpload();
 
-    if (syncAll && this.collab.isCollaborating()) {
-      await Promise.all([
-        broadcastPromise,
-        this.collab.saveCollabRoomToFirebase(syncableElements),
-      ]);
-    } else {
-      await broadcastPromise;
-    }
+    await this._broadcastSocketData(data as SocketUpdateData);
   };
 
   broadcastIdleChange = (userState: UserIdleState) => {

src/excalidraw-app/collab/reconciliation.ts

@@ -78,8 +78,14 @@ export const reconcileElements = (
       continue;
     }
 
+    // Mark duplicate for removal as it'll be replaced with the remote element
     if (local) {
-      // mark for removal since it'll be replaced with the remote element
+      // Unless the ramote and local elements are the same element in which case
+      // we need to keep it as we'd otherwise discard it from the resulting
+      // array.
+      if (local[0] === remoteElement) {
+        continue;
+      }
       duplicates.set(local[0], true);
     }
 

src/excalidraw-app/data/firebase.ts

@@ -2,11 +2,17 @@ import { ExcalidrawElement, FileId } from "../../element/types";
 import { getSceneVersion } from "../../element";
 import Portal from "../collab/Portal";
 import { restoreElements } from "../../data/restore";
-import { BinaryFileData, BinaryFileMetadata, DataURL } from "../../types";
+import {
+  AppState,
+  BinaryFileData,
+  BinaryFileMetadata,
+  DataURL,
+} from "../../types";
 import { FILE_CACHE_MAX_AGE_SEC } from "../app_constants";
 import { decompressData } from "../../data/encode";
 import { encryptData, decryptData } from "../../data/encryption";
 import { MIME_TYPES } from "../../constants";
+import { reconcileElements } from "../collab/reconciliation";
 
 // private
 // -----------------------------------------------------------------------------
@@ -108,11 +114,13 @@ const encryptElements = async (
 };
 
 const decryptElements = async (
-  key: string,
-  iv: Uint8Array,
-  ciphertext: ArrayBuffer | Uint8Array,
+  data: FirebaseStoredScene,
+  roomKey: string,
 ): Promise<readonly ExcalidrawElement[]> => {
-  const decrypted = await decryptData(iv, ciphertext, key);
+  const ciphertext = data.ciphertext.toUint8Array();
+  const iv = data.iv.toUint8Array();
+
+  const decrypted = await decryptData(iv, ciphertext, roomKey);
   const decodedData = new TextDecoder("utf-8").decode(
     new Uint8Array(decrypted),
   );
@@ -171,57 +179,86 @@ export const saveFilesToFirebase = async ({
   return { savedFiles, erroredFiles };
 };
 
+const createFirebaseSceneDocument = async (
+  firebase: ResolutionType<typeof loadFirestore>,
+  elements: readonly ExcalidrawElement[],
+  roomKey: string,
+) => {
+  const sceneVersion = getSceneVersion(elements);
+  const { ciphertext, iv } = await encryptElements(roomKey, elements);
+  return {
+    sceneVersion,
+    ciphertext: firebase.firestore.Blob.fromUint8Array(
+      new Uint8Array(ciphertext),
+    ),
+    iv: firebase.firestore.Blob.fromUint8Array(iv),
+  } as FirebaseStoredScene;
+};
+
 export const saveToFirebase = async (
   portal: Portal,
   elements: readonly ExcalidrawElement[],
+  appState: AppState,
 ) => {
   const { roomId, roomKey, socket } = portal;
   if (
-    // if no room exists, consider the room saved because there's nothing we can
-    // do at this point
+    // bail if no room exists as there's nothing we can do at this point
     !roomId ||
     !roomKey ||
     !socket ||
     isSavedToFirebase(portal, elements)
   ) {
-    return true;
+    return false;
   }
 
   const firebase = await loadFirestore();
-  const sceneVersion = getSceneVersion(elements);
-  const { ciphertext, iv } = await encryptElements(roomKey, elements);
+  const firestore = firebase.firestore();
 
-  const nextDocData = {
-    sceneVersion,
-    ciphertext: firebase.firestore.Blob.fromUint8Array(
-      new Uint8Array(ciphertext),
-    ),
-    iv: firebase.firestore.Blob.fromUint8Array(iv),
-  } as FirebaseStoredScene;
+  const docRef = firestore.collection("scenes").doc(roomId);
 
-  const db = firebase.firestore();
-  const docRef = db.collection("scenes").doc(roomId);
-  const didUpdate = await db.runTransaction(async (transaction) => {
-    const doc = await transaction.get(docRef);
-    if (!doc.exists) {
-      transaction.set(docRef, nextDocData);
-      return true;
-    }
+  const savedData = await firestore.runTransaction(async (transaction) => {
+    const snapshot = await transaction.get(docRef);
 
-    const prevDocData = doc.data() as FirebaseStoredScene;
-    if (prevDocData.sceneVersion >= nextDocData.sceneVersion) {
-      return false;
+    if (!snapshot.exists) {
+      const sceneDocument = await createFirebaseSceneDocument(
+        firebase,
+        elements,
+        roomKey,
+      );
+
+      transaction.set(docRef, sceneDocument);
+
+      return {
+        sceneVersion: sceneDocument.sceneVersion,
+        reconciledElements: null,
+      };
     }
 
-    transaction.update(docRef, nextDocData);
-    return true;
+    const prevDocData = snapshot.data() as FirebaseStoredScene;
+    const prevElements = await decryptElements(prevDocData, roomKey);
+
+    const reconciledElements = reconcileElements(
+      elements,
+      prevElements,
+      appState,
+    );
+
+    const sceneDocument = await createFirebaseSceneDocument(
+      firebase,
+      reconciledElements,
+      roomKey,
+    );
+
+    transaction.update(docRef, sceneDocument);
+    return {
+      reconciledElements,
+      sceneVersion: sceneDocument.sceneVersion,
+    };
   });
 
-  if (didUpdate) {
-    firebaseSceneVersionCache.set(socket, sceneVersion);
-  }
+  firebaseSceneVersionCache.set(socket, savedData.sceneVersion);
 
-  return didUpdate;
+  return savedData;
 };
 
 export const loadFromFirebase = async (
@@ -238,10 +275,7 @@ export const loadFromFirebase = async (
     return null;
   }
   const storedScene = doc.data() as FirebaseStoredScene;
-  const ciphertext = storedScene.ciphertext.toUint8Array();
-  const iv = storedScene.iv.toUint8Array();
-
-  const elements = await decryptElements(roomKey, iv, ciphertext);
+  const elements = await decryptElements(storedScene, roomKey);
 
   if (socket) {
     firebaseSceneVersionCache.set(socket, getSceneVersion(elements));

src/excalidraw-app/index.tsx

@@ -455,7 +455,7 @@ const ExcalidrawWrapper = () => {
     files: BinaryFiles,
   ) => {
     if (collabAPI?.isCollaborating()) {
-      collabAPI.broadcastElements(elements);
+      collabAPI.syncElements(elements);
     }
 
     // this check is redundant, but since this is a hot path, it's best

src/tests/reconciliation.test.ts

@@ -9,36 +9,60 @@ import { randomInteger } from "../random";
 import { AppState } from "../types";
 
 type Id = string;
-type Ids = Id[];
+type ElementLike = {
+  id: string;
+  version: number;
+  versionNonce: number;
+  parent?: string | null;
+};
 
 type Cache = Record<string, ExcalidrawElement | undefined>;
 
-const parseId = (uid: string) => {
-  const [, parent, id, version] = uid.match(
-    /^(?:\((\^|\w+)\))?(\w+)(?::(\d+))?(?:\((\w+)\))?$/,
-  )!;
+const createElement = (opts: { uid: string } | ElementLike) => {
+  let uid: string;
+  let id: string;
+  let version: number | null;
+  let parent: string | null = null;
+  let versionNonce: number | null = null;
+  if ("uid" in opts) {
+    const match = opts.uid.match(
+      /^(?:\((\^|\w+)\))?(\w+)(?::(\d+))?(?:\((\w+)\))?$/,
+    )!;
+    parent = match[1];
+    id = match[2];
+    version = match[3] ? parseInt(match[3]) : null;
+    uid = version ? `${id}:${version}` : id;
+  } else {
+    ({ id, version, versionNonce } = opts);
+    parent = parent || null;
+    uid = id;
+  }
   return {
-    uid: version ? `${id}:${version}` : id,
+    uid,
     id,
-    version: version ? parseInt(version) : null,
+    version,
+    versionNonce: versionNonce || randomInteger(),
     parent: parent || null,
   };
 };
 
 const idsToElements = (
-  ids: Ids,
+  ids: (Id | ElementLike)[],
   cache: Cache = {},
 ): readonly ExcalidrawElement[] => {
   return ids.reduce((acc, _uid, idx) => {
-    const { uid, id, version, parent } = parseId(_uid);
+    const { uid, id, version, parent, versionNonce } = createElement(
+      typeof _uid === "string" ? { uid: _uid } : _uid,
+    );
     const cached = cache[uid];
     const elem = {
       id,
       version: version ?? 0,
-      versionNonce: randomInteger(),
+      versionNonce,
       ...cached,
       parent,
     } as BroadcastedExcalidrawElement;
+    // @ts-ignore
     cache[uid] = elem;
     acc.push(elem);
     return acc;
@@ -67,8 +91,8 @@ const cleanElements = (elements: ReconciledElements) => {
 const cloneDeep = (data: any) => JSON.parse(JSON.stringify(data));
 
 const test = <U extends `${string}:${"L" | "R"}`>(
-  local: Ids,
-  remote: Ids,
+  local: (Id | ElementLike)[],
+  remote: (Id | ElementLike)[],
   target: U[],
   bidirectional = true,
 ) => {
@@ -80,6 +104,7 @@ const test = <U extends `${string}:${"L" | "R"}`>(
     return (source === "L" ? _local : _remote).find((e) => e.id === id)!;
   }) as any as ReconciledElements;
   const remoteReconciled = reconcileElements(_local, _remote, {} as AppState);
+  expect(target.length).equal(remoteReconciled.length);
   expect(cleanElements(remoteReconciled)).deep.equal(
     cleanElements(_target),
     "remote reconciliation",
@@ -301,4 +326,92 @@ describe("elements reconciliation", () => {
     test(["A:2", "B:2"], ["(A)C", "B:1"], ["A:L", "C:R", "B:L"]);
     test(["A:2", "B:2"], ["(A)C", "B:1"], ["A:L", "C:R", "B:L"]);
   });
+
+  it("test identical elements reconciliation", () => {
+    const testIdentical = (
+      local: ElementLike[],
+      remote: ElementLike[],
+      expected: Id[],
+    ) => {
+      const ret = reconcileElements(
+        local as any as ExcalidrawElement[],
+        remote as any as ExcalidrawElement[],
+        {} as AppState,
+      );
+
+      if (new Set(ret.map((x) => x.id)).size !== ret.length) {
+        throw new Error("reconcileElements: duplicate elements found");
+      }
+
+      expect(ret.map((x) => x.id)).to.deep.equal(expected);
+    };
+
+    // identical id/version/versionNonce
+    // -------------------------------------------------------------------------
+
+    testIdentical(
+      [{ id: "A", version: 1, versionNonce: 1 }],
+      [{ id: "A", version: 1, versionNonce: 1 }],
+      ["A"],
+    );
+    testIdentical(
+      [
+        { id: "A", version: 1, versionNonce: 1 },
+        { id: "B", version: 1, versionNonce: 1 },
+      ],
+      [
+        { id: "B", version: 1, versionNonce: 1 },
+        { id: "A", version: 1, versionNonce: 1 },
+      ],
+      ["B", "A"],
+    );
+    testIdentical(
+      [
+        { id: "A", version: 1, versionNonce: 1 },
+        { id: "B", version: 1, versionNonce: 1 },
+      ],
+      [
+        { id: "B", version: 1, versionNonce: 1 },
+        { id: "A", version: 1, versionNonce: 1 },
+      ],
+      ["B", "A"],
+    );
+
+    // actually identical (arrays and element objects)
+    // -------------------------------------------------------------------------
+
+    const elements1 = [
+      {
+        id: "A",
+        version: 1,
+        versionNonce: 1,
+        parent: null,
+      },
+      {
+        id: "B",
+        version: 1,
+        versionNonce: 1,
+        parent: null,
+      },
+    ];
+
+    testIdentical(elements1, elements1, ["A", "B"]);
+    testIdentical(elements1, elements1.slice(), ["A", "B"]);
+    testIdentical(elements1.slice(), elements1, ["A", "B"]);
+    testIdentical(elements1.slice(), elements1.slice(), ["A", "B"]);
+
+    const el1 = {
+      id: "A",
+      version: 1,
+      versionNonce: 1,
+      parent: null,
+    };
+    const el2 = {
+      id: "B",
+      version: 1,
+      versionNonce: 1,
+      parent: null,
+    };
+    testIdentical([el1, el2], [el2, el1], ["A", "B"]);
+  });
 });