Make retried Onyx writes idempotent on storage retry

Follow-up to Expensify#787. Splits each retriable collection write into an
orchestrator (cache + subscriber notify + storage prep) and a small
file-private write helper (Storage.multiSet/multiMerge + retryOperation
+ devtools log). retryOperation re-enters the write helper, not the
orchestrator, which fixes two pre-existing bugs flagged in Expensify#787 review:

1. New keys no longer get silently downgraded to multiMerge on retry.
   The existing/new key split was previously re-derived from getAllKeys
   against a cache that the first attempt had already mutated, routing
   brand-new keys through multiMerge instead of multiSet on retry.
   Benign on IDB but wrong semantics and a crash on backends that
   require multiSet for missing keys.

2. keysChanged/keyChanged subscribers (notably waitForCollectionCallback
   ones, which re-fire on every call by contract) are now notified
   exactly once per logical operation rather than once per retry attempt.

Affected paths: mergeCollectionWithPatches, multiSetWithRetry,
setCollectionWithRetry, partialSetCollection. setWithRetry was already
safe via its hasValueChanged guard. partialSetCollection shares the
new persistCollectionWrite helper with setCollectionWithRetry since
their storage tails are identical.

Adds a 'retry side-effect idempotency' test block covering all four
paths. Tests confirmed to fail without the lib changes.

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

Author: elirangoshen

lib/OnyxUtils.ts

@@ -1333,6 +1333,21 @@ function setWithRetry<TKey extends OnyxKey>({key, value, options}: SetParams<TKe
         });
 }
 
+/**
+ * Storage-write tail of multiSetWithRetry, isolated so that retryOperation re-enters only the
+ * 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> {
+    const {keyValuePairsToStore, newData} = params;
+
+    return Storage.multiSet(keyValuePairsToStore)
+        .catch((error) => OnyxUtils.retryOperation(error, persistMultiSetWrite, params, retryAttempt))
+        .then(() => {
+            OnyxUtils.sendActionToDevTools(OnyxUtils.METHOD.MULTI_SET, undefined, newData);
+        });
+}
+
 /**
  * Sets multiple keys and values.
  * Serves as core implementation for `Onyx.multiSet()` public function, the difference being
@@ -1411,10 +1426,34 @@ function multiSetWithRetry(data: OnyxMultiSetInput, retryAttempt?: number): Prom
         return !OnyxKeys.isRamOnlyKey(key);
     });
 
-    return Storage.multiSet(keyValuePairsToStore)
-        .catch((error) => OnyxUtils.retryOperation(error, multiSetWithRetry, newData, retryAttempt))
+    return persistMultiSetWrite({keyValuePairsToStore, newData}, retryAttempt);
+}
+
+/**
+ * Storage-write tail of setCollectionWithRetry, isolated so that retryOperation re-enters only the
+ * storage step. Cache and subscriber notifications already happened in the orchestrator, so
+ * retries no longer re-fire `waitForCollectionCallback` subscribers with the same payload.
+ */
+function persistCollectionWrite<TKey extends CollectionKeyBase>(
+    params: {
+        collectionKey: TKey;
+        keyValuePairs: StorageKeyValuePair[];
+        mutableCollection: OnyxInputKeyValueMapping;
+    },
+    retryAttempt?: number,
+): Promise<void> {
+    const {collectionKey, keyValuePairs, mutableCollection} = params;
+
+    // RAM-only keys never persist to storage.
+    if (OnyxKeys.isRamOnlyKey(collectionKey)) {
+        OnyxUtils.sendActionToDevTools(OnyxUtils.METHOD.SET_COLLECTION, undefined, mutableCollection);
+        return Promise.resolve();
+    }
+
+    return Storage.multiSet(keyValuePairs)
+        .catch((error) => OnyxUtils.retryOperation(error, persistCollectionWrite, params, retryAttempt))
         .then(() => {
-            OnyxUtils.sendActionToDevTools(OnyxUtils.METHOD.MULTI_SET, undefined, newData);
+            OnyxUtils.sendActionToDevTools(OnyxUtils.METHOD.SET_COLLECTION, undefined, mutableCollection);
         });
 }
 
@@ -1478,20 +1517,46 @@ function setCollectionWithRetry<TKey extends CollectionKeyBase>({collectionKey,
 
         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;
-        }
-
-        return Storage.multiSet(keyValuePairs)
-            .catch((error) => OnyxUtils.retryOperation(error, setCollectionWithRetry, {collectionKey, collection}, retryAttempt))
-            .then(() => {
-                OnyxUtils.sendActionToDevTools(OnyxUtils.METHOD.SET_COLLECTION, undefined, mutableCollection);
-            });
+        return persistCollectionWrite({collectionKey, keyValuePairs, mutableCollection}, retryAttempt);
     });
 }
 
+/**
+ * Storage-write tail of mergeCollectionWithPatches, isolated so that retryOperation re-enters only
+ * the storage step. Cache and subscriber notifications already happened in the orchestrator, and
+ * the existing/new key split is captured in `params` — so retries don't re-fire subscribers and
+ * don't reclassify keys against a cache that was already mutated on the first attempt.
+ */
+function persistMergedCollectionWrite<TKey extends CollectionKeyBase>(
+    params: {
+        collectionKey: TKey;
+        keyValuePairsForExistingCollection: StorageKeyValuePair[];
+        keyValuePairsForNewCollection: StorageKeyValuePair[];
+        resultCollection: OnyxInputKeyValueMapping;
+    },
+    retryAttempt?: number,
+): Promise<void> {
+    const {collectionKey, keyValuePairsForExistingCollection, keyValuePairsForNewCollection, resultCollection} = params;
+
+    const promises = [];
+
+    // New keys are added via multiSet while existing keys are updated using multiMerge; using
+    // multiMerge on a key that doesn't exist yet throws on some storage backends. RAM-only keys
+    // never persist to storage.
+    if (!OnyxKeys.isRamOnlyKey(collectionKey) && keyValuePairsForExistingCollection.length > 0) {
+        promises.push(Storage.multiMerge(keyValuePairsForExistingCollection));
+    }
+    if (!OnyxKeys.isRamOnlyKey(collectionKey) && keyValuePairsForNewCollection.length > 0) {
+        promises.push(Storage.multiSet(keyValuePairsForNewCollection));
+    }
+
+    return Promise.all(promises)
+        .catch((error) => retryOperation(error, persistMergedCollectionWrite, params, retryAttempt))
+        .then(() => {
+            sendActionToDevTools(METHOD.MERGE_COLLECTION, undefined, resultCollection);
+        });
+}
+
 /**
  * Merges a collection based on their keys.
  * Serves as core implementation for `Onyx.mergeCollection()` public function, the difference being
@@ -1613,32 +1678,7 @@ function mergeCollectionWithPatches<TKey extends CollectionKeyBase>(
                 cache.merge(finalMergedCollection);
                 keysChanged(collectionKey, finalMergedCollection, previousCollection);
 
-                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
-                // 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));
-                }
-
-                // We can skip this step for RAM-only keys as they should never be saved to storage
-                if (!OnyxKeys.isRamOnlyKey(collectionKey) && keyValuePairsForNewCollection.length > 0) {
-                    promises.push(Storage.multiSet(keyValuePairsForNewCollection));
-                }
-
-                return Promise.all(promises)
-                    .catch((error) =>
-                        retryOperation(
-                            error,
-                            mergeCollectionWithPatches,
-                            {collectionKey, collection: resultCollection as OnyxMergeCollectionInput<TKey>, mergeReplaceNullPatches, isProcessingCollectionUpdate},
-                            retryAttempt,
-                        ),
-                    )
-                    .then(() => {
-                        sendActionToDevTools(METHOD.MERGE_COLLECTION, undefined, resultCollection);
-                    });
+                return persistMergedCollectionWrite({collectionKey, keyValuePairsForExistingCollection, keyValuePairsForNewCollection, resultCollection}, retryAttempt);
             });
         })
         .then(() => undefined);
@@ -1692,16 +1732,7 @@ function partialSetCollection<TKey extends CollectionKeyBase>({collectionKey, co
 
         keysChanged(collectionKey, mutableCollection, previousCollection);
 
-        if (OnyxKeys.isRamOnlyKey(collectionKey)) {
-            sendActionToDevTools(METHOD.SET_COLLECTION, undefined, mutableCollection);
-            return;
-        }
-
-        return Storage.multiSet(keyValuePairs)
-            .catch((error) => retryOperation(error, partialSetCollection, {collectionKey, collection}, retryAttempt))
-            .then(() => {
-                sendActionToDevTools(METHOD.SET_COLLECTION, undefined, mutableCollection);
-            });
+        return persistCollectionWrite({collectionKey, keyValuePairs, mutableCollection}, retryAttempt);
     });
 }
 
@@ -1766,12 +1797,15 @@ const OnyxUtils = {
     reduceCollectionWithSelector,
     updateSnapshots,
     mergeCollectionWithPatches,
+    persistMergedCollectionWrite,
     partialSetCollection,
     logKeyChanged,
     logKeyRemoved,
     setWithRetry,
     multiSetWithRetry,
+    persistMultiSetWrite,
     setCollectionWithRetry,
+    persistCollectionWrite,
 };
 
 export type {OnyxMethod};

lib/types.ts

@@ -371,8 +371,11 @@ type MergeCollectionWithPatchesParams<TKey extends CollectionKeyBase> = {
 type RetriableOnyxOperation =
     | typeof OnyxUtils.setWithRetry
     | typeof OnyxUtils.multiSetWithRetry
+    | typeof OnyxUtils.persistMultiSetWrite
     | typeof OnyxUtils.setCollectionWithRetry
+    | typeof OnyxUtils.persistCollectionWrite
     | typeof OnyxUtils.mergeCollectionWithPatches
+    | typeof OnyxUtils.persistMergedCollectionWrite
     | typeof OnyxUtils.partialSetCollection;
 
 /**

tests/unit/onyxUtilsTest.ts

@@ -920,6 +920,156 @@ describe('OnyxUtils', () => {
         });
     });
 
+    describe('retry side-effect idempotency', () => {
+        // Save originals so each test can replace StorageMock.multiMerge / StorageMock.multiSet
+        // with a mock that rejects once and then resolves, exercising the retryOperation path
+        // without burning the full retry budget. Restoring keeps mocks from leaking into the
+        // storage-eviction describe block below (which depends on these storage methods).
+        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 — new keys still route through multiSet on retry, not downgraded to multiMerge', async () => {
+            const collectionKey = ONYXKEYS.COLLECTION.TEST_KEY;
+            const existingMemberKey = `${collectionKey}1`;
+            const newMemberKey = `${collectionKey}2`;
+
+            // Seed an existing member via setItem so the multiMerge mock below doesn't intercept seeding.
+            await Onyx.set(existingMemberKey, {value: 'initial'});
+
+            const multiMergeSpy = jest.fn(originalMultiMerge).mockRejectedValueOnce(transientError);
+            StorageMock.multiMerge = multiMergeSpy;
+
+            await Onyx.mergeCollection(collectionKey, {
+                [existingMemberKey]: {value: 'merged'},
+                [newMemberKey]: {value: 'new'},
+            } as GenericCollection);
+
+            // Before this fix, the retry attempt re-derived the existing/new split against a cache
+            // already mutated by the first attempt, which silently routed `newMemberKey` through
+            // multiMerge. Benign on IDB (multiMerge on a missing key behaves like set) but wrong
+            // semantically and a crash on storage backends that require multiSet for new keys.
+            const allMultiMergeKeys = multiMergeSpy.mock.calls.flatMap((args) => (args[0] as Array<[string, unknown]>).map(([key]) => key));
+            expect(allMultiMergeKeys).not.toContain(newMemberKey);
+
+            // Sanity: the retry actually happened.
+            expect(multiMergeSpy).toHaveBeenCalledTimes(2);
+        });
+
+        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 re-enter only the storage-write helper, 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');