Update batching to be applied to useOnyx if they come from Onyx.update

lib/Onyx.ts

@@ -162,6 +162,13 @@ function disconnect(connection: Connection): void {
  * @param options optional configuration object
  */
 function set<TKey extends OnyxKey>(key: TKey, value: OnyxSetInput<TKey>, options?: SetOptions): Promise<void> {
+    return setInternal(key, value, options);
+}
+/**
+ * @param isFromUpdate - Whether this call originates from Onyx.update()
+ * When isFromUpdate = true (called from Onyx.update()), useOnyx hook subscribers are batched to escape excessive re-renders in components
+ */
+function setInternal<TKey extends OnyxKey>(key: TKey, value: OnyxSetInput<TKey>, options?: SetOptions, isFromUpdate = false): Promise<void> {
     // When we use Onyx.set to set a key we want to clear the current delta changes from Onyx.merge that were queued
     // before the value was set. If Onyx.merge is currently reading the old value from storage, it will then not apply the changes.
     if (OnyxUtils.hasPendingMergeForKey(key)) {
@@ -214,15 +221,15 @@ function set<TKey extends OnyxKey>(key: TKey, value: OnyxSetInput<TKey>, options
     OnyxUtils.logKeyChanged(OnyxUtils.METHOD.SET, key, value, hasChanged);
 
     // This approach prioritizes fast UI changes without waiting for data to be stored in device storage.
-    const updatePromise = OnyxUtils.broadcastUpdate(key, valueWithoutNestedNullValues, hasChanged);
+    const updatePromise = OnyxUtils.broadcastUpdate(key, valueWithoutNestedNullValues, hasChanged, isFromUpdate);
 
     // If the value has not changed or the key got removed, calling Storage.setItem() would be redundant and a waste of performance, so return early instead.
     if (!hasChanged) {
         return updatePromise;
     }
 
     return Storage.setItem(key, valueWithoutNestedNullValues)
-        .catch((error) => OnyxUtils.evictStorageAndRetry(error, set, key, valueWithoutNestedNullValues))
+        .catch((error) => OnyxUtils.evictStorageAndRetry(error, setInternal, key, valueWithoutNestedNullValues, undefined, isFromUpdate))
         .then(() => {
             OnyxUtils.sendActionToDevTools(OnyxUtils.METHOD.SET, key, valueWithoutNestedNullValues);
             return updatePromise;
@@ -237,6 +244,13 @@ function set<TKey extends OnyxKey>(key: TKey, value: OnyxSetInput<TKey>, options
  * @param data object keyed by ONYXKEYS and the values to set
  */
 function multiSet(data: OnyxMultiSetInput): Promise<void> {
+    return multiSetInternal(data);
+}
+/**
+ * @param isFromUpdate - Whether this call originates from Onyx.update()
+ * When isFromUpdate = true (called from Onyx.update()), useOnyx hook subscribers are batched to escape excessive re-renders in components
+ */
+function multiSetInternal(data: OnyxMultiSetInput, isFromUpdate = false): Promise<void> {
     let newData = data;
 
     const skippableCollectionMemberIDs = OnyxUtils.getSkippableCollectionMemberIDs();
@@ -269,11 +283,11 @@ function multiSet(data: OnyxMultiSetInput): Promise<void> {
 
         // Update cache and optimistically inform subscribers on the next tick
         cache.set(key, value);
-        return OnyxUtils.scheduleSubscriberUpdate(key, value, prevValue);
+        return OnyxUtils.scheduleSubscriberUpdate(key, value, prevValue, undefined, isFromUpdate);
     });
 
     return Storage.multiSet(keyValuePairsToSet)
-        .catch((error) => OnyxUtils.evictStorageAndRetry(error, multiSet, newData))
+        .catch((error) => OnyxUtils.evictStorageAndRetry(error, multiSetInternal, newData, isFromUpdate))
         .then(() => {
             OnyxUtils.sendActionToDevTools(OnyxUtils.METHOD.MULTI_SET, undefined, newData);
             return Promise.all(updatePromises);
@@ -298,6 +312,13 @@ function multiSet(data: OnyxMultiSetInput): Promise<void> {
  * Onyx.merge(ONYXKEYS.POLICY, {name: 'My Workspace'}); // -> {id: 1, name: 'My Workspace'}
  */
 function merge<TKey extends OnyxKey>(key: TKey, changes: OnyxMergeInput<TKey>): Promise<void> {
+    return mergeInternal(key, changes);
+}
+/**
+ * @param isFromUpdate - Whether this call originates from Onyx.update()
+ * When isFromUpdate = true (called from Onyx.update()), useOnyx hook subscribers are batched to escape excessive re-renders in components
+ */
+function mergeInternal<TKey extends OnyxKey>(key: TKey, changes: OnyxMergeInput<TKey>, isFromUpdate = false): Promise<void> {
     const skippableCollectionMemberIDs = OnyxUtils.getSkippableCollectionMemberIDs();
     if (skippableCollectionMemberIDs.size) {
         try {
@@ -360,7 +381,7 @@ function merge<TKey extends OnyxKey>(key: TKey, changes: OnyxMergeInput<TKey>):
                 return Promise.resolve();
             }
 
-            return OnyxMerge.applyMerge(key, existingValue, validChanges).then(({mergedValue, updatePromise}) => {
+            return OnyxMerge.applyMerge(key, existingValue, validChanges, isFromUpdate).then(({mergedValue, updatePromise}) => {
                 OnyxUtils.sendActionToDevTools(OnyxUtils.METHOD.MERGE, key, changes, mergedValue);
                 return updatePromise;
             });
@@ -387,7 +408,14 @@ 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> {
-    return OnyxUtils.mergeCollectionWithPatches(collectionKey, collection);
+    return mergeCollectionInternal(collectionKey, collection);
+}
+/**
+ * @param isFromUpdate - Whether this call originates from Onyx.update()
+ * When isFromUpdate = true (called from Onyx.update()), useOnyx hook subscribers are batched to escape excessive re-renders in components
+ */
+function mergeCollectionInternal<TKey extends CollectionKeyBase, TMap>(collectionKey: TKey, collection: OnyxMergeCollectionInput<TKey, TMap>, isFromUpdate = false): Promise<void> {
+    return OnyxUtils.mergeCollectionWithPatches(collectionKey, collection, undefined, isFromUpdate);
 }
 
 /**
@@ -476,10 +504,10 @@ function clear(keysToPreserve: OnyxKey[] = []): Promise<void> {
 
             // Notify the subscribers for each key/value group so they can receive the new values
             Object.entries(keyValuesToResetIndividually).forEach(([key, value]) => {
-                updatePromises.push(OnyxUtils.scheduleSubscriberUpdate(key, value, cache.get(key, false)));
+                updatePromises.push(OnyxUtils.scheduleSubscriberUpdate(key, value, cache.get(key, false), undefined, false));
             });
             Object.entries(keyValuesToResetAsCollection).forEach(([key, value]) => {
-                updatePromises.push(OnyxUtils.scheduleNotifyCollectionSubscribers(key, value));
+                updatePromises.push(OnyxUtils.scheduleNotifyCollectionSubscribers(key, value, undefined, false));
             });
 
             const defaultKeyValuePairs = Object.entries(
@@ -570,7 +598,7 @@ function update(data: OnyxUpdate[]): Promise<void> {
                     collectionKeys.forEach((collectionKey) => enqueueMergeOperation(collectionKey, mergedCollection[collectionKey]));
                 }
             },
-            [OnyxUtils.METHOD.SET_COLLECTION]: (k, v) => promises.push(() => setCollection(k, v as Collection<CollectionKey, unknown, unknown>)),
+            [OnyxUtils.METHOD.SET_COLLECTION]: (k, v) => promises.push(() => setCollectionInternal(k, v as Collection<CollectionKey, unknown, unknown>, true)),
             [OnyxUtils.METHOD.MULTI_SET]: (k, v) => Object.entries(v as Partial<OnyxInputKeyValueMapping>).forEach(([entryKey, entryValue]) => enqueueSetOperation(entryKey, entryValue)),
             [OnyxUtils.METHOD.CLEAR]: () => {
                 clearPromise = clear();
@@ -625,27 +653,28 @@ function update(data: OnyxUpdate[]): Promise<void> {
                     collectionKey,
                     batchedCollectionUpdates.merge as Collection<CollectionKey, unknown, unknown>,
                     batchedCollectionUpdates.mergeReplaceNullPatches,
+                    true,
                 ),
             );
         }
         if (!utils.isEmptyObject(batchedCollectionUpdates.set)) {
-            promises.push(() => multiSet(batchedCollectionUpdates.set));
+            promises.push(() => multiSetInternal(batchedCollectionUpdates.set, true));
         }
     });
 
     Object.entries(updateQueue).forEach(([key, operations]) => {
         if (operations[0] === null) {
             const batchedChanges = OnyxUtils.mergeChanges(operations).result;
-            promises.push(() => set(key, batchedChanges));
+            promises.push(() => setInternal(key, batchedChanges, undefined, true));
             return;
         }
 
         operations.forEach((operation) => {
-            promises.push(() => merge(key, operation));
+            promises.push(() => mergeInternal(key, operation, true));
         });
     });
 
-    const snapshotPromises = OnyxUtils.updateSnapshots(data, merge);
+    const snapshotPromises = OnyxUtils.updateSnapshots(data, (key, changes) => mergeInternal(key, changes, true));
 
     // We need to run the snapshot updates before the other updates so the snapshot data can be updated before the loading state in the snapshot
     const finalPromises = snapshotPromises.concat(promises);
@@ -667,6 +696,13 @@ function update(data: OnyxUpdate[]): Promise<void> {
  * @param collection Object collection keyed by individual collection member keys and values
  */
 function setCollection<TKey extends CollectionKeyBase, TMap>(collectionKey: TKey, collection: OnyxMergeCollectionInput<TKey, TMap>): Promise<void> {
+    return setCollectionInternal(collectionKey, collection);
+}
+/**
+ * @param isFromUpdate - Whether this call originates from Onyx.update()
+ * When isFromUpdate = true (called from Onyx.update()), useOnyx hook subscribers are batched to escape excessive re-renders in components
+ */
+function setCollectionInternal<TKey extends CollectionKeyBase, TMap>(collectionKey: TKey, collection: OnyxMergeCollectionInput<TKey, TMap>, isFromUpdate = false): Promise<void> {
     let resultCollection: OnyxInputKeyValueMapping = collection;
     let resultCollectionKeys = Object.keys(resultCollection);
 
@@ -714,10 +750,10 @@ function setCollection<TKey extends CollectionKeyBase, TMap>(collectionKey: TKey
 
         keyValuePairs.forEach(([key, value]) => cache.set(key, value));
 
-        const updatePromise = OnyxUtils.scheduleNotifyCollectionSubscribers(collectionKey, mutableCollection, previousCollection);
+        const updatePromise = OnyxUtils.scheduleNotifyCollectionSubscribers(collectionKey, mutableCollection, previousCollection, isFromUpdate);
 
         return Storage.multiSet(keyValuePairs)
-            .catch((error) => OnyxUtils.evictStorageAndRetry(error, setCollection, collectionKey, collection))
+            .catch((error) => OnyxUtils.evictStorageAndRetry(error, setCollectionInternal, collectionKey, collection, isFromUpdate))
             .then(() => {
                 OnyxUtils.sendActionToDevTools(OnyxUtils.METHOD.SET_COLLECTION, undefined, mutableCollection);
                 return updatePromise;

lib/OnyxMerge/index.native.ts

@@ -8,6 +8,7 @@ const applyMerge: ApplyMerge = <TKey extends OnyxKey, TValue extends OnyxInput<T
     key: TKey,
     existingValue: TValue,
     validChanges: TChange[],
+    isFromUpdate = false,
 ) => {
     // If any of the changes is null, we need to discard the existing value.
     const baseValue = validChanges.includes(null as TChange) ? undefined : existingValue;
@@ -26,7 +27,7 @@ const applyMerge: ApplyMerge = <TKey extends OnyxKey, TValue extends OnyxInput<T
     OnyxUtils.logKeyChanged(OnyxUtils.METHOD.MERGE, key, mergedValue, hasChanged);
 
     // This approach prioritizes fast UI changes without waiting for data to be stored in device storage.
-    const updatePromise = OnyxUtils.broadcastUpdate(key, mergedValue as OnyxValue<TKey>, hasChanged);
+    const updatePromise = OnyxUtils.broadcastUpdate(key, mergedValue as OnyxValue<TKey>, hasChanged, isFromUpdate);
 
     // If the value has not changed, calling Storage.setItem() would be redundant and a waste of performance, so return early instead.
     if (!hasChanged) {

lib/OnyxMerge/index.ts

@@ -8,6 +8,7 @@ const applyMerge: ApplyMerge = <TKey extends OnyxKey, TValue extends OnyxInput<T
     key: TKey,
     existingValue: TValue,
     validChanges: TChange[],
+    isFromUpdate = false,
 ) => {
     const {result: mergedValue} = OnyxUtils.mergeChanges(validChanges, existingValue);
 
@@ -18,7 +19,7 @@ const applyMerge: ApplyMerge = <TKey extends OnyxKey, TValue extends OnyxInput<T
     OnyxUtils.logKeyChanged(OnyxUtils.METHOD.MERGE, key, mergedValue, hasChanged);
 
     // This approach prioritizes fast UI changes without waiting for data to be stored in device storage.
-    const updatePromise = OnyxUtils.broadcastUpdate(key, mergedValue as OnyxValue<TKey>, hasChanged);
+    const updatePromise = OnyxUtils.broadcastUpdate(key, mergedValue as OnyxValue<TKey>, hasChanged, isFromUpdate);
 
     // If the value has not changed, calling Storage.setItem() would be redundant and a waste of performance, so return early instead.
     if (!hasChanged) {

lib/OnyxMerge/types.ts

@@ -9,6 +9,7 @@ type ApplyMerge = <TKey extends OnyxKey, TValue extends OnyxInput<OnyxKey> | und
     key: TKey,
     existingValue: TValue,
     validChanges: TChange[],
+    isFromUpdate?: boolean,
 ) => Promise<ApplyMergeResult<TChange>>;
 
 export type {ApplyMerge, ApplyMergeResult};

lib/OnyxUtils.ts

@@ -622,6 +622,7 @@ function keysChanged<TKey extends CollectionKeyBase>(
     partialPreviousCollection: OnyxCollection<KeyValueMapping[TKey]> | undefined,
     notifyConnectSubscribers = true,
     notifyWithOnyxSubscribers = true,
+    notifyUseOnyxHookSubscribers = true,
 ): void {
     // We prepare the "cached collection" which is the entire collection + the new partial data that
     // was merged in via mergeCollection().
@@ -657,7 +658,8 @@ function keysChanged<TKey extends CollectionKeyBase>(
 
         // Regular Onyx.connect() subscriber found.
         if (typeof subscriber.callback === 'function') {
-            if (!notifyConnectSubscribers) {
+            // Check if it's a useOnyx or a regular Onyx.connect() subscriber
+            if ((subscriber.isUseOnyxSubscriber && !notifyUseOnyxHookSubscribers) || (!subscriber.isUseOnyxSubscriber && !notifyConnectSubscribers)) {
                 continue;
             }
 
@@ -816,6 +818,7 @@ function keyChanged<TKey extends OnyxKey>(
     canUpdateSubscriber: (subscriber?: Mapping<OnyxKey>) => boolean = () => true,
     notifyConnectSubscribers = true,
     notifyWithOnyxSubscribers = true,
+    notifyUseOnyxHookSubscribers = true,
 ): void {
     // Add or remove this key from the recentlyAccessedKeys lists
     if (value !== null) {
@@ -857,7 +860,8 @@ function keyChanged<TKey extends OnyxKey>(
 
         // Subscriber is a regular call to connect() and provided a callback
         if (typeof subscriber.callback === 'function') {
-            if (!notifyConnectSubscribers) {
+            // Check if it's a useOnyx or a regular Onyx.connect() subscriber
+            if ((subscriber.isUseOnyxSubscriber && !notifyUseOnyxHookSubscribers) || (!subscriber.isUseOnyxSubscriber && !notifyConnectSubscribers)) {
                 continue;
             }
             if (lastConnectionCallbackData.has(subscriber.subscriptionID) && lastConnectionCallbackData.get(subscriber.subscriptionID) === value) {
@@ -1077,9 +1081,10 @@ function scheduleSubscriberUpdate<TKey extends OnyxKey>(
     value: OnyxValue<TKey>,
     previousValue: OnyxValue<TKey>,
     canUpdateSubscriber: (subscriber?: Mapping<OnyxKey>) => boolean = () => true,
+    isFromUpdate = false,
 ): Promise<void> {
-    const promise = Promise.resolve().then(() => keyChanged(key, value, previousValue, canUpdateSubscriber, true, false));
-    batchUpdates(() => keyChanged(key, value, previousValue, canUpdateSubscriber, false, true));
+    const promise = Promise.resolve().then(() => keyChanged(key, value, previousValue, canUpdateSubscriber, true, false, !isFromUpdate));
+    batchUpdates(() => keyChanged(key, value, previousValue, canUpdateSubscriber, false, true, isFromUpdate));
     return Promise.all([maybeFlushBatchUpdates(), promise]).then(() => undefined);
 }
 
@@ -1092,9 +1097,10 @@ function scheduleNotifyCollectionSubscribers<TKey extends OnyxKey>(
     key: TKey,
     value: OnyxCollection<KeyValueMapping[TKey]>,
     previousValue?: OnyxCollection<KeyValueMapping[TKey]>,
+    isFromUpdate = false,
 ): Promise<void> {
-    const promise = Promise.resolve().then(() => keysChanged(key, value, previousValue, true, false));
-    batchUpdates(() => keysChanged(key, value, previousValue, false, true));
+    const promise = Promise.resolve().then(() => keysChanged(key, value, previousValue, true, false, !isFromUpdate));
+    batchUpdates(() => keysChanged(key, value, previousValue, false, true, isFromUpdate));
     return Promise.all([maybeFlushBatchUpdates(), promise]).then(() => undefined);
 }
 
@@ -1156,7 +1162,7 @@ function evictStorageAndRetry<TMethod extends typeof Onyx.set | typeof Onyx.mult
 /**
  * Notifies subscribers and writes current value to cache
  */
-function broadcastUpdate<TKey extends OnyxKey>(key: TKey, value: OnyxValue<TKey>, hasChanged?: boolean): Promise<void> {
+function broadcastUpdate<TKey extends OnyxKey>(key: TKey, value: OnyxValue<TKey>, hasChanged?: boolean, isFromUpdate = false): Promise<void> {
     const prevValue = cache.get(key, false) as OnyxValue<TKey>;
 
     // Update subscribers if the cached value has changed, or when the subscriber specifically requires
@@ -1167,7 +1173,7 @@ function broadcastUpdate<TKey extends OnyxKey>(key: TKey, value: OnyxValue<TKey>
         cache.addToAccessedKeys(key);
     }
 
-    return scheduleSubscriberUpdate(key, value, prevValue, (subscriber) => hasChanged || subscriber?.initWithStoredValues === false).then(() => undefined);
+    return scheduleSubscriberUpdate(key, value, prevValue, (subscriber) => hasChanged || subscriber?.initWithStoredValues === false, isFromUpdate).then(() => undefined);
 }
 
 function hasPendingMergeForKey(key: OnyxKey): boolean {
@@ -1514,6 +1520,7 @@ function mergeCollectionWithPatches<TKey extends CollectionKeyBase, TMap>(
     collectionKey: TKey,
     collection: OnyxMergeCollectionInput<TKey, TMap>,
     mergeReplaceNullPatches?: MultiMergeReplaceNullPatches,
+    isFromUpdate = false,
 ): Promise<void> {
     if (!isValidNonEmptyCollectionForMerge(collection)) {
         Logger.logInfo('mergeCollection() called with invalid or empty value. Skipping this update.');
@@ -1614,7 +1621,7 @@ function mergeCollectionWithPatches<TKey extends CollectionKeyBase, TMap>(
             // and update all subscribers
             const promiseUpdate = previousCollectionPromise.then((previousCollection) => {
                 cache.merge(finalMergedCollection);
-                return scheduleNotifyCollectionSubscribers(collectionKey, finalMergedCollection, previousCollection);
+                return scheduleNotifyCollectionSubscribers(collectionKey, finalMergedCollection, previousCollection, isFromUpdate);
             });
 
             return Promise.all(promises)

lib/types.ts

@@ -273,6 +273,9 @@ type BaseConnectOptions = {
      * with the same connect configurations.
      */
     reuseConnection?: boolean;
+
+    /** Indicates whether this subscriber is created from the useOnyx hook. */
+    isUseOnyxSubscriber?: boolean;
 };
 
 /** Represents the callback function used in `Onyx.connect()` method with a regular key. */

lib/useOnyx.ts

@@ -359,6 +359,7 @@ function useOnyx<TKey extends OnyxKey, TReturnValue = OnyxValue<TKey>>(
                 initWithStoredValues: options?.initWithStoredValues,
                 waitForCollectionCallback: OnyxUtils.isCollectionKey(key) as true,
                 reuseConnection: options?.reuseConnection,
+                isUseOnyxSubscriber: true,
             });
 
             checkEvictableKey();