fix: don't evict in-flight keys on retry; correct stale multiMerge-new-key comment

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: elirangoshen

lib/OnyxCache.ts

@@ -362,11 +362,18 @@ class OnyxCache {
 
     /**
      * Finds the least recently accessed key that can be safely evicted from storage.
+     * `excludeKeys` skips keys that must not be evicted (e.g. the in-flight write's own keys,
+     * whose cache value is the merge base the retry depends on).
      */
-    getKeyForEviction(): OnyxKey | undefined {
+    getKeyForEviction(excludeKeys?: Set<OnyxKey>): OnyxKey | undefined {
         // recentlyAccessedKeys is ordered from least to most recently accessed,
-        // so the first element is the best candidate for eviction.
-        return this.recentlyAccessedKeys.values().next().value;
+        // so the first non-excluded key is the best candidate for eviction.
+        for (const key of this.recentlyAccessedKeys) {
+            if (!excludeKeys?.has(key)) {
+                return key;
+            }
+        }
+        return undefined;
     }
 
     /**

lib/OnyxUtils.ts

@@ -777,14 +777,9 @@ function getCollectionDataAndSendAsObject<TKey extends OnyxKey>(matchingKeys: Co
 /**
  * Remove a key from Onyx and update the subscribers
  */
-function remove<TKey extends OnyxKey>(key: TKey, isProcessingCollectionUpdate?: boolean, skipNotify?: boolean): Promise<void> {
+function remove<TKey extends OnyxKey>(key: TKey, isProcessingCollectionUpdate?: boolean): Promise<void> {
     cache.drop(key);
-    // skipNotify is used by retryOperation's eviction branch — the imminent retry's cache.set
-    // will re-populate cache, so firing keyChanged(undefined) here would only strand subscribers
-    // in the "removed" state across the retry.
-    if (!skipNotify) {
-        keyChanged(key, undefined as OnyxValue<TKey>, undefined, isProcessingCollectionUpdate);
-    }
+    keyChanged(key, undefined as OnyxValue<TKey>, undefined, isProcessingCollectionUpdate);
 
     if (OnyxKeys.isRamOnlyKey(key)) {
         return Promise.resolve();
@@ -851,8 +846,10 @@ function retryOperation<TMethod extends RetriableOnyxOperation>(
         return onyxMethod(defaultParams, nextRetryAttempt);
     }
 
-    // Find the least recently accessed evictable key that we can remove
-    const keyForRemoval = cache.getKeyForEviction();
+    // Find the least recently accessed evictable key that we can remove. Never evict an in-flight
+    // key — its cache value is the merge base this retry depends on, so dropping it would truncate
+    // the write to just the delta and diverge cache from storage.
+    const keyForRemoval = cache.getKeyForEviction(inFlightKeys);
     if (!keyForRemoval) {
         // If we have no acceptable keys to remove then we are possibly trying to save mission critical data. If this is the case,
         // then we should stop retrying as there is not much the user can do to fix this. Instead of getting them stuck in an infinite loop we
@@ -865,12 +862,10 @@ function retryOperation<TMethod extends RetriableOnyxOperation>(
     Logger.logInfo(`Out of storage. Evicting least recently accessed key (${keyForRemoval}) and retrying. Error: ${error}`);
     reportStorageQuota(error);
 
-    // Only suppress keyChanged(undefined) when the evicted key is part of the in-flight
-    // write — then cache.set on retry will restore it. For unrelated keys, eviction is a
-    // genuine loss and subscribers must see the removed state.
-    const willBeRestored = inFlightKeys?.has(keyForRemoval) ?? false;
+    // The evicted key is never in-flight (excluded above), so this is a genuine removal —
+    // subscribers must see the removed state.
     // @ts-expect-error No overload matches this call.
-    return remove(keyForRemoval, undefined, willBeRestored).then(() => onyxMethod(defaultParams, nextRetryAttempt));
+    return remove(keyForRemoval).then(() => onyxMethod(defaultParams, nextRetryAttempt));
 }
 
 /**
@@ -1666,8 +1661,9 @@ function mergeCollectionWithPatches<TKey extends CollectionKeyBase>(
 
                 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
+                // New keys go through multiSet and existing keys through multiMerge. multiMerge on a
+                // missing key stores the value just like multiSet across all backends; splitting them lets
+                // multiSet strip nested nulls (the merge layer keeps them to delete nested storage keys).
                 // 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));

tests/unit/onyxUtilsTest.ts

@@ -1566,13 +1566,66 @@ describe('OnyxUtils', () => {
 
             await LocalOnyx.multiSet({[memberKey]: {value: 'updated'}});
 
-            // The in-flight key was evicted then restored by the retry's cache.set. Subscriber's
-            // last value must be the new value, never a transient undefined from the eviction.
+            // The in-flight key is excluded from eviction, so its cache value (the merge base) is
+            // never dropped. Subscriber's last value is the new value, never a transient undefined.
             expect(LocalOnyxCache.get(memberKey)).toEqual({value: 'updated'});
             expect(subscriberCalls.at(-1)).toEqual({value: 'updated'});
             // Subscriber should never have seen undefined in the middle of the eviction-retry cycle.
             expect(subscriberCalls).not.toContain(undefined);
         });
+
+        it('mergeCollection — evicts an unrelated key, not the in-flight key, so its fields survive', async () => {
+            const collectionKey = ONYXKEYS.COLLECTION.TEST_KEY;
+            const memberKey = `${collectionKey}1`;
+            const unrelatedKey = `${collectionKey}2`;
+
+            // Seed the in-flight member with extra fields, plus a separate evictable key. The merge
+            // only touches memberKey; the unrelated key is the genuine eviction target.
+            await LocalOnyx.set(memberKey, {id: 1, value: 'orig'});
+            await LocalOnyx.set(unrelatedKey, {value: 'evict-me'});
+
+            const memberCalls: unknown[] = [];
+            LocalOnyx.connect({key: memberKey, callback: (value) => memberCalls.push(value)});
+            await waitForPromisesToResolve();
+            memberCalls.length = 0;
+
+            // Storage.multiMerge rejects once with disk-full, then succeeds on retry.
+            LocalStorageMock.multiMerge = jest.fn(LocalStorageMock.multiMerge).mockRejectedValueOnce(diskFullError).mockImplementation(LocalStorageMock.multiMerge);
+
+            await LocalOnyx.mergeCollection(collectionKey, {[memberKey]: {value: 'merged'}} as GenericCollection);
+
+            // The old code evicted the in-flight key and re-ran the merge against an empty cache,
+            // collapsing {id: 1, value: 'orig'} + {value: 'merged'} to just {value: 'merged'}. Now
+            // the in-flight key is protected, so its pre-existing {id: 1} survives.
+            expect(LocalOnyxCache.get(memberKey)).toEqual({id: 1, value: 'merged'});
+            expect(memberCalls.at(-1)).toEqual({id: 1, value: 'merged'});
+            expect(memberCalls).not.toContain(undefined);
+            // The unrelated key was the genuine eviction target.
+            expect(LocalOnyxCache.hasCacheForKey(unrelatedKey)).toBe(false);
+        });
+
+        it('mergeCollection — does not truncate the in-flight key when it is the only evictable key', async () => {
+            const collectionKey = ONYXKEYS.COLLECTION.TEST_KEY;
+            const memberKey = `${collectionKey}1`;
+
+            await LocalOnyx.set(memberKey, {id: 1, value: 'orig'});
+            expect(LocalOnyxCache.getKeyForEviction()).toBe(memberKey);
+
+            const memberCalls: unknown[] = [];
+            LocalOnyx.connect({key: memberKey, callback: (value) => memberCalls.push(value)});
+            await waitForPromisesToResolve();
+            memberCalls.length = 0;
+
+            // The only evictable key is the in-flight one, which is now excluded — so retryOperation
+            // finds no acceptable key and reports the quota instead of dropping (and truncating) it.
+            LocalStorageMock.multiMerge = jest.fn(LocalStorageMock.multiMerge).mockRejectedValue(diskFullError);
+
+            await LocalOnyx.mergeCollection(collectionKey, {[memberKey]: {value: 'merged'}} as GenericCollection);
+
+            expect(LocalOnyxCache.get(memberKey)).toEqual({id: 1, value: 'merged'});
+            expect(memberCalls.at(-1)).toEqual({id: 1, value: 'merged'});
+            expect(memberCalls).not.toContain(undefined);
+        });
     });
 
     describe('afterInit', () => {

tests/unit/storage/providers/IDBKeyvalProviderTest.ts

@@ -172,6 +172,11 @@ describe('IDBKeyValProvider', () => {
                 ),
             ).toEqual(expectedEntries.map((e) => (e[1] === null ? undefined : e[1])));
         });
+
+        it('should insert a new record when key does not exist', async () => {
+            await IDBKeyValProvider.multiMerge([[ONYXKEYS.TEST_KEY_2, {fresh: true}]]);
+            expect(await IDBKeyValProvider.getItem(ONYXKEYS.TEST_KEY_2)).toEqual({fresh: true});
+        });
     });
 
     describe('mergeItem', () => {