[POC] Remove macro/micro tasks during subscriber update

lib/Onyx.ts

@@ -254,9 +254,8 @@ 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).then(({mergedValue}) => {
                     OnyxUtils.sendActionToDevTools(OnyxUtils.METHOD.MERGE, key, changes, mergedValue);
-                    return updatePromise;
                 });
             } catch (error) {
                 Logger.logAlert(`An error occurred while applying merge for key: ${key}, Error: ${error}`);
@@ -374,16 +373,6 @@ function clear(keysToPreserve: OnyxKey[] = []): Promise<void> {
                     keysToBeClearedFromStorage.push(key);
                 }
 
-                const updatePromises: Array<Promise<void>> = [];
-
-                // Notify the subscribers for each key/value group so they can receive the new values
-                for (const [key, value] of Object.entries(keyValuesToResetIndividually)) {
-                    updatePromises.push(OnyxUtils.scheduleSubscriberUpdate(key, value));
-                }
-                for (const [key, value] of Object.entries(keyValuesToResetAsCollection)) {
-                    updatePromises.push(OnyxUtils.scheduleNotifyCollectionSubscribers(key, value.newValues, value.oldValues));
-                }
-
                 // Exclude RAM-only keys to prevent them from being saved to storage
                 const defaultKeyValuePairs = Object.entries(
                     Object.keys(defaultKeyStates)
@@ -402,7 +391,14 @@ function clear(keysToPreserve: OnyxKey[] = []): Promise<void> {
                     .then(() => Storage.multiSet(defaultKeyValuePairs))
                     .then(() => {
                         DevTools.clearState(keysToPreserve);
-                        return Promise.all(updatePromises);
+
+                        // Notify the subscribers for each key/value group so they can receive the new values
+                        for (const [key, value] of Object.entries(keyValuesToResetIndividually)) {
+                            OnyxUtils.keyChanged(key, value);
+                        }
+                        for (const [key, value] of Object.entries(keyValuesToResetAsCollection)) {
+                            OnyxUtils.keysChanged(key, value.newValues, value.oldValues);
+                        }
                     });
             })
             .then(() => undefined);

lib/OnyxMerge/index.native.ts

@@ -26,21 +26,20 @@ 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);
+    OnyxUtils.broadcastUpdate(key, mergedValue as OnyxValue<TKey>, hasChanged);
 
     const shouldSkipStorageOperations = !hasChanged || OnyxUtils.isRamOnlyKey(key);
 
     // If the value has not changed, calling Storage.setItem() would be redundant and a waste of performance, so return early instead.
     // If the key is marked as RAM-only, it should not be saved nor updated in the storage.
     if (shouldSkipStorageOperations) {
-        return Promise.resolve({mergedValue, updatePromise});
+        return Promise.resolve({mergedValue});
     }
 
     // For native platforms we use `mergeItem` that will take advantage of JSON_PATCH and JSON_REPLACE SQL operations to
     // merge the object in a performant way.
     return Storage.mergeItem(key, batchedChanges as OnyxValue<TKey>, replaceNullPatches).then(() => ({
         mergedValue,
-        updatePromise,
     }));
 };
 

lib/OnyxMerge/index.ts

@@ -18,20 +18,19 @@ 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);
+    OnyxUtils.broadcastUpdate(key, mergedValue as OnyxValue<TKey>, hasChanged);
 
     const shouldSkipStorageOperations = !hasChanged || OnyxUtils.isRamOnlyKey(key);
 
     // If the value has not changed, calling Storage.setItem() would be redundant and a waste of performance, so return early instead.
     // If the key is marked as RAM-only, it should not be saved nor updated in the storage.
     if (shouldSkipStorageOperations) {
-        return Promise.resolve({mergedValue, updatePromise});
+        return Promise.resolve({mergedValue});
     }
 
     // For web platforms we use `setItem` since the object was already merged with its changes before.
     return Storage.setItem(key, mergedValue as OnyxValue<TKey>).then(() => ({
         mergedValue,
-        updatePromise,
     }));
 };
 

