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> = (
   currentValue?: 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,27 +169,27 @@ export type Mutator<Data = any> = (
   cache: Cache,
   key: Key,
   data?: Data | Promise<Data> | MutatorCallback<Data>,
-  opts?: boolean | MutatorOptions
+  opts?: boolean | MutatorOptions<Data>
 ) => Promise<Data | undefined>
 
 export interface ScopedMutator<Data = any> {
   /** This is used for bound mutator */
   (
     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] || []
+  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, 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, , keyErr] = serialize(_key)
@@ -53,6 +55,14 @@ export const internalMutate = async <Data>(
   // Update global timestamps.
   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 (hasOptimisticData) {
+    cache.set(key, optimisticData)
+    broadcastState(cache, key, optimisticData)
+  }
 
   if (isFunction(data)) {
     // `data` is a function, call it passing current cache value.
@@ -78,6 +88,11 @@ export const internalMutate = async <Data>(
     if (beforeMutationTs !== MUTATION_TS[key]) {
       if (error) throw error
       return 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])
+  })
 })