refactor(core): introduce TracingService for snapshot-based tracing (angular#58771)

This commit introduces a private API, the `TracingService` DI token. By
providing this token, Angular can be configured to capture tracing snapshots
for certain operations such as change detection notifications, and to run
downstream operations within the context of those snapshots.
`TracingService` abstracts this context propagation and makes it pluggable.

PR Close angular#58771

Author: alxhub

packages/core/src/application/application_ref.ts

@@ -45,6 +45,7 @@ import {isPromise} from '../util/lang';
 import {NgZone} from '../zone/ng_zone';
 
 import {ApplicationInitStatus} from './application_init';
+import {TracingService} from './tracing';
 import {EffectScheduler} from '../render3/reactivity/root_effect_scheduler';
 
 /**
@@ -311,6 +312,7 @@ export class ApplicationRef {
   private readonly afterRenderManager = inject(AfterRenderManager);
   private readonly zonelessEnabled = inject(ZONELESS_ENABLED);
   private readonly rootEffectScheduler = inject(EffectScheduler);
+  private readonly tracing = inject(TracingService, {optional: true});
 
   /**
    * Current dirty state of the application across a number of dimensions (views, afterRender hooks,
@@ -329,6 +331,16 @@ export class ApplicationRef {
    */
   deferredDirtyFlags = ApplicationRefDirtyFlags.None;
 
+  /**
+   * Most recent snapshot from the `TracingService`, if any.
+   *
+   * This snapshot attempts to capture the context when `tick()` was first scheduled. It then runs
+   * wrapped in this context.
+   *
+   * @internal
+   */
+  tracingSnapshot: unknown | undefined = undefined;
+
   // Needed for ComponentFixture temporarily during migration of autoDetect behavior
   // Eventually the hostView of the fixture should just attach to ApplicationRef.
   private externalTestViews: Set<InternalViewRef<unknown>> = new Set();
@@ -577,11 +589,15 @@ export class ApplicationRef {
     if (!this.zonelessEnabled) {
       this.dirtyFlags |= ApplicationRefDirtyFlags.ViewTreeGlobal;
     }
-    this._tick();
+
+    // Run `_tick()` in the context of the most recent snapshot, if one exists.
+    this.tracing?.run(this._tick, this.tracingSnapshot) ?? this._tick();
   }
 
   /** @internal */
-  _tick(): void {
+  _tick = (): void => {
+    this.tracingSnapshot = undefined;
+
     (typeof ngDevMode === 'undefined' || ngDevMode) && this.warnIfDestroyed();
     if (this._runningTick) {
       throw new RuntimeError(
@@ -608,7 +624,7 @@ export class ApplicationRef {
       setActiveConsumer(prevConsumer);
       this.afterTick.next();
     }
-  }
+  };
 
   /**
    * Performs the core work of synchronizing the application state with the UI, resolving any

packages/core/src/application/tracing.ts

@@ -0,0 +1,34 @@
+/**
+ * @license
+ * Copyright Google LLC All Rights Reserved.
+ *
+ * Use of this source code is governed by an MIT-style license that can be
+ * found in the LICENSE file at https://angular.dev/license
+ */
+
+import {InjectionToken} from '../di/injection_token';
+
+/**
+ * Injection token for a `TracingService`, optionally provided.
+ */
+export const TracingService = new InjectionToken<TracingService<unknown>>('');
+
+/**
+ * Tracing mechanism which can associate causes (snapshots) with runs of subsequent operations.
+ *
+ * Not defined by Angular directly, but defined in contexts where tracing is desired.
+ */
+export interface TracingService<TSnapshot> {
+  /**
+   * Take a snapshot of the current context which will be stored by Angular and used when additional
+   * work is performed that was scheduled in this context.
+   */
+  snapshot(): TSnapshot;
+
+  /**
+   * Invoke `fn` within the given tracing snapshot, which may be `undefined`.
+   *
+   * This _must_ return the result of the function invocation.
+   */
+  run<T>(fn: () => T, snapshot: TSnapshot | undefined): T;
+}

packages/core/src/change_detection/scheduling/zoneless_scheduling_impl.ts

@@ -30,6 +30,7 @@ import {
   ZONELESS_SCHEDULER_DISABLED,
   SCHEDULE_IN_ROOT_ZONE,
 } from './zoneless_scheduling';
+import {TracingService} from '../../application/tracing';
 
 const CONSECUTIVE_MICROTASK_NOTIFICATION_LIMIT = 100;
 let consecutiveMicrotaskNotifications = 0;
@@ -60,6 +61,7 @@ export class ChangeDetectionSchedulerImpl implements ChangeDetectionScheduler {
   private readonly taskService = inject(PendingTasksInternal);
   private readonly ngZone = inject(NgZone);
   private readonly zonelessEnabled = inject(ZONELESS_ENABLED);
+  private readonly tracing = inject(TracingService, {optional: true});
   private readonly disableScheduling =
     inject(ZONELESS_SCHEDULER_DISABLED, {optional: true}) ?? false;
   private readonly zoneIsDefined = typeof Zone !== 'undefined' && !!Zone.root.run;
@@ -192,6 +194,10 @@ export class ChangeDetectionSchedulerImpl implements ChangeDetectionScheduler {
       }
     }
 
+    // If not already defined, attempt to capture a tracing snapshot of this notification so that
+    // the resulting CD run can be attributed to the context which produced the notification.
+    this.appRef.tracingSnapshot ??= this.tracing?.snapshot();
+
     if (!this.shouldScheduleTick(force)) {
       return;
     }

packages/core/src/core_private_export.ts

@@ -17,6 +17,7 @@ export {
   IMAGE_CONFIG_DEFAULTS as ɵIMAGE_CONFIG_DEFAULTS,
   ImageConfig as ɵImageConfig,
 } from './application/application_tokens';
+export {TracingService as ɵTracingService} from './application/tracing';
 export {internalCreateApplication as ɵinternalCreateApplication} from './application/create_application';
 export {
   defaultIterableDiffers as ɵdefaultIterableDiffers,

packages/core/src/render3/after_render/hooks.ts

@@ -6,6 +6,7 @@
  * found in the LICENSE file at https://angular.dev/license
  */
 
+import {TracingService} from '../../application/tracing';
 import {assertInInjectionContext} from '../../di';
 import {Injector} from '../../di/injector';
 import {inject} from '../../di/injector_compatibility';
@@ -454,13 +455,16 @@ function afterRenderImpl(
   // tree-shaken if `afterRender` and `afterNextRender` aren't used.
   manager.impl ??= injector.get(AfterRenderImpl);
 
+  const tracing = injector.get(TracingService, null, {optional: true});
+
   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,
     destroyRef,
+    tracing?.snapshot(),
   );
   manager.impl.register(sequence);
   return sequence;

packages/core/src/render3/after_render/manager.ts

@@ -16,6 +16,7 @@ import {
   NotificationSource,
 } from '../../change_detection/scheduling/zoneless_scheduling';
 import {type DestroyRef} from '../../linker/destroy_ref';
+import {TracingService} from '../../application/tracing';
 
 export class AfterRenderManager {
   impl: AfterRenderImpl | null = null;
@@ -44,6 +45,7 @@ export class AfterRenderImpl {
   private readonly ngZone = inject(NgZone);
   private readonly scheduler = inject(ChangeDetectionScheduler);
   private readonly errorHandler = inject(ErrorHandler, {optional: true});
+  private readonly tracing = inject(TracingService, {optional: true});
 
   /** Current set of active sequences. */
   private readonly sequences = new Set<AfterRenderSequence>();
@@ -68,7 +70,10 @@ export class AfterRenderImpl {
 
         try {
           sequence.pipelinedValue = this.ngZone.runOutsideAngular(() =>
-            sequence.hooks[phase]!(sequence.pipelinedValue),
+            this.maybeTrace(
+              () => sequence.hooks[phase]!(sequence.pipelinedValue),
+              sequence.snapshot,
+            ),
           );
         } catch (err) {
           sequence.erroredOrDestroyed = true;
@@ -124,6 +129,11 @@ export class AfterRenderImpl {
     }
   }
 
+  protected maybeTrace<T>(fn: () => T, snapshot: unknown): T {
+    // Only trace the execution if the snapshot is defined.
+    return this.tracing && snapshot ? this.tracing.run(fn, snapshot) : fn();
+  }
+
   /** @nocollapse */
   static ɵprov = /** @pureOrBreakMyCode */ /* @__PURE__ */ ɵɵdefineInjectable({
     token: AfterRenderImpl,
@@ -160,13 +170,19 @@ export class AfterRenderSequence implements AfterRenderRef {
     readonly hooks: AfterRenderHooks,
     public once: boolean,
     destroyRef: DestroyRef | null,
+    public snapshot: unknown,
   ) {
     this.unregisterOnDestroy = destroyRef?.onDestroy(() => this.destroy());
   }
 
   afterRun(): void {
     this.erroredOrDestroyed = false;
     this.pipelinedValue = undefined;
+
+    // Clear the tracing snapshot after the initial run. This snapshot only 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 = undefined;
   }
 
   destroy(): void {

packages/core/src/render3/reactivity/after_render_effect.ts

@@ -36,6 +36,7 @@ import {NOOP_AFTER_RENDER_REF, type AfterRenderOptions} from '../after_render/ho
 import {DestroyRef} from '../../linker/destroy_ref';
 import {assertNotInReactiveContext} from './asserts';
 import {assertInInjectionContext} from '../../di/contextual';
+import {TracingService} from '../../application/tracing';
 
 const NOT_SET = Symbol('NOT_SET');
 const EMPTY_CLEANUP_SET = new Set<() => void>();
@@ -175,10 +176,11 @@ class AfterRenderEffectSequence extends AfterRenderSequence {
     effectHooks: Array<AfterRenderPhaseEffectHook | undefined>,
     readonly scheduler: ChangeDetectionScheduler,
     destroyRef: DestroyRef,
+    snapshot: unknown,
   ) {
     // Note that we also initialize the underlying `AfterRenderSequence` hooks to `undefined` and
     // populate them as we create reactive nodes below.
-    super(impl, [undefined, undefined, undefined, undefined], false, destroyRef);
+    super(impl, [undefined, undefined, undefined, undefined], false, destroyRef, snapshot);
 
     // Setup a reactive node for each phase.
     for (const phase of AFTER_RENDER_PHASES) {
@@ -370,6 +372,7 @@ export function afterRenderEffect<E = never, W = never, M = never>(
   const injector = options?.injector ?? inject(Injector);
   const scheduler = injector.get(ChangeDetectionScheduler);
   const manager = injector.get(AfterRenderManager);
+  const tracing = injector.get(TracingService, null, {optional: true});
   manager.impl ??= injector.get(AfterRenderImpl);
 
   let spec = callbackOrSpec;
@@ -382,6 +385,7 @@ export function afterRenderEffect<E = never, W = never, M = never>(
     [spec.earlyRead, spec.write, spec.mixedReadWrite, spec.read] as AfterRenderPhaseEffectHook[],
     scheduler,
     injector.get(DestroyRef),
+    tracing?.snapshot(),
   );
   manager.impl.register(sequence);
   return sequence;