[Do Not Merge] Avoid subscriber re-fire when storage writes retry

lib/OnyxUtils.ts

@@ -769,9 +769,14 @@ 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): Promise<void> {
+function remove<TKey extends OnyxKey>(key: TKey, isProcessingCollectionUpdate?: boolean, skipNotify?: boolean): Promise<void> {
     cache.drop(key);
-    keyChanged(key, undefined as OnyxValue<TKey>, undefined, isProcessingCollectionUpdate);
+    // 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);
+    }
 
     if (OnyxKeys.isRamOnlyKey(key)) {
         return Promise.resolve();
@@ -797,7 +802,13 @@ function reportStorageQuota(error?: Error): Promise<void> {
  * - Non-retriable errors: logs an alert and resolves without retrying
  * - Other errors: retries the operation
  */
-function retryOperation<TMethod extends RetriableOnyxOperation>(error: Error, onyxMethod: TMethod, defaultParams: Parameters<TMethod>[0], retryAttempt: number | undefined): Promise<void> {
+function retryOperation<TMethod extends RetriableOnyxOperation>(
+    error: Error,
+    onyxMethod: TMethod,
+    defaultParams: Parameters<TMethod>[0],
+    retryAttempt: number | undefined,
+    inFlightKeys?: Set<OnyxKey>,
+): Promise<void> {
     const currentRetryAttempt = retryAttempt ?? 0;
     const nextRetryAttempt = currentRetryAttempt + 1;
 
@@ -842,8 +853,12 @@ 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);
 
+    // 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;
     // @ts-expect-error No overload matches this call.
-    return remove(keyForRemoval).then(() => onyxMethod(defaultParams, nextRetryAttempt));
+    return remove(keyForRemoval, undefined, willBeRestored).then(() => onyxMethod(defaultParams, nextRetryAttempt));
 }
 
 /**
@@ -1395,14 +1410,21 @@ function multiSetWithRetry(data: OnyxMultiSetInput, retryAttempt?: number): Prom
             // so re-entrant callbacks (e.g. Onyx.set inside a callback) see consistent cache
             // and subscriber state, matching the original per-key notification semantics.
             cache.set(key, value);
-            keyChanged(key, value);
+            // Skip subscriber notification on retry — already notified on attempt 0.
+            // waitForCollectionCallback subscribers re-fire on every keyChanged by contract.
+            if (!retryAttempt) {
+                keyChanged(key, value);
+            }
         }
     }
 
     // One keysChanged() per collection — fires each collection-level subscriber once and lets
     // keysChanged() internally decide which individual member subscribers need notification.
-    for (const [collectionKey, batch] of collectionBatches) {
-        keysChanged(collectionKey as CollectionKeyBase, batch.partial, batch.previous);
+    // Skip on retry — already notified on attempt 0 (see same-reason comment above).
+    if (!retryAttempt) {
+        for (const [collectionKey, batch] of collectionBatches) {
+            keysChanged(collectionKey as CollectionKeyBase, batch.partial, batch.previous);
+        }
     }
 
     const keyValuePairsToStore = keyValuePairsToSet.filter((keyValuePair) => {
@@ -1411,8 +1433,10 @@ function multiSetWithRetry(data: OnyxMultiSetInput, retryAttempt?: number): Prom
         return !OnyxKeys.isRamOnlyKey(key);
     });
 
+    const inFlightKeys = new Set<OnyxKey>(keyValuePairsToSet.map(([key]) => key));
+
     return Storage.multiSet(keyValuePairsToStore)
-        .catch((error) => OnyxUtils.retryOperation(error, multiSetWithRetry, newData, retryAttempt))
+        .catch((error) => OnyxUtils.retryOperation(error, multiSetWithRetry, newData, retryAttempt, inFlightKeys))
         .then(() => {
             OnyxUtils.sendActionToDevTools(OnyxUtils.METHOD.MULTI_SET, undefined, newData);
         });
@@ -1476,16 +1500,22 @@ function setCollectionWithRetry<TKey extends CollectionKeyBase>({collectionKey,
 
         for (const [key, value] of keyValuePairs) cache.set(key, value);
 
-        keysChanged(collectionKey, mutableCollection, previousCollection);
+        // Skip subscriber notification on retry — already notified on attempt 0.
+        // waitForCollectionCallback subscribers re-fire on every keysChanged by contract.
+        if (!retryAttempt) {
+            keysChanged(collectionKey, mutableCollection, previousCollection);
+        }
 
         // RAM-only keys are not supposed to be saved to storage
         if (OnyxKeys.isRamOnlyKey(collectionKey)) {
             OnyxUtils.sendActionToDevTools(OnyxUtils.METHOD.SET_COLLECTION, undefined, mutableCollection);
             return;
         }
 
+        const inFlightKeys = new Set<OnyxKey>(keyValuePairs.map(([key]) => key));
+
         return Storage.multiSet(keyValuePairs)
-            .catch((error) => OnyxUtils.retryOperation(error, setCollectionWithRetry, {collectionKey, collection}, retryAttempt))
+            .catch((error) => OnyxUtils.retryOperation(error, setCollectionWithRetry, {collectionKey, collection}, retryAttempt, inFlightKeys))
             .then(() => {
                 OnyxUtils.sendActionToDevTools(OnyxUtils.METHOD.SET_COLLECTION, undefined, mutableCollection);
             });
@@ -1611,7 +1641,11 @@ function mergeCollectionWithPatches<TKey extends CollectionKeyBase>(
                 // write fails.
                 const previousCollection = getCachedCollection(collectionKey, existingKeys);
                 cache.merge(finalMergedCollection);
-                keysChanged(collectionKey, finalMergedCollection, previousCollection);
+                // Skip subscriber notification on retry — already notified on attempt 0.
+                // waitForCollectionCallback subscribers re-fire on every keysChanged by contract.
+                if (!retryAttempt) {
+                    keysChanged(collectionKey, finalMergedCollection, previousCollection);
+                }
 
                 const promises = [];
 
@@ -1627,13 +1661,16 @@ function mergeCollectionWithPatches<TKey extends CollectionKeyBase>(
                     promises.push(Storage.multiSet(keyValuePairsForNewCollection));
                 }
 
+                const inFlightKeys = new Set<OnyxKey>(Object.keys(finalMergedCollection));
+
                 return Promise.all(promises)
                     .catch((error) =>
                         retryOperation(
                             error,
                             mergeCollectionWithPatches,
                             {collectionKey, collection: resultCollection as OnyxMergeCollectionInput<TKey>, mergeReplaceNullPatches, isProcessingCollectionUpdate},
                             retryAttempt,
+                            inFlightKeys,
                         ),
                     )
                     .then(() => {
@@ -1690,15 +1727,21 @@ function partialSetCollection<TKey extends CollectionKeyBase>({collectionKey, co
 
         for (const [key, value] of keyValuePairs) cache.set(key, value);
 
-        keysChanged(collectionKey, mutableCollection, previousCollection);
+        // Skip subscriber notification on retry — already notified on attempt 0.
+        // waitForCollectionCallback subscribers re-fire on every keysChanged by contract.
+        if (!retryAttempt) {
+            keysChanged(collectionKey, mutableCollection, previousCollection);
+        }
 
         if (OnyxKeys.isRamOnlyKey(collectionKey)) {
             sendActionToDevTools(METHOD.SET_COLLECTION, undefined, mutableCollection);
             return;
         }
 
+        const inFlightKeys = new Set<OnyxKey>(keyValuePairs.map(([key]) => key));
+
         return Storage.multiSet(keyValuePairs)
-            .catch((error) => retryOperation(error, partialSetCollection, {collectionKey, collection}, retryAttempt))
+            .catch((error) => retryOperation(error, partialSetCollection, {collectionKey, collection}, retryAttempt, inFlightKeys))
             .then(() => {
                 sendActionToDevTools(METHOD.SET_COLLECTION, undefined, mutableCollection);
             });

tests/unit/onyxUtilsTest.ts

@@ -920,6 +920,128 @@ describe('OnyxUtils', () => {
         });
     });
 
+    describe('retry side-effect idempotency', () => {
+        // Save originals so each test can replace StorageMock.multiMerge / StorageMock.multiSet
+        // with a one-shot rejecting mock that triggers retryOperation's transient-error path.
+        // Restoring keeps mocks from leaking into the storage-eviction describe block below.
+        const originalMultiMerge = StorageMock.multiMerge;
+        const originalMultiSet = StorageMock.multiSet;
+
+        afterEach(() => {
+            StorageMock.multiMerge = originalMultiMerge;
+            StorageMock.multiSet = originalMultiSet;
+        });
+
+        // A retriable error: not in NON_RETRIABLE_ERRORS, not in STORAGE_ERRORS, so retryOperation
+        // re-enters the failing method on the next attempt.
+        const transientError = new Error('Transient storage error');
+
+        it('mergeCollection — waitForCollectionCallback subscriber fires once across retries', async () => {
+            const collectionKey = ONYXKEYS.COLLECTION.TEST_KEY;
+            const existingMemberKey = `${collectionKey}1`;
+            const newMemberKey = `${collectionKey}2`;
+
+            await Onyx.set(existingMemberKey, {value: 'initial'});
+
+            const collectionCallback = jest.fn();
+            Onyx.connect({
+                key: collectionKey,
+                waitForCollectionCallback: true,
+                callback: collectionCallback,
+            });
+            await waitForPromisesToResolve();
+            collectionCallback.mockClear();
+
+            StorageMock.multiMerge = jest.fn(originalMultiMerge).mockRejectedValueOnce(transientError);
+
+            await Onyx.mergeCollection(collectionKey, {
+                [existingMemberKey]: {value: 'merged'},
+                [newMemberKey]: {value: 'new'},
+            } as GenericCollection);
+
+            // Before this fix, every retry attempt re-fired keysChanged() — and
+            // waitForCollectionCallback subscribers fire on every keysChanged() call by contract.
+            // After the fix, retries skip the keysChanged re-fire, so subscribers are notified
+            // exactly once per logical operation.
+            expect(collectionCallback).toHaveBeenCalledTimes(1);
+        });
+
+        it('Onyx.multiSet — collection subscriber fires once across retries', async () => {
+            const collectionKey = ONYXKEYS.COLLECTION.TEST_KEY;
+            const memberKey1 = `${collectionKey}1`;
+            const memberKey2 = `${collectionKey}2`;
+
+            const collectionCallback = jest.fn();
+            Onyx.connect({
+                key: collectionKey,
+                waitForCollectionCallback: true,
+                callback: collectionCallback,
+            });
+            await waitForPromisesToResolve();
+            collectionCallback.mockClear();
+
+            StorageMock.multiSet = jest.fn(originalMultiSet).mockRejectedValueOnce(transientError);
+
+            await Onyx.multiSet({
+                [memberKey1]: {value: 'first'},
+                [memberKey2]: {value: 'second'},
+            });
+
+            expect(collectionCallback).toHaveBeenCalledTimes(1);
+        });
+
+        it('Onyx.setCollection — collection subscriber fires once across retries', async () => {
+            const collectionKey = ONYXKEYS.COLLECTION.TEST_KEY;
+            const memberKey1 = `${collectionKey}1`;
+            const memberKey2 = `${collectionKey}2`;
+
+            const collectionCallback = jest.fn();
+            Onyx.connect({
+                key: collectionKey,
+                waitForCollectionCallback: true,
+                callback: collectionCallback,
+            });
+            await waitForPromisesToResolve();
+            collectionCallback.mockClear();
+
+            StorageMock.multiSet = jest.fn(originalMultiSet).mockRejectedValueOnce(transientError);
+
+            await Onyx.setCollection(collectionKey, {
+                [memberKey1]: {value: 'first'},
+                [memberKey2]: {value: 'second'},
+            } as GenericCollection);
+
+            expect(collectionCallback).toHaveBeenCalledTimes(1);
+        });
+
+        it('OnyxUtils.partialSetCollection — collection subscriber fires once across retries', async () => {
+            const collectionKey = ONYXKEYS.COLLECTION.TEST_KEY;
+            const memberKey1 = `${collectionKey}1`;
+            const memberKey2 = `${collectionKey}2`;
+
+            const collectionCallback = jest.fn();
+            Onyx.connect({
+                key: collectionKey,
+                waitForCollectionCallback: true,
+                callback: collectionCallback,
+            });
+            await waitForPromisesToResolve();
+            collectionCallback.mockClear();
+
+            StorageMock.multiSet = jest.fn(originalMultiSet).mockRejectedValueOnce(transientError);
+
+            await OnyxUtils.partialSetCollection({
+                collectionKey,
+                collection: {
+                    [memberKey1]: {value: 'first'},
+                    [memberKey2]: {value: 'second'},
+                } as GenericCollection,
+            });
+
+            expect(collectionCallback).toHaveBeenCalledTimes(1);
+        });
+    });
+
     describe('storage eviction', () => {
         const diskFullError = new Error('database or disk is full');
 
@@ -1059,6 +1181,63 @@ 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}`);
         });
+
+        it('multiSet — eviction of an UNRELATED key still notifies its subscribers (codex regression guard)', async () => {
+            const evictableKey = `${ONYXKEYS.COLLECTION.TEST_KEY}1`;
+            const writeKey = `${ONYXKEYS.COLLECTION.TEST_KEY}2`;
+
+            // Seed the evictable key first so it becomes the LRU evictable. The subsequent multiSet
+            // writes a DIFFERENT key, so the evicted key is unrelated to the in-flight write.
+            await LocalOnyx.set(evictableKey, {value: 'will-be-evicted'});
+            expect(LocalOnyxCache.getKeyForEviction()).toBe(evictableKey);
+
+            const subscriberCalls: unknown[] = [];
+            LocalOnyx.connect({
+                key: evictableKey,
+                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({[writeKey]: {value: 'new'}});
+
+            // evictableKey was the LRU evictable, so retryOperation evicted it. It's not in the
+            // in-flight write's keys, so the retry's cache.set won't restore it — subscribers MUST
+            // see keyChanged(undefined) so they reflect the genuine removal (not stale value).
+            expect(LocalOnyxCache.hasCacheForKey(evictableKey)).toBe(false);
+            expect(subscriberCalls.at(-1)).toBeUndefined();
+        });
+
+        it('multiSet — eviction of an IN-FLIGHT key does not strand its subscriber', async () => {
+            const memberKey = `${ONYXKEYS.COLLECTION.TEST_KEY}1`;
+
+            // Seed memberKey so it becomes the LRU evictable. The multiSet below writes to the SAME
+            // key, so eviction picks an in-flight key.
+            await LocalOnyx.set(memberKey, {value: 'original'});
+            expect(LocalOnyxCache.getKeyForEviction()).toBe(memberKey);
+
+            const subscriberCalls: unknown[] = [];
+            LocalOnyx.connect({
+                key: memberKey,
+                callback: (value) => subscriberCalls.push(value),
+            });
+            await waitForPromisesToResolve();
+            subscriberCalls.length = 0;
+
+            LocalStorageMock.multiSet = jest.fn(LocalStorageMock.multiSet).mockRejectedValueOnce(diskFullError).mockImplementation(LocalStorageMock.multiSet);
+
+            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.
+            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);
+        });
     });
 
     describe('afterInit', () => {