From 7976f7044e66b3b7adac4c72a392935704658b10 Mon Sep 17 00:00:00 2001
From: Evan You <yyx990803@gmail.com>
Date: Tue, 9 Jan 2024 07:22:13 +0800
Subject: [PATCH] refactor(runtime-core): safer currentInstance reset

---
 packages/runtime-core/src/apiLifecycle.ts   |  5 ++---
 packages/runtime-core/src/apiWatch.ts       | 10 ++--------
 packages/runtime-core/src/component.ts      | 13 +++++++++----
 packages/runtime-core/src/componentProps.ts |  5 ++---
 4 files changed, 15 insertions(+), 18 deletions(-)

diff --git a/packages/runtime-core/src/apiLifecycle.ts b/packages/runtime-core/src/apiLifecycle.ts
index 9de0887250d..741d43ec45c 100644
--- a/packages/runtime-core/src/apiLifecycle.ts
+++ b/packages/runtime-core/src/apiLifecycle.ts
@@ -3,7 +3,6 @@ import {
   currentInstance,
   isInSSRComponentSetup,
   setCurrentInstance,
-  unsetCurrentInstance,
 } from './component'
 import type { ComponentPublicInstance } from './componentPublicInstance'
 import { ErrorTypeStrings, callWithAsyncErrorHandling } from './errorHandling'
@@ -41,9 +40,9 @@ export function injectHook(
         // Set currentInstance during hook invocation.
         // This assumes the hook does not synchronously trigger other hooks, which
         // can only be false when the user does something really funky.
-        setCurrentInstance(target)
+        const reset = setCurrentInstance(target)
         const res = callWithAsyncErrorHandling(hook, target, type, args)
-        unsetCurrentInstance()
+        reset()
         resetTracking()
         return res
       })
diff --git a/packages/runtime-core/src/apiWatch.ts b/packages/runtime-core/src/apiWatch.ts
index 3a2d9e46c33..bc10547824e 100644
--- a/packages/runtime-core/src/apiWatch.ts
+++ b/packages/runtime-core/src/apiWatch.ts
@@ -30,7 +30,6 @@ import {
   currentInstance,
   isInSSRComponentSetup,
   setCurrentInstance,
-  unsetCurrentInstance,
 } from './component'
 import {
   ErrorCodes,
@@ -448,14 +447,9 @@ export function instanceWatch(
     cb = value.handler as Function
     options = value
   }
-  const cur = currentInstance
-  setCurrentInstance(this)
+  const reset = setCurrentInstance(this)
   const res = doWatch(getter, cb.bind(publicThis), options)
-  if (cur) {
-    setCurrentInstance(cur)
-  } else {
-    unsetCurrentInstance()
-  }
+  reset()
   return res
 }
 
diff --git a/packages/runtime-core/src/component.ts b/packages/runtime-core/src/component.ts
index c77a17ff80b..a8edbb9e228 100644
--- a/packages/runtime-core/src/component.ts
+++ b/packages/runtime-core/src/component.ts
@@ -696,8 +696,13 @@ if (__SSR__) {
 }
 
 export const setCurrentInstance = (instance: ComponentInternalInstance) => {
+  const prev = currentInstance
   internalSetCurrentInstance(instance)
   instance.scope.on()
+  return () => {
+    instance.scope.off()
+    internalSetCurrentInstance(prev)
+  }
 }
 
 export const unsetCurrentInstance = () => {
@@ -785,7 +790,7 @@ function setupStatefulComponent(
     const setupContext = (instance.setupContext =
       setup.length > 1 ? createSetupContext(instance) : null)
 
-    setCurrentInstance(instance)
+    const reset = setCurrentInstance(instance)
     pauseTracking()
     const setupResult = callWithErrorHandling(
       setup,
@@ -797,7 +802,7 @@ function setupStatefulComponent(
       ],
     )
     resetTracking()
-    unsetCurrentInstance()
+    reset()
 
     if (isPromise(setupResult)) {
       setupResult.then(unsetCurrentInstance, unsetCurrentInstance)
@@ -972,13 +977,13 @@ export function finishComponentSetup(
 
   // support for 2.x options
   if (__FEATURE_OPTIONS_API__ && !(__COMPAT__ && skipOptions)) {
-    setCurrentInstance(instance)
+    const reset = setCurrentInstance(instance)
     pauseTracking()
     try {
       applyOptions(instance)
     } finally {
       resetTracking()
-      unsetCurrentInstance()
+      reset()
     }
   }
 
diff --git a/packages/runtime-core/src/componentProps.ts b/packages/runtime-core/src/componentProps.ts
index 2124853f123..088f4da3c6e 100644
--- a/packages/runtime-core/src/componentProps.ts
+++ b/packages/runtime-core/src/componentProps.ts
@@ -32,7 +32,6 @@ import {
   type ConcreteComponent,
   type Data,
   setCurrentInstance,
-  unsetCurrentInstance,
 } from './component'
 import { isEmitListener } from './componentEmits'
 import { InternalObjectKey } from './vnode'
@@ -470,7 +469,7 @@ function resolvePropValue(
         if (key in propsDefaults) {
           value = propsDefaults[key]
         } else {
-          setCurrentInstance(instance)
+          const reset = setCurrentInstance(instance)
           value = propsDefaults[key] = defaultValue.call(
             __COMPAT__ &&
               isCompatEnabled(DeprecationTypes.PROPS_DEFAULT_THIS, instance)
@@ -478,7 +477,7 @@ function resolvePropValue(
               : null,
             props,
           )
-          unsetCurrentInstance()
+          reset()
         }
       } else {
         value = defaultValue