feat: refactor local persistence & fix race condition on SW reload (#5032)

src/excalidraw-app/collab/CollabWrapper.tsx

@@ -68,6 +68,7 @@ import {
 } from "./reconciliation";
 import { decryptData } from "../../data/encryption";
 import { resetBrowserStateVersions } from "../data/tabSync";
+import { LocalData } from "../data/LocalData";
 
 interface CollabState {
   modalIsShown: boolean;
@@ -109,10 +110,11 @@ class CollabWrapper extends PureComponent<Props, CollabState> {
   portal: Portal;
   fileManager: FileManager;
   excalidrawAPI: Props["excalidrawAPI"];
-  isCollaborating: boolean = false;
   activeIntervalId: number | null;
   idleTimeoutId: number | null;
 
+  // marked as private to ensure we don't change it outside this class
+  private _isCollaborating: boolean = false;
   private socketInitializationTimer?: number;
   private lastBroadcastedOrReceivedSceneVersion: number = -1;
   private collaborators = new Map<string, Collaborator>();
@@ -193,6 +195,8 @@ class CollabWrapper extends PureComponent<Props, CollabState> {
     }
   }
 
+  isCollaborating = () => this._isCollaborating;
+
   private onUnload = () => {
     this.destroySocketClient({ isUnload: true });
   };
@@ -203,7 +207,7 @@ class CollabWrapper extends PureComponent<Props, CollabState> {
     );
 
     if (
-      this.isCollaborating &&
+      this._isCollaborating &&
       (this.fileManager.shouldPreventUnload(syncableElements) ||
         !isSavedToFirebase(this.portal, syncableElements))
     ) {
@@ -285,7 +289,8 @@ class CollabWrapper extends PureComponent<Props, CollabState> {
       this.setState({
         activeRoomLink: "",
       });
-      this.isCollaborating = false;
+      this._isCollaborating = false;
+      LocalData.resumeSave("collaboration");
     }
     this.lastBroadcastedOrReceivedSceneVersion = -1;
     this.portal.close();
@@ -353,7 +358,8 @@ class CollabWrapper extends PureComponent<Props, CollabState> {
 
     const scenePromise = resolvablePromise<ImportedDataState | null>();
 
-    this.isCollaborating = true;
+    this._isCollaborating = true;
+    LocalData.pauseSave("collaboration");
 
     const { default: socketIOClient } = await import(
       /* webpackChunkName: "socketIoClient" */ "socket.io-client"
@@ -760,7 +766,7 @@ class CollabWrapper extends PureComponent<Props, CollabState> {
       this.contextValue = {} as CollabAPI;
     }
 
-    this.contextValue.isCollaborating = () => this.isCollaborating;
+    this.contextValue.isCollaborating = this.isCollaborating;
     this.contextValue.username = this.state.username;
     this.contextValue.onPointerUpdate = this.onPointerUpdate;
     this.contextValue.initializeSocketClient = this.initializeSocketClient;

src/excalidraw-app/collab/Portal.tsx

@@ -172,7 +172,7 @@ class Portal {
 
     this.queueFileUpload();
 
-    if (syncAll && this.collab.isCollaborating) {
+    if (syncAll && this.collab.isCollaborating()) {
       await Promise.all([
         broadcastPromise,
         this.collab.saveCollabRoomToFirebase(syncableElements),

src/excalidraw-app/data/LocalData.ts

@@ -0,0 +1,154 @@
+/**
+ * This file deals with saving data state (appState, elements, images, ...)
+ * locally to the browser.
+ *
+ * Notes:
+ *
+ * - DataState refers to full state of the app: appState, elements, images,
+ *   though some state is saved separately (collab username, library) for one
+ *   reason or another. We also save different data to different sotrage
+ *   (localStorage, indexedDB).
+ */
+
+import { createStore, keys, del, getMany, set } from "idb-keyval";
+import { clearAppStateForLocalStorage } from "../../appState";
+import { clearElementsForLocalStorage } from "../../element";
+import { ExcalidrawElement, FileId } from "../../element/types";
+import { AppState, BinaryFileData, BinaryFiles } from "../../types";
+import { debounce } from "../../utils";
+import { SAVE_TO_LOCAL_STORAGE_TIMEOUT, STORAGE_KEYS } from "../app_constants";
+import { FileManager } from "./FileManager";
+import { Locker } from "./Locker";
+import { updateBrowserStateVersion } from "./tabSync";
+
+const filesStore = createStore("files-db", "files-store");
+
+class LocalFileManager extends FileManager {
+  clearObsoleteFiles = async (opts: { currentFileIds: FileId[] }) => {
+    const allIds = await keys(filesStore);
+    for (const id of allIds) {
+      if (!opts.currentFileIds.includes(id as FileId)) {
+        del(id, filesStore);
+      }
+    }
+  };
+}
+
+const saveDataStateToLocalStorage = (
+  elements: readonly ExcalidrawElement[],
+  appState: AppState,
+) => {
+  try {
+    localStorage.setItem(
+      STORAGE_KEYS.LOCAL_STORAGE_ELEMENTS,
+      JSON.stringify(clearElementsForLocalStorage(elements)),
+    );
+    localStorage.setItem(
+      STORAGE_KEYS.LOCAL_STORAGE_APP_STATE,
+      JSON.stringify(clearAppStateForLocalStorage(appState)),
+    );
+    updateBrowserStateVersion(STORAGE_KEYS.VERSION_DATA_STATE);
+  } catch (error: any) {
+    // Unable to access window.localStorage
+    console.error(error);
+  }
+};
+
+type SavingLockTypes = "collaboration";
+
+export class LocalData {
+  private static _save = debounce(
+    async (
+      elements: readonly ExcalidrawElement[],
+      appState: AppState,
+      files: BinaryFiles,
+      onFilesSaved: () => void,
+    ) => {
+      saveDataStateToLocalStorage(elements, appState);
+
+      await this.fileStorage.saveFiles({
+        elements,
+        files,
+      });
+      onFilesSaved();
+    },
+    SAVE_TO_LOCAL_STORAGE_TIMEOUT,
+  );
+
+  /** Saves DataState, including files. Bails if saving is paused */
+  static save = (
+    elements: readonly ExcalidrawElement[],
+    appState: AppState,
+    files: BinaryFiles,
+    onFilesSaved: () => void,
+  ) => {
+    // we need to make the `isSavePaused` check synchronously (undebounced)
+    if (!this.isSavePaused()) {
+      this._save(elements, appState, files, onFilesSaved);
+    }
+  };
+
+  static flushSave = () => {
+    this._save.flush();
+  };
+
+  private static locker = new Locker<SavingLockTypes>();
+
+  static pauseSave = (lockType: SavingLockTypes) => {
+    this.locker.lock(lockType);
+  };
+
+  static resumeSave = (lockType: SavingLockTypes) => {
+    this.locker.unlock(lockType);
+  };
+
+  static isSavePaused = () => {
+    return document.hidden || this.locker.isLocked();
+  };
+
+  // ---------------------------------------------------------------------------
+
+  static fileStorage = new LocalFileManager({
+    getFiles(ids) {
+      return getMany(ids, filesStore).then(
+        (filesData: (BinaryFileData | undefined)[]) => {
+          const loadedFiles: BinaryFileData[] = [];
+          const erroredFiles = new Map<FileId, true>();
+          filesData.forEach((data, index) => {
+            const id = ids[index];
+            if (data) {
+              loadedFiles.push(data);
+            } else {
+              erroredFiles.set(id, true);
+            }
+          });
+
+          return { loadedFiles, erroredFiles };
+        },
+      );
+    },
+    async saveFiles({ addedFiles }) {
+      const savedFiles = new Map<FileId, true>();
+      const erroredFiles = new Map<FileId, true>();
+
+      // before we use `storage` event synchronization, let's update the flag
+      // optimistically. Hopefully nothing fails, and an IDB read executed
+      // before an IDB write finishes will read the latest value.
+      updateBrowserStateVersion(STORAGE_KEYS.VERSION_FILES);
+
+      await Promise.all(
+        [...addedFiles].map(async ([id, fileData]) => {
+          try {
+            await set(id, fileData, filesStore);
+            savedFiles.set(id, true);
+          } catch (error: any) {
+            console.error(error);
+            erroredFiles.set(id, true);
+          }
+        }),
+      );
+
+      return { savedFiles, erroredFiles };
+    },
+  });
+}

src/excalidraw-app/data/Locker.ts

@@ -0,0 +1,18 @@
+export class Locker<T extends string> {
+  private locks = new Map<T, true>();
+
+  lock = (lockType: T) => {
+    this.locks.set(lockType, true);
+  };
+
+  /** @returns whether no locks remaining */
+  unlock = (lockType: T) => {
+    this.locks.delete(lockType);
+    return !this.isLocked();
+  };
+
+  /** @returns whether some (or specific) locks are present */
+  isLocked(lockType?: T) {
+    return lockType ? this.locks.has(lockType) : !!this.locks.size;
+  }
+}

src/excalidraw-app/data/localStorage.ts

@@ -5,7 +5,6 @@ import {
   getDefaultAppState,
 } from "../../appState";
 import { clearElementsForLocalStorage } from "../../element";
-import { updateBrowserStateVersion } from "./tabSync";
 import { STORAGE_KEYS } from "../app_constants";
 
 export const saveUsernameToLocalStorage = (username: string) => {
@@ -34,26 +33,6 @@ export const importUsernameFromLocalStorage = (): string | null => {
   return null;
 };
 
-export const saveToLocalStorage = (
-  elements: readonly ExcalidrawElement[],
-  appState: AppState,
-) => {
-  try {
-    localStorage.setItem(
-      STORAGE_KEYS.LOCAL_STORAGE_ELEMENTS,
-      JSON.stringify(clearElementsForLocalStorage(elements)),
-    );
-    localStorage.setItem(
-      STORAGE_KEYS.LOCAL_STORAGE_APP_STATE,
-      JSON.stringify(clearAppStateForLocalStorage(appState)),
-    );
-    updateBrowserStateVersion(STORAGE_KEYS.VERSION_DATA_STATE);
-  } catch (error: any) {
-    // Unable to access window.localStorage
-    console.error(error);
-  }
-};
-
 export const importFromLocalStorage = () => {
   let savedElements = null;
   let savedState = null;

src/excalidraw-app/index.tsx

@@ -28,7 +28,6 @@ import {
   AppState,
   LibraryItems,
   ExcalidrawImperativeAPI,
-  BinaryFileData,
   BinaryFiles,
 } from "../types";
 import {
@@ -42,7 +41,6 @@ import {
 } from "../utils";
 import {
   FIREBASE_STORAGE_PREFIXES,
-  SAVE_TO_LOCAL_STORAGE_TIMEOUT,
   STORAGE_KEYS,
   SYNC_BROWSER_TABS_TIMEOUT,
 } from "./app_constants";
@@ -57,7 +55,6 @@ import {
   getLibraryItemsFromStorage,
   importFromLocalStorage,
   importUsernameFromLocalStorage,
-  saveToLocalStorage,
 } from "./data/localStorage";
 import CustomStats from "./CustomStats";
 import { restoreAppState, RestoredDataState } from "../data/restore";
@@ -67,72 +64,12 @@ import { shield } from "../components/icons";
 import "./index.scss";
 import { ExportToExcalidrawPlus } from "./components/ExportToExcalidrawPlus";
 
-import { getMany, set, del, keys, createStore } from "idb-keyval";
-import { FileManager, updateStaleImageStatuses } from "./data/FileManager";
+import { updateStaleImageStatuses } from "./data/FileManager";
 import { newElementWith } from "../element/mutateElement";
 import { isInitializedImageElement } from "../element/typeChecks";
 import { loadFilesFromFirebase } from "./data/firebase";
-import {
-  isBrowserStorageStateNewer,
-  updateBrowserStateVersion,
-} from "./data/tabSync";
-
-const filesStore = createStore("files-db", "files-store");
-
-const clearObsoleteFilesFromIndexedDB = async (opts: {
-  currentFileIds: FileId[];
-}) => {
-  const allIds = await keys(filesStore);
-  for (const id of allIds) {
-    if (!opts.currentFileIds.includes(id as FileId)) {
-      del(id, filesStore);
-    }
-  }
-};
-
-const localFileStorage = new FileManager({
-  getFiles(ids) {
-    return getMany(ids, filesStore).then(
-      (filesData: (BinaryFileData | undefined)[]) => {
-        const loadedFiles: BinaryFileData[] = [];
-        const erroredFiles = new Map<FileId, true>();
-        filesData.forEach((data, index) => {
-          const id = ids[index];
-          if (data) {
-            loadedFiles.push(data);
-          } else {
-            erroredFiles.set(id, true);
-          }
-        });
-
-        return { loadedFiles, erroredFiles };
-      },
-    );
-  },
-  async saveFiles({ addedFiles }) {
-    const savedFiles = new Map<FileId, true>();
-    const erroredFiles = new Map<FileId, true>();
-
-    // before we use `storage` event synchronization, let's update the flag
-    // optimistically. Hopefully nothing fails, and an IDB read executed
-    // before an IDB write finishes will read the latest value.
-    updateBrowserStateVersion(STORAGE_KEYS.VERSION_FILES);
-
-    await Promise.all(
-      [...addedFiles].map(async ([id, fileData]) => {
-        try {
-          await set(id, fileData, filesStore);
-          savedFiles.set(id, true);
-        } catch (error: any) {
-          console.error(error);
-          erroredFiles.set(id, true);
-        }
-      }),
-    );
-
-    return { savedFiles, erroredFiles };
-  },
-});
+import { LocalData } from "./data/LocalData";
+import { isBrowserStorageStateNewer } from "./data/tabSync";
 
 const languageDetector = new LanguageDetector();
 languageDetector.init({
@@ -143,28 +80,6 @@ languageDetector.init({
   checkWhitelist: false,
 });
 
-const saveDebounced = debounce(
-  async (
-    elements: readonly ExcalidrawElement[],
-    appState: AppState,
-    files: BinaryFiles,
-    onFilesSaved: () => void,
-  ) => {
-    saveToLocalStorage(elements, appState);
-
-    await localFileStorage.saveFiles({
-      elements,
-      files,
-    });
-    onFilesSaved();
-  },
-  SAVE_TO_LOCAL_STORAGE_TIMEOUT,
-);
-
-const onBlur = () => {
-  saveDebounced.flush();
-};
-
 const initializeScene = async (opts: {
   collabAPI: CollabAPI;
 }): Promise<
@@ -366,7 +281,7 @@ const ExcalidrawWrapper = () => {
           });
         } else if (isInitialLoad) {
           if (fileIds.length) {
-            localFileStorage
+            LocalData.fileStorage
               .getFiles(fileIds)
               .then(({ loadedFiles, erroredFiles }) => {
                 if (loadedFiles.length) {
@@ -381,7 +296,7 @@ const ExcalidrawWrapper = () => {
           }
           // on fresh load, clear unused files from IDB (from previous
           // session)
-          clearObsoleteFilesFromIndexedDB({ currentFileIds: fileIds });
+          LocalData.fileStorage.clearObsoleteFiles({ currentFileIds: fileIds });
         }
       }
 
@@ -458,7 +373,7 @@ const ExcalidrawWrapper = () => {
               return acc;
             }, [] as FileId[]) || [];
           if (fileIds.length) {
-            localFileStorage
+            LocalData.fileStorage
               .getFiles(fileIds)
               .then(({ loadedFiles, erroredFiles }) => {
                 if (loadedFiles.length) {
@@ -475,28 +390,50 @@ const ExcalidrawWrapper = () => {
       }
     }, SYNC_BROWSER_TABS_TIMEOUT);
 
+    const onUnload = () => {
+      LocalData.flushSave();
+    };
+
+    const visibilityChange = (event: FocusEvent | Event) => {
+      if (event.type === EVENT.BLUR || document.hidden) {
+        LocalData.flushSave();
+      }
+      if (
+        event.type === EVENT.VISIBILITY_CHANGE ||
+        event.type === EVENT.FOCUS
+      ) {
+        syncData();
+      }
+    };
+
     window.addEventListener(EVENT.HASHCHANGE, onHashChange, false);
-    window.addEventListener(EVENT.UNLOAD, onBlur, false);
-    window.addEventListener(EVENT.BLUR, onBlur, false);
-    document.addEventListener(EVENT.VISIBILITY_CHANGE, syncData, false);
-    window.addEventListener(EVENT.FOCUS, syncData, false);
+    window.addEventListener(EVENT.UNLOAD, onUnload, false);
+    window.addEventListener(EVENT.BLUR, visibilityChange, false);
+    document.addEventListener(EVENT.VISIBILITY_CHANGE, visibilityChange, false);
+    window.addEventListener(EVENT.FOCUS, visibilityChange, false);
     return () => {
       window.removeEventListener(EVENT.HASHCHANGE, onHashChange, false);
-      window.removeEventListener(EVENT.UNLOAD, onBlur, false);
-      window.removeEventListener(EVENT.BLUR, onBlur, false);
-      window.removeEventListener(EVENT.FOCUS, syncData, false);
-      document.removeEventListener(EVENT.VISIBILITY_CHANGE, syncData, false);
+      window.removeEventListener(EVENT.UNLOAD, onUnload, false);
+      window.removeEventListener(EVENT.BLUR, visibilityChange, false);
+      window.removeEventListener(EVENT.FOCUS, visibilityChange, false);
+      document.removeEventListener(
+        EVENT.VISIBILITY_CHANGE,
+        visibilityChange,
+        false,
+      );
       clearTimeout(titleTimeout);
     };
   }, [collabAPI, excalidrawAPI]);
 
   useEffect(() => {
     const unloadHandler = (event: BeforeUnloadEvent) => {
-      saveDebounced.flush();
+      LocalData.flushSave();
 
       if (
         excalidrawAPI &&
-        localFileStorage.shouldPreventUnload(excalidrawAPI.getSceneElements())
+        LocalData.fileStorage.shouldPreventUnload(
+          excalidrawAPI.getSceneElements(),
+        )
       ) {
         preventUnload(event);
       }
@@ -518,8 +455,12 @@ const ExcalidrawWrapper = () => {
   ) => {
     if (collabAPI?.isCollaborating()) {
       collabAPI.broadcastElements(elements);
-    } else {
-      saveDebounced(elements, appState, files, () => {
+    }
+
+    // this check is redundant, but since this is a hot path, it's best
+    // not to evaludate the nested expression every time
+    if (!LocalData.isSavePaused()) {
+      LocalData.save(elements, appState, files, () => {
         if (excalidrawAPI) {
           let didChange = false;
 
@@ -527,7 +468,9 @@ const ExcalidrawWrapper = () => {
           const elements = excalidrawAPI
             .getSceneElementsIncludingDeleted()
             .map((element) => {
-              if (localFileStorage.shouldUpdateImageElementStatus(element)) {
+              if (
+                LocalData.fileStorage.shouldUpdateImageElementStatus(element)
+              ) {
                 didChange = true;
                 const newEl = newElementWith(element, { status: "saved" });
                 if (pendingImageElement === element) {
@@ -687,7 +630,7 @@ const ExcalidrawWrapper = () => {
   };
 
   const onRoomClose = useCallback(() => {
-    localFileStorage.reset();
+    LocalData.fileStorage.reset();
   }, []);
 
   return (