From 3b765367f31b6d1bb32406505f18151acdf1f2b2 Mon Sep 17 00:00:00 2001
From: arielbackenroth <arielb@alumni.rice.edu>
Date: Wed, 27 Nov 2024 12:52:34 +0000
Subject: [PATCH] fix(core): Explicitly manage TracingSnapshot lifecycle and
 dispose of it once it's been used. (#58929)

Provide a callback to the TracingService implementation when a Snapshot can be disposed.
The underlying tracing implementation may use refcounting and needs to release resources
to enable the trace to complete.

While change detection uses the snapshot for exactly one callback, after render runs
multiple hooks in the sequence so we need a more predictable way to indicate that the snapshot
can be finalized.s

PR Close #58929
---
 packages/core/src/application/application_ref.ts  | 1 +
 packages/core/src/application/tracing.ts          | 7 +++++++
 packages/core/src/render3/after_render/manager.ts | 1 +
 packages/core/test/acceptance/tracing_spec.ts     | 1 +
 4 files changed, 10 insertions(+)

diff --git a/packages/core/src/application/application_ref.ts b/packages/core/src/application/application_ref.ts
index 7514d772625..06fe1c2c7a2 100644
--- a/packages/core/src/application/application_ref.ts
+++ b/packages/core/src/application/application_ref.ts
@@ -604,6 +604,7 @@ export class ApplicationRef {
       // if one exists. Snapshots may be reference counted by the implementation so
       // we want to ensure that if we request a snapshot that we use it.
       snapshot.run(TracingAction.CHANGE_DETECTION, this._tick);
+      snapshot.dispose();
       return;
     }
 
diff --git a/packages/core/src/application/tracing.ts b/packages/core/src/application/tracing.ts
index 663fd6e9d7b..d8c30b44c78 100644
--- a/packages/core/src/application/tracing.ts
+++ b/packages/core/src/application/tracing.ts
@@ -17,6 +17,9 @@ export enum TracingAction {
 /** A single tracing snapshot. */
 export interface TracingSnapshot {
   run<T>(action: TracingAction, fn: () => T): T;
+
+  /** Disposes of the tracing snapshot. Must be run exactly once per TracingSnapshot. */
+  dispose(): void;
 }
 
 /**
@@ -39,6 +42,10 @@ export interface TracingService<T extends TracingSnapshot> {
    * used when additional work is performed that was scheduled in this context.
    *
    * @param linkedSnapshot Optional snapshot to use link to the current context.
+   * The caller is no longer responsible for calling dispose on the linkedSnapshot.
+   *
+   * @return The tracing snapshot. The caller is responsible for diposing of the
+   * snapshot.
    */
   snapshot(linkedSnapshot: T | null): T;
 }
diff --git a/packages/core/src/render3/after_render/manager.ts b/packages/core/src/render3/after_render/manager.ts
index 6680357856c..dba8bf55932 100644
--- a/packages/core/src/render3/after_render/manager.ts
+++ b/packages/core/src/render3/after_render/manager.ts
@@ -187,6 +187,7 @@ export class AfterRenderSequence implements AfterRenderRef {
     // associates the initial run of the hook with the context that created it.
     // Follow-up runs are independent of that initial context and have different
     // triggers.
+    this.snapshot?.dispose();
     this.snapshot = null;
   }
 
diff --git a/packages/core/test/acceptance/tracing_spec.ts b/packages/core/test/acceptance/tracing_spec.ts
index d4a97f19682..66c9c6883a7 100644
--- a/packages/core/test/acceptance/tracing_spec.ts
+++ b/packages/core/test/acceptance/tracing_spec.ts
@@ -27,6 +27,7 @@ describe('TracingService', () => {
         actions.push(action);
         return fn();
       },
+      dispose() {},
     };
     mockTracingService = {
       snapshot: jasmine.createSpy('snapshot').and.returnValue(fakeSnapshot),