lib/OnyxMerge/types.ts

@@ -2,7 +2,6 @@ import type {OnyxInput, OnyxKey} from '../types';
 
 type ApplyMergeResult<TValue> = {
     mergedValue: TValue;
-    updatePromise: Promise<void>;
 };
 
 type ApplyMerge = <TKey extends OnyxKey, TValue extends OnyxInput<OnyxKey> | undefined, TChange extends OnyxInput<OnyxKey> | null>(

lib/OnyxUtils.ts

@@ -74,9 +74,6 @@ type OnyxMethod = ValueOf<typeof METHOD>;
 let mergeQueue: Record<OnyxKey, Array<OnyxValue<OnyxKey>>> = {};
 let mergeQueuePromise: Record<OnyxKey, Promise<void>> = {};
 
-// Used to schedule subscriber update to the macro task queue
-let nextMacrotaskPromise: Promise<void> | null = null;
-
 // Holds a mapping of all the React components that want their state subscribed to a store key
 let callbackToStateMapping: Record<string, CallbackToStateMapping<OnyxKey>> = {};
 
@@ -798,6 +795,7 @@ function keyChanged<TKey extends OnyxKey>(
                 }
 
                 cachedCollection[key] = value;
+                lastConnectionCallbackData.set(subscriber.subscriptionID, cachedCollection);
                 subscriber.callback(cachedCollection, subscriber.key, {[key]: value});
                 continue;
             }
@@ -825,11 +823,10 @@ function sendDataToConnection<TKey extends OnyxKey>(mapping: CallbackToStateMapp
 
     // For regular callbacks, we never want to pass null values, but always just undefined if a value is not set in cache or storage.
     const valueToPass = value === null ? undefined : value;
-    const lastValue = lastConnectionCallbackData.get(mapping.subscriptionID);
-    lastConnectionCallbackData.get(mapping.subscriptionID);
 
-    // If the value has not changed we do not need to trigger the callback
-    if (lastConnectionCallbackData.has(mapping.subscriptionID) && valueToPass === lastValue) {
+    // If the subscriber was already notified (e.g. by a synchronous keyChanged call),
+    // skip the initial data delivery to prevent duplicate callbacks.
+    if (lastConnectionCallbackData.has(mapping.subscriptionID)) {
         return;
     }
 
@@ -862,57 +859,12 @@ function getCollectionDataAndSendAsObject<TKey extends OnyxKey>(matchingKeys: Co
     });
 }
 
