fix(core): fixes issues with control flow and incremental hydration (#58644)

If a defer block is nested inside control flow while also being nested
underneath a defer block all using incremental hydration, timing issues
prevented the child nodes from being properly hydrated. This ensures
hydration happens on next render.

PR Close #58644

packages/core/src/defer/registry.ts

@@ -35,9 +35,11 @@ export class DehydratedBlockRegistry {
   private cleanupFns = new Map<string, Function[]>();
   private jsActionMap: Map<string, Set<Element>> = inject(JSACTION_BLOCK_ELEMENT_MAP);
   private contract: EventContractDetails = inject(JSACTION_EVENT_CONTRACT);
+
   add(blockId: string, info: DehydratedDeferBlock) {
     this.registry.set(blockId, info);
   }
+
   get(blockId: string): DehydratedDeferBlock | null {
     return this.registry.get(blockId) ?? null;
   }

packages/core/src/defer/triggering.ts

@@ -6,6 +6,7 @@
  * found in the LICENSE file at https://angular.dev/license
  */
 
+import {afterNextRender} from '../render3/after_render/hooks';
 import {Injector} from '../di';
 import {internalImportProvidersFrom} from '../di/provider_collection';
 import {RuntimeError, RuntimeErrorCode} from '../errors';
@@ -46,7 +47,7 @@ import {
   TDeferBlockDetails,
   TriggerType,
 } from './interfaces';
-import {DEHYDRATED_BLOCK_REGISTRY} from './registry';
+import {DEHYDRATED_BLOCK_REGISTRY, DehydratedBlockRegistry} from './registry';
 import {
   DEFER_BLOCK_CONFIG,
   DEFER_BLOCK_DEPENDENCY_INTERCEPTOR,
@@ -354,12 +355,43 @@ export async function triggerHydrationFromBlockName(
  * Triggers the resource loading for a defer block and passes back a promise
  * to handle cleanup on completion
  */
-export function triggerAndWaitForCompletion(deferBlock: DehydratedDeferBlock): Promise<void> {
-  const lDetails = getLDeferBlockDetails(deferBlock.lView, deferBlock.tNode);
-  const promise = new Promise<void>((resolve) => {
-    onDeferBlockCompletion(lDetails, resolve);
+export function triggerAndWaitForCompletion(
+  dehydratedBlockId: string,
+  dehydratedBlockRegistry: DehydratedBlockRegistry,
+  injector: Injector,
+): Promise<void> {
+  // TODO(incremental-hydration): This is a temporary fix to resolve control flow
+  // cases where nested defer blocks are inside control flow. We wait for each nested
+  // defer block to load and render before triggering the next one in a sequence. This is
+  // needed to ensure that corresponding LViews & LContainers are available for a block
+  // before we trigger it. We need to investigate how to get rid of the `afterNextRender`
+  // calls (in the nearest future) and do loading of all dependencies of nested defer blocks
+  // in parallel (later).
+
+  let resolve: VoidFunction;
+  const promise = new Promise<void>((resolveFn) => {
+    resolve = resolveFn;
   });
-  triggerDeferBlock(deferBlock.lView, deferBlock.tNode);
+
+  afterNextRender(
+    () => {
+      const deferBlock = dehydratedBlockRegistry.get(dehydratedBlockId);
+      // Since we trigger hydration for nested defer blocks in a sequence (parent -> child),
+      // there is a chance that a defer block may not be present at hydration time. For example,
+      // when a nested block was in an `@if` condition, which has changed.
+      // TODO(incremental-hydration): add tests to verify the behavior mentioned above.
+      if (deferBlock !== null) {
+        const {tNode, lView} = deferBlock;
+        const lDetails = getLDeferBlockDetails(lView, tNode);
+        onDeferBlockCompletion(lDetails, resolve);
+        triggerDeferBlock(lView, tNode);
+        // TODO(incremental-hydration): handle the cleanup for cases when
+        // defer block is no longer present during hydration (e.g. `@if` condition
+        // has changed during hydration/rendering).
+      }
+    },
+    {injector},
+  );
   return promise;
 }
 
@@ -401,12 +433,9 @@ async function triggerBlockTreeHydrationByName(
 
   // Step 3: hydrate each block in the queue. It will be in descending order from the top down.
   for (const dehydratedBlockId of hydrationQueue) {
-    // The registry will have the item in the queue after each loop.
-    const deferBlock = dehydratedBlockRegistry.get(dehydratedBlockId)!;
-
     // Step 4: Run the actual trigger function to fetch dependencies.
     // Triggering a block adds any of its child defer blocks to the registry.
-    await triggerAndWaitForCompletion(deferBlock);
+    await triggerAndWaitForCompletion(dehydratedBlockId, dehydratedBlockRegistry, injector);
   }
 
   const hydratedBlocks = new Set<string>(hydrationQueue);

packages/platform-server/test/incremental_hydration_spec.ts

@@ -28,12 +28,9 @@ import {
 } from '@angular/platform-browser';
 import {TestBed} from '@angular/core/testing';
 import {PLATFORM_BROWSER_ID} from '@angular/common/src/platform_id';
-import {DEHYDRATED_BLOCK_REGISTRY, DehydratedBlockRegistry} from '@angular/core/src/defer/registry';
+import {DEHYDRATED_BLOCK_REGISTRY} from '@angular/core/src/defer/registry';
 import {JSACTION_BLOCK_ELEMENT_MAP} from '@angular/core/src/hydration/tokens';
-import {
-  EventContractDetails,
-  JSACTION_EVENT_CONTRACT,
-} from '@angular/core/src/event_delegation_utils';
+import {JSACTION_EVENT_CONTRACT} from '@angular/core/src/event_delegation_utils';
 
 describe('platform-server partial hydration integration', () => {
   const originalWindow = globalThis.window;
@@ -1445,6 +1442,88 @@ describe('platform-server partial hydration integration', () => {
     });
   });
 
+  describe('control flow', () => {
+    it('should support hydration for all items in a for loop', async () => {
+      @Component({
+        standalone: true,
+        selector: 'app',
+        template: `
+          <main>
+            @defer (on interaction; hydrate on interaction) {
+              <div id="main" (click)="fnA()">
+                <p>Main defer block rendered!</p>
+                @for (item of items; track $index) {
+                  @defer (on interaction; hydrate on interaction) {
+                    <article id="item-{{item}}">
+                      defer block {{item}} rendered!
+                      <span (click)="fnB()">{{value()}}</span>
+                    </article>
+                  } @placeholder {
+                    <span>Outer block placeholder</span>
+                  }
+                }
+              </div>
+            } @placeholder {
+              <span>Outer block placeholder</span>
+            }
+          </main>
+        `,
+      })
+      class SimpleComponent {
+        value = signal('start');
+        items = [1, 2, 3, 4, 5, 6];
+        fnA() {}
+        fnB() {
+          this.value.set('end');
+        }
+      }
+
+      const appId = 'custom-app-id';
+      const providers = [{provide: APP_ID, useValue: appId}];
+      const hydrationFeatures = () => [withIncrementalHydration()];
+
+      const html = await ssr(SimpleComponent, {envProviders: providers, hydrationFeatures});
+      const ssrContents = getAppContents(html);
+
+      // <main> uses "eager" `custom-app-id` namespace.
+      // <div>s inside a defer block have `d0` as a namespace.
+      expect(ssrContents).toContain('<article id="item-1" jsaction="click:;keydown:;"');
+      // Outer defer block is rendered.
+      expect(ssrContents).toContain('defer block 1 rendered');
+
+      // Internal cleanup before we do server->client transition in this test.
+      resetTViewsFor(SimpleComponent);
+
+      ////////////////////////////////
+      const doc = getDocument();
+      const appRef = await prepareEnvironmentAndHydrate(doc, html, SimpleComponent, {
+        envProviders: [...providers, {provide: PLATFORM_ID, useValue: 'browser'}],
+        hydrationFeatures,
+      });
+      const compRef = getComponentRef<SimpleComponent>(appRef);
+      appRef.tick();
+      await whenStable(appRef);
+
+      const appHostNode = compRef.location.nativeElement;
+
+      expect(appHostNode.outerHTML).toContain('<article id="item-1" jsaction="click:;keydown:;"');
+
+      // Emit an event inside of a defer block, which should result
+      // in triggering the defer block (start loading deps, etc) and
+      // subsequent hydration.
+      const article = doc.getElementById('item-1')!;
+      const clickEvent = new CustomEvent('click', {bubbles: true});
+      article.dispatchEvent(clickEvent);
+      await timeout(1000); // wait for defer blocks to resolve
+
+      appRef.tick();
+      expect(appHostNode.outerHTML).not.toContain(
+        '<article id="item-1" jsaction="click:;keydown:;"',
+      );
+      expect(appHostNode.outerHTML).not.toContain('<span>Outer block placeholder</span>');
+    });
+  });
+
   describe('cleanup', () => {
     it('should cleanup partial hydration blocks appropriately', async () => {
       @Component({
@@ -1585,7 +1664,6 @@ describe('platform-server partial hydration integration', () => {
 
       await timeout(1000); // wait for defer blocks to resolve
 
-      appRef.tick();
       expect(registry.size).toBe(1);
       expect(registry.has('d0')).toBeFalsy();
       expect(jsActionMap.size).toBe(1);