Merge pull request Expensify#792 from callstack-internal/elirangoshen/fix/retry-side-effects-idempotency

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

Author: Julesssss

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

@@ -809,7 +809,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;
 
@@ -840,8 +846,10 @@ function retryOperation<TMethod extends RetriableOnyxOperation>(error: Error, on
         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
@@ -1407,14 +1415,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) => {
@@ -1423,8 +1438,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);
         });
@@ -1488,16 +1505,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);
             });
@@ -1628,12 +1651,17 @@ 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 = [];
 
-                // 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));
@@ -1644,13 +1672,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(() => {
@@ -1707,15 +1738,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);
             });