@milaboratories/pl-tree 1.10.6 → 1.11.1

package/src/sync.ts

@@ -1,7 +1,9 @@
 import type {
   FieldData,
+  Filter,
   OptionalSignedResourceId,
   PlTransaction,
+  ResourceTreeFrame,
   SignedResourceId,
 } from "@milaboratories/pl-client";
 import Denque from "denque";
@@ -23,19 +25,32 @@ export interface TreeLoadingRequest {
    * be retrieved for them. */
   readonly finalResources: Set<SignedResourceId>;
 
-  /** This function is applied to each resource data field list, before
-   * using it continue traversal. This modification also is applied to
-   * output data to make result self-consistent in terms that it will contain
-   * all referenced resources, this is required to be able to pass it to tree
-   * to update the state. */
+  /** Applied to each resource field list in fallback BFS mode and to streamed results. */
   readonly pruningFunction?: PruningFunction;
+
+  /** ResourceTree field filter passed to the backend when supported. */
+  readonly fieldFilter?: Filter;
+
+  /** ResourceTree traversal stop rules passed to the backend when supported. */
+  readonly traverseStopRules?: Filter;
 }
 
+const CapabilityTreeFilter = "treeFilter:v1";
+
+/** Controls which tree-loading path is used.
+ * - `"auto"` (default): use backend streaming when the backend advertises `treeFilter:v1`,
+ *   fall back to client-side BFS otherwise.
+ * - `"client-bfs"`: always use client-side BFS, even on capable backends.
+ * - `"backend-streaming"`: always prefer backend streaming; if the capability is absent,
+ *   logs a warning and falls back to BFS (never throws).
+ */
+export type TraversalMode = "auto" | "client-bfs" | "backend-streaming";
+
 /** Given the current tree state, build the request object to pass to
  * {@link loadTreeState} to load updated state. */
 export function constructTreeLoadingRequest(
   tree: PlTreeState,
-  pruningFunction?: PruningFunction,
+  options: Pick<TreeLoadingRequest, "pruningFunction" | "fieldFilter" | "traverseStopRules"> = {},
 ): TreeLoadingRequest {
   const seedResources: SignedResourceId[] = [];
   const finalResources = new Set<SignedResourceId>();
@@ -47,7 +62,13 @@ export function constructTreeLoadingRequest(
   // if tree is empty, seeding tree reconstruction from the specified root
   if (seedResources.length === 0 && finalResources.size === 0) seedResources.push(tree.root);
 
-  return { seedResources, finalResources, pruningFunction };
+  return {
+    seedResources,
+    finalResources,
+    pruningFunction: options.pruningFunction,
+    fieldFilter: options.fieldFilter,
+    traverseStopRules: options.traverseStopRules,
+  };
 }
 
 export type TreeLoadingStat = {
@@ -61,6 +82,10 @@ export type TreeLoadingStat = {
   prunedFields: number;
   finalResourcesSkipped: number;
   millisSpent: number;
+  /** Stop-marker frames whose id was already final locally and were skipped. */
+  stopMarkersSkipped: number;
+  /** Number of follow-up resourceTree() calls issued to resolve unknown stop markers. */
+  stopMarkerFollowUpRoundTrips: number;
 };
 
 export function initialTreeLoadingStat(): TreeLoadingStat {
@@ -75,6 +100,8 @@ export function initialTreeLoadingStat(): TreeLoadingStat {
     prunedFields: 0,
     finalResourcesSkipped: 0,
     millisSpent: 0,
+    stopMarkersSkipped: 0,
+    stopMarkerFollowUpRoundTrips: 0,
   };
 }
 
@@ -88,24 +115,30 @@ export function formatTreeLoadingStat(stat: TreeLoadingStat): string {
   result += `Data Bytes: ${stat.retrievedResourceDataBytes}\n`;
   result += `KV Bytes: ${stat.retrievedKeyValueBytes}\n`;
   result += `Pruned fields: ${stat.prunedFields}\n`;
-  result += `Final resources skipped: ${stat.finalResourcesSkipped}`;
+  result += `Final resources skipped: ${stat.finalResourcesSkipped}\n`;
+  result += `Stop markers skipped: ${stat.stopMarkersSkipped}\n`;
+  result += `Stop marker follow-up round-trips: ${stat.stopMarkerFollowUpRoundTrips}`;
   return result;
 }
 
-/** Given the transaction (preferably read-only) and loading request, executes
- * the tree traversal algorithm, and collects fresh states of resources
- * to update the tree state. */
-export async function loadTreeState(
+function supportsResourceTreeTraversal(capabilities: readonly string[] = []): boolean {
+  return capabilities.includes(CapabilityTreeFilter);
+}
+
+function collectStatsForResource(resource: ExtendedResourceData, stats?: TreeLoadingStat) {
+  if (!stats) return;
+  stats.retrievedResources++;
+  stats.retrievedFields += resource.fields.length;
+  stats.retrievedKeyValues += resource.kv.length;
+  stats.retrievedResourceDataBytes += resource.data?.length ?? 0;
+  for (const kv of resource.kv) stats.retrievedKeyValueBytes += kv.value.length;
+}
+
+async function loadTreeStateViaBfs(
   tx: PlTransaction,
   loadingRequest: TreeLoadingRequest,
   stats?: TreeLoadingStat,
 ): Promise<ExtendedResourceData[]> {
-  // saving start timestamp to add time spent in this function to the stats at the end of the method
-  const startTimestamp = Date.now();
-
-  // counting the request
-  if (stats) stats.requests++;
-
   const { seedResources, finalResources, pruningFunction } = loadingRequest;
 
   // Limits the number of concurrent gRPC fetches to bound peak memory
@@ -187,23 +220,163 @@ export async function loadTreeState(
     }
 
     // collecting stats
-    if (stats) {
-      stats.retrievedResources++;
-      stats.retrievedFields += nextResource.fields.length;
-      stats.retrievedKeyValues += nextResource.kv.length;
-      stats.retrievedResourceDataBytes += nextResource.data?.length ?? 0;
-      for (const kv of nextResource.kv) stats.retrievedKeyValueBytes += kv.value.length;
-    }
+    collectStatsForResource(nextResource, stats);
 
     // aggregating the state
     result.push(nextResource);
   }
 
-  // adding the time we spent in this method to stats
-  if (stats) {
-    stats.millisSpent += Date.now() - startTimestamp;
-    stats.roundTrips += numberOfRoundTrips;
+  if (stats) stats.roundTrips += numberOfRoundTrips;
+
+  return result;
+}
+
+async function processResourceTreeStream(
+  treeItems: AsyncIterable<ResourceTreeFrame>,
+  finalResources: Set<SignedResourceId>,
+  pruningFunction: PruningFunction | undefined,
+  stats: TreeLoadingStat | undefined,
+): Promise<{ result: ExtendedResourceData[]; followUpSeeds: SignedResourceId[] }> {
+  const result: ExtendedResourceData[] = [];
+  const followUpSeeds: SignedResourceId[] = [];
+
+  // backend returns two types of frames:
+  // - 'resource' frames contain the resource state and are processed normally
+  // - 'stopMarker' frames indicate resources that are ignored due stop rules fired
+  //
+  // Usually stop rules indicates the resources with final state. In that case middle layer
+  // should make a decision: has it already loaded the resource or should it be requested for get the latest state?
+  for await (const frame of treeItems) {
+    if (frame.frameKind === "stopMarker") {
+      if (finalResources.has(frame.id)) {
+        if (stats) stats.stopMarkersSkipped++;
+        continue;
+      }
+      followUpSeeds.push(frame.id);
+      continue;
+    }
+
+    // Normal resource frame.
+    if (finalResources.has(frame.id)) {
+      if (stats) stats.finalResourcesSkipped++;
+      continue;
+    }
+
+    let nextResource: ExtendedResourceData = {
+      id: frame.id,
+      type: frame.type,
+      kind: frame.kind,
+      data: frame.data,
+      resourceReady: frame.resourceReady,
+      error: frame.error,
+      originalResourceId: frame.originalResourceId,
+      // traverseWasStopped: backend matched traverse stop rules — children were not streamed.
+      // Mark as terminal; fields are resolved below.
+      final: frame.final || frame.traverseWasStopped,
+      inputsLocked: frame.inputsLocked,
+      outputsLocked: frame.outputsLocked,
+      fields: frame.fields,
+      kv: frame.kv,
+    };
+
+    // Apply field rules: traverseWasStopped drops all fields to keep the refCount
+    // invariant; pruning function further filters the remaining fields.
+    const rawFields = frame.traverseWasStopped ? [] : nextResource.fields;
+    const resolvedFields =
+      pruningFunction !== undefined
+        ? pruningFunction({ ...nextResource, fields: rawFields })
+        : rawFields;
+    if (stats) stats.prunedFields += nextResource.fields.length - resolvedFields.length;
+    nextResource = { ...nextResource, fields: resolvedFields };
+
+    collectStatsForResource(nextResource, stats);
+    result.push(nextResource);
+  }
+
+  return { result, followUpSeeds };
+}
+
+async function loadTreeStateViaResourceTree(
+  tx: PlTransaction,
+  loadingRequest: TreeLoadingRequest,
+  stats?: TreeLoadingStat,
+  logger?: { warn: (msg: string) => void; info?: (msg: unknown) => void },
+): Promise<ExtendedResourceData[]> {
+  const { seedResources, finalResources, pruningFunction, fieldFilter, traverseStopRules } =
+    loadingRequest;
+
+  // Round 0: initial tree traversal.
+  const treeItems = tx.resourceTree(seedResources, {
+    includeKv: true,
+    fieldFilter,
+    traverseStopRules,
+  });
+
+  const { result, followUpSeeds } = await processResourceTreeStream(
+    treeItems,
+    finalResources,
+    pruningFunction,
+    stats,
+  );
+  if (stats) stats.roundTrips++;
+
+  // Client must request full resource tree in case when stop-marker seeds are returned,
+  // to ensure all resources are loaded and stop-marker frames are processed.
+  if (followUpSeeds.length > 0) {
+    const followUpItems = tx.resourceTree(followUpSeeds, {
+      includeKv: true,
+      fieldFilter,
+    });
+    const { result: followUpResult } = await processResourceTreeStream(
+      followUpItems,
+      finalResources,
+      pruningFunction,
+      stats,
+    );
+    result.push(...followUpResult);
+    if (stats) {
+      logger?.info?.(
+        `loadTreeStateViaResourceTree: follow-up request for ${followUpSeeds.length} stop-marker seeds: ${JSON.stringify(followUpSeeds)}`,
+      );
+      stats.roundTrips++;
+      stats.stopMarkerFollowUpRoundTrips++;
+    }
   }
 
   return result;
 }
+
+/** Given the transaction (preferably read-only) and loading request, executes
+ * the tree traversal algorithm, and collects fresh states of resources
+ * to update the tree state. */
+export async function loadTreeState(
+  tx: PlTransaction,
+  loadingRequest: TreeLoadingRequest,
+  stats?: TreeLoadingStat,
+  capabilities: readonly string[] = [],
+  mode: TraversalMode = "auto",
+  logger?: { warn: (msg: string) => void; info?: (msg: unknown) => void },
+): Promise<ExtendedResourceData[]> {
+  const startTimestamp = Date.now();
+  if (stats) stats.requests++;
+
+  try {
+    const wantsStreaming =
+      mode === "backend-streaming" ||
+      (mode === "auto" && supportsResourceTreeTraversal(capabilities));
+
+    if (wantsStreaming && !supportsResourceTreeTraversal(capabilities)) {
+      const msg =
+        "traversalMode=backend-streaming but backend lacks treeFilter:v1 capability; falling back to BFS";
+      if (logger) logger.warn(msg);
+      else console.warn(msg);
+      return await loadTreeStateViaBfs(tx, loadingRequest, stats);
+    }
+
+    return wantsStreaming
+      ? await loadTreeStateViaResourceTree(tx, loadingRequest, stats, logger)
+      : await loadTreeStateViaBfs(tx, loadingRequest, stats);
+  } finally {
+    if (stats) stats.millisSpent += Date.now() - startTimestamp;
+  }
+}

package/src/synchronized_tree.ts

@@ -6,12 +6,18 @@ import type {
   SignedResourceId,
   TxOps,
 } from "@milaboratories/pl-client";
+import type { Filter } from "@milaboratories/pl-client";
 import { isTimeoutOrCancelError } from "@milaboratories/pl-client";
 import type { ExtendedResourceData } from "./state";
 import { PlTreeState, TreeStateUpdateError } from "./state";
-import type { PruningFunction, TreeLoadingStat } from "./sync";
+import type { PruningFunction, TraversalMode, TreeLoadingStat } from "./sync";
 import { constructTreeLoadingRequest, initialTreeLoadingStat, loadTreeState } from "./sync";
 import * as tp from "node:timers/promises";
+
+/** Hard floor between consecutive tree-refresh calls.
+ * Applies even when {@link scheduleOnNextState} has woken the loop early,
+ * preventing tight polling loops during rapid state transitions. */
+const MIN_POLLING_INTERVAL_MS = 100;
 import type { MiLogger } from "@milaboratories/ts-helpers";
 
 type StatLoggingMode = "cumulative" | "per-request";
@@ -20,10 +26,15 @@ export type SynchronizedTreeOps = {
   /** Override final predicate from the PlClient */
   finalPredicateOverride?: FinalResourceDataPredicate;
 
-  /** Pruning function to limit set of fields through which tree will
-   * traverse during state synchronization */
+  /** Pruning function for legacy fallback path. */
   pruning?: PruningFunction;
 
+  /** ResourceTree field filter for modern backend path. */
+  fieldFilter?: Filter;
+
+  /** ResourceTree traversal stop rules for modern backend path. */
+  traverseStopRules?: Filter;
+
   /** Interval after last sync to sleep before the next one */
   pollingInterval: number;
   /** For how long to continue polling after the last derived value access */
@@ -34,6 +45,9 @@ export type SynchronizedTreeOps = {
 
   /** Timeout for initial tree loading. If not specified, will use default for RO tx from pl-client. */
   initialTreeLoadingTimeout?: number;
+
+  /** Controls which tree-loading path to use.  Default `"auto"`. */
+  traversalMode?: TraversalMode;
 };
 
 type ScheduledRefresh = {
@@ -46,6 +60,9 @@ export class SynchronizedTreeState {
   private state: PlTreeState;
   private readonly pollingInterval: number;
   private readonly pruning?: PruningFunction;
+  private readonly fieldFilter?: Filter;
+  private readonly traverseStopRules?: Filter;
+  private readonly traversalMode: TraversalMode;
   private readonly logStat?: StatLoggingMode;
   private readonly hooks: PollingComputableHooks;
   private readonly abortController = new AbortController();
@@ -56,8 +73,20 @@ export class SynchronizedTreeState {
     ops: SynchronizedTreeOps,
     private readonly logger?: MiLogger,
   ) {
-    const { finalPredicateOverride, pruning, pollingInterval, stopPollingDelay, logStat } = ops;
+    const {
+      finalPredicateOverride,
+      pruning,
+      fieldFilter,
+      traverseStopRules,
+      traversalMode,
+      pollingInterval,
+      stopPollingDelay,
+      logStat,
+    } = ops;
     this.pruning = pruning;
+    this.fieldFilter = fieldFilter;
+    this.traverseStopRules = traverseStopRules;
+    this.traversalMode = traversalMode ?? "auto";
     this.pollingInterval = pollingInterval;
     this.finalPredicate = finalPredicateOverride ?? pl.finalPredicate;
     this.logStat = logStat;
@@ -123,11 +152,22 @@ export class SynchronizedTreeState {
   /** Executed from the main loop, and initialization procedure. */
   private async refresh(stats?: TreeLoadingStat, txOps?: TxOps): Promise<void> {
     if (this.terminated) throw new Error("tree synchronization is terminated");
-    const request = constructTreeLoadingRequest(this.state, this.pruning);
+    const request = constructTreeLoadingRequest(this.state, {
+      pruningFunction: this.pruning,
+      fieldFilter: this.fieldFilter,
+      traverseStopRules: this.traverseStopRules,
+    });
     const data = await this.pl.withReadTx(
       "ReadingTree",
       async (tx) => {
-        return await loadTreeState(tx, request, stats);
+        return await loadTreeState(
+          tx,
+          request,
+          stats,
+          this.pl.serverInfo.capabilities ?? [],
+          this.traversalMode,
+          this.logger,
+        );
       },
       txOps,
     );
@@ -204,23 +244,40 @@ export class SynchronizedTreeState {
 
       if (!this.keepRunning || this.terminated) break;
 
+      // Phase 1: mandatory floor — always wait at least MIN_POLLING_INTERVAL_MS.
+      // Not interruptible by scheduleOnNextState; only termination aborts it.
+      try {
+        await tp.setTimeout(MIN_POLLING_INTERVAL_MS, undefined, {
+          signal: this.abortController.signal,
+        });
+      } catch (e: unknown) {
+        if (!isTimeoutOrCancelError(e)) throw new Error("Unexpected error", { cause: e });
+        if (this.abortController.signal.aborted) break;
+      }
+
+      if (!this.keepRunning || this.terminated) break;
+
+      // Phase 2: optional remainder up to pollingInterval — interruptible by
+      // scheduleOnNextState so that an external nudge wakes the loop promptly.
       if (this.scheduledOnNextState.length === 0) {
-        try {
-          this.currentLoopDelayInterrupt = new AbortController();
-          await tp.setTimeout(this.pollingInterval, undefined, {
-            signal: AbortSignal.any([
-              this.abortController.signal,
-              this.currentLoopDelayInterrupt.signal,
-            ]),
-          });
-        } catch (e: unknown) {
-          if (!isTimeoutOrCancelError(e)) throw new Error("Unexpected error", { cause: e });
-          // If the main abort controller fired, this is a permanent termination
-          if (this.abortController.signal.aborted) break;
-          // Otherwise it was just the loop delay interrupt (scheduleOnNextState),
-          // continue to the next iteration
-        } finally {
-          this.currentLoopDelayInterrupt = undefined;
+        const remaining = Math.max(0, this.pollingInterval - MIN_POLLING_INTERVAL_MS);
+        if (remaining > 0) {
+          try {
+            this.currentLoopDelayInterrupt = new AbortController();
+            await tp.setTimeout(remaining, undefined, {
+              signal: AbortSignal.any([
+                this.abortController.signal,
+                this.currentLoopDelayInterrupt.signal,
+              ]),
+            });
+          } catch (e: unknown) {
+            if (!isTimeoutOrCancelError(e)) throw new Error("Unexpected error", { cause: e });
+            if (this.abortController.signal.aborted) break;
+            // Otherwise it was just the loop delay interrupt (scheduleOnNextState),
+            // continue to the next iteration
+          } finally {
+            this.currentLoopDelayInterrupt = undefined;
+          }
         }
       }
     }