refactor(core): Additional cleanup for incremental hydration (angular#58394)

This adds comments, removes some duplicate logic, and eliminates some unnecessary complexity of the incremental hydration core logic.

PR Close angular#58394

Author: thePunderWoman

packages/core/src/defer/utils.ts

@@ -166,3 +166,18 @@ export function isTDeferBlockDetails(value: unknown): value is TDeferBlockDetail
     typeof (value as TDeferBlockDetails).primaryTmplIndex === 'number'
   );
 }
+
+/**
+ * Whether a given TNode represents a defer block.
+ */
+export function isDeferBlock(tView: TView, tNode: TNode): boolean {
+  let tDetails: TDeferBlockDetails | null = null;
+  const slotIndex = getDeferBlockDataIndex(tNode.index);
+  // Check if a slot index is in the reasonable range.
+  // Note: we do `-1` on the right border, since defer block details are stored
+  // in the `n+1` slot, see `getDeferBlockDataIndex` for more info.
+  if (HEADER_OFFSET < slotIndex && slotIndex < tView.bindingStartIndex) {
+    tDetails = getTDeferBlockDetails(tView, tNode);
+  }
+  return !!tDetails && isTDeferBlockDetails(tDetails);
+}

packages/core/src/hydration/annotate.ts

@@ -14,7 +14,7 @@ import {
   HydrateTriggerDetails,
   TDeferBlockDetails,
 } from '../defer/interfaces';
-import {getLDeferBlockDetails, getTDeferBlockDetails} from '../defer/utils';
+import {getLDeferBlockDetails, getTDeferBlockDetails, isDeferBlock} from '../defer/utils';
 import {isDetachedByI18n} from '../i18n/utils';
 import {ViewEncapsulation} from '../metadata';
 import {Renderer2} from '../render';
@@ -91,7 +91,6 @@ import {
   TextNodeMarker,
 } from './utils';
 import {Injector} from '../di';
