Try to synchronize operations within one update

lib/Onyx.ts

@@ -9,11 +9,13 @@ import utils from './utils';
 import DevTools from './DevTools';
 import type {
     Collection,
+    CollectionKey,
     CollectionKeyBase,
     ConnectOptions,
     InitOptions,
     KeyValueMapping,
     Mapping,
+    MixedOperationsQueue,
     NonUndefined,
     NullableKeyValueMapping,
     NullishDeep,
@@ -74,21 +76,21 @@ function init({
  *     callback: onSessionChange,
  * });
  *
- * @param mapping the mapping information to connect Onyx to the components state
- * @param mapping.key ONYXKEY to subscribe to
- * @param [mapping.statePropertyName] the name of the property in the state to connect the data to
- * @param [mapping.withOnyxInstance] whose setState() method will be called with any changed data
+ * @param connectOptions the mapping information to connect Onyx to the components state
+ * @param connectOptions.key ONYXKEY to subscribe to
+ * @param [connectOptions.statePropertyName] the name of the property in the state to connect the data to
+ * @param [connectOptions.withOnyxInstance] whose setState() method will be called with any changed data
  *      This is used by React components to connect to Onyx
- * @param [mapping.callback] a method that will be called with changed data
+ * @param [connectOptions.callback] a method that will be called with changed data
  *      This is used by any non-React code to connect to Onyx
- * @param [mapping.initWithStoredValues] If set to false, then no data will be prefilled into the
+ * @param [connectOptions.initWithStoredValues] If set to false, then no data will be prefilled into the
  *  component
- * @param [mapping.waitForCollectionCallback] If set to true, it will return the entire collection to the callback as a single object
- * @param [mapping.selector] THIS PARAM IS ONLY USED WITH withOnyx(). If included, this will be used to subscribe to a subset of an Onyx key's data.
+ * @param [connectOptions.waitForCollectionCallback] If set to true, it will return the entire collection to the callback as a single object
+ * @param [connectOptions.selector] THIS PARAM IS ONLY USED WITH withOnyx(). If included, this will be used to subscribe to a subset of an Onyx key's data.
  *       The sourceData and withOnyx state are passed to the selector and should return the simplified data. Using this setting on `withOnyx` can have very positive
  *       performance benefits because the component will only re-render when the subset of data changes. Otherwise, any change of data on any property would normally
  *       cause the component to re-render (and that can be expensive from a performance standpoint).
- * @param [mapping.initialValue] THIS PARAM IS ONLY USED WITH withOnyx().
+ * @param [connectOptions.initialValue] THIS PARAM IS ONLY USED WITH withOnyx().
  * If included, this will be passed to the component so that something can be rendered while data is being fetched from the DB.
  * Note that it will not cause the component to have the loading prop set to true.
  * @returns an ID to use when calling disconnect
@@ -182,7 +184,7 @@ function connect<TKey extends OnyxKey>(connectOptions: ConnectOptions<TKey>): nu
 }
 
 /**
- * Remove the listener for a react component
+ * Remove the listener for a React component
  * @example
  * Onyx.disconnect(connectionID);
  *
@@ -389,32 +391,13 @@ function merge<TKey extends OnyxKey>(key: TKey, changes: NonUndefined<OnyxEntry<
  * @param collection Object collection keyed by individual collection member keys and values
  */
 function mergeCollection<TKey extends CollectionKeyBase, TMap>(collectionKey: TKey, collection: Collection<TKey, TMap, NullishDeep<KeyValueMapping[TKey]>>): Promise<void> {
-    if (typeof collection !== 'object' || Array.isArray(collection) || utils.isEmptyObject(collection)) {
-        Logger.logInfo('mergeCollection() called with invalid or empty value. Skipping this update.');
-        return Promise.resolve();
-    }
-    const mergedCollection: NullableKeyValueMapping = collection;
-
-    // Confirm all the collection keys belong to the same parent
-    let hasCollectionKeyCheckFailed = false;
-    Object.keys(mergedCollection).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.isValidMergeCollection(collectionKey, collection)) {
         return Promise.resolve();
     }
 
+    const mergedCollection: NullableKeyValueMapping = collection;
+
     return OnyxUtils.getAllKeys()
         .then((persistedKeys) => {
             // Split to keys that exist in storage and keys that don't
@@ -616,23 +599,44 @@ 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));
+                if (OnyxUtils.isValidMergeCollection(key, value as Collection<CollectionKey, unknown, unknown>)) {
+                    Object.entries(value).forEach(([entryKey, entryValue]) => enqueueMergeOperation(entryKey, entryValue));
+                }
                 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();
@@ -642,6 +646,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(() => undefined);
 }
 

lib/OnyxUtils.ts

@@ -27,6 +27,8 @@ import type {
     OnyxEntry,
     KeyValueMapping,
     DefaultConnectCallback,
+    Collection,
+    NullishDeep,
 } from './types';
 import type Onyx from './Onyx';
 
@@ -45,11 +47,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 list of keys that have been directly subscribed to or recently modified from least to most recent
 let recentlyAccessedKeys: OnyxKey[] = [];
@@ -106,11 +108,11 @@ function getDefaultKeyStates(): Record<OnyxKey, OnyxValue<OnyxKey>> {
 function initStoreValues(keys: DeepRecord<string, OnyxKey>, initialKeyStates: Partial<NullableKeyValueMapping>, 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;
@@ -256,11 +258,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}` {
@@ -1072,6 +1081,33 @@ function initializeWithDefaultKeyStates(): Promise<void> {
     });
 }
 
+/**
+ * Verify if the collection is valid for merging into the collection key using mergeCollection()
+ */
+function isValidMergeCollection<TKey extends CollectionKeyBase, TMap>(collectionKey: TKey, collection: Collection<TKey, TMap, NullishDeep<KeyValueMapping[TKey]>>): boolean {
+    if (typeof collection !== 'object' || Array.isArray(collection) || utils.isEmptyObject(collection)) {
+        Logger.logInfo('mergeCollection() called with invalid or empty value. Skipping this update.');
+        return false;
+    }
+
+    // Confirm all the collection keys belong to the same parent
+    let hasCollectionKeyCheckFailed = false;
+    Object.keys(collection).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,
@@ -1084,6 +1120,7 @@ const OnyxUtils = {
     batchUpdates,
     get,
     getAllKeys,
+    getCollectionKeys,
     isCollectionKey,
     isCollectionMemberKey,
     splitCollectionMemberKey,
@@ -1112,6 +1149,7 @@ const OnyxUtils = {
     prepareKeyValuePairsForStorage,
     applyMerge,
     initializeWithDefaultKeyStates,
+    isValidMergeCollection,
 };
 
 export default OnyxUtils;

lib/types.ts

@@ -400,6 +400,14 @@ type InitOptions = {
     debugSetState?: boolean;
 };
 
+/**
+ * Represents a combination of Merge and Set operations that should be executed in Onyx
+ */
+type MixedOperationsQueue = {
+    merge: NullableKeyValueMapping;
+    set: NullableKeyValueMapping;
+};
+
 export type {
     BaseConnectOptions,
     Collection,
@@ -430,4 +438,5 @@ export type {
     WithOnyxConnectOptions,
     WithOnyxInstance,
     WithOnyxInstanceState,
+    MixedOperationsQueue,
 };