angular · pkozlowski-opensource · Apr 11, 2023
diff --git a/goldens/public-api/core/index.md b/goldens/public-api/core/index.md
@@ -512,7 +512,7 @@ export interface DoCheck {
 }
 
 // @public
-export function effect(effectFn: () => EffectCleanupFn | void, options?: CreateEffectOptions): EffectRef;
+export function effect(effectFn: (onCleanup: EffectCleanupRegisterFn) => void, options?: CreateEffectOptions): EffectRef;
 
 // @public
 export type EffectCleanupFn = () => void;

diff --git a/packages/core/src/render3/reactivity/effect.ts b/packages/core/src/render3/reactivity/effect.ts
@@ -14,22 +14,29 @@ import {DestroyRef} from '../../linker/destroy_ref';
 import {Watch} from '../../signals';
 
 /**
- * An effect can, optionally, return a cleanup function. If returned, the cleanup is executed before
- * the next effect run. The cleanup function makes it possible to "cancel" any work that the
+ * An effect can, optionally, register a cleanup function. If registered, the cleanup is executed
+ * before the next effect run. The cleanup function makes it possible to "cancel" any work that the
  * previous effect run might have started.
  *
  * @developerPreview
  */
 export type EffectCleanupFn = () => void;
 
+/**
+ * A callback passed to the effect function that makes it possible to register cleanup logic.
+ */
+export type EffectCleanupRegisterFn = (cleanupFn: EffectCleanupFn) => void;
+
 /**
  * Tracks all effects registered within a given application and runs them via `flush`.
  */
 export class EffectManager {
   private all = new Set<Watch>();
   private queue = new Map<Watch, Zone>();
 
-  create(effectFn: () => void, destroyRef: DestroyRef|null, allowSignalWrites: boolean): EffectRef {
+  create(
+      effectFn: (onCleanup: (cleanupFn: EffectCleanupFn) => void) => void,
+      destroyRef: DestroyRef|null, allowSignalWrites: boolean): EffectRef {
     const zone = Zone.current;
     const watch = new Watch(effectFn, (watch) => {
       if (!this.all.has(watch)) {
@@ -131,7 +138,8 @@ export interface CreateEffectOptions {
  * @developerPreview
  */
 export function effect(
-    effectFn: () => EffectCleanupFn | void, options?: CreateEffectOptions): EffectRef {
+    effectFn: (onCleanup: EffectCleanupRegisterFn) => void,
+    options?: CreateEffectOptions): EffectRef {
   !options?.injector && assertInInjectionContext(effect);
   const injector = options?.injector ?? inject(Injector);
   const effectManager = injector.get(EffectManager);

diff --git a/packages/core/src/signals/src/watch.ts b/packages/core/src/signals/src/watch.ts
@@ -9,11 +9,16 @@
 import {ReactiveNode, setActiveConsumer} from './graph';
 
 /**
- * A cleanup function that can be optionally returned from the watch logic. When returned, the
+ * A cleanup function that can be optionally registered from the watch logic. If registered, the
  * cleanup logic runs before the next watch execution.
  */
 export type WatchCleanupFn = () => void;
 
+/**
+ * A callback passed to the watch function that makes it possible to register cleanup logic.
+ */
+export type WatchCleanupRegisterFn = (cleanupFn: WatchCleanupFn) => void;
+
 const NOOP_CLEANUP_FN: WatchCleanupFn = () => {};
 
 /**
@@ -27,10 +32,14 @@ export class Watch extends ReactiveNode {
   protected override readonly consumerAllowSignalWrites: boolean;
   private dirty = false;
   private cleanupFn = NOOP_CLEANUP_FN;
+  private registerOnCleanup =
+      (cleanupFn: WatchCleanupFn) => {
+        this.cleanupFn = cleanupFn;
+      }
 
   constructor(
-      private watch: () => void|WatchCleanupFn, private schedule: (watch: Watch) => void,
-      allowSignalWrites: boolean) {
+      private watch: (onCleanup: WatchCleanupRegisterFn) => void,
+      private schedule: (watch: Watch) => void, allowSignalWrites: boolean) {
     super();
     this.consumerAllowSignalWrites = allowSignalWrites;
   }
@@ -66,7 +75,8 @@ export class Watch extends ReactiveNode {
     this.trackingVersion++;
     try {
       this.cleanupFn();
-      this.cleanupFn = this.watch() ?? NOOP_CLEANUP_FN;
+      this.cleanupFn = NOOP_CLEANUP_FN;
+      this.watch(this.registerOnCleanup);
     } finally {
       setActiveConsumer(prevConsumer);
     }

diff --git a/packages/core/test/render3/reactivity_spec.ts b/packages/core/test/render3/reactivity_spec.ts
@@ -130,11 +130,11 @@ describe('effects', () => {
     })
     class Cmp {
       counter = signal(0);
-      effectRef = effect(() => {
+      effectRef = effect((onCleanup) => {
         counterLog.push(this.counter());
-        return () => {
+        onCleanup(() => {
           cleanupCount++;
-        };
+        });
       });
     }
 
@@ -179,6 +179,7 @@ describe('effects', () => {
 
     expect(didRun).toBeTrue();
   });
+
   it('should disallow writing to signals within effects by default',
      withBody('<test-cmp></test-cmp>', async () => {
        @Component({

diff --git a/packages/core/test/signals/effect_util.ts b/packages/core/test/signals/effect_util.ts
@@ -6,14 +6,15 @@
  * found in the LICENSE file at https://angular.io/license
  */
 
-import {Watch} from '@angular/core/src/signals';
+import {Watch, WatchCleanupFn} from '@angular/core/src/signals';
 
 let queue = new Set<Watch>();
 
 /**
  * A wrapper around `Watch` that emulates the `effect` API and allows for more streamlined testing.
  */
-export function testingEffect(effectFn: () => void): void {
+export function testingEffect(effectFn: (onCleanup: (cleanupFn: WatchCleanupFn) => void) => void):
+    void {
   const watch = new Watch(effectFn, queue.add.bind(queue), true);
 
   // Effects start dirty.

diff --git a/packages/core/test/signals/watch_spec.ts b/packages/core/test/signals/watch_spec.ts
@@ -100,19 +100,19 @@ describe('watchers', () => {
     expect(updateCounter).toEqual(3);
   });
 
-  it('should allow returning cleanup function from the watch logic', () => {
+  it('should allow registering cleanup function from the watch logic', () => {
     const source = signal(0);
 
     const seenCounterValues: number[] = [];
-    testingEffect(() => {
+    testingEffect((onCleanup) => {
       seenCounterValues.push(source());
 
-      // return a cleanup function that is executed every time an effect re-runs
-      return () => {
+      // register a cleanup function that is executed every time an effect re-runs
+      onCleanup(() => {
         if (seenCounterValues.length === 2) {
           seenCounterValues.length = 0;
         }
-      };
+      });
     });
 
     flushEffects();
@@ -128,6 +128,37 @@ describe('watchers', () => {
     expect(seenCounterValues).toEqual([2]);
   });
 
+  it('should forget previously registered cleanup function when effect re-runs', () => {
+    const source = signal(0);
+
+    const seenCounterValues: number[] = [];
+    testingEffect((onCleanup) => {
+      const value = source();
+
+      seenCounterValues.push(value);
+
+      // register a cleanup function that is executed next time an effect re-runs
+      if (value === 0) {
+        onCleanup(() => {
+          seenCounterValues.length = 0;
+        });
+      }
+    });
+
+    flushEffects();
+    expect(seenCounterValues).toEqual([0]);
+
+    source.set(2);
+    flushEffects();
+    // cleanup (array trim) should have run before executing effect
+    expect(seenCounterValues).toEqual([2]);
+
+    source.set(3);
+    flushEffects();
+    // cleanup (array trim) should _not_ be registered again
+    expect(seenCounterValues).toEqual([2, 3]);
+  });
+
   it('should throw an error when reading a signal during the notification phase', () => {
     const source = signal(0);
     let ranScheduler = false;