fix(core): ensures immediate trigger fires properly with lazy loaded routes (#60203)

thePunderWoman · AndrewKushnir · commit c82384104fd9 · 2025-03-06T12:47:46.000-08:00
In the case that a route was lazy loaded, some triggers would never properly finish hydrating due to things firing before the route finished resolving. This will find the topmost parent defer block and ensure the registry knows about it before trying to hydrate. In the case that the registry does not yet know, just the affected triggers await app stability before initializing. fixes #59997 PR Close #60203
diff --git a/packages/core/src/defer/registry.ts b/packages/core/src/defer/registry.ts
@@ -38,6 +38,16 @@ export class DehydratedBlockRegistry {
 
   add(blockId: string, info: DehydratedDeferBlock) {
     this.registry.set(blockId, info);
+    // It's possible that hydration is queued that's waiting for the
+    // resolution of a lazy loaded route. In this case, we ensure
+    // the callback function is called to continue the hydration process
+    // for the queued block set.
+    if (this.awaitingCallbacks.has(blockId)) {
+      const awaitingCallbacks = this.awaitingCallbacks.get(blockId)!;
+      for (const cb of awaitingCallbacks) {
+        cb();
+      }
+    }
   }
 
   get(blockId: string): DehydratedDeferBlock | null {
@@ -55,6 +65,7 @@ export class DehydratedBlockRegistry {
       this.jsActionMap.delete(blockId);
       this.invokeTriggerCleanupFns(blockId);
       this.hydrating.delete(blockId);
+      this.awaitingCallbacks.delete(blockId);
     }
     if (this.size === 0) {
       this.contract.instance?.cleanUp();
@@ -87,6 +98,15 @@ export class DehydratedBlockRegistry {
   // Blocks that are being hydrated.
   hydrating = new Map<string, PromiseWithResolvers<void>>();
 
+  // Blocks that are awaiting a defer instruction finish.
+  private awaitingCallbacks = new Map<string, Function[]>();
+
+  awaitParentBlock(topmostParentBlock: string, callback: Function) {
+    const parentBlockAwaitCallbacks = this.awaitingCallbacks.get(topmostParentBlock) ?? [];
+    parentBlockAwaitCallbacks.push(callback);
+    this.awaitingCallbacks.set(topmostParentBlock, parentBlockAwaitCallbacks);
+  }
+
   /** @nocollapse */
   static ɵprov = /** @pureOrBreakMyCode */ /* @__PURE__ */ ɵɵdefineInjectable({
     token: DehydratedBlockRegistry,
diff --git a/packages/core/src/defer/triggering.ts b/packages/core/src/defer/triggering.ts
@@ -359,8 +359,8 @@ export function triggerDeferBlock(triggerType: TriggerType, lView: LView, tNode:
 }
 
 /**
- * The core mechanism for incremental hydration. This triggers
- * hydration for all the blocks in the tree that need to be hydrated
+ * The core mechanism for incremental hydration. This triggers or
+ * queues hydration for all the blocks in the tree that need to be hydrated
  * and keeps track of all those blocks that were hydrated along the way.
  *
  * Note: the `replayQueuedEventsFn` is only provided when hydration is invoked
@@ -381,31 +381,70 @@ export async function triggerHydrationFromBlockName(
     return;
   }
 
-  // The parent promise is the possible case of a list of defer blocks already being queued
-  // If it is queued, it'll exist; otherwise it'll be null. The hydration queue will contain all
-  // elements that need to be hydrated, sans any that have promises already
+  // Trigger resource loading and hydration for the blocks in the queue in the order of highest block
+  // to lowest block. Once a block has finished resource loading, after next render fires after hydration
+  // finishes. The new block will have its defer instruction called and will be in the registry.
+  // Due to timing related to potential nested control flow, this has to be scheduled after the next render.
   const {parentBlockPromise, hydrationQueue} = getParentBlockHydrationQueue(blockName, injector);
+  if (hydrationQueue.length === 0) return;
+
+  // It's possible that the hydrationQueue topmost item is actually in the process of hydrating and has
+  // a promise already. In that case, we don't want to destroy that promise and queue it again.
+  if (parentBlockPromise !== null) {
+    hydrationQueue.shift();
+  }
 
   // The hydrating map in the registry prevents re-triggering hydration for a block that's already in
   // the hydration queue. Here we generate promises for each of the blocks about to be hydrated
   populateHydratingStateForQueue(dehydratedBlockRegistry, hydrationQueue);
 
-  // Trigger resource loading and hydration for the blocks in the queue in the order of highest block
-  // to lowest block. Once a block has finished resource loading, after next render fires after hydration
-  // finishes. The new block will have its defer instruction called and will be in the registry.
-  // Due to timing related to potential nested control flow, this has to be scheduled after the next render.
+  // We await this after populating the hydration state so we can prevent re-triggering hydration for
+  // the same blocks while this promise is being awaited.
+  if (parentBlockPromise !== null) {
+    await parentBlockPromise;
+  }
+
+  const topmostParentBlock = hydrationQueue[0];
+  if (dehydratedBlockRegistry.has(topmostParentBlock)) {
+    // the topmost parent block is already in the registry and we can proceed
+    // with hydration.
+    await triggerHydrationForBlockQueue(injector, hydrationQueue, replayQueuedEventsFn);
+  } else {
+    // the topmost parent block is not yet in the registry, which may mean
+    // a lazy loaded route, a control flow branch was taken, a route has
+    // been navigated, etc. So we need to queue up the hydration process
+    // so that it can be finished after the top block has had its defer
+    // instruction executed.
+    dehydratedBlockRegistry.awaitParentBlock(
+      topmostParentBlock,
+      async () =>
+        await triggerHydrationForBlockQueue(injector, hydrationQueue, replayQueuedEventsFn),
+    );
+  }
+}
+
+/**
+ * The core mechanism for incremental hydration. This triggers
+ * hydration for all the blocks in the tree that need to be hydrated
+ * and keeps track of all those blocks that were hydrated along the way.
+ *
+ * Note: the `replayQueuedEventsFn` is only provided when hydration is invoked
+ * as a result of an event replay (via JsAction). When hydration is invoked from
+ * an instruction set (e.g. `deferOnImmediate`) - there is no need to replay any
+ * events.
+ */
+export async function triggerHydrationForBlockQueue(
+  injector: Injector,
+  hydrationQueue: string[],
+  replayQueuedEventsFn?: Function,
+): Promise<void> {
+  const dehydratedBlockRegistry = injector.get(DEHYDRATED_BLOCK_REGISTRY);
+  const blocksBeingHydrated = dehydratedBlockRegistry.hydrating;
 
   // Indicate that we have some pending async work.
   const pendingTasks = injector.get(PendingTasksInternal);
   const taskId = pendingTasks.add();
 
-  // If the parent block was being hydrated, but the process has
-  // not yet complete, wait until parent block promise settles before
-  // going over dehydrated blocks from the queue.
-  if (parentBlockPromise !== null) {
-    await parentBlockPromise;
-  }
-
   // Actually do the triggering and hydration of the queue of blocks
   for (let blockQueueIdx = 0; blockQueueIdx < hydrationQueue.length; blockQueueIdx++) {
     const dehydratedBlockId = hydrationQueue[blockQueueIdx];
@@ -443,8 +482,10 @@ export async function triggerHydrationFromBlockName(
     }
   }
 
-  // Await hydration completion for the requested block.
-  await blocksBeingHydrated.get(blockName)?.promise;
+  const lastBlockName = hydrationQueue[hydrationQueue.length - 1];
+
+  // Await hydration completion for the last block.
+  await blocksBeingHydrated.get(lastBlockName)?.promise;
 
   // All async work is done, remove the taskId from the registry.
   pendingTasks.remove(taskId);
@@ -456,7 +497,7 @@ export async function triggerHydrationFromBlockName(
 
   // Cleanup after hydration of all affected defer blocks.
   cleanupHydratedDeferBlocks(
-    dehydratedBlockRegistry.get(blockName),
+    dehydratedBlockRegistry.get(lastBlockName),
     hydrationQueue,
     dehydratedBlockRegistry,
     injector.get(ApplicationRef),
@@ -683,7 +724,7 @@ function setTimerTriggers(injector: Injector, elementTriggers: ElementTrigger[])
 function setImmediateTriggers(injector: Injector, elementTriggers: ElementTrigger[]) {
   for (const elementTrigger of elementTriggers) {
     // Note: we intentionally avoid awaiting each call and instead kick off
-    // th hydration process simultaneously for all defer blocks with this trigger;
+    // the hydration process simultaneously for all defer blocks with this trigger;
     triggerHydrationFromBlockName(injector, elementTrigger.blockName);
   }
 }
diff --git a/packages/core/src/hydration/utils.ts b/packages/core/src/hydration/utils.ts
@@ -586,7 +586,6 @@ export function getParentBlockHydrationQueue(
     isTopMostDeferBlock = dehydratedBlockRegistry.has(currentBlockId);
     const hydratingParentBlock = dehydratedBlockRegistry.hydrating.get(currentBlockId);
     if (parentBlockPromise === null && hydratingParentBlock != null) {
-      // TODO: add an ngDevMode asset that `hydratingParentBlock.promise` exists and is of type Promise.
       parentBlockPromise = hydratingParentBlock.promise;
       break;
     }
diff --git a/packages/platform-server/test/incremental_hydration_spec.ts b/packages/platform-server/test/incremental_hydration_spec.ts
@@ -1202,7 +1202,7 @@ describe('platform-server partial hydration integration', () => {
         appRef.tick();
 
         expect(appHostNode.outerHTML).toContain('<span id="test">end</span>');
-      }, 100_000);
+      });
 
       describe('idle', () => {
         /**
@@ -2619,5 +2619,102 @@ describe('platform-server partial hydration integration', () => {
 
       expect(appHostNode.outerHTML).toContain('<p>OtherCmp content</p>');
     });
+
+    it('should trigger immediate with a lazy loaded route', async () => {
+      @Component({
+        selector: 'nested-more',
+        template: `
+          <div>
+            @defer(hydrate on immediate) {
+              <button id="click-me" (click)="clickMe()">Click me I'm dehydrated?</button>
+              <p id="hydrated">{{hydrated()}}</p>
+            }
+          </div>
+        `,
+      })
+      class NestedMoreCmp {
+        hydrated = signal('nope');
+        constructor() {
+          if (!isPlatformServer(inject(PLATFORM_ID))) {
+            this.hydrated.set('yup');
+          }
+        }
+      }
+      @Component({
+        selector: 'nested',
+        imports: [NestedMoreCmp],
+        template: `
+          <div>
+            @defer(hydrate on interaction) {
+              <nested-more />
+            }
+          </div>
+        `,
+      })
+      class NestedCmp {}
+
+      @Component({
+        selector: 'lazy',
+        imports: [NestedCmp],
+        template: `
+          @defer (hydrate on interaction) {
+            <nested />
+          }
+        `,
+      })
+      class LazyCmp {}
+
+      const routes: Routes = [
+        {
+          path: '',
+          loadComponent: () => dynamicImportOf(LazyCmp, 50),
+        },
+      ];
+
+      @Component({
+        selector: 'app',
+        imports: [RouterOutlet],
+        template: `
+          Works!
+          <router-outlet />
+        `,
+      })
+      class SimpleComponent {
+        location = inject(Location);
+      }
+
+      const appId = 'custom-app-id';
+      const providers = [
+        {provide: APP_ID, useValue: appId},
+        {provide: PlatformLocation, useClass: MockPlatformLocation},
+        provideRouter(routes),
+      ] as unknown as Provider[];
+      const hydrationFeatures = () => [withIncrementalHydration()];
+
+      const html = await ssr(SimpleComponent, {envProviders: providers, hydrationFeatures});
+      const ssrContents = getAppContents(html);
+
+      expect(ssrContents).toContain(
+        `<button id="click-me" jsaction="click:;" ngb="d2">Click me I'm dehydrated?</button>`,
+      );
+      expect(ssrContents).toContain(`<p id="hydrated">nope</p>`);
+
+      resetTViewsFor(SimpleComponent, LazyCmp);
+
+      const doc = getDocument();
+      const appRef = await prepareEnvironmentAndHydrate(doc, html, SimpleComponent, {
+        envProviders: [...providers],
+        hydrationFeatures,
+      });
+      const compRef = getComponentRef<SimpleComponent>(appRef);
+      await appRef.whenStable();
+      await allPendingDynamicImports();
+      const appHostNode = compRef.location.nativeElement;
+
+      expect(appHostNode.outerHTML).toContain(
+        `<button id="click-me">Click me I'm dehydrated?</button>`,
+      );
+      expect(appHostNode.outerHTML).toContain(`<p id="hydrated">yup</p>`);
+    });
   });
 });

Original file line number	Diff line number	Diff line change
`@@ -586,7 +586,6 @@ export function getParentBlockHydrationQueue(`
`586`	`586`	`isTopMostDeferBlock = dehydratedBlockRegistry.has(currentBlockId);`
`587`	`587`	`const hydratingParentBlock = dehydratedBlockRegistry.hydrating.get(currentBlockId);`
`588`	`588`	`if (parentBlockPromise === null && hydratingParentBlock != null) {`
`589`		- // TODO: add an ngDevMode asset that `hydratingParentBlock.promise` exists and is of type Promise.
`590`	`589`	`parentBlockPromise = hydratingParentBlock.promise;`
`591`	`590`	`break;`
`592`	`591`	`}`