feat: Optimistic mutation with error rollback (vercel#1745)

* add optimisticData and rollbackOnError

* add test cases

* add test for rollback

Author: shuding

src/types.ts

@@ -142,9 +142,11 @@ export type MutatorCallback<Data = any> = (
   currentData?: Data
 ) => Promise<undefined | Data> | undefined | Data
 
-export type MutatorOptions = {
+export type MutatorOptions<Data = any> = {
   revalidate?: boolean
   populateCache?: boolean
+  optimisticData?: Data
+  rollbackOnError?: boolean
 }
 
 export type Broadcaster<Data = any, Error = any> = (
@@ -167,7 +169,7 @@ export type MutatorFn<Data = any> = (
   cache: Cache,
   key: Key,
   data?: Data | Promise<Data> | MutatorCallback<Data>,
-  opts?: boolean | MutatorOptions
+  opts?: boolean | MutatorOptions<Data>
 ) => Promise<Data | undefined>
 
 export type MutatorWrapper<Fn> = Fn extends (
@@ -189,19 +191,19 @@ export interface ScopedMutator<Data = any> {
   (
     key: Key,
     data?: Data | Promise<Data> | MutatorCallback<Data>,
-    opts?: boolean | MutatorOptions
+    opts?: boolean | MutatorOptions<Data>
   ): Promise<Data | undefined>
   /** This is used for global mutator */
   <T = any>(
     key: Key,
     data?: T | Promise<T> | MutatorCallback<T>,
-    opts?: boolean | MutatorOptions
+    opts?: boolean | MutatorOptions<Data>
   ): Promise<T | undefined>
 }
 
 export type KeyedMutator<Data> = (
   data?: Data | Promise<Data> | MutatorCallback<Data>,
-  opts?: boolean | MutatorOptions
+  opts?: boolean | MutatorOptions<Data>
 ) => Promise<Data | undefined>
 
 // Public types

src/utils/broadcast-state.ts

@@ -9,15 +9,15 @@ export const broadcastState: Broadcaster = (
   error,
   isValidating,
   revalidate,
-  populateCache = true
+  broadcast = true
 ) => {
   const [EVENT_REVALIDATORS, STATE_UPDATERS, , , CONCURRENT_REQUESTS] =
     SWRGlobalState.get(cache) as GlobalState
   const revalidators = EVENT_REVALIDATORS[key]
   const updaters = STATE_UPDATERS[key]
 
   // Cache was populated, update states of all hooks.
-  if (populateCache && updaters) {
+  if (broadcast && updaters) {
     for (let i = 0; i < updaters.length; ++i) {
       updaters[i](data, error, isValidating)
     }

src/utils/mutate.ts

@@ -1,5 +1,5 @@
 import { serialize } from './serialize'
-import { isFunction, isUndefined, mergeObjects, UNDEFINED } from './helper'
+import { isFunction, isUndefined, UNDEFINED } from './helper'
 import { SWRGlobalState, GlobalState } from './global-state'
 import { broadcastState } from './broadcast-state'
 import { getTimestamp } from './timestamp'
@@ -11,7 +11,7 @@ export const internalMutate = async <Data>(
     Cache,
     Key,
     undefined | Data | Promise<Data | undefined> | MutatorCallback<Data>,
-    undefined | boolean | MutatorOptions
+    undefined | boolean | MutatorOptions<Data>
   ]
 ) => {
   const [cache, _key, _data, _opts] = args
@@ -22,8 +22,10 @@ export const internalMutate = async <Data>(
     typeof _opts === 'boolean' ? { revalidate: _opts } : _opts || {}
 
   // Fallback to `true` if it's not explicitly set to `false`
+  let populateCache = options.populateCache !== false
   const revalidate = options.revalidate !== false
-  const populateCache = options.populateCache !== false
+  const rollbackOnError = options.rollbackOnError !== false
+  const optimisticData = options.optimisticData
 
   // Serilaize key
   const [key, , keyInfo] = serialize(_key)
@@ -48,16 +50,13 @@ export const internalMutate = async <Data>(
   let error: unknown
 
   // Update global timestamps.
-  const beforeMutationTs = getTimestamp()
-  MUTATION[key] = [beforeMutationTs, 0]
-  const hasCustomOptimisticData = !isUndefined(customOptimisticData)
+  const beforeMutationTs = (MUTATION_TS[key] = getTimestamp())
+  MUTATION_END_TS[key] = 0
+  const hasOptimisticData = !isUndefined(optimisticData)
   const rollbackData = cache.get(key)
 
   // Do optimistic data update.
-  if (hasCustomOptimisticData) {
-    const optimisticData = isFunction(customOptimisticData)
-      ? customOptimisticData(rollbackData)
-      : customOptimisticData
+  if (hasOptimisticData) {
     cache.set(key, optimisticData)
     broadcastState(cache, key, optimisticData)
   }
@@ -86,9 +85,8 @@ export const internalMutate = async <Data>(
     if (beforeMutationTs !== MUTATION[key][0]) {
       if (error) throw error
       return data
-    } else if (error && hasCustomOptimisticData && rollbackOnError) {
-      // Rollback. Always populate the cache in this case but without
-      // transforming the data.
+    } else if (error && hasOptimisticData && rollbackOnError) {
+      // Rollback. Always populate the cache in this case.
       populateCache = true
       data = rollbackData
       cache.set(key, rollbackData)

test/use-swr-local-mutation.test.tsx

@@ -1011,4 +1011,105 @@ describe('useSWR - local mutation', () => {
     await sleep(30)
     await screen.findByText('data: foo')
   })
+
+  it('should support optimistic updates via `optimisticData`', async () => {
+    const key = createKey()
+    const renderedData = []
+    let mutate
+
+    function Page() {
+      const { data, mutate: boundMutate } = useSWR(key, () =>
+        createResponse('foo', { delay: 20 })
+      )
+      mutate = boundMutate
+      renderedData.push(data)
+      return <div>data: {String(data)}</div>
+    }
+
+    renderWithConfig(<Page />)
+    await screen.findByText('data: foo')
+
+    await act(() =>
+      mutate(createResponse('baz', { delay: 20 }), {
+        optimisticData: 'bar'
+      })
+    )
+    await sleep(30)
+    expect(renderedData).toEqual([undefined, 'foo', 'bar', 'baz', 'foo'])
+  })
+
+  it('should rollback optimistic updates when mutation fails', async () => {
+    const key = createKey()
+    const renderedData = []
+    let mutate
+    let cnt = 0
+
+    function Page() {
+      const { data, mutate: boundMutate } = useSWR(key, () =>
+        createResponse(cnt++, { 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')
+
+    try {
+      await act(() =>
+        mutate(createResponse(new Error('baz'), { delay: 20 }), {
+          optimisticData: 'bar'
+        })
+      )
+    } catch (e) {
+      expect(e.message).toEqual('baz')
+    }
+
+    await sleep(30)
+    expect(renderedData).toEqual([undefined, 0, 'bar', 0, 1])
+  })
+
+  it('should not rollback optimistic updates if `rollbackOnError`', async () => {
+    const key = createKey()
+    const renderedData = []
+    let mutate
+    let cnt = 0
+
+    function Page() {
+      const { data, mutate: boundMutate } = useSWR(key, () =>
+        createResponse(cnt++, { 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')
+
+    try {
+      await act(() =>
+        mutate(createResponse(new Error('baz'), { delay: 20 }), {
+          optimisticData: 'bar',
+          rollbackOnError: false
+        })
+      )
+    } catch (e) {
+      expect(e.message).toEqual('baz')
+    }
+
+    await sleep(30)
+    expect(renderedData).toEqual([undefined, 0, 'bar', 1])
+  })
 })