Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
elirangoshen wants to merge 7 commits into
base: main
Choose a base branch
from
+299 −9
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
3295637
perf: batch cold-cache pre-warm in mergeCollectionWithPatches via mul…
elirangoshen May 21, 2026
d02382f
test: cover mergeCollectionWithPatches multiGet pre-warm fast/slow paths
elirangoshen May 26, 2026
d66152a
fix: align multiGet cache hit check with cache.hasCacheForKey
elirangoshen May 26, 2026
7fb447d
fix: preserve cache-first merge when multiGet pre-warm rejects
elirangoshen May 26, 2026
1764d2c
Address PR #793 review comments
elirangoshen May 27, 2026
2f9fb8a
Tighten comments to <= 3 lines
elirangoshen May 27, 2026
98ddd13
multiGet: prefer cache over stale storage on concurrent writes
elirangoshen May 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 22 additions & 9 deletions lib/OnyxUtils.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -347,9 +347,10 @@ function multiGet<TKey extends OnyxKey>(keys: CollectionKeyBase[]): Promise<Map<
continue;
}

const cacheValue = cache.get(key) as OnyxValue<TKey>;
if (cacheValue) {
dataMap.set(key, cacheValue);
// hasCacheForKey catches cached falsy values (0, '', false, null) as cache hits, which
// a truthy check on the value would miss.
if (cache.hasCacheForKey(key)) {
dataMap.set(key, cache.get(key) as OnyxValue<TKey>);
continue;
}

Expand Down Expand Up @@ -401,6 +402,13 @@ function multiGet<TKey extends OnyxKey>(keys: CollectionKeyBase[]): Promise<Map<
}
}

// Prefer cache over stale storage if a concurrent write populated it during
// the read — otherwise cache.merge(temp) below would resurrect dropped fields.
if (cache.hasCacheForKey(key)) {
dataMap.set(key, cache.get(key) as OnyxValue<TKey>);
continue;
}
Copy link
Copy Markdown
Contributor Author

@elirangoshen elirangoshen May 27, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @fabioh8010 here I addressed codex comment, this change might be better to do in a new pr, because it affect all multiGet calls, so I pushed it here to show the change I want to do, but I want to ask what you think if its ok to try it on this pr to open new one for that ?


dataMap.set(key, value as OnyxValue<TKey>);
temp[key] = value as OnyxValue<TKey>;
}
Expand Down Expand Up @@ -1597,12 +1605,17 @@ 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 cache-miss existingKeys so cache.merge() merges the new delta into
// the real previous storage value. Fast path (all warm) skips the pre-warm to preserve
// promise-chain depth; slow path batches the misses into one Storage.multiGet.
const hasColdExistingKey = existingKeys.some((key) => !cache.hasCacheForKey(key));
// Swallow pre-warm read failures so a transient Storage.multiGet rejection doesn't
// skip the cache.merge() + keysChanged() below. Subscribers still see the merge even
// when storage reads fail.
const prewarmPromise = hasColdExistingKey
? multiGet(existingKeys).catch((err) => Logger.logInfo(`mergeCollectionWithPatches pre-warm failed; proceeding with cache-only merge. Error: ${err}`))
Copy link
Copy Markdown

@chatgpt-codex-connector chatgpt-codex-connector Bot May 27, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Preserve concurrent cache writes during batched pre-warm

When this slow path runs for a cold existing key, multiGet() writes the storage result back with cache.merge() without re-checking whether another Onyx write populated the same key while the storage read was in flight. The previous get()-based pre-warm explicitly preferred a value that appeared in cache after the read, so a concurrent set/merge during a cold-cache mergeCollection would keep the newer cache value; with this change, stale persisted fields can be merged over the newer cache state before the final delta is applied and broadcast.

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Contributor Author

@elirangoshen elirangoshen May 27, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

https://github.com/Expensify/react-native-onyx/pull/793/changes#r3310600265

: 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
Expand Down
277 changes: 277 additions & 0 deletions tests/unit/onyxUtilsTest.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -920,6 +920,283 @@ describe('OnyxUtils', () => {
});
});

describe('mergeCollection pre-warm', () => {
// retryOperation tests above replace StorageMock methods without restoring them, leaving
// rejecting mocks behind. Capture pristine refs at file-load time and restore in beforeEach
// so our Onyx.set seeding 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 still tracked as persisted. OnyxCache.drop
// also removes the key from `storageKeys`, so we re-register it afterwards to reliably hit
// the cold-but-persisted state regardless of getAllKeys()'s fallback path.
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 () => {
Copy link
Copy Markdown
Contributor

@fabioh8010 fabioh8010 May 27, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
it('preserves cache-first invariant when Storage.multiGet rejects on the slow path', async () => {
it('preserves cache-first invariant when Storage.multiGet rejects on the slow path', async () => {

Copy link
Copy Markdown
Contributor Author

@elirangoshen elirangoshen May 27, 2026 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done

// A Storage.multiGet rejection during pre-warm must not skip the cache.merge() +
// keysChanged() that follow. Without the .catch() at the pre-warm call site,
// subscribers would miss the merge and Onyx.mergeCollection would reject.
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 and flush the subscriber's initial load BEFORE installing the rejecting mock —
// otherwise the connect's own multiGet (no .catch) consumes the mockRejectedValueOnce and
// leaks an unhandled rejection instead of exercising the merge pre-warm path.
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 fire so subscribers see the merge. Use
// toMatchObject because a concurrent read may have re-populated the persisted value;
// what matters is that the new {merged: true} delta is applied on top.
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
// refs at file-load time and restore them in beforeEach, since retryOperation tests leak.
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);
});

it('prefers cache when a concurrent write lands during the storage read', async () => {
// Concurrent write during multiGet's storage read must not be overwritten by the
// stale snapshot via cache.merge.
const key = `${ONYXKEYS.COLLECTION.TEST_KEY}race`;

OnyxCache.drop(key);
OnyxCache.addKey(key);

// Set cache inside the mock so it lands before Storage.multiGet's promise resolves —
// multiGet's .then() then sees a populated cache and skips writing the stale value.
StorageMock.multiGet = jest.fn().mockImplementation(() => {
OnyxCache.set(key, {fresh: 'data'});
return Promise.resolve([[key, {stale: 'a', alsoStale: 'b'}]]);
});

const result = await OnyxUtils.multiGet([key]);

expect(OnyxCache.get(key)).toEqual({fresh: 'data'});
expect(result.get(key)).toEqual({fresh: 'data'});
});
});

describe('storage eviction', () => {
const diskFullError = new Error('database or disk is full');

Expand Down
Loading