Make retried Onyx writes idempotent on storage retry

API-INTERNAL.md

@@ -125,17 +125,33 @@ It will also mark deep nested objects that need to be entirely replaced during t
 Serves as core implementation for <code>Onyx.set()</code> public function, the difference being
 that this internal function allows passing an additional <code>retryAttempt</code> parameter to retry on failure.</p>
 </dd>
+<dt><a href="#persistMultiSetWrite">persistMultiSetWrite()</a></dt>
+<dd><p>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 <code>waitForCollectionCallback</code> subscribers with the same payload.</p>
+</dd>
 <dt><a href="#multiSetWithRetry">multiSetWithRetry(data, retryAttempt)</a></dt>
 <dd><p>Sets multiple keys and values.
 Serves as core implementation for <code>Onyx.multiSet()</code> public function, the difference being
 that this internal function allows passing an additional <code>retryAttempt</code> parameter to retry on failure.</p>
 </dd>
+<dt><a href="#persistCollectionWrite">persistCollectionWrite()</a></dt>
+<dd><p>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 <code>waitForCollectionCallback</code> subscribers with the same payload.</p>
+</dd>
 <dt><a href="#setCollectionWithRetry">setCollectionWithRetry(params, retryAttempt)</a></dt>
 <dd><p>Sets a collection by replacing all existing collection members with new values.
 Any existing collection members not included in the new data will be removed.
 Serves as core implementation for <code>Onyx.setCollection()</code> public function, the difference being
 that this internal function allows passing an additional <code>retryAttempt</code> parameter to retry on failure.</p>
 </dd>
+<dt><a href="#persistMergedCollectionWrite">persistMergedCollectionWrite()</a></dt>
+<dd><p>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 <code>params</code> — so retries don&#39;t re-fire subscribers and
+don&#39;t reclassify keys against a cache that was already mutated on the first attempt.</p>
+</dd>
 <dt><a href="#mergeCollectionWithPatches">mergeCollectionWithPatches(params, retryAttempt)</a></dt>
 <dd><p>Merges a collection based on their keys.
 Serves as core implementation for <code>Onyx.mergeCollection()</code> public function, the difference being
@@ -435,6 +451,14 @@ that this internal function allows passing an additional `retryAttempt` paramete
 | params.options | optional configuration object |
 | retryAttempt | retry attempt |
 
+<a name="persistMultiSetWrite"></a>
+
+## persistMultiSetWrite()
+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.
+
+**Kind**: global function  
 <a name="multiSetWithRetry"></a>
 
 ## multiSetWithRetry(data, retryAttempt)
@@ -449,6 +473,14 @@ that this internal function allows passing an additional `retryAttempt` paramete
 | data | object keyed by ONYXKEYS and the values to set |
 | retryAttempt | retry attempt |
 
+<a name="persistCollectionWrite"></a>
+
+## persistCollectionWrite()
+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.
+
+**Kind**: global function  
 <a name="setCollectionWithRetry"></a>
 
 ## setCollectionWithRetry(params, retryAttempt)
@@ -466,6 +498,15 @@ that this internal function allows passing an additional `retryAttempt` paramete
 | params.collection | Object collection keyed by individual collection member keys and values |
 | retryAttempt | retry attempt |
 
+<a name="persistMergedCollectionWrite"></a>
+
+## persistMergedCollectionWrite()
+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.
+
+**Kind**: global function  
 <a name="mergeCollectionWithPatches"></a>
 
 ## mergeCollectionWithPatches(params, retryAttempt)

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);
+    });
 }
 
 /**
@@ -1333,6 +1345,24 @@ 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; restoreEvictedKey?: (evictedKey: OnyxKey) => void},
+    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 +1441,59 @@ function multiSetWithRetry(data: OnyxMultiSetInput, retryAttempt?: number): Prom
         return !OnyxKeys.isRamOnlyKey(key);
     });
 
-    return Storage.multiSet(keyValuePairsToStore)
-        .catch((error) => OnyxUtils.retryOperation(error, multiSetWithRetry, 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);
+}
+
+/**
+ * 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;
+        restoreEvictedKey?: (evictedKey: OnyxKey) => void;
+    },
+    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 +1557,62 @@ 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;
+        // 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 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, restoreEvictedKey}, 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;
+        restoreEvictedKey?: (evictedKey: OnyxKey) => void;
+    },
+    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 +1734,29 @@ 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));
+                // 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));
                 }
 
-                // 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));
-                }
+                // 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 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, restoreEvictedKey}, retryAttempt);
             });
         })
         .then(() => undefined);
@@ -1692,16 +1810,22 @@ function partialSetCollection<TKey extends CollectionKeyBase>({collectionKey, co
 
         keysChanged(collectionKey, mutableCollection, previousCollection);
 
-        if (OnyxKeys.isRamOnlyKey(collectionKey)) {
-            sendActionToDevTools(METHOD.SET_COLLECTION, undefined, mutableCollection);
-            return;
+        // 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 Storage.multiSet(keyValuePairs)
-            .catch((error) => retryOperation(error, partialSetCollection, {collectionKey, collection}, retryAttempt))
-            .then(() => {
-                sendActionToDevTools(METHOD.SET_COLLECTION, undefined, mutableCollection);
-            });
+        return persistCollectionWrite({collectionKey, keyValuePairs, mutableCollection, restoreEvictedKey}, retryAttempt);
     });
 }
 
@@ -1766,12 +1890,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;
 
 /**