Restore cache state for evicted in-flight keys on retry

Addresses Codex P1 on Expensify#792 review:
Expensify#792 (comment)

When `retryOperation()` evicts an LRU key via `remove()` on a storage
capacity error, `remove()` immediately drops the key from cache and
fires `keyChanged(undefined)`. With the orchestrator/write-helper split
introduced in this PR, the retry only re-attempts the storage write --
it does not re-apply cache state. So if the evicted key is part of the
in-flight write, cache + subscribers stay in the "removed" state even
though the storage retry eventually succeeds, leaving the App's view
of the world inconsistent with persisted data.

Fix: each orchestrator now constructs a `restoreEvictedKey(key)`
closure that re-applies its cache write + subscriber notification for
a single in-flight key, and passes it via the write helper's params.
`retryOperation` invokes the closure after `remove(keyForRemoval)`,
before re-entering the write helper. The closure is a no-op when the
evicted key isn't part of this write (the typical case where some
unrelated evictable key gets selected).

For mergeCollection specifically, the closure snapshots the cache's
post-merge state for in-flight keys (rather than just the merge delta)
so that previously-merged-in fields are preserved when the eviction's
cache.drop forces us to re-populate cache from scratch.

Adds 3 regression tests under `storage eviction`:
- `multiSet` -- LRU evictable is the in-flight key (persistMultiSetWrite path)
- `setCollection` -- same scenario via persistCollectionWrite
- `mergeCollection` -- same scenario via persistMergedCollectionWrite,
  asserts the previously-existing `id` field survives the eviction

All 3 tests fail with `Received: undefined` / missing fields if the
lib changes are reverted; pass with them in place.

All 459 unit tests pass (456 existing + 3 new).

Author: elirangoshen

lib/OnyxUtils.ts

@@ -842,8 +842,20 @@ function retryOperation<TMethod extends RetriableOnyxOperation>(error: Error, on
     Logger.logInfo(`Out of storage. Evicting least recently accessed key (${keyForRemoval}) and retrying. Error: ${error}`);
     reportStorageQuota(error);
 
-    // @ts-expect-error No overload matches this call.
-    return remove(keyForRemoval).then(() => onyxMethod(defaultParams, nextRetryAttempt));
+    return remove(keyForRemoval).then(() => {
+        // remove() drops `keyForRemoval` from cache and fires keyChanged(undefined). If that key is
+        // part of this in-flight write, the upcoming retry would land the value in storage but
+        // cache + subscribers would stay in the "removed" state. Each orchestrator passes a
+        // restoreEvictedKey closure via params that re-applies the orchestrator's cache write +
+        // subscriber notification for the evicted key only — preserving the cache-first invariant
+        // across eviction-driven retries.
+        const restore = (defaultParams as {restoreEvictedKey?: (key: OnyxKey) => void}).restoreEvictedKey;
+        if (typeof restore === 'function') {
+            restore(keyForRemoval);
+        }
+        // @ts-expect-error No overload matches this call.
+        return onyxMethod(defaultParams, nextRetryAttempt);
+    });
 }
 
 /**
@@ -1338,7 +1350,10 @@ function setWithRetry<TKey extends OnyxKey>({key, value, options}: SetParams<TKe
  * storage step. Cache and subscriber notifications already happened in the orchestrator, so
  * retries no longer re-fire `waitForCollectionCallback` subscribers with the same payload.
  */
-function persistMultiSetWrite(params: {keyValuePairsToStore: StorageKeyValuePair[]; newData: OnyxMultiSetInput}, retryAttempt?: number): Promise<void> {
+function persistMultiSetWrite(
+    params: {keyValuePairsToStore: StorageKeyValuePair[]; newData: OnyxMultiSetInput; restoreEvictedKey?: (evictedKey: OnyxKey) => void},
+    retryAttempt?: number,
+): Promise<void> {
     const {keyValuePairsToStore, newData} = params;
 
     return Storage.multiSet(keyValuePairsToStore)
@@ -1426,7 +1441,31 @@ function multiSetWithRetry(data: OnyxMultiSetInput, retryAttempt?: number): Prom
         return !OnyxKeys.isRamOnlyKey(key);
     });
 
