From 2dd94517f2cb60cef500aadd39d469706784f346 Mon Sep 17 00:00:00 2001
From: Shu Ding <g@shud.in>
Date: Sun, 15 May 2022 17:09:42 +0200
Subject: [PATCH] Fix race conditions related to optimistic UI (#1970)

* fix optimistic ui race conditions

* fix test

* fix
---
 _internal/utils/mutate.ts            | 27 ++++++++++----
 package.json                         |  2 +-
 test/use-swr-local-mutation.test.tsx | 56 ++++++++++++++++++++++++++++
 3 files changed, 76 insertions(+), 9 deletions(-)

diff --git a/_internal/utils/mutate.ts b/_internal/utils/mutate.ts
index 74cbca98b..cc2544956 100644
--- a/_internal/utils/mutate.ts
+++ b/_internal/utils/mutate.ts
@@ -1,5 +1,5 @@
 import { serialize } from './serialize'
-import { createCacheHelper, isFunction, isUndefined } from './helper'
+import { createCacheHelper, isFunction, isUndefined, UNDEFINED } from './helper'
 import { SWRGlobalState } from './global-state'
 import { getTimestamp } from './timestamp'
 import * as revalidateEvents from '../constants'
@@ -8,7 +8,8 @@ import {
   Cache,
   MutatorCallback,
   MutatorOptions,
-  GlobalState
+  GlobalState,
+  State
 } from '../types'
 
 export const internalMutate = async <Data>(
@@ -38,12 +39,19 @@ export const internalMutate = async <Data>(
   const [key] = serialize(_key)
   if (!key) return
 
-  const [get, set] = createCacheHelper<Data>(cache, key)
+  const [get, set] = createCacheHelper<
+    Data,
+    State<Data, any> & {
+      // The original data.
+      _o?: Data
+    }
+  >(cache, key)
   const [EVENT_REVALIDATORS, MUTATION, FETCH] = SWRGlobalState.get(
     cache
   ) as GlobalState
+
   const revalidators = EVENT_REVALIDATORS[key]
-  const startRevalidate = async () => {
+  const startRevalidate = () => {
     if (revalidate) {
       // Invalidate the key by deleting the concurrent request markers so new
       // requests will not be deduped.
@@ -56,6 +64,7 @@ export const internalMutate = async <Data>(
     }
     return get().data
   }
+
   // If there is no new data provided, revalidate the key with current state.
   if (args.length < 3) {
     // Revalidate and broadcast state.
@@ -70,14 +79,16 @@ export const internalMutate = async <Data>(
   MUTATION[key] = [beforeMutationTs, 0]
 
   const hasOptimisticData = !isUndefined(optimisticData)
-  const originalData = get().data
+  const state = get()
+  const currentData = state.data
+  const originalData = isUndefined(state._o) ? currentData : state._o
 
   // Do optimistic data update.
   if (hasOptimisticData) {
     optimisticData = isFunction(optimisticData)
       ? optimisticData(originalData)
       : optimisticData
-    set({ data: optimisticData })
+    set({ data: optimisticData, _o: originalData })
   }
 
   if (isFunction(data)) {
@@ -125,8 +136,8 @@ export const internalMutate = async <Data>(
       set({ data })
     }
 
-    // Always update or reset the error.
-    set({ error })
+    // Always update or reset the error and original data field.
+    set({ error, _o: UNDEFINED })
   }
 
   // Reset the timestamp to mark the mutation has ended.
diff --git a/package.json b/package.json
index 765d93734..cfd0fd8fd 100644
--- a/package.json
+++ b/package.json
@@ -82,7 +82,7 @@
     "lint": "eslint . --ext .ts,.tsx --cache",
     "lint:fix": "yarn lint --fix",
     "coverage": "jest --coverage",
-    "test": "jest && tsc --noEmit -p test"
+    "test": "tsc --noEmit -p test && jest"
   },
   "husky": {
     "hooks": {
diff --git a/test/use-swr-local-mutation.test.tsx b/test/use-swr-local-mutation.test.tsx
index c81776717..c149081f8 100644
--- a/test/use-swr-local-mutation.test.tsx
+++ b/test/use-swr-local-mutation.test.tsx
@@ -1142,6 +1142,62 @@ describe('useSWR - local mutation', () => {
     expect(renderedData).toEqual([undefined, 0, 'bar', 0, 1])
   })
 
+  it('should not revert to optimistic data when rolling back', async () => {
+    const key = createKey()
+    const renderedData = []
+    let mutate
+    let previousValue
+    let previousValue2
+
+    function Page() {
+      const { data, mutate: boundMutate } = useSWR(key, () =>
+        createResponse(0, { delay: 20 })
+      )
+      mutate = boundMutate
+
+      if (
+        !renderedData.length ||
+        renderedData[renderedData.length - 1] !== data
+      ) {
+        renderedData.push(data)
+      }
+
+      return <div>data: {String(data)}</div>
+    }
+
+    renderWithConfig(<Page />)
+    await screen.findByText('data: 0')
+
+    await executeWithoutBatching(async () => {
+      const p1 = mutate(createResponse(new Error(), { delay: 20 }), {
+        optimisticData: 1
+      })
+      await sleep(10)
+      const p2 = mutate(
+        v => {
+          previousValue = v
+          return createResponse(new Error(), { delay: 20 })
+        },
+        {
+          optimisticData: v => {
+            previousValue2 = v
+            return 2
+          }
+        }
+      )
+      return Promise.all([p1, p2])
+    }).catch(_e => {})
+
+    await sleep(30)
+
+    // It should revert to `0` instead of `1` at the end.
+    expect(renderedData).toEqual([undefined, 0, 1, 2, 0])
+
+    // It should receive the original displayed data instead of current displayed data.
+    expect(previousValue).toBe(0)
+    expect(previousValue2).toBe(0)
+  })
+
   it('should not rollback optimistic updates if `rollbackOnError`', async () => {
     const key = createKey()
     const renderedData = []