From 2cc5615590de77126e8df46136de0240dbde5004 Mon Sep 17 00:00:00 2001
From: Evan You <yyx990803@gmail.com>
Date: Tue, 5 Mar 2024 22:53:08 +0800
Subject: [PATCH] feat(reactivity): `onEffectCleanup` API

ref #10173

Instead of exposing `getCurrentEffect`, this version accepts a second
argument to suppress the no-active-effect warning.
---
 packages/reactivity/__tests__/effect.spec.ts | 39 +++++++++++++++++
 packages/reactivity/src/effect.ts            | 44 ++++++++++++++++++++
 packages/reactivity/src/index.ts             |  1 +
 3 files changed, 84 insertions(+)

diff --git a/packages/reactivity/__tests__/effect.spec.ts b/packages/reactivity/__tests__/effect.spec.ts
index 99453d35d87..3936216ae61 100644
--- a/packages/reactivity/__tests__/effect.spec.ts
+++ b/packages/reactivity/__tests__/effect.spec.ts
@@ -23,6 +23,7 @@ import {
 } from '@vue/runtime-test'
 import {
   endBatch,
+  onEffectCleanup,
   pauseTracking,
   resetTracking,
   startBatch,
@@ -1131,4 +1132,42 @@ describe('reactivity/effect', () => {
       expect(getSubCount(depC)).toBe(1)
     })
   })
+
+  describe('onEffectCleanup', () => {
+    it('should get called correctly', async () => {
+      const count = ref(0)
+      const cleanupEffect = vi.fn()
+
+      const e = effect(() => {
+        onEffectCleanup(cleanupEffect)
+        count.value
+      })
+
+      count.value++
+      await nextTick()
+      expect(cleanupEffect).toHaveBeenCalledTimes(1)
+
+      count.value++
+      await nextTick()
+      expect(cleanupEffect).toHaveBeenCalledTimes(2)
+
+      // call it on stop
+      e.effect.stop()
+      expect(cleanupEffect).toHaveBeenCalledTimes(3)
+    })
+
+    it('should warn if called without active effect', () => {
+      onEffectCleanup(() => {})
+      expect(
+        `onEffectCleanup() was called when there was no active effect`,
+      ).toHaveBeenWarned()
+    })
+
+    it('should not warn without active effect when failSilently argument is passed', () => {
+      onEffectCleanup(() => {}, true)
+      expect(
+        `onEffectCleanup() was called when there was no active effect`,
+      ).not.toHaveBeenWarned()
+    })
+  })
 })
diff --git a/packages/reactivity/src/effect.ts b/packages/reactivity/src/effect.ts
index b48463d3bf7..1a8673a9a5a 100644
--- a/packages/reactivity/src/effect.ts
+++ b/packages/reactivity/src/effect.ts
@@ -126,6 +126,10 @@ export class ReactiveEffect<T = any>
    * @internal
    */
   nextEffect?: ReactiveEffect = undefined
+  /**
+   * @internal
+   */
+  cleanup?: () => void = undefined
 
   scheduler?: EffectScheduler = undefined
   onStop?: () => void
@@ -165,6 +169,7 @@ export class ReactiveEffect<T = any>
     }
 
     this.flags |= EffectFlags.RUNNING
+    cleanupEffect(this)
     prepareDeps(this)
     const prevEffect = activeSub
     const prevShouldTrack = shouldTrack
@@ -193,6 +198,7 @@ export class ReactiveEffect<T = any>
         removeSub(link)
       }
       this.deps = this.depsTail = undefined
+      cleanupEffect(this)
       this.onStop && this.onStop()
       this.flags &= ~EffectFlags.ACTIVE
     }
@@ -479,3 +485,41 @@ export function resetTracking() {
   const last = trackStack.pop()
   shouldTrack = last === undefined ? true : last
 }
+
+/**
+ * Registers a cleanup function for the current active effect.
+ * The cleanup function is called right before the next effect run, or when the
+ * effect is stopped.
+ *
+ * Throws a warning iff there is no currenct active effect. The warning can be
+ * suppressed by passing `true` to the second argument.
+ *
+ * @param fn - the cleanup function to be registered
+ * @param failSilently - if `true`, will not throw warning when called without
+ * an active effect.
+ */
+export function onEffectCleanup(fn: () => void, failSilently = false) {
+  if (activeSub instanceof ReactiveEffect) {
+    activeSub.cleanup = fn
+  } else if (__DEV__ && !failSilently) {
+    warn(
+      `onEffectCleanup() was called when there was no active effect` +
+        ` to associate with.`,
+    )
+  }
+}
+
+function cleanupEffect(e: ReactiveEffect) {
+  const { cleanup } = e
+  e.cleanup = undefined
+  if (cleanup) {
+    // run cleanup without active effect
+    const prevSub = activeSub
+    activeSub = undefined
+    try {
+      cleanup()
+    } finally {
+      activeSub = prevSub
+    }
+  }
+}
diff --git a/packages/reactivity/src/index.ts b/packages/reactivity/src/index.ts
index 609afc05f8a..97815fdcde2 100644
--- a/packages/reactivity/src/index.ts
+++ b/packages/reactivity/src/index.ts
@@ -53,6 +53,7 @@ export {
   enableTracking,
   pauseTracking,
   resetTracking,
+  onEffectCleanup,
   ReactiveEffect,
   EffectFlags,
   type ReactiveEffectRunner,