From 03e26845e2c220b1350a35179acf3435e2711282 Mon Sep 17 00:00:00 2001
From: Evan You <yyx990803@gmail.com>
Date: Tue, 6 Jul 2021 14:31:53 -0400
Subject: [PATCH] fix(sfc): fix `<script setup>` async context preservation
 logic

fix #4050
---
 .../__tests__/compileScript.spec.ts           | 20 +++-
 packages/compiler-sfc/src/compileScript.ts    | 39 ++++++--
 .../__tests__/apiSetupHelpers.spec.ts         | 94 ++++++++++++++++---
 packages/runtime-core/src/apiSetupHelpers.ts  | 47 ++++++----
 4 files changed, 152 insertions(+), 48 deletions(-)

diff --git a/packages/compiler-sfc/__tests__/compileScript.spec.ts b/packages/compiler-sfc/__tests__/compileScript.spec.ts
index 1da416b36fb..d75927d33bb 100644
--- a/packages/compiler-sfc/__tests__/compileScript.spec.ts
+++ b/packages/compiler-sfc/__tests__/compileScript.spec.ts
@@ -847,6 +847,9 @@ const emit = defineEmits(['a', 'b'])
       const { content } = compile(`<script setup>${code}</script>`, {
         refSugar: true
       })
+      if (shouldAsync) {
+        expect(content).toMatch(`let __temp, __restore`)
+      }
       expect(content).toMatch(`${shouldAsync ? `async ` : ``}setup(`)
       if (typeof expected === 'string') {
         expect(content).toMatch(expected)
@@ -856,28 +859,35 @@ const emit = defineEmits(['a', 'b'])
     }
 
     test('expression statement', () => {
-      assertAwaitDetection(`await foo`, `await _withAsyncContext(foo)`)
+      assertAwaitDetection(
+        `await foo`,
+        `;(([__temp,__restore]=_withAsyncContext(()=>(foo))),__temp=await __temp,__restore())`
+      )
     })
 
     test('variable', () => {
       assertAwaitDetection(
         `const a = 1 + (await foo)`,
-        `1 + (await _withAsyncContext(foo))`
+        `1 + ((([__temp,__restore]=_withAsyncContext(()=>(foo))),__temp=await __temp,__restore(),__temp))`
       )
     })
 
     test('ref', () => {
       assertAwaitDetection(
         `ref: a = 1 + (await foo)`,
-        `1 + (await _withAsyncContext(foo))`
+        `1 + ((([__temp,__restore]=_withAsyncContext(()=>(foo))),__temp=await __temp,__restore(),__temp))`
       )
     })
 
     test('nested statements', () => {
       assertAwaitDetection(`if (ok) { await foo } else { await bar }`, code => {
         return (
-          code.includes(`await _withAsyncContext(foo)`) &&
-          code.includes(`await _withAsyncContext(bar)`)
+          code.includes(
+            `;(([__temp,__restore]=_withAsyncContext(()=>(foo))),__temp=await __temp,__restore())`
+          ) &&
+          code.includes(
+            `;(([__temp,__restore]=_withAsyncContext(()=>(bar))),__temp=await __temp,__restore())`
+          )
         )
       })
     })
diff --git a/packages/compiler-sfc/src/compileScript.ts b/packages/compiler-sfc/src/compileScript.ts
index 94d351ab03c..e5a792d77c3 100644
--- a/packages/compiler-sfc/src/compileScript.ts
+++ b/packages/compiler-sfc/src/compileScript.ts
@@ -32,7 +32,8 @@ import {
   LabeledStatement,
   CallExpression,
   RestElement,
-  TSInterfaceBody
+  TSInterfaceBody,
+  AwaitExpression
 } from '@babel/types'
 import { walk } from 'estree-walker'
 import { RawSourceMap } from 'source-map'
@@ -487,6 +488,25 @@ export function compileScript(
     })
   }
 
+  /**
+   * await foo()
+   * -->
+   * (([__temp, __restore] = withAsyncContext(() => foo())),__temp=await __temp,__restore(),__temp)
+   */
+  function processAwait(node: AwaitExpression, isStatement: boolean) {
+    s.overwrite(
+      node.start! + startOffset,
+      node.argument.start! + startOffset,
+      `${isStatement ? `;` : ``}(([__temp,__restore]=${helper(
+        `withAsyncContext`
+      )}(()=>(`
+    )
+    s.appendLeft(
+      node.end! + startOffset,
+      `))),__temp=await __temp,__restore()${isStatement ? `` : `,__temp`})`
+    )
+  }
+
   function processRefExpression(exp: Expression, statement: LabeledStatement) {
     if (exp.type === 'AssignmentExpression') {
       const { left, right } = exp
@@ -949,17 +969,13 @@ export function compileScript(
       node.type.endsWith('Statement')
     ) {
       ;(walk as any)(node, {
-        enter(node: Node) {
-          if (isFunction(node)) {
+        enter(child: Node, parent: Node) {
+          if (isFunction(child)) {
             this.skip()
           }
-          if (node.type === 'AwaitExpression') {
+          if (child.type === 'AwaitExpression') {
             hasAwait = true
-            s.prependRight(
-              node.argument.start! + startOffset,
-              helper(`withAsyncContext`) + `(`
-            )
-            s.appendLeft(node.argument.end! + startOffset, `)`)
+            processAwait(child, parent.type === 'ExpressionStatement')
           }
         }
       })
@@ -1151,6 +1167,11 @@ export function compileScript(
   if (propsIdentifier) {
     s.prependRight(startOffset, `\nconst ${propsIdentifier} = __props`)
   }
+  // inject temp variables for async context preservation
+  if (hasAwait) {
+    const any = isTS ? `:any` : ``
+    s.prependRight(startOffset, `\nlet __temp${any}, __restore${any}\n`)
+  }
 
   const destructureElements =
     hasDefineExposeCall || !options.inlineTemplate ? [`expose`] : []
diff --git a/packages/runtime-core/__tests__/apiSetupHelpers.spec.ts b/packages/runtime-core/__tests__/apiSetupHelpers.spec.ts
index 5244a9a87b8..3dc22e67ce5 100644
--- a/packages/runtime-core/__tests__/apiSetupHelpers.spec.ts
+++ b/packages/runtime-core/__tests__/apiSetupHelpers.spec.ts
@@ -119,12 +119,20 @@ describe('SFC <script setup> helpers', () => {
 
       const Comp = defineComponent({
         async setup() {
+          let __temp: any, __restore: any
+
           beforeInstance = getCurrentInstance()
-          const msg = await withAsyncContext(
-            new Promise(r => {
-              resolve = r
-            })
-          )
+
+          const msg = (([__temp, __restore] = withAsyncContext(
+            () =>
+              new Promise(r => {
+                resolve = r
+              })
+          )),
+          (__temp = await __temp),
+          __restore(),
+          __temp)
+
           // register the lifecycle after an await statement
           onMounted(spy)
           afterInstance = getCurrentInstance()
@@ -155,13 +163,18 @@ describe('SFC <script setup> helpers', () => {
 
       const Comp = defineComponent({
         async setup() {
+          let __temp: any, __restore: any
+
           beforeInstance = getCurrentInstance()
           try {
-            await withAsyncContext(
-              new Promise((r, rj) => {
-                reject = rj
-              })
+            ;[__temp, __restore] = withAsyncContext(
+              () =>
+                new Promise((_, rj) => {
+                  reject = rj
+                })
             )
+            __temp = await __temp
+            __restore()
           } catch (e) {
             // ignore
           }
@@ -206,11 +219,20 @@ describe('SFC <script setup> helpers', () => {
 
       const Comp = defineComponent({
         async setup() {
+          let __temp: any, __restore: any
+
           beforeInstance = getCurrentInstance()
+
           // first await
-          await withAsyncContext(Promise.resolve())
+          ;[__temp, __restore] = withAsyncContext(() => Promise.resolve())
+          __temp = await __temp
+          __restore()
+
           // setup exit, instance set to null, then resumed
-          await withAsyncContext(doAsyncWork())
+          ;[__temp, __restore] = withAsyncContext(() => doAsyncWork())
+          __temp = await __temp
+          __restore()
+
           afterInstance = getCurrentInstance()
           return () => {
             resolve()
@@ -237,8 +259,13 @@ describe('SFC <script setup> helpers', () => {
 
       const Comp = defineComponent({
         async setup() {
-          await withAsyncContext(Promise.resolve())
-          await withAsyncContext(Promise.reject())
+          let __temp: any, __restore: any
+          ;[__temp, __restore] = withAsyncContext(() => Promise.resolve())
+          __temp = await __temp
+          __restore()
+          ;[__temp, __restore] = withAsyncContext(() => Promise.reject())
+          __temp = await __temp
+          __restore()
         },
         render() {}
       })
@@ -256,6 +283,42 @@ describe('SFC <script setup> helpers', () => {
       expect(getCurrentInstance()).toBeNull()
     })
 
+    // #4050
+    test('race conditions', async () => {
+      const uids = {
+        one: { before: NaN, after: NaN },
+        two: { before: NaN, after: NaN }
+      }
+
+      const Comp = defineComponent({
+        props: ['name'],
+        async setup(props: { name: 'one' | 'two' }) {
+          let __temp: any, __restore: any
+
+          uids[props.name].before = getCurrentInstance()!.uid
+          ;[__temp, __restore] = withAsyncContext(() => Promise.resolve())
+          __temp = await __temp
+          __restore()
+
+          uids[props.name].after = getCurrentInstance()!.uid
+          return () => ''
+        }
+      })
+
+      const app = createApp(() =>
+        h(Suspense, () =>
+          h('div', [h(Comp, { name: 'one' }), h(Comp, { name: 'two' })])
+        )
+      )
+      const root = nodeOps.createElement('div')
+      app.mount(root)
+
+      await new Promise(r => setTimeout(r))
+      expect(uids.one.before).not.toBe(uids.two.before)
+      expect(uids.one.before).toBe(uids.one.after)
+      expect(uids.two.before).toBe(uids.two.after)
+    })
+
     test('should teardown in-scope effects', async () => {
       let resolve: (val?: any) => void
       const ready = new Promise(r => {
@@ -266,7 +329,10 @@ describe('SFC <script setup> helpers', () => {
 
       const Comp = defineComponent({
         async setup() {
-          await withAsyncContext(Promise.resolve())
+          let __temp: any, __restore: any
+          ;[__temp, __restore] = withAsyncContext(() => Promise.resolve())
+          __temp = await __temp
+          __restore()
 
           c = computed(() => {})
           // register the lifecycle after an await statement
diff --git a/packages/runtime-core/src/apiSetupHelpers.ts b/packages/runtime-core/src/apiSetupHelpers.ts
index 5bbcffd9575..4dcadbabdde 100644
--- a/packages/runtime-core/src/apiSetupHelpers.ts
+++ b/packages/runtime-core/src/apiSetupHelpers.ts
@@ -1,9 +1,9 @@
 import { isPromise } from '../../shared/src'
 import {
   getCurrentInstance,
+  setCurrentInstance,
   SetupContext,
-  createSetupContext,
-  setCurrentInstance
+  createSetupContext
 } from './component'
 import { EmitFn, EmitsOptions } from './componentEmits'
 import {
@@ -230,25 +230,32 @@ export function mergeDefaults(
 }
 
 /**
- * Runtime helper for storing and resuming current instance context in
- * async setup().
+ * `<script setup>` helper for persisting the current instance context over
+ * async/await flows.
+ *
+ * `@vue/compiler-sfc` converts the following:
+ *
+ * ```ts
+ * const x = await foo()
+ * ```
+ *
+ * into:
+ *
+ * ```ts
+ * let __temp, __restore
+ * const x = (([__temp, __restore] = withAsyncContext(() => foo())),__temp=await __temp,__restore(),__temp)
+ * ```
+ * @internal
  */
-export function withAsyncContext<T>(awaitable: T | Promise<T>): Promise<T> {
+export function withAsyncContext(getAwaitable: () => any) {
   const ctx = getCurrentInstance()
-  setCurrentInstance(null) // unset after storing instance
-  if (__DEV__ && !ctx) {
-    warn(`withAsyncContext() called when there is no active context instance.`)
+  let awaitable = getAwaitable()
+  setCurrentInstance(null)
+  if (isPromise(awaitable)) {
+    awaitable = awaitable.catch(e => {
+      setCurrentInstance(ctx)
+      throw e
+    })
   }
-  return isPromise<T>(awaitable)
-    ? awaitable.then(
-        res => {
-          setCurrentInstance(ctx)
-          return res
-        },
-        err => {
-          setCurrentInstance(ctx)
-          throw err
-        }
-      )
-    : (awaitable as any)
+  return [awaitable, () => setCurrentInstance(ctx)]
 }