-    return persistMultiSetWrite({keyValuePairsToStore, newData}, retryAttempt);
+    // Capture the in-flight key/value map so retryOperation can re-apply cache state for any
+    // evicted-then-retried key (see retryOperation's eviction branch for context).
+    const inFlightValueByKey = new Map<OnyxKey, OnyxValue<OnyxKey>>();
+    for (const [inFlightKey, inFlightValue] of keyValuePairsToSet) {
+        inFlightValueByKey.set(inFlightKey, inFlightValue);
+    }
+    const restoreEvictedKey = (evictedKey: OnyxKey): void => {
+        if (!inFlightValueByKey.has(evictedKey)) {
+            return;
+        }
+        const value = inFlightValueByKey.get(evictedKey) as OnyxValue<OnyxKey>;
+        const evictedCollectionKey = OnyxKeys.getCollectionKey(evictedKey);
+        cache.set(evictedKey, value);
+        if (evictedCollectionKey && OnyxKeys.isCollectionMemberKey(evictedCollectionKey, evictedKey)) {
+            keysChanged(
+                evictedCollectionKey as CollectionKeyBase,
+                {[evictedKey]: value} as Record<string, OnyxValue<OnyxKey>>,
+                {[evictedKey]: undefined} as Record<string, OnyxValue<OnyxKey>>,
+            );
+        } else {
+            keyChanged(evictedKey, value);
+        }
+    };
+
+    return persistMultiSetWrite({keyValuePairsToStore, newData, restoreEvictedKey}, retryAttempt);
 }
 
 /**
@@ -1439,6 +1478,7 @@ function persistCollectionWrite<TKey extends CollectionKeyBase>(
         collectionKey: TKey;
         keyValuePairs: StorageKeyValuePair[];
         mutableCollection: OnyxInputKeyValueMapping;
+        restoreEvictedKey?: (evictedKey: OnyxKey) => void;
     },
     retryAttempt?: number,
 ): Promise<void> {
@@ -1517,7 +1557,22 @@ function setCollectionWithRetry<TKey extends CollectionKeyBase>({collectionKey,
 
         keysChanged(collectionKey, mutableCollection, previousCollection);
 
-        return persistCollectionWrite({collectionKey, keyValuePairs, mutableCollection}, retryAttempt);
+        // Capture the in-flight key/value map so retryOperation can re-apply cache state for any
+        // evicted-then-retried key (see retryOperation's eviction branch for context).
+        const inFlightValueByKey = new Map<OnyxKey, OnyxValue<OnyxKey>>();
+        for (const [inFlightKey, inFlightValue] of keyValuePairs) {
+            inFlightValueByKey.set(inFlightKey, inFlightValue);
+        }
+        const restoreEvictedKey = (evictedKey: OnyxKey): void => {
+            if (!inFlightValueByKey.has(evictedKey)) {
+                return;
+            }
+            const value = inFlightValueByKey.get(evictedKey) as OnyxValue<OnyxKey>;
+            cache.set(evictedKey, value);
+            keysChanged(collectionKey, {[evictedKey]: value} as Record<string, OnyxValue<OnyxKey>>, {[evictedKey]: undefined} as Record<string, OnyxValue<OnyxKey>>);
+        };
+
+        return persistCollectionWrite({collectionKey, keyValuePairs, mutableCollection, restoreEvictedKey}, retryAttempt);
     });
 }
 
@@ -1533,6 +1588,7 @@ function persistMergedCollectionWrite<TKey extends CollectionKeyBase>(
         keyValuePairsForExistingCollection: StorageKeyValuePair[];
         keyValuePairsForNewCollection: StorageKeyValuePair[];
         resultCollection: OnyxInputKeyValueMapping;
+        restoreEvictedKey?: (evictedKey: OnyxKey) => void;
     },
     retryAttempt?: number,
 ): Promise<void> {
@@ -1678,7 +1734,29 @@ function mergeCollectionWithPatches<TKey extends CollectionKeyBase>(
                 cache.merge(finalMergedCollection);
                 keysChanged(collectionKey, finalMergedCollection, previousCollection);
 
-                return persistMergedCollectionWrite({collectionKey, keyValuePairsForExistingCollection, keyValuePairsForNewCollection, resultCollection}, retryAttempt);
+                // Snapshot the post-merge cache values for the in-flight keys. The orchestrator's
+                // cache.merge already merged the deltas into the previous storage values, so the cache
+                // now holds the *full* merged value for each key. Using the snapshot (rather than the
+                // delta in `finalMergedCollection`) preserves previously-merged-in fields when the
+                // eviction's cache.drop forces us to re-populate cache from scratch on retry.
+                const inFlightMergedSnapshot = new Map<OnyxKey, OnyxValue<OnyxKey>>();
+                for (const inFlightKey of Object.keys(finalMergedCollection)) {
+                    inFlightMergedSnapshot.set(inFlightKey, cache.get(inFlightKey));
+                }
+
+                // Closure invoked by retryOperation after an eviction picks a key that's part of this
+                // merge. We re-apply the full merged value + a keysChanged notification for that one
+                // key so subscribers don't stay in the "removed" state across the imminent storage retry.
+                const restoreEvictedKey = (evictedKey: OnyxKey): void => {
+                    if (!inFlightMergedSnapshot.has(evictedKey)) {
+                        return;
+                    }
+                    const value = inFlightMergedSnapshot.get(evictedKey) as OnyxValue<OnyxKey>;
+                    cache.set(evictedKey, value);
+                    keysChanged(collectionKey, {[evictedKey]: value} as Record<string, OnyxValue<OnyxKey>>, {[evictedKey]: undefined} as Record<string, OnyxValue<OnyxKey>>);
+                };
+
+                return persistMergedCollectionWrite({collectionKey, keyValuePairsForExistingCollection, keyValuePairsForNewCollection, resultCollection, restoreEvictedKey}, retryAttempt);
             });
         })
         .then(() => undefined);
@@ -1732,7 +1810,22 @@ function partialSetCollection<TKey extends CollectionKeyBase>({collectionKey, co
 
         keysChanged(collectionKey, mutableCollection, previousCollection);
 
-        return persistCollectionWrite({collectionKey, keyValuePairs, mutableCollection}, retryAttempt);
+        // Capture the in-flight key/value map so retryOperation can re-apply cache state for any
+        // evicted-then-retried key (see retryOperation's eviction branch for context).
+        const inFlightValueByKey = new Map<OnyxKey, OnyxValue<OnyxKey>>();
+        for (const [inFlightKey, inFlightValue] of keyValuePairs) {
+            inFlightValueByKey.set(inFlightKey, inFlightValue);
+        }
+        const restoreEvictedKey = (evictedKey: OnyxKey): void => {
+            if (!inFlightValueByKey.has(evictedKey)) {
+                return;
+            }
+            const value = inFlightValueByKey.get(evictedKey) as OnyxValue<OnyxKey>;
+            cache.set(evictedKey, value);
+            keysChanged(collectionKey, {[evictedKey]: value} as Record<string, OnyxValue<OnyxKey>>, {[evictedKey]: undefined} as Record<string, OnyxValue<OnyxKey>>);
+        };
+
+        return persistCollectionWrite({collectionKey, keyValuePairs, mutableCollection, restoreEvictedKey}, retryAttempt);
     });
 }
 

tests/unit/onyxUtilsTest.ts

@@ -1209,6 +1209,93 @@ describe('OnyxUtils', () => {
             expect(logInfoSpy).toHaveBeenCalledWith(`Out of storage. Evicting least recently accessed key (${key1}) and retrying. Error: ${diskFullError}`);
             expect(logInfoSpy).toHaveBeenCalledWith(`Storage Quota Check -- bytesUsed: 0 bytesRemaining: Infinity. Original error: ${diskFullError}`);
         });
+
+        // Eviction-restoration tests: when retryOperation evicts a key via remove() that happens to be
+        // part of the in-flight write, the orchestrator's restoreEvictedKey closure must re-apply
+        // cache state for that key. Otherwise cache + subscribers stay in the "removed" state from
+        // the eviction's keyChanged(undefined) even after the retry succeeds in storage.
+        it('multiSet — restores cache + subscribers when the in-flight key is the LRU evictable', async () => {
+            const memberKey = `${ONYXKEYS.COLLECTION.TEST_KEY}1`;
+
+            // Seed memberKey so it enters the recentlyAccessedKeys set. It's the only evictable key,
+            // so getKeyForEviction() will return memberKey at the time of the multiSet retry.
+            await LocalOnyx.set(memberKey, {id: 1, value: 'original'});
+            expect(LocalOnyxCache.getKeyForEviction()).toBe(memberKey);
+
+            const subscriberCalls: Array<unknown> = [];
+            LocalOnyx.connect({
+                key: memberKey,
+                callback: (value) => subscriberCalls.push(value),
+            });
+            await waitForPromisesToResolve();
+            subscriberCalls.length = 0;
+
+            // Storage.multiSet rejects once with disk-full, then succeeds on retry.
+            LocalStorageMock.multiSet = jest.fn(LocalStorageMock.multiSet).mockRejectedValueOnce(diskFullError).mockImplementation(LocalStorageMock.multiSet);
+
+            await LocalOnyx.multiSet({[memberKey]: {id: 1, value: 'updated'}});
+
+            // Cache must reflect the new value, not undefined from the eviction's remove().
+            expect(LocalOnyxCache.get(memberKey)).toEqual({id: 1, value: 'updated'});
+
+            // Subscriber's last value must be the new value, not undefined.
+            expect(subscriberCalls.at(-1)).toEqual({id: 1, value: 'updated'});
+        });
+
+        it('setCollection — restores cache + subscribers when the in-flight key is the LRU evictable', async () => {
+            const memberKey = `${ONYXKEYS.COLLECTION.TEST_KEY}1`;
+
+            await LocalOnyx.set(memberKey, {id: 1, value: 'original'});
+            expect(LocalOnyxCache.getKeyForEviction()).toBe(memberKey);
+
+            const collectionCalls: Array<unknown> = [];
+            LocalOnyx.connect({
+                key: ONYXKEYS.COLLECTION.TEST_KEY,
+                waitForCollectionCallback: true,
+                callback: (value) => collectionCalls.push(value),
+            });
+            await waitForPromisesToResolve();
+            collectionCalls.length = 0;
+
+            LocalStorageMock.multiSet = jest.fn(LocalStorageMock.multiSet).mockRejectedValueOnce(diskFullError).mockImplementation(LocalStorageMock.multiSet);
+
+            await LocalOnyx.setCollection(ONYXKEYS.COLLECTION.TEST_KEY, {
+                [memberKey]: {id: 1, value: 'updated'},
+            } as GenericCollection);
+
+            expect(LocalOnyxCache.get(memberKey)).toEqual({id: 1, value: 'updated'});
+            // Last collection callback must include the restored member with its new value.
+            const lastBroadcast = collectionCalls.at(-1) as Record<string, unknown> | undefined;
+            expect(lastBroadcast?.[memberKey]).toEqual({id: 1, value: 'updated'});
+        });
+
+        it('mergeCollection — restores cache + subscribers when the in-flight key is the LRU evictable', async () => {
+            const memberKey = `${ONYXKEYS.COLLECTION.TEST_KEY}1`;
+
+            await LocalOnyx.set(memberKey, {id: 1, value: 'original'});
+            expect(LocalOnyxCache.getKeyForEviction()).toBe(memberKey);
+
+            const collectionCalls: Array<unknown> = [];
+            LocalOnyx.connect({
+                key: ONYXKEYS.COLLECTION.TEST_KEY,
+                waitForCollectionCallback: true,
+                callback: (value) => collectionCalls.push(value),
+            });
+            await waitForPromisesToResolve();
+            collectionCalls.length = 0;
+
+            // mergeCollection routes existing keys through multiMerge — fail it once then succeed.
+            LocalStorageMock.multiMerge = jest.fn(LocalStorageMock.multiMerge).mockRejectedValueOnce(diskFullError).mockImplementation(LocalStorageMock.multiMerge);
+
+            await LocalOnyx.mergeCollection(ONYXKEYS.COLLECTION.TEST_KEY, {
+                [memberKey]: {value: 'merged'},
+            } as GenericCollection);
+
+            // Cache must reflect the merged value, not undefined from the eviction's remove().
+            expect(LocalOnyxCache.get(memberKey)).toEqual({id: 1, value: 'merged'});
+            const lastBroadcast = collectionCalls.at(-1) as Record<string, unknown> | undefined;
+            expect(lastBroadcast?.[memberKey]).toEqual({id: 1, value: 'merged'});
+        });
     });
 
     describe('afterInit', () => {