fix: use WeakRef for dep to avoid ReactiveEffect memory leak (vue…

…js#9233)
johnsoncodehk · Oct 20, 2023 · 66da31a · 66da31a
1 parent fc76e81
commit 66da31a
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 7 deletions.
diff --git a/packages/reactivity/src/dep.ts b/packages/reactivity/src/dep.ts
@@ -1,7 +1,7 @@
 import type { ReactiveEffect } from './effect'
 import type { ComputedRefImpl } from './computed'
 
-export type Dep = Map<ReactiveEffect, number> & {
+export type Dep = Map<WeakRef<ReactiveEffect>, number> & {
   cleanup: () => void
   computed?: ComputedRefImpl<any>
 }
@@ -10,7 +10,7 @@ export const createDep = (
   cleanup: () => void,
   computed?: ComputedRefImpl<any>
 ): Dep => {
-  const dep: Dep = new Map() as Dep
+  const dep = new Map() as Dep
   dep.cleanup = cleanup
   dep.computed = computed
   return dep

diff --git a/packages/reactivity/src/effect.ts b/packages/reactivity/src/effect.ts
@@ -46,6 +46,7 @@ export class ReactiveEffect<T = any> {
 
   _dirtyLevel = DirtyLevels.Dirty
   _queryingDirty = false
+  _trackToken = new WeakRef(this)
   _trackId = 0
   _runnings = 0
   _depsLength = 0
@@ -131,9 +132,9 @@ function postCleanupEffect(effect: ReactiveEffect) {
 }
 
 function cleanupDepEffect(dep: Dep, effect: ReactiveEffect) {
-  const trackId = dep.get(effect)
+  const trackId = dep.get(effect._trackToken)
   if (trackId !== undefined && effect._trackId !== trackId) {
-    dep.delete(effect)
+    dep.delete(effect._trackToken)
     if (dep.size === 0) {
       dep.cleanup()
     }
@@ -255,8 +256,8 @@ export function trackEffect(
   dep: Dep,
   debuggerEventExtraInfo?: DebuggerEventExtraInfo
 ) {
-  if (dep.get(effect) !== effect._trackId) {
-    dep.set(effect, effect._trackId)
+  if (dep.get(effect._trackToken) !== effect._trackId) {
+    dep.set(effect._trackToken, effect._trackId)
     const oldDep = effect.deps[effect._depsLength]
     if (oldDep !== dep) {
       if (oldDep) {
@@ -281,7 +282,14 @@ export function triggerEffects(
   debuggerEventExtraInfo?: DebuggerEventExtraInfo
 ) {
   pauseScheduling()
-  for (const effect of dep.keys()) {
+  let reclaimedTokens: WeakRef<ReactiveEffect>[] | undefined
+  for (const trackToken of dep.keys()) {
+    const effect = trackToken.deref()
+    if (!effect) {
+      reclaimedTokens ??= []
+      reclaimedTokens.push(trackToken)
+      continue
+    }
     if (!effect.allowRecurse && effect._runnings) {
       continue
     }
@@ -303,5 +311,10 @@ export function triggerEffects(
       }
     }
   }
+  if (reclaimedTokens) {
+    for (const token of reclaimedTokens) {
+      dep.delete(token)
+    }
+  }
   resetScheduling()
 }