fix(core): provide flag to opt into manual cleanup for after render hooks (#57917)

crisbeto · alxhub · commit 2f347ef8fcef · 2024-09-26T14:20:55.000-07:00
Adds a `manualCleanup` flag to `afterRender` and `afterNextRender`, similarly to `effect`. The reason is that currently if the hook is created outside of an injection context, it requires an injector to be passed in. In some cases that injector might be an injector that is never destroyed (e.g. `EnvironmentInjector`) which can give a false sense of security users thinking that the hook will be cleaned up automatically. We fell into this in the CDK which caused a memory leak (see angular/components#29709). With the `manualCleanup` option users explicitly opt into cleaning the hook up themselves. PR Close #57917
diff --git a/goldens/public-api/core/index.api.md b/goldens/public-api/core/index.api.md
@@ -52,6 +52,7 @@ export function afterRender(callback: VoidFunction, options?: AfterRenderOptions
 // @public
 export interface AfterRenderOptions {
     injector?: Injector;
+    manualCleanup?: boolean;
     // @deprecated
     phase?: AfterRenderPhase;
 }
diff --git a/packages/core/src/render3/after_render/hooks.ts b/packages/core/src/render3/after_render/hooks.ts
@@ -44,6 +44,14 @@ export interface AfterRenderOptions {
    */
   injector?: Injector;
 
+  /**
+   * Whether the hook should require manual cleanup.
+   *
+   * If this is `false` (the default) the hook will automatically register itself to be cleaned up
+   * with the current `DestroyRef`.
+   */
+  manualCleanup?: boolean;
+
   /**
    * The phase the callback should be invoked in.
    *
@@ -448,11 +456,12 @@ function afterRenderImpl(
   manager.impl ??= injector.get(AfterRenderImpl);
 
   const hooks = options?.phase ?? AfterRenderPhase.MixedReadWrite;
+  const destroyRef = options?.manualCleanup !== true ? injector.get(DestroyRef) : null;
   const sequence = new AfterRenderSequence(
     manager.impl,
     getHooks(callbackOrSpec, hooks),
     once,
-    injector.get(DestroyRef),
+    destroyRef,
   );
   manager.impl.register(sequence);
   return sequence;
diff --git a/packages/core/src/render3/after_render/manager.ts b/packages/core/src/render3/after_render/manager.ts
@@ -149,15 +149,15 @@ export class AfterRenderSequence implements AfterRenderRef {
    */
   pipelinedValue: unknown = undefined;
 
-  private unregisterOnDestroy: () => void;
+  private unregisterOnDestroy: (() => void) | undefined;
 
   constructor(
     readonly impl: AfterRenderImpl,
     readonly hooks: AfterRenderHooks,
     public once: boolean,
-    destroyRef: DestroyRef,
+    destroyRef: DestroyRef | null,
   ) {
-    this.unregisterOnDestroy = destroyRef.onDestroy(() => this.destroy());
+    this.unregisterOnDestroy = destroyRef?.onDestroy(() => this.destroy());
   }
 
   afterRun(): void {
@@ -167,6 +167,6 @@ export class AfterRenderSequence implements AfterRenderRef {
 
   destroy(): void {
     this.impl.unregister(this);
-    this.unregisterOnDestroy();
+    this.unregisterOnDestroy?.();
   }
 }
diff --git a/packages/core/test/acceptance/after_render_hook_spec.ts b/packages/core/test/acceptance/after_render_hook_spec.ts
@@ -708,6 +708,53 @@ describe('after render hooks', () => {
           );
         });
       });
+
+      it('should not destroy automatically if manualCleanup is set', () => {
+        let afterRenderRef: AfterRenderRef | null = null;
+        let count = 0;
+
+        @Component({selector: 'comp', template: ''})
+        class Comp {
+          constructor() {
+            afterRenderRef = afterRender(() => count++, {manualCleanup: true});
+          }
+        }
+
+        @Component({
+          imports: [Comp],
+          template: `
+            @if (shouldShow) {
+              <comp/>
+            }
+          `,
+        })
+        class App {
+          shouldShow = true;
+        }
+
+        TestBed.configureTestingModule({
+          declarations: [App, Comp],
+          ...COMMON_CONFIGURATION,
+        });
+        const component = createAndAttachComponent(App);
+        const appRef = TestBed.inject(ApplicationRef);
+        expect(count).toBe(0);
+
+        appRef.tick();
+        expect(count).toBe(1);
+
+        component.instance.shouldShow = false;
+        component.changeDetectorRef.detectChanges();
+        appRef.tick();
+        expect(count).toBe(2);
+        appRef.tick();
+        expect(count).toBe(3);
+
+        // Ensure that manual destruction still works.
+        afterRenderRef!.destroy();
+        appRef.tick();
+        expect(count).toBe(3);
+      });
     });
 
     describe('afterNextRender', () => {

Original file line number	Diff line number	Diff line change
`@@ -52,6 +52,7 @@ export function afterRender(callback: VoidFunction, options?: AfterRenderOptions`
`52`	`52`	`// @public`
`53`	`53`	`export interface AfterRenderOptions {`
`54`	`54`	`injector?: Injector;`
	`55`	`+ manualCleanup?: boolean;`
`55`	`56`	`// @deprecated`
`56`	`57`	`phase?: AfterRenderPhase;`
`57`	`58`	`}`