Merge pull request #490 from paultsimura/fix/update-order

tgolen · web-flow · commit 6317b21afae1 · 2024-07-16T08:29:18.000-06:00
Try to synchronize operations within one update
diff --git a/lib/Onyx.ts b/lib/Onyx.ts
@@ -9,13 +9,16 @@ import Storage from './storage';
 import utils from './utils';
 import DevTools from './DevTools';
 import type {
+    Collection,
+    CollectionKey,
     CollectionKeyBase,
     ConnectOptions,
     InitOptions,
     KeyValueMapping,
     Mapping,
     OnyxInputKeyValueMapping,
     OnyxCollection,
+    MixedOperationsQueue,
     OnyxKey,
     OnyxMergeCollectionInput,
     OnyxMergeInput,
@@ -437,30 +440,16 @@ function merge<TKey extends OnyxKey>(key: TKey, changes: OnyxMergeInput<TKey>):
  * @param collection Object collection keyed by individual collection member keys and values
  */
 function mergeCollection<TKey extends CollectionKeyBase, TMap>(collectionKey: TKey, collection: OnyxMergeCollectionInput<TKey, TMap>): Promise<void> {
-    if (typeof collection !== 'object' || Array.isArray(collection) || utils.isEmptyObject(collection)) {
+    if (!OnyxUtils.isValidNonEmptyCollectionForMerge(collection)) {
         Logger.logInfo('mergeCollection() called with invalid or empty value. Skipping this update.');
         return Promise.resolve();
     }
+
     const mergedCollection: OnyxInputKeyValueMapping = collection;
 
     // Confirm all the collection keys belong to the same parent
-    let hasCollectionKeyCheckFailed = false;
     const mergedCollectionKeys = Object.keys(mergedCollection);
-    mergedCollectionKeys.forEach((dataKey) => {
-        if (OnyxUtils.isKeyMatch(collectionKey, dataKey)) {
-            return;
-        }
-
-        if (process.env.NODE_ENV === 'development') {
-            throw new Error(`Provided collection doesn't have all its data belonging to the same parent. CollectionKey: ${collectionKey}, DataKey: ${dataKey}`);
-        }
-
-        hasCollectionKeyCheckFailed = true;
-        Logger.logAlert(`Provided collection doesn't have all its data belonging to the same parent. CollectionKey: ${collectionKey}, DataKey: ${dataKey}`);
-    });
-
-    // Gracefully handle bad mergeCollection updates so it doesn't block the merge queue
-    if (hasCollectionKeyCheckFailed) {
+    if (!OnyxUtils.doAllCollectionItemsBelongToSameParent(collectionKey, mergedCollectionKeys)) {
         return Promise.resolve();
     }
 
@@ -712,23 +701,55 @@ function update(data: OnyxUpdate[]): Promise<void> {
         }
     });
 
+    // The queue of operations within a single `update` call in the format of <item key - list of operations updating the item>.
+    // This allows us to batch the operations per item and merge them into one operation in the order they were requested.
+    const updateQueue: Record<OnyxKey, Array<OnyxValue<OnyxKey>>> = {};
+    const enqueueSetOperation = (key: OnyxKey, value: OnyxValue<OnyxKey>) => {
+        // If a `set` operation is enqueued, we should clear the whole queue.
+        // Since the `set` operation replaces the value entirely, there's no need to perform any previous operations.
+        // To do this, we first put `null` in the queue, which removes the existing value, and then merge the new value.
+        updateQueue[key] = [null, value];
+    };
+    const enqueueMergeOperation = (key: OnyxKey, value: OnyxValue<OnyxKey>) => {
+        if (value === null) {
+            // If we merge `null`, the value is removed and all the previous operations are discarded.
+            updateQueue[key] = [null];
+        } else if (!updateQueue[key]) {
+            updateQueue[key] = [value];
+        } else {
+            updateQueue[key].push(value);
+        }
+    };
+
     const promises: Array<() => Promise<void>> = [];
     let clearPromise: Promise<void> = Promise.resolve();
 
     data.forEach(({onyxMethod, key, value}) => {
         switch (onyxMethod) {
             case OnyxUtils.METHOD.SET:
-                promises.push(() => set(key, value));
+                enqueueSetOperation(key, value);
                 break;
             case OnyxUtils.METHOD.MERGE:
-                promises.push(() => merge(key, value));
+                enqueueMergeOperation(key, value);
                 break;
-            case OnyxUtils.METHOD.MERGE_COLLECTION:
-                // eslint-disable-next-line @typescript-eslint/no-explicit-any -- We validated that the value is a collection
-                promises.push(() => mergeCollection(key, value as any));
+            case OnyxUtils.METHOD.MERGE_COLLECTION: {
+                const collection = value as Collection<CollectionKey, unknown, unknown>;
+                if (!OnyxUtils.isValidNonEmptyCollectionForMerge(collection)) {
+                    Logger.logInfo('mergeCollection enqueued within update() with invalid or empty value. Skipping this operation.');
+                    break;
+                }
+
+                // Confirm all the collection keys belong to the same parent
+                const collectionKeys = Object.keys(collection);
+                if (OnyxUtils.doAllCollectionItemsBelongToSameParent(key, collectionKeys)) {
+                    const mergedCollection: OnyxInputKeyValueMapping = collection;
+                    collectionKeys.forEach((collectionKey) => enqueueMergeOperation(collectionKey, mergedCollection[collectionKey]));
+                }
+
                 break;
+            }
             case OnyxUtils.METHOD.MULTI_SET:
-                promises.push(() => multiSet(value));
+                Object.entries(value).forEach(([entryKey, entryValue]) => enqueueSetOperation(entryKey, entryValue));
                 break;
             case OnyxUtils.METHOD.CLEAR:
                 clearPromise = clear();
@@ -738,6 +759,58 @@ function update(data: OnyxUpdate[]): Promise<void> {
         }
     });
 
+    // Group all the collection-related keys and update each collection in a single `mergeCollection` call.
+    // This is needed to prevent multiple `mergeCollection` calls for the same collection and `merge` calls for the individual items of the said collection.
+    // This way, we ensure there is no race condition in the queued updates of the same key.
+    OnyxUtils.getCollectionKeys().forEach((collectionKey) => {
+        const collectionItemKeys = Object.keys(updateQueue).filter((key) => OnyxUtils.isKeyMatch(collectionKey, key));
+        if (collectionItemKeys.length <= 1) {
+            // If there are no items of this collection in the updateQueue, we should skip it.
+            // If there is only one item, we should update it individually, therefore retain it in the updateQueue.
+            return;
+        }
+
+        const batchedCollectionUpdates = collectionItemKeys.reduce(
+            (queue: MixedOperationsQueue, key: string) => {
+                const operations = updateQueue[key];
+
+                // Remove the collection-related key from the updateQueue so that it won't be processed individually.
+                delete updateQueue[key];
+
+                const updatedValue = OnyxUtils.applyMerge(undefined, operations, false);
+                if (operations[0] === null) {
+                    // eslint-disable-next-line no-param-reassign
+                    queue.set[key] = updatedValue;
+                } else {
+                    // eslint-disable-next-line no-param-reassign
+                    queue.merge[key] = updatedValue;
+                }
+                return queue;
+            },
+            {
+                merge: {},
+                set: {},
+            },
+        );
+
+        if (!utils.isEmptyObject(batchedCollectionUpdates.merge)) {
+            promises.push(() => mergeCollection(collectionKey, batchedCollectionUpdates.merge as Collection<CollectionKey, unknown, unknown>));
+        }
+        if (!utils.isEmptyObject(batchedCollectionUpdates.set)) {
+            promises.push(() => multiSet(batchedCollectionUpdates.set));
+        }
+    });
+
+    Object.entries(updateQueue).forEach(([key, operations]) => {
+        const batchedChanges = OnyxUtils.applyMerge(undefined, operations, false);
+
+        if (operations[0] === null) {
+            promises.push(() => set(key, batchedChanges));
+        } else {
+            promises.push(() => merge(key, batchedChanges));
+        }
+    });
+
     return clearPromise
         .then(() => Promise.all(promises.map((p) => p())))
         .then(() => updateSnapshots(data))
diff --git a/lib/OnyxUtils.ts b/lib/OnyxUtils.ts
@@ -23,6 +23,7 @@ import type {
     OnyxEntry,
     OnyxInput,
     OnyxKey,
+    OnyxMergeCollectionInput,
     OnyxValue,
     Selector,
     WithOnyxConnectOptions,
@@ -45,11 +46,11 @@ type OnyxMethod = ValueOf<typeof METHOD>;
 const mergeQueue: Record<OnyxKey, Array<OnyxValue<OnyxKey>>> = {};
 const mergeQueuePromise: Record<OnyxKey, Promise<void>> = {};
 
-// Holds a mapping of all the react components that want their state subscribed to a store key
+// Holds a mapping of all the React components that want their state subscribed to a store key
 const callbackToStateMapping: Record<string, Mapping<OnyxKey>> = {};
 
 // Keeps a copy of the values of the onyx collection keys as a map for faster lookups
-let onyxCollectionKeyMap = new Map<OnyxKey, OnyxValue<OnyxKey>>();
+let onyxCollectionKeySet = new Set<OnyxKey>();
 
 // Holds a mapping of the connected key to the connectionID for faster lookups
 const onyxKeyToConnectionIDs = new Map();
@@ -115,11 +116,11 @@ function getDefaultKeyStates(): Record<OnyxKey, OnyxValue<OnyxKey>> {
 function initStoreValues(keys: DeepRecord<string, OnyxKey>, initialKeyStates: Partial<KeyValueMapping>, safeEvictionKeys: OnyxKey[]): void {
     // We need the value of the collection keys later for checking if a
     // key is a collection. We store it in a map for faster lookup.
-    const collectionValues = Object.values(keys.COLLECTION ?? {});
-    onyxCollectionKeyMap = collectionValues.reduce((acc, val) => {
-        acc.set(val, true);
+    const collectionValues = Object.values(keys.COLLECTION ?? {}) as string[];
+    onyxCollectionKeySet = collectionValues.reduce((acc, val) => {
+        acc.add(val);
         return acc;
-    }, new Map());
+    }, new Set<OnyxKey>());
 
     // Set our default key states to use when initializing and clearing Onyx data
     defaultKeyStates = initialKeyStates;
@@ -376,11 +377,18 @@ function getAllKeys(): Promise<Set<OnyxKey>> {
 }
 
 /**
- * Checks to see if the a subscriber's supplied key
+ * Returns set of all registered collection keys
+ */
+function getCollectionKeys(): Set<OnyxKey> {
+    return onyxCollectionKeySet;
+}
+
+/**
+ * Checks to see if the subscriber's supplied key
  * is associated with a collection of keys.
  */
 function isCollectionKey(key: OnyxKey): key is CollectionKeyBase {
-    return onyxCollectionKeyMap.has(key);
+    return onyxCollectionKeySet.has(key);
 }
 
 function isCollectionMemberKey<TCollectionKey extends CollectionKeyBase>(collectionKey: TCollectionKey, key: string): key is `${TCollectionKey}${string}` {
@@ -1182,6 +1190,34 @@ function initializeWithDefaultKeyStates(): Promise<void> {
     });
 }
 
+/**
+ * Validate the collection is not empty and has a correct type before applying mergeCollection()
+ */
+function isValidNonEmptyCollectionForMerge<TKey extends CollectionKeyBase, TMap>(collection: OnyxMergeCollectionInput<TKey, TMap>): boolean {
+    return typeof collection === 'object' && !Array.isArray(collection) && !utils.isEmptyObject(collection);
+}
+
+/**
+ * Verify if all the collection keys belong to the same parent
+ */
+function doAllCollectionItemsBelongToSameParent<TKey extends CollectionKeyBase>(collectionKey: TKey, collectionKeys: string[]): boolean {
+    let hasCollectionKeyCheckFailed = false;
+    collectionKeys.forEach((dataKey) => {
+        if (OnyxUtils.isKeyMatch(collectionKey, dataKey)) {
+            return;
+        }
+
+        if (process.env.NODE_ENV === 'development') {
+            throw new Error(`Provided collection doesn't have all its data belonging to the same parent. CollectionKey: ${collectionKey}, DataKey: ${dataKey}`);
+        }
+
+        hasCollectionKeyCheckFailed = true;
+        Logger.logAlert(`Provided collection doesn't have all its data belonging to the same parent. CollectionKey: ${collectionKey}, DataKey: ${dataKey}`);
+    });
+
+    return !hasCollectionKeyCheckFailed;
+}
+
 const OnyxUtils = {
     METHOD,
     getMergeQueue,
@@ -1194,6 +1230,7 @@ const OnyxUtils = {
     batchUpdates,
     get,
     getAllKeys,
+    getCollectionKeys,
     isCollectionKey,
     isCollectionMemberKey,
     splitCollectionMemberKey,
@@ -1227,6 +1264,8 @@ const OnyxUtils = {
     deleteKeyByConnections,
     getSnapshotKey,
     multiGet,
+    isValidNonEmptyCollectionForMerge,
+    doAllCollectionItemsBelongToSameParent,
 };
 
 export default OnyxUtils;
diff --git a/lib/types.ts b/lib/types.ts
@@ -432,6 +432,14 @@ type InitOptions = {
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 type GenericFunction = (...args: any[]) => any;
 
+/**
+ * Represents a combination of Merge and Set operations that should be executed in Onyx
+ */
+type MixedOperationsQueue = {
+    merge: OnyxInputKeyValueMapping;
+    set: OnyxInputKeyValueMapping;
+};
+
 export type {
     BaseConnectOptions,
     Collection,
@@ -468,4 +476,5 @@ export type {
     OnyxValue,
     Selector,
     WithOnyxConnectOptions,
+    MixedOperationsQueue,
 };
diff --git a/tests/unit/onyxTest.ts b/tests/unit/onyxTest.ts
