[POC] Remove macro/micro tasks during subscriber update

lib/Onyx.ts

@@ -248,9 +248,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}`);
@@ -366,16 +365,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));
-            }
-
             const defaultKeyValuePairs = Object.entries(
                 Object.keys(defaultKeyStates)
                     .filter((key) => !keysToPreserve.includes(key))
@@ -393,7 +382,13 @@ 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.scheduleSubscriberUpdate(key, value);
+                    }
+                    for (const [key, value] of Object.entries(keyValuesToResetAsCollection)) {
+                        OnyxUtils.scheduleNotifyCollectionSubscribers(key, value.newValues, value.oldValues);
+                    }
                 });
         })
         .then(() => undefined);

lib/OnyxMerge/index.native.ts

@@ -26,18 +26,17 @@ 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);
 
     // If the value has not changed, calling Storage.setItem() would be redundant and a waste of performance, so return early instead.
     if (!hasChanged) {
-        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/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

@@ -75,9 +75,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>> = {};
 
@@ -739,6 +736,7 @@ function keyChanged<TKey extends OnyxKey>(
                 }
 
                 cachedCollection[key] = value;
+                lastConnectionCallbackData.set(subscriber.subscriptionID, cachedCollection);
                 subscriber.callback(cachedCollection, subscriber.key, {[key]: value});
                 continue;
             }
@@ -770,7 +768,7 @@ function sendDataToConnection<TKey extends OnyxKey>(mapping: CallbackToStateMapp
     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 (lastConnectionCallbackData.has(mapping.subscriptionID) && lastValue) {
         return;
     }
 
@@ -803,23 +801,6 @@ 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).
  *
@@ -831,21 +812,17 @@ function scheduleSubscriberUpdate<TKey extends OnyxKey>(
     value: OnyxValue<TKey>,
     canUpdateSubscriber: (subscriber?: CallbackToStateMapping<OnyxKey>) => boolean = () => true,
     isProcessingCollectionUpdate = false,
-): Promise<void> {
-    return prepareSubscriberUpdate(() => keyChanged(key, value, canUpdateSubscriber, isProcessingCollectionUpdate));
+): void {
+    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));
+function scheduleNotifyCollectionSubscribers<TKey extends OnyxKey>(key: TKey, value: OnyxCollection<KeyValueMapping[TKey]>, previousValue?: OnyxCollection<KeyValueMapping[TKey]>): void {
+    keysChanged(key, value, previousValue);
 }
 
 /**
@@ -919,7 +896,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) {
@@ -928,7 +905,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);
+    scheduleSubscriberUpdate(key, value, (subscriber) => hasChanged || subscriber?.initWithStoredValues === false);
 }
 
 function hasPendingMergeForKey(key: OnyxKey): boolean {
@@ -1314,18 +1291,17 @@ 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();
     }
 
     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;
         });
 }
 

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

@@ -452,7 +452,7 @@ describe('OnyxUtils', () => {
                 Object.entries(mockedReportActionsMap).map(([k, v]) => [k, createRandomReportAction(Number(v.reportActionID))] as const),
             ) as GenericCollection;
 
-            await measureAsyncFunction(() => OnyxUtils.scheduleNotifyCollectionSubscribers(collectionKey, changedReportActions, mockedReportActionsMap), {
+            await measureFunction(() => OnyxUtils.scheduleNotifyCollectionSubscribers(collectionKey, changedReportActions, mockedReportActionsMap), {
                 beforeEach: async () => {
                     await Onyx.multiSet(mockedReportActionsMap);
                     for (const key of mockedReportActionsKeys) {
@@ -517,7 +517,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

@@ -1535,10 +1535,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);