Merge pull request Expensify#787 from callstack-internal/elirangoshen/fix/90634-mergeCollectionWithPatches-cache-first

Make mergeCollectionWithPatches cache-first

Author: Gonals

lib/OnyxUtils.ts

@@ -1594,47 +1594,52 @@ function mergeCollectionWithPatches<TKey extends CollectionKeyBase>(
             // because we will simply overwrite the existing values in storage.
             const keyValuePairsForNewCollection = prepareKeyValuePairsForStorage(newCollection, true);
 
-            const promises = [];
-
-            // We need to get the previously existing values so we can compare the new ones
-            // against them, to avoid unnecessary subscriber updates.
-            const previousCollectionPromise = Promise.all(existingKeys.map((key) => get(key).then((value) => [key, value]))).then(Object.fromEntries);
-
-            // New keys will be added via multiSet while existing keys will be updated using multiMerge
-            // This is because setting a key that doesn't exist yet with multiMerge will throw errors
-            // We can skip this step for RAM-only keys as they should never be saved to storage
-            if (!OnyxKeys.isRamOnlyKey(collectionKey) && keyValuePairsForExistingCollection.length > 0) {
-                promises.push(Storage.multiMerge(keyValuePairsForExistingCollection));
-            }
-
-            // We can skip this step for RAM-only keys as they should never be saved to storage
-            if (!OnyxKeys.isRamOnlyKey(collectionKey) && keyValuePairsForNewCollection.length > 0) {
-                promises.push(Storage.multiSet(keyValuePairsForNewCollection));
-            }
-
             // finalMergedCollection contains all the keys that were merged, without the keys of incompatible updates
             const finalMergedCollection = {...existingKeyCollection, ...newCollection};
 
-            // Prefill cache if necessary by calling get() on any existing keys and then merge original data to cache
-            // and update all subscribers
-            const promiseUpdate = previousCollectionPromise.then((previousCollection) => {
+            // Pre-warm cache for any existing storage keys that aren't yet in cache. get() is a no-op
+            // (sync-resolved) for cache hits, and on a cache miss it reads from storage and writes the
+            // value back to cache. This is required so the subsequent cache.merge() merges the new delta
+            // into the real previous storage value (rather than starting from `undefined` and dropping
+            // the existing keys).
+            return Promise.all(existingKeys.map((key) => get(key))).then(() => {
+                // Snapshot previous values from the (now-warm) cache for keysChanged's diff, then update
+                // cache and notify subscribers synchronously BEFORE issuing storage writes. This matches
+                // the cache-first / storage-second invariant followed by every other Onyx write method
+                // (setWithRetry, applyMerge, setCollectionWithRetry, partialSetCollection, clear),
+                // ensuring subscribers still reflect the merged data even if the subsequent storage
+                // write fails.
+                const previousCollection = getCachedCollection(collectionKey, existingKeys);
                 cache.merge(finalMergedCollection);
                 keysChanged(collectionKey, finalMergedCollection, previousCollection);
-            });
 
-            return Promise.all(promises)
-                .catch((error) =>
-                    retryOperation(
-                        error,
-                        mergeCollectionWithPatches,
-                        {collectionKey, collection: resultCollection as OnyxMergeCollectionInput<TKey>, mergeReplaceNullPatches, isProcessingCollectionUpdate},
-                        retryAttempt,
-                    ),
-                )
-                .then(() => {
-                    sendActionToDevTools(METHOD.MERGE_COLLECTION, undefined, resultCollection);
-                    return promiseUpdate;
-                });
+                const promises = [];
+
+                // New keys will be added via multiSet while existing keys will be updated using multiMerge
+                // This is because setting a key that doesn't exist yet with multiMerge will throw errors
+                // We can skip this step for RAM-only keys as they should never be saved to storage
+                if (!OnyxKeys.isRamOnlyKey(collectionKey) && keyValuePairsForExistingCollection.length > 0) {
+                    promises.push(Storage.multiMerge(keyValuePairsForExistingCollection));
+                }
+
+                // We can skip this step for RAM-only keys as they should never be saved to storage
+                if (!OnyxKeys.isRamOnlyKey(collectionKey) && keyValuePairsForNewCollection.length > 0) {
+                    promises.push(Storage.multiSet(keyValuePairsForNewCollection));
+                }
+
+                return Promise.all(promises)
+                    .catch((error) =>
+                        retryOperation(
+                            error,
+                            mergeCollectionWithPatches,
+                            {collectionKey, collection: resultCollection as OnyxMergeCollectionInput<TKey>, mergeReplaceNullPatches, isProcessingCollectionUpdate},
+                            retryAttempt,
+                        ),
+                    )
+                    .then(() => {
+                        sendActionToDevTools(METHOD.MERGE_COLLECTION, undefined, resultCollection);
+                    });
+            });
         })
         .then(() => undefined);
 }

tests/unit/onyxUtilsTest.ts

@@ -825,6 +825,101 @@ describe('OnyxUtils', () => {
         });
     });
 
+    describe('mergeCollection cache-first ordering', () => {
+        // Save originals so we can restore them after each test. The tests below replace
+        // StorageMock.multiMerge / StorageMock.multiSet with rejecting mocks; without
+        // restoring, the mock leaks into later describe blocks (e.g. eviction tests) whose
+        // setup relies on these storage methods working normally.
+        const originalMultiMerge = StorageMock.multiMerge;
+        const originalMultiSet = StorageMock.multiSet;
+
+        afterEach(() => {
+            StorageMock.multiMerge = originalMultiMerge;
+            StorageMock.multiSet = originalMultiSet;
+        });
+
+        it('updates cache and notifies subscribers even when Storage.multiMerge rejects', async () => {
+            const collectionKey = ONYXKEYS.COLLECTION.TEST_KEY;
+            const existingMemberKey = `${collectionKey}1`;
+            const newMemberKey = `${collectionKey}2`;
+
+            // Seed an existing member so the merge path exercises multiMerge (existing) + multiSet (new)
+            await Onyx.set(existingMemberKey, {value: 'initial'});
+
+            const collectionCallback = jest.fn();
+            Onyx.connect({
+                key: collectionKey,
+                waitForCollectionCallback: true,
+                callback: collectionCallback,
+            });
+            await waitForPromisesToResolve();
+            collectionCallback.mockClear();
+
+            // Force Storage.multiMerge to reject with a non-retriable IDB error so the failure
+            // path is taken without burning the full retry budget and without rejecting the
+            // outer Onyx.mergeCollection promise.
+            const nonRetriableIdbError = Object.assign(new Error('Internal error opening backing store for indexedDB.open.'), {name: 'UnknownError'});
+            StorageMock.multiMerge = jest.fn().mockRejectedValue(nonRetriableIdbError);
+
+            await Onyx.mergeCollection(collectionKey, {
+                [existingMemberKey]: {value: 'merged'},
+                [newMemberKey]: {value: 'new'},
+            } as GenericCollection);
+
+            // Cache must reflect the merge regardless of the multiMerge rejection. This is the
+            // cache-first / storage-second invariant that mergeCollectionWithPatches must honor.
+            const cachedCollection = OnyxCache.getCollectionData(collectionKey);
+            expect(cachedCollection?.[existingMemberKey]).toEqual({value: 'merged'});
+            expect(cachedCollection?.[newMemberKey]).toEqual({value: 'new'});
+
+            // Subscribers must have been notified with the merged values.
+            expect(collectionCallback).toHaveBeenCalled();
+            const lastBroadcast = collectionCallback.mock.calls.at(-1)?.[0] as Record<string, unknown> | undefined;
+            expect(lastBroadcast?.[existingMemberKey]).toEqual({value: 'merged'});
+            expect(lastBroadcast?.[newMemberKey]).toEqual({value: 'new'});
+        });
+
+        it('updates cache and notifies subscribers even when Storage.multiSet rejects', async () => {
+            const collectionKey = ONYXKEYS.COLLECTION.TEST_KEY;
+            const newMemberKey1 = `${collectionKey}1`;
+            const newMemberKey2 = `${collectionKey}2`;
+
+            // No keys are seeded, so every merged key is a "new" key. This forces the merge path
+            // to use Storage.multiSet (existing keys would go through Storage.multiMerge).
+            const collectionCallback = jest.fn();
+            Onyx.connect({
+                key: collectionKey,
+                waitForCollectionCallback: true,
+                callback: collectionCallback,
+            });
+            await waitForPromisesToResolve();
+            collectionCallback.mockClear();
+
+            // Force Storage.multiSet to reject with a non-retriable IDB error so the failure
+            // path is taken without burning the full retry budget and without rejecting the
+            // outer Onyx.mergeCollection promise.
+            const nonRetriableIdbError = Object.assign(new Error('Internal error opening backing store for indexedDB.open.'), {name: 'UnknownError'});
+            StorageMock.multiSet = jest.fn().mockRejectedValue(nonRetriableIdbError);
+
+            await Onyx.mergeCollection(collectionKey, {
+                [newMemberKey1]: {value: 'first'},
+                [newMemberKey2]: {value: 'second'},
+            } as GenericCollection);
+
+            // Cache must reflect the merge regardless of the multiSet rejection. This is the
+            // cache-first / storage-second invariant that mergeCollectionWithPatches must honor.
+            const cachedCollection = OnyxCache.getCollectionData(collectionKey);
+            expect(cachedCollection?.[newMemberKey1]).toEqual({value: 'first'});
+            expect(cachedCollection?.[newMemberKey2]).toEqual({value: 'second'});
+
+            // Subscribers must have been notified with the merged values.
+            expect(collectionCallback).toHaveBeenCalled();
+            const lastBroadcast = collectionCallback.mock.calls.at(-1)?.[0] as Record<string, unknown> | undefined;
+            expect(lastBroadcast?.[newMemberKey1]).toEqual({value: 'first'});
+            expect(lastBroadcast?.[newMemberKey2]).toEqual({value: 'second'});
+        });
+    });
+
     describe('storage eviction', () => {
         const diskFullError = new Error('database or disk is full');