-import {isDeferBlock} from './blocks';
 
 /**
  * A collection that tracks all serialized views (`ngh` DOM annotations)

packages/core/src/hydration/api.ts

@@ -46,7 +46,6 @@ import {
 import {enableRetrieveHydrationInfoImpl, NGH_DATA_KEY, SSR_CONTENT_INTEGRITY_MARKER} from './utils';
 import {enableFindMatchingDehydratedViewImpl} from './views';
 import {bootstrapIncrementalHydration, enableRetrieveDeferBlockDataImpl} from './incremental';
-import {enableHydrateFromBlockNameImpl} from './blocks';
 
 /**
  * Indicates whether the hydration-related code was added,
@@ -124,7 +123,6 @@ function enableIncrementalHydrationRuntimeSupport() {
   if (!isIncrementalHydrationRuntimeSupportEnabled) {
     isIncrementalHydrationRuntimeSupportEnabled = true;
     enableRetrieveDeferBlockDataImpl();
-    enableHydrateFromBlockNameImpl();
   }
 }
 

packages/core/src/hydration/blocks.ts

@@ -13,25 +13,26 @@ import {Injector} from '../di';
 import {TransferState} from '../transfer_state';
 import {removeListenersFromBlocks} from '../event_delegation_utils';
 import {cleanupLContainer} from './cleanup';
-import {TNode} from '../render3/interfaces/node';
-import {HEADER_OFFSET, TView} from '../render3/interfaces/view';
-import {DeferBlock, TDeferBlockDetails} from '../defer/interfaces';
-import {getDeferBlockDataIndex, getTDeferBlockDetails, isTDeferBlockDetails} from '../defer/utils';
+import {DeferBlock} from '../defer/interfaces';
 import {whenStable, ApplicationRef} from '../application/application_ref';
 
 /**
  * Finds first hydrated parent `@defer` block for a given block id.
  * If there are any dehydrated `@defer` blocks found along the way,
  * they are also stored and returned from the function (as a list of ids).
+ * Note: This is utilizing serialized information to navigate up the tree
  */
 export function findFirstHydratedParentDeferBlock(deferBlockId: string, injector: Injector) {
   const deferBlockRegistry = injector.get(DeferBlockRegistry);
   const transferState = injector.get(TransferState);
   const deferBlockParents = transferState.get(NGH_DEFER_BLOCKS_KEY, {});
   const dehydratedBlocks: string[] = [];
 
-  let deferBlock = deferBlockRegistry.get(deferBlockId) ?? null;
+  let deferBlock = deferBlockRegistry.get(deferBlockId);
   let currentBlockId: string | null = deferBlockId;
+  // at each level we check if the registry has the given defer block id
+  // - if it does, we know it was already hydrated and can stop here
+  // - if it does not, we continue on
   while (!deferBlock) {
     dehydratedBlocks.unshift(currentBlockId);
     currentBlockId = deferBlockParents[currentBlockId][DEFER_PARENT_BLOCK_ID];
@@ -42,24 +43,21 @@ export function findFirstHydratedParentDeferBlock(deferBlockId: string, injector
 }
 
 /**
- * Hydrates a defer block by block name through jsaction code paths
- */
-let _hydrateFromBlockNameImpl: typeof hydrateFromBlockNameImpl = () => {
-  return Promise.resolve({deferBlock: null, hydratedBlocks: new Set<string>()});
-};
-
-/**
- * Hydrates a defer block by block name using non jsaction code paths
+ * The core mechanism for incremental hydration. This recursively 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.
+ *
+ * @param injector
+ * @param blockName
+ * @param onTriggerFn The function that triggers the block and fetches deps
+ * @param hydratedBlocks The set of blocks currently being hydrated in the tree
+ * @returns
  */
-let _incrementallyHydrateFromBlockNameImpl: typeof incrementallyHydrateFromBlockNameImpl = () => {
-  return Promise.resolve();
-};
-
-async function hydrateFromBlockNameImpl(
+export async function hydrateFromBlockName(
   injector: Injector,
   blockName: string,
-  onTriggerFn: (deferBlock: any) => void,
-  hydratedBlocks: Set<string>,
+  onTriggerFn: (deferBlock: DeferBlock) => void,
+  hydratedBlocks: Set<string> = new Set(),
 ): Promise<{
   deferBlock: DeferBlock | null;
   hydratedBlocks: Set<string>;
@@ -74,80 +72,48 @@ async function hydrateFromBlockNameImpl(
     injector,
   );
   if (deferBlock && blockId) {
+    // Step 2: Add the current block to the tracking sets to prevent
+    // attempting to trigger hydration on a block more than once
+    // simulataneously.
     hydratedBlocks.add(blockId);
     deferBlockRegistry.hydrating.add(blockId);
 
+    // Step 3: Run the actual trigger function to fetch dependencies
     await onTriggerFn(deferBlock);
+
+    // Step 4: Recursively trigger, fetch, and hydrate from the top of the hierarchy down
     let hydratedBlock: DeferBlock | null = deferBlock;
     for (const dehydratedBlock of dehydratedBlocks) {
-      const hydratedInfo = await hydrateFromBlockNameImpl(
+      const hydratedInfo = await hydrateFromBlockName(
         injector,
         dehydratedBlock,
         onTriggerFn,
         hydratedBlocks,
       );
       hydratedBlock = hydratedInfo.deferBlock;
     }
-    // this is going to be the wrong defer block. We need to get the final one.
+    // TODO(incremental-hydration): this is likely where we want to do Step 5: some cleanup work in the
+    // DeferBlockRegistry.
     return {deferBlock: hydratedBlock, hydratedBlocks};
   } else {
     // TODO(incremental-hydration): this is likely an error, consider producing a `console.error`.
     return {deferBlock: null, hydratedBlocks};
   }
 }
 
-/**
- * Sets the implementation for the `retrieveDeferBlockData` function.
- */
-export function enableHydrateFromBlockNameImpl() {
-  _hydrateFromBlockNameImpl = hydrateFromBlockNameImpl;
-  _incrementallyHydrateFromBlockNameImpl = incrementallyHydrateFromBlockNameImpl;
-}
-
-export async function hydrateFromBlockName(
-  injector: Injector,
-  blockName: string,
-  onTriggerFn: (deferBlock: any) => void,
-): Promise<{
-  deferBlock: DeferBlock | null;
-  hydratedBlocks: Set<string>;
-}> {
-  return await _hydrateFromBlockNameImpl(injector, blockName, onTriggerFn, new Set<string>());
-}
-
-async function incrementallyHydrateFromBlockNameImpl(
+export async function incrementallyHydrateFromBlockName(
   injector: Injector,
   blockName: string,
-  triggerFn: (deferBlock: any) => void,
+  triggerFn: (deferBlock: DeferBlock) => void,
 ): Promise<void> {
   const {deferBlock, hydratedBlocks} = await hydrateFromBlockName(injector, blockName, triggerFn);
-  removeListenersFromBlocks([...hydratedBlocks], injector);
   if (deferBlock !== null) {
+    // hydratedBlocks is a set, and needs to be converted to an array
+    // for removing listeners
+    removeListenersFromBlocks([...hydratedBlocks], injector);
     cleanupLContainer(deferBlock.lContainer);
-    const appRef = injector.get(ApplicationRef);
-    await whenStable(appRef);
-  }
-}
-
-export function incrementallyHydrateFromBlockName(
-  injector: Injector,
-  blockName: string,
-  triggerFn: (deferBlock: any) => void,
-): Promise<void> {
-  return _incrementallyHydrateFromBlockNameImpl(injector, blockName, triggerFn);
-}
-
-/**
- * Whether a given TNode represents a defer block.
- */
-export function isDeferBlock(tView: TView, tNode: TNode): boolean {
-  let tDetails: TDeferBlockDetails | null = null;
-  const slotIndex = getDeferBlockDataIndex(tNode.index);
-  // Check if a slot index is in the reasonable range.
-  // Note: we do `-1` on the right border, since defer block details are stored
-  // in the `n+1` slot, see `getDeferBlockDataIndex` for more info.
-  if (HEADER_OFFSET < slotIndex && slotIndex < tView.bindingStartIndex) {
-    tDetails = getTDeferBlockDetails(tView, tNode);
+    // we need to wait for app stability here so we don't continue before
+    // the hydration process has finished, which could result in problems
+    await whenStable(injector.get(ApplicationRef));
   }
-  return !!tDetails && isTDeferBlockDetails(tDetails);
 }

packages/core/src/hydration/event_replay.ts

@@ -44,17 +44,13 @@ import {DeferBlock, DeferBlockTrigger, HydrateTriggerDetails} from '../defer/int
 import {triggerAndWaitForCompletion} from '../defer/instructions';
 import {cleanupDehydratedViews, cleanupLContainer} from './cleanup';
 import {hoverEventNames, interactionEventNames} from '../defer/dom_triggers';
+import {DeferBlockRegistry} from '../defer/registry';
 
 /** Apps in which we've enabled event replay.
  *  This is to prevent initializing event replay more than once per app.
  */
 const appsWithEventReplay = new WeakSet<ApplicationRef>();
 
-/**
- * A set of in progress hydrating blocks
- */
-let hydratingBlocks = new Set<string>();
-
 /**
  * A list of block events that need to be replayed
  */
@@ -94,6 +90,9 @@ export function withEventReplay(): Provider[] {
       provide: ENVIRONMENT_INITIALIZER,
       useValue: () => {
         const injector = inject(Injector);
+        // We have to check for the appRef here due to the possibility of multiple apps
+        // being present on the same page. We only want to enable event replay for the
+        // apps that actually want it.
         const appRef = injector.get(ApplicationRef);
         if (!appsWithEventReplay.has(appRef)) {
           const jsActionMap = inject(BLOCK_ELEMENT_MAP);
@@ -118,6 +117,9 @@ export function withEventReplay(): Provider[] {
               return;
             }
 
+            // We have to check for the appRef here due to the possibility of multiple apps
+            // being present on the same page. We only want to enable event replay for the
+            // apps that actually want it.
             if (!appsWithEventReplay.has(appRef)) {
               appsWithEventReplay.add(appRef);
               appRef.onDestroy(() => appsWithEventReplay.delete(appRef));
@@ -242,30 +244,16 @@ async function hydrateAndInvokeBlockListeners(
   currentTarget: Element,
 ) {
   blockEventQueue.push({event, currentTarget});
-  if (!hydratingBlocks.has(blockName)) {
-    hydratingBlocks.add(blockName);
-    await triggerBlockHydration(injector, blockName, fetchAndRenderDeferBlock);
-    hydratingBlocks.delete(blockName);
-  }
-}
-
-export async function fetchAndRenderDeferBlock(deferBlock: DeferBlock): Promise<DeferBlock> {
-  await triggerAndWaitForCompletion(deferBlock);
-  return deferBlock;
-}
-
-async function triggerBlockHydration(
-  injector: Injector,
-  blockName: string,
-  onTriggerFn: (deferBlock: any) => void,
-) {
-  // grab the list of dehydrated blocks and queue them up
-  const {dehydratedBlocks} = findFirstHydratedParentDeferBlock(blockName, injector);
-  for (let block of dehydratedBlocks) {
-    hydratingBlocks.add(block);
-  }
-  const {deferBlock, hydratedBlocks} = await hydrateFromBlockName(injector, blockName, onTriggerFn);
+  const {deferBlock, hydratedBlocks} = await hydrateFromBlockName(
+    injector,
+    blockName,
+    fetchAndRenderDeferBlock,
+  );
   if (deferBlock !== null) {
+    // TODO(incremental-hydration): extract this work into a post
+    // hydration cleanup function
+    const deferBlockRegistry = injector.get(DeferBlockRegistry);
+    deferBlockRegistry.hydrating.delete(blockName);
     hydratedBlocks.add(blockName);
     const appRef = injector.get(ApplicationRef);
     await appRef.whenStable();
@@ -274,6 +262,11 @@ async function triggerBlockHydration(
   }
 }
 
+export async function fetchAndRenderDeferBlock(deferBlock: DeferBlock): Promise<DeferBlock> {
+  await triggerAndWaitForCompletion(deferBlock);
+  return deferBlock;
+}
+
 function replayQueuedBlockEvents(hydratedBlocks: Set<string>, injector: Injector) {
   // clone the queue
   const queue = [...blockEventQueue];