-/**
- * Delays promise resolution until the next macrotask to prevent race condition if the key subscription is in progress.
- *
- * @param callback The keyChanged/keysChanged callback
- * */
-function prepareSubscriberUpdate(callback: () => void): Promise<void> {
-    if (!nextMacrotaskPromise) {
-        nextMacrotaskPromise = new Promise<void>((resolve) => {
-            setTimeout(() => {
-                nextMacrotaskPromise = null;
-                resolve();
-            }, 0);
-        });
-    }
-    return Promise.all([nextMacrotaskPromise, Promise.resolve().then(callback)]).then();
-}
-
-/**
- * Schedules an update that will be appended to the macro task queue (so it doesn't update the subscribers immediately).
- *
- * @example
- * scheduleSubscriberUpdate(key, value, subscriber => subscriber.initWithStoredValues === false)
- */
-function scheduleSubscriberUpdate<TKey extends OnyxKey>(
-    key: TKey,
-    value: OnyxValue<TKey>,
-    canUpdateSubscriber: (subscriber?: CallbackToStateMapping<OnyxKey>) => boolean = () => true,
-    isProcessingCollectionUpdate = false,
-): Promise<void> {
-    return prepareSubscriberUpdate(() => keyChanged(key, value, canUpdateSubscriber, isProcessingCollectionUpdate));
-}
-
-/**
- * This method is similar to scheduleSubscriberUpdate but it is built for working specifically with collections
- * so that keysChanged() is triggered for the collection and not keyChanged(). If this was not done, then the
- * subscriber callbacks receive the data in a different format than they normally expect and it breaks code.
- */
-function scheduleNotifyCollectionSubscribers<TKey extends OnyxKey>(
-    key: TKey,
-    value: OnyxCollection<KeyValueMapping[TKey]>,
-    previousValue?: OnyxCollection<KeyValueMapping[TKey]>,
-): Promise<void> {
-    return prepareSubscriberUpdate(() => keysChanged(key, value, previousValue));
-}
-
 /**
  * Remove a key from Onyx and update the subscribers
  */
 function remove<TKey extends OnyxKey>(key: TKey, isProcessingCollectionUpdate?: boolean): Promise<void> {
     cache.drop(key);
-    scheduleSubscriberUpdate(key, undefined as OnyxValue<TKey>, undefined, isProcessingCollectionUpdate);
+    keyChanged(key, undefined as OnyxValue<TKey>, undefined, isProcessingCollectionUpdate);
 
     if (isRamOnlyKey(key)) {
         return Promise.resolve();
@@ -983,7 +935,7 @@ function retryOperation<TMethod extends RetriableOnyxOperation>(error: Error, on
 /**
  * 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): void {
     // Update subscribers if the cached value has changed, or when the subscriber specifically requires
     // all updates regardless of value changes (indicated by initWithStoredValues set to false).
     if (hasChanged) {
@@ -992,7 +944,7 @@ function broadcastUpdate<TKey extends OnyxKey>(key: TKey, value: OnyxValue<TKey>
         cache.addToAccessedKeys(key);
     }
 
-    return scheduleSubscriberUpdate(key, value, (subscriber) => hasChanged || subscriber?.initWithStoredValues === false).then(() => undefined);
+    keyChanged(key, value, (subscriber) => hasChanged || subscriber?.initWithStoredValues === false);
 }
 
 function hasPendingMergeForKey(key: OnyxKey): boolean {
@@ -1386,24 +1338,23 @@ function setWithRetry<TKey extends OnyxKey>({key, value, options}: SetParams<TKe
     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);
+    OnyxUtils.broadcastUpdate(key, valueWithoutNestedNullValues, hasChanged);
 
     // If the value has not changed and this isn't a retry attempt, calling Storage.setItem() would be redundant and a waste of performance, so return early instead.
     if (!hasChanged && !retryAttempt) {
-        return updatePromise;
+        return Promise.resolve();
     }
 
     // If a key is a RAM-only key or a member of RAM-only collection, we skip the step that modifies the storage
     if (isRamOnlyKey(key)) {
         OnyxUtils.sendActionToDevTools(OnyxUtils.METHOD.SET, key, valueWithoutNestedNullValues);
-        return updatePromise;
+        return Promise.resolve();
     }
 
     return Storage.setItem(key, valueWithoutNestedNullValues)
         .catch((error) => OnyxUtils.retryOperation(error, setWithRetry, {key, value: valueWithoutNestedNullValues, options}, retryAttempt))
         .then(() => {
             OnyxUtils.sendActionToDevTools(OnyxUtils.METHOD.SET, key, valueWithoutNestedNullValues);
-            return updatePromise;
         });
 }
 
@@ -1446,7 +1397,7 @@ function multiSetWithRetry(data: OnyxMultiSetInput, retryAttempt?: number): Prom
 
         // Update cache and optimistically inform subscribers on the next tick
         cache.set(key, value);
-        return OnyxUtils.scheduleSubscriberUpdate(key, value);
+        return OnyxUtils.keyChanged(key, value);
     });
 
     const keyValuePairsToStore = keyValuePairsToSet.filter((keyValuePair) => {
@@ -1522,7 +1473,7 @@ function setCollectionWithRetry<TKey extends CollectionKeyBase>({collectionKey,
 
         for (const [key, value] of keyValuePairs) cache.set(key, value);
 
-        const updatePromise = OnyxUtils.scheduleNotifyCollectionSubscribers(collectionKey, mutableCollection, previousCollection);
+        const updatePromise = OnyxUtils.keysChanged(collectionKey, mutableCollection, previousCollection);
 
         // RAM-only keys are not supposed to be saved to storage
         if (isRamOnlyKey(collectionKey)) {
@@ -1657,7 +1608,7 @@ function mergeCollectionWithPatches<TKey extends CollectionKeyBase>(
             // and update all subscribers
             const promiseUpdate = previousCollectionPromise.then((previousCollection) => {
                 cache.merge(finalMergedCollection);
-                return scheduleNotifyCollectionSubscribers(collectionKey, finalMergedCollection, previousCollection);
+                return keysChanged(collectionKey, finalMergedCollection, previousCollection);
             });
 
             return Promise.all(promises)
@@ -1723,7 +1674,7 @@ function partialSetCollection<TKey extends CollectionKeyBase>({collectionKey, co
 
         for (const [key, value] of keyValuePairs) cache.set(key, value);
 
-        const updatePromise = scheduleNotifyCollectionSubscribers(collectionKey, mutableCollection, previousCollection);
+        const updatePromise = keysChanged(collectionKey, mutableCollection, previousCollection);
 
         if (isRamOnlyKey(collectionKey)) {
             sendActionToDevTools(METHOD.SET_COLLECTION, undefined, mutableCollection);
@@ -1782,8 +1733,6 @@ const OnyxUtils = {
     sendDataToConnection,
     getCollectionKey,
     getCollectionDataAndSendAsObject,
-    scheduleSubscriberUpdate,
-    scheduleNotifyCollectionSubscribers,
     remove,
     reportStorageQuota,
     retryOperation,
@@ -1840,10 +1789,6 @@ GlobalSettings.addGlobalSettingsChangeListener(({enablePerformanceMetrics}) => {
     // @ts-expect-error Reassign
     sendDataToConnection = decorateWithMetrics(sendDataToConnection, 'OnyxUtils.sendDataToConnection');
     // @ts-expect-error Reassign
-    scheduleSubscriberUpdate = decorateWithMetrics(scheduleSubscriberUpdate, 'OnyxUtils.scheduleSubscriberUpdate');
-    // @ts-expect-error Reassign
-    scheduleNotifyCollectionSubscribers = decorateWithMetrics(scheduleNotifyCollectionSubscribers, 'OnyxUtils.scheduleNotifyCollectionSubscribers');
-    // @ts-expect-error Reassign
     remove = decorateWithMetrics(remove, 'OnyxUtils.remove');
     // @ts-expect-error Reassign
     reportStorageQuota = decorateWithMetrics(reportStorageQuota, 'OnyxUtils.reportStorageQuota');

tests/perf-test/OnyxUtils.perf-test.ts

@@ -431,66 +431,6 @@ describe('OnyxUtils', () => {
         });
     });
 
-    describe('scheduleSubscriberUpdate', () => {
-        test('10k calls scheduling updates', async () => {
-            const subscriptionMap = new Map<string, number>();
-
-            const changedReportActions = Object.fromEntries(
-                Object.entries(mockedReportActionsMap).map(([k, v]) => [k, createRandomReportAction(Number(v.reportActionID))] as const),
-            ) as GenericCollection;
-
-            await measureAsyncFunction(() => Promise.all(Object.entries(changedReportActions).map(([key, value]) => OnyxUtils.scheduleSubscriberUpdate(key, value))), {
-                beforeEach: async () => {
-                    await Onyx.multiSet(mockedReportActionsMap);
-                    for (const key of mockedReportActionsKeys) {
-                        const id = OnyxUtils.subscribeToKey({key, callback: jest.fn(), initWithStoredValues: false});
-                        subscriptionMap.set(key, id);
-                    }
-                },
-                afterEach: async () => {
-                    for (const key of mockedReportActionsKeys) {
-                        const id = subscriptionMap.get(key);
-                        if (id) {
-                            OnyxUtils.unsubscribeFromKey(id);
-                        }
-                    }
-                    subscriptionMap.clear();
-                    await clearOnyxAfterEachMeasure();
-                },
-            });
-        });
-    });
-
-    describe('scheduleNotifyCollectionSubscribers', () => {
-        test('one call with 10k heavy objects to update 10k subscribers', async () => {
-            const subscriptionMap = new Map<string, number>();
-
-            const changedReportActions = Object.fromEntries(
-                Object.entries(mockedReportActionsMap).map(([k, v]) => [k, createRandomReportAction(Number(v.reportActionID))] as const),
-            ) as GenericCollection;
-
-            await measureAsyncFunction(() => OnyxUtils.scheduleNotifyCollectionSubscribers(collectionKey, changedReportActions, mockedReportActionsMap), {
-                beforeEach: async () => {
-                    await Onyx.multiSet(mockedReportActionsMap);
-                    for (const key of mockedReportActionsKeys) {
-                        const id = OnyxUtils.subscribeToKey({key, callback: jest.fn(), initWithStoredValues: false});
-                        subscriptionMap.set(key, id);
-                    }
-                },
-                afterEach: async () => {
-                    for (const key of mockedReportActionsKeys) {
-                        const id = subscriptionMap.get(key);
-                        if (id) {
-                            OnyxUtils.unsubscribeFromKey(id);
-                        }
-                    }
-                    subscriptionMap.clear();
-                    await clearOnyxAfterEachMeasure();
-                },
-            });
-        });
-    });
-
     describe('remove', () => {
         test('10k calls', async () => {
             await measureAsyncFunction(() => Promise.all(mockedReportActionsKeys.map((key) => OnyxUtils.remove(key))), {
@@ -534,7 +474,7 @@ describe('OnyxUtils', () => {
             const reportAction = mockedReportActionsMap[`${collectionKey}0`];
             const changedReportAction = createRandomReportAction(Number(reportAction.reportActionID));
 
-            await measureAsyncFunction(() => OnyxUtils.broadcastUpdate(key, changedReportAction, true), {
+            await measureFunction(() => OnyxUtils.broadcastUpdate(key, changedReportAction, true), {
                 beforeEach: async () => {
                     await Onyx.set(key, reportAction);
                 },

tests/unit/onyxTest.ts

@@ -1482,10 +1482,9 @@ describe('Onyx', () => {
                 return waitForPromisesToResolve();
             })
             .then(() => {
-                expect(collectionCallback).toHaveBeenCalledTimes(3);
+                expect(collectionCallback).toHaveBeenCalledTimes(2);
                 expect(collectionCallback).toHaveBeenNthCalledWith(1, {[cat]: initialValue}, ONYX_KEYS.COLLECTION.ANIMALS, {[cat]: initialValue});
-                expect(collectionCallback).toHaveBeenNthCalledWith(2, {[cat]: initialValue}, ONYX_KEYS.COLLECTION.ANIMALS, undefined);
-                expect(collectionCallback).toHaveBeenNthCalledWith(3, collectionDiff, ONYX_KEYS.COLLECTION.ANIMALS, {[cat]: initialValue, [dog]: {name: 'Rex'}});
+                expect(collectionCallback).toHaveBeenNthCalledWith(2, collectionDiff, ONYX_KEYS.COLLECTION.ANIMALS, {[cat]: initialValue, [dog]: {name: 'Rex'}});
 
                 // Cat hasn't changed from its original value, expect only the initial connect callback
                 expect(catCallback).toHaveBeenCalledTimes(1);