From 0240e82a38e2e0c5f0b63c228fd02b059a19073d Mon Sep 17 00:00:00 2001
From: Evan You <yyx990803@gmail.com>
Date: Tue, 29 Jun 2021 09:24:12 -0400
Subject: [PATCH] feat(sfc): auto restore current instance after await
 statements in async setup()

---
 .../__tests__/compileScript.spec.ts           | 49 ++++++++++++++++---
 packages/compiler-sfc/src/compileScript.ts    |  5 ++
 .../__tests__/apiSetupHelpers.spec.ts         | 44 ++++++++++++++++-
 packages/runtime-core/src/apiSetupHelpers.ts  | 17 ++++++-
 packages/runtime-core/src/index.ts            |  2 +
 5 files changed, 106 insertions(+), 11 deletions(-)

diff --git a/packages/compiler-sfc/__tests__/compileScript.spec.ts b/packages/compiler-sfc/__tests__/compileScript.spec.ts
index eae846ee517..54a4b698b47 100644
--- a/packages/compiler-sfc/__tests__/compileScript.spec.ts
+++ b/packages/compiler-sfc/__tests__/compileScript.spec.ts
@@ -824,37 +824,70 @@ const emit = defineEmits(['a', 'b'])
   })
 
   describe('async/await detection', () => {
-    function assertAwaitDetection(code: string, shouldAsync = true) {
+    function assertAwaitDetection(
+      code: string,
+      expected: string | ((content: string) => boolean),
+      shouldAsync = true
+    ) {
       const { content } = compile(`<script setup>${code}</script>`)
       expect(content).toMatch(`${shouldAsync ? `async ` : ``}setup(`)
+      if (typeof expected === 'string') {
+        expect(content).toMatch(expected)
+      } else {
+        expect(expected(content)).toBe(true)
+      }
     }
 
     test('expression statement', () => {
-      assertAwaitDetection(`await foo`)
+      assertAwaitDetection(`await foo`, `await _withAsyncContext(foo)`)
     })
 
     test('variable', () => {
-      assertAwaitDetection(`const a = 1 + (await foo)`)
+      assertAwaitDetection(
+        `const a = 1 + (await foo)`,
+        `1 + (await _withAsyncContext(foo))`
+      )
     })
 
     test('ref', () => {
-      assertAwaitDetection(`ref: a = 1 + (await foo)`)
+      assertAwaitDetection(
+        `ref: a = 1 + (await foo)`,
+        `1 + (await _withAsyncContext(foo))`
+      )
     })
 
     test('nested statements', () => {
-      assertAwaitDetection(`if (ok) { await foo } else { await bar }`)
+      assertAwaitDetection(`if (ok) { await foo } else { await bar }`, code => {
+        return (
+          code.includes(`await _withAsyncContext(foo)`) &&
+          code.includes(`await _withAsyncContext(bar)`)
+        )
+      })
     })
 
     test('should ignore await inside functions', () => {
       // function declaration
-      assertAwaitDetection(`async function foo() { await bar }`, false)
+      assertAwaitDetection(
+        `async function foo() { await bar }`,
+        `await bar`,
+        false
+      )
       // function expression
-      assertAwaitDetection(`const foo = async () => { await bar }`, false)
+      assertAwaitDetection(
+        `const foo = async () => { await bar }`,
+        `await bar`,
+        false
+      )
       // object method
-      assertAwaitDetection(`const obj = { async method() { await bar }}`, false)
+      assertAwaitDetection(
+        `const obj = { async method() { await bar }}`,
+        `await bar`,
+        false
+      )
       // class method
       assertAwaitDetection(
         `const cls = class Foo { async method() { await bar }}`,
+        `await bar`,
         false
       )
     })
diff --git a/packages/compiler-sfc/src/compileScript.ts b/packages/compiler-sfc/src/compileScript.ts
index cfc5f8021ac..9347e996b50 100644
--- a/packages/compiler-sfc/src/compileScript.ts
+++ b/packages/compiler-sfc/src/compileScript.ts
@@ -900,6 +900,11 @@ export function compileScript(
           }
           if (node.type === 'AwaitExpression') {
             hasAwait = true
+            s.prependRight(
+              node.argument.start! + startOffset,
+              helper(`withAsyncContext`) + `(`
+            )
+            s.appendLeft(node.argument.end! + startOffset, `)`)
           }
         }
       })
diff --git a/packages/runtime-core/__tests__/apiSetupHelpers.spec.ts b/packages/runtime-core/__tests__/apiSetupHelpers.spec.ts
index c1e72ee4e6b..17d4e2988e8 100644
--- a/packages/runtime-core/__tests__/apiSetupHelpers.spec.ts
+++ b/packages/runtime-core/__tests__/apiSetupHelpers.spec.ts
@@ -1,9 +1,13 @@
 import {
+  ComponentInternalInstance,
   defineComponent,
+  getCurrentInstance,
   h,
   nodeOps,
+  onMounted,
   render,
-  SetupContext
+  SetupContext,
+  Suspense
 } from '@vue/runtime-test'
 import {
   defineEmits,
@@ -12,7 +16,8 @@ import {
   withDefaults,
   useAttrs,
   useSlots,
-  mergeDefaults
+  mergeDefaults,
+  withAsyncContext
 } from '../src/apiSetupHelpers'
 
 describe('SFC <script setup> helpers', () => {
@@ -89,4 +94,39 @@ describe('SFC <script setup> helpers', () => {
       `props default key "foo" has no corresponding declaration`
     ).toHaveBeenWarned()
   })
+
+  test('withAsyncContext', async () => {
+    const spy = jest.fn()
+
+    let beforeInstance: ComponentInternalInstance | null = null
+    let afterInstance: ComponentInternalInstance | null = null
+    let resolve: (msg: string) => void
+
+    const Comp = defineComponent({
+      async setup() {
+        beforeInstance = getCurrentInstance()
+        const msg = await withAsyncContext(
+          new Promise(r => {
+            resolve = r
+          })
+        )
+        // register the lifecycle after an await statement
+        onMounted(spy)
+        afterInstance = getCurrentInstance()
+        return () => msg
+      }
+    })
+
+    const root = nodeOps.createElement('div')
+    render(h(() => h(Suspense, () => h(Comp))), root)
+
+    expect(spy).not.toHaveBeenCalled()
+    resolve!('hello')
+    // wait a macro task tick for all micro ticks to resolve
+    await new Promise(r => setTimeout(r))
+    // mount hook should have been called
+    expect(spy).toHaveBeenCalled()
+    // should retain same instance before/after the await call
+    expect(beforeInstance).toBe(afterInstance)
+  })
 })
diff --git a/packages/runtime-core/src/apiSetupHelpers.ts b/packages/runtime-core/src/apiSetupHelpers.ts
index 493cb713959..b967414349a 100644
--- a/packages/runtime-core/src/apiSetupHelpers.ts
+++ b/packages/runtime-core/src/apiSetupHelpers.ts
@@ -1,7 +1,8 @@
 import {
   getCurrentInstance,
   SetupContext,
-  createSetupContext
+  createSetupContext,
+  setCurrentInstance
 } from './component'
 import { EmitFn, EmitsOptions } from './componentEmits'
 import {
@@ -226,3 +227,17 @@ export function mergeDefaults(
   }
   return props
 }
+
+/**
+ * Runtime helper for storing and resuming current instance context in
+ * async setup().
+ * @internal
+ */
+export async function withAsyncContext<T>(
+  awaitable: T | Promise<T>
+): Promise<T> {
+  const ctx = getCurrentInstance()
+  const res = await awaitable
+  setCurrentInstance(ctx)
+  return res
+}
diff --git a/packages/runtime-core/src/index.ts b/packages/runtime-core/src/index.ts
index e7ef359d6e3..d16313b2a44 100644
--- a/packages/runtime-core/src/index.ts
+++ b/packages/runtime-core/src/index.ts
@@ -48,12 +48,14 @@ export { defineAsyncComponent } from './apiAsyncComponent'
 // <script setup> API ----------------------------------------------------------
 
 export {
+  // macros runtime, for warnings only
   defineProps,
   defineEmits,
   defineExpose,
   withDefaults,
   // internal
   mergeDefaults,
+  withAsyncContext,
   // deprecated
   defineEmit,
   useContext