perf: batch cold-cache pre-warm in mergeCollectionWithPatches via multiGet

lib/OnyxUtils.ts

@@ -347,9 +347,11 @@ function multiGet<TKey extends OnyxKey>(keys: CollectionKeyBase[]): Promise<Map<
             continue;
         }
 
-        const cacheValue = cache.get(key) as OnyxValue<TKey>;
-        if (cacheValue) {
-            dataMap.set(key, cacheValue);
+        // Use hasCacheForKey, not a truthy check on the value — otherwise cached falsy values
+        // (0, '', false, null) get treated as cache misses and re-fetched from storage, which
+        // can overwrite the warm cached value with a stale storage value via cache.merge().
+        if (cache.hasCacheForKey(key)) {
+            dataMap.set(key, cache.get(key) as OnyxValue<TKey>);
             continue;
         }
 
@@ -1597,12 +1599,27 @@ function mergeCollectionWithPatches<TKey extends CollectionKeyBase>(
             // finalMergedCollection contains all the keys that were merged, without the keys of incompatible updates
             const finalMergedCollection = {...existingKeyCollection, ...newCollection};
 
-            // 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(() => {
+            // Pre-warm cache for any existing storage keys that aren't yet in cache 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).
+            //
+            // Fast path: when every existingKey is already in cache, skip the pre-warm entirely. This
+            // preserves the original promise-chain depth and the subscriber-callback timing that
+            // dependent tests rely on.
+            //
+            // Slow path: when at least one existingKey is a cache miss, use multiGet — it batches the
+            // missing keys into a single Storage.multiGet call (vs. N parallel get() invocations) and
+            // writes the storage values back to cache before resolving.
+            const hasColdExistingKey = existingKeys.some((key) => !cache.hasCacheForKey(key));
+            // Swallow pre-warm read failures the same way the previous get()-based pre-warm did
+            // (see get() catch at the bottom of its definition). Without this, a transient
+            // Storage.multiGet rejection would skip cache.merge() + keysChanged() below and
+            // regress the cache-first invariant established in #787 — subscribers would miss
+            // the merge and Onyx.mergeCollection would reject up to the caller.
+            const prewarmPromise = hasColdExistingKey
+                ? multiGet(existingKeys).catch((err) => Logger.logInfo(`mergeCollectionWithPatches pre-warm failed; proceeding with cache-only merge. Error: ${err}`))
+                : Promise.resolve();
+            return prewarmPromise.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

tests/unit/onyxUtilsTest.ts

@@ -920,6 +920,271 @@ describe('OnyxUtils', () => {
         });
     });
 
+    describe('mergeCollection pre-warm', () => {
+        // The retryOperation describe block above mutates StorageMock.setItem (and never restores it),
+        // so by the time we run, setItem may be a rejecting mock from a prior test. Capture pristine
+        // references at file-load time and restore them in beforeEach so our seeding via Onyx.set
+        // actually reaches the in-memory storage provider.
+        const pristineSetItem = StorageMock.setItem;
+        const pristineMultiSet = StorageMock.multiSet;
+        const pristineMultiGet = StorageMock.multiGet;
+        const pristineGetItem = StorageMock.getItem;
+        const pristineMultiMerge = StorageMock.multiMerge;
+
+        beforeEach(() => {
+            StorageMock.setItem = pristineSetItem;
+            StorageMock.multiSet = pristineMultiSet;
+            StorageMock.multiGet = pristineMultiGet;
+            StorageMock.getItem = pristineGetItem;
+            StorageMock.multiMerge = pristineMultiMerge;
+        });
+
+        // Make a key "cold" — value evicted from cache but the key is still tracked as persisted.
+        // OnyxCache.drop also removes the key from `storageKeys`, which would cause getAllKeys() to
+        // miss it (unless the entire storageKeys set is empty, in which case the function falls back
+        // to Storage.getAllKeys). To reliably hit the cold-but-persisted state regardless of how many
+        // other keys remain in cache, we re-register the key as known after dropping its value.
+        const evictFromCache = (...keys: string[]) => {
+            for (const key of keys) {
+                OnyxCache.drop(key);
+                OnyxCache.addKey(key);
+            }
+        };
+
+        it('fast path: skips storage reads entirely when every existing key is warm in cache', async () => {
+            const collectionKey = ONYXKEYS.COLLECTION.TEST_KEY;
+            const existingKey1 = `${collectionKey}1`;
+            const existingKey2 = `${collectionKey}2`;
+
+            // Seed both members so they are both warm in cache and present in storage.
+            await Onyx.set(existingKey1, {value: 'initial-1'});
+            await Onyx.set(existingKey2, {value: 'initial-2'});
+
+            const multiGetSpy = jest.spyOn(StorageMock, 'multiGet');
+            const getItemSpy = jest.spyOn(StorageMock, 'getItem');
+
+            await Onyx.mergeCollection(collectionKey, {
+                [existingKey1]: {value: 'merged-1'},
+                [existingKey2]: {value: 'merged-2'},
+            } as GenericCollection);
+
+            // With every existingKey warm, the diff swaps Promise.all(get) for Promise.resolve(),
+            // so no storage reads should happen during the pre-warm.
+            expect(multiGetSpy).not.toHaveBeenCalled();
+            expect(getItemSpy).not.toHaveBeenCalled();
+
+            // Cache still reflects the merge.
+            const cached = OnyxCache.getCollectionData(collectionKey);
+            expect(cached?.[existingKey1]).toEqual({value: 'merged-1'});
+            expect(cached?.[existingKey2]).toEqual({value: 'merged-2'});
+        });
+
+        it('slow path: batches cold existing keys into a single Storage.multiGet, with no individual getItem calls', async () => {
+            const collectionKey = ONYXKEYS.COLLECTION.TEST_KEY;
+            const coldKey1 = `${collectionKey}1`;
+            const coldKey2 = `${collectionKey}2`;
+            const warmKey = `${collectionKey}3`;
+
+            // Seed all three in storage, then evict two from cache so they are cold-but-persisted.
+            await Onyx.set(coldKey1, {value: 'persisted-1'});
+            await Onyx.set(coldKey2, {value: 'persisted-2'});
+            await Onyx.set(warmKey, {value: 'persisted-3'});
+            evictFromCache(coldKey1, coldKey2);
+
+            // Reset spies AFTER seeding so we only count calls made during mergeCollection itself.
+            const multiGetSpy = jest.spyOn(StorageMock, 'multiGet').mockClear();
+            const getItemSpy = jest.spyOn(StorageMock, 'getItem').mockClear();
+
+            await Onyx.mergeCollection(collectionKey, {
+                [coldKey1]: {value: 'merged-1'},
+                [coldKey2]: {value: 'merged-2'},
+                [warmKey]: {value: 'merged-3'},
+            } as GenericCollection);
+
+            // OnyxUtils.multiGet filters to cache-missing keys before issuing Storage.multiGet, so we
+            // expect exactly one batched read containing only the cold keys (the warm key is skipped).
+            expect(multiGetSpy).toHaveBeenCalledTimes(1);
+            const requestedKeys = multiGetSpy.mock.calls[0][0] as string[];
+            expect(requestedKeys.sort()).toEqual([coldKey1, coldKey2].sort());
+
+            // No individual Storage.getItem calls during pre-warm. Old code path would have fired one
+            // get() per existing key, each potentially landing in Storage.getItem on cache miss.
+            expect(getItemSpy).not.toHaveBeenCalled();
+        });
+
+        it('slow path: cold-cache merge layers the new delta on top of existing storage data (no field drops)', async () => {
+            const collectionKey = ONYXKEYS.COLLECTION.TEST_KEY;
+            const coldKey = `${collectionKey}1`;
+
+            // Seed an object with multiple fields in storage, then evict from cache so the merge base
+            // must come from a storage read — not from `undefined`.
+            await Onyx.set(coldKey, {a: 1, b: 2});
+            evictFromCache(coldKey);
+
+            await Onyx.mergeCollection(collectionKey, {
+                [coldKey]: {c: 3},
+            } as GenericCollection);
+
+            // If the pre-warm did NOT populate the cache from storage, fastMerge would treat the
+            // previous value as undefined and the result would drop {a:1, b:2}. With the pre-warm
+            // running multiGet on the cold key, the merge layers {c:3} on top of {a:1, b:2}.
+            const cached = OnyxCache.getCollectionData(collectionKey);
+            expect(cached?.[coldKey]).toEqual({a: 1, b: 2, c: 3});
+        });
+
+        it('warm cache: subscriber receives a single merged broadcast for an Onyx.update batch (no transient undefined)', async () => {
+            const collectionKey = ONYXKEYS.COLLECTION.TEST_KEY;
+            const existingKey = `${collectionKey}1`;
+
+            await Onyx.set(existingKey, {value: 'initial'});
+
+            const collectionCallback = jest.fn();
+            Onyx.connect({
+                key: collectionKey,
+                waitForCollectionCallback: true,
+                callback: collectionCallback,
+            });
+            await waitForPromisesToResolve();
+            collectionCallback.mockClear();
+
+            await Onyx.update([
+                {
+                    onyxMethod: Onyx.METHOD.MERGE_COLLECTION,
+                    key: collectionKey,
+                    value: {
+                        [existingKey]: {value: 'merged'},
+                    } as GenericCollection,
+                },
+            ]);
+
+            // The fast path resolves the pre-warm synchronously (Promise.resolve()), preserving the
+            // original promise-chain depth. The Onyx.update batch must therefore broadcast exactly
+            // one merged value — not undefined first and the merged value on a later microtask.
+            const broadcasts = collectionCallback.mock.calls.map((c) => c[0]);
+            expect(broadcasts).toHaveLength(1);
+            expect(broadcasts[0]?.[existingKey]).toEqual({value: 'merged'});
+        });
+
+        it('equivalence: warm-path and cold-path produce the same final cache state for the same merge', async () => {
+            const collectionKey = ONYXKEYS.COLLECTION.TEST_KEY;
+            const memberKey = `${collectionKey}1`;
+            const delta = {value: 'after'} as const;
+
+            // Warm-path run.
+            await Onyx.set(memberKey, {value: 'before', extra: 'kept'});
+            await Onyx.mergeCollection(collectionKey, {
+                [memberKey]: delta,
+            } as GenericCollection);
+            const warmResult = OnyxCache.getCollectionData(collectionKey)?.[memberKey];
+
+            // Reset and replay with a cold cache before the merge.
+            await Onyx.clear();
+            await Onyx.set(memberKey, {value: 'before', extra: 'kept'});
+            evictFromCache(memberKey);
+            await Onyx.mergeCollection(collectionKey, {
+                [memberKey]: delta,
+            } as GenericCollection);
+            const coldResult = OnyxCache.getCollectionData(collectionKey)?.[memberKey];
+
+            expect(warmResult).toEqual(coldResult);
+            expect(coldResult).toEqual({value: 'after', extra: 'kept'});
+        });
+        it('preserves cache-first invariant when Storage.multiGet rejects on the slow path', async () => {
+            // Regression for codex review #3302454987 / PR #793. Before the .catch() at the
+            // pre-warm call site, a Storage.multiGet rejection would propagate up and skip
+            // cache.merge() + keysChanged() entirely — subscribers would miss the merge and
+            // Onyx.mergeCollection would reject. The previous get()-based pre-warm swallowed
+            // these errors per-key inside get()'s own .catch(), so the cache-first invariant
+            // from #787 held even on a flaky read.
+            const collectionKey = ONYXKEYS.COLLECTION.TEST_KEY;
+            const coldMemberKey = `${collectionKey}1`;
+            const newMemberKey = `${collectionKey}2`;
+
+            // Seed an existing member, then evict it from cache so it's "tracked but unloaded" —
+            // the slow path will try to multiGet it.
+            await Onyx.set(coldMemberKey, {value: 'persisted'});
+            evictFromCache(coldMemberKey);
+
+            // Connect the subscriber and flush its initial data load FIRST, before installing
+            // the rejecting mock. Otherwise the connect's getCollectionDataAndSendAsObject path
+            // (whose multiGet call has no .catch) would consume the mockRejectedValueOnce and
+            // leak an unhandled rejection — not the merge pre-warm we're actually exercising.
+            const collectionCallback = jest.fn();
+            Onyx.connect({
+                key: collectionKey,
+                waitForCollectionCallback: true,
+                callback: collectionCallback,
+            });
+            await waitForPromisesToResolve();
+            collectionCallback.mockClear();
+
+            // The subscriber's connect re-populated cache, so re-evict to force the merge into
+            // the slow (cold-key) path. Then reject the next Storage.multiGet so the pre-warm
+            // read fails.
+            evictFromCache(coldMemberKey);
+            const transientError = new Error('Transient IndexedDB read error');
+            StorageMock.multiGet = jest.fn(pristineMultiGet).mockRejectedValueOnce(transientError);
+
+            // Outer promise must resolve, not reject, even when the pre-warm read fails.
+            let outerRejected: unknown = null;
+            const result = await Onyx.mergeCollection(collectionKey, {
+                [coldMemberKey]: {merged: true},
+                [newMemberKey]: {value: 'new'},
+            } as GenericCollection).catch((e: unknown) => {
+                outerRejected = e;
+            });
+            expect(outerRejected).toBeNull();
+            expect(result).toBeUndefined();
+
+            // cache.merge() + keysChanged() must still have fired so subscribers see the merge.
+            // Use toMatchObject for the cold key because a successful concurrent read may have
+            // re-populated the persisted value into cache; what we care about is that the new
+            // {merged: true} delta is applied.
+            expect(collectionCallback).toHaveBeenCalled();
+            const lastBroadcast = collectionCallback.mock.calls.at(-1)?.[0] as Record<string, unknown> | undefined;
+            expect(lastBroadcast?.[coldMemberKey]).toMatchObject({merged: true});
+            expect(lastBroadcast?.[newMemberKey]).toEqual({value: 'new'});
+        });
+    });
+
+    describe('multiGet cache hit consistency', () => {
+        // Same suite-pollution guard as the pre-warm block above: capture pristine StorageMock
+        // references at file-load time and restore them in beforeEach. The retryOperation
+        // describe block above mutates StorageMock.setItem (and never restores it), so by the
+        // time we run, setItem may be a rejecting mock from a prior test.
+        const pristineSetItem = StorageMock.setItem;
+        const pristineMultiGet = StorageMock.multiGet;
+        const pristineGetItem = StorageMock.getItem;
+
+        beforeEach(() => {
+            StorageMock.setItem = pristineSetItem;
+            StorageMock.multiGet = pristineMultiGet;
+            StorageMock.getItem = pristineGetItem;
+        });
+
+        it('does not re-fetch a cached falsy value from storage', async () => {
+            const falsyKey = ONYXKEYS.TEST_KEY;
+
+            // Seed cache with the falsy value 0 (a number, but the same logic applies to '',
+            // false, and null). Using `Onyx.set` ensures the value lands in cache and storage.
+            await Onyx.set(falsyKey, 0);
+
+            // Spy on Storage methods to confirm multiGet does NOT round-trip to storage for
+            // the cached falsy value.
+            const multiGetSpy = jest.spyOn(StorageMock, 'multiGet');
+            const getItemSpy = jest.spyOn(StorageMock, 'getItem');
+
+            const result = await OnyxUtils.multiGet([falsyKey]);
+
+            // The cached value must be returned without any storage read. Before this fix,
+            // `if (cacheValue)` treated the cached 0 as a miss and triggered Storage.multiGet,
+            // which would then overwrite the warm value via cache.merge().
+            expect(multiGetSpy).not.toHaveBeenCalled();
+            expect(getItemSpy).not.toHaveBeenCalled();
+            expect(result.get(falsyKey)).toBe(0);
+        });
+    });
+
     describe('storage eviction', () => {
         const diskFullError = new Error('database or disk is full');