diff --git a/packages/dds/tree/src/feature-libraries/flex-tree/lazyField.ts b/packages/dds/tree/src/feature-libraries/flex-tree/lazyField.ts index 4e44e7da0d21..370fee17b88c 100644 --- a/packages/dds/tree/src/feature-libraries/flex-tree/lazyField.ts +++ b/packages/dds/tree/src/feature-libraries/flex-tree/lazyField.ts @@ -23,6 +23,7 @@ import { assertValidRangeIndices, disposeSymbol, fail, + getOrCreate, } from "../../util/index.js"; // TODO: stop depending on contextuallyTyped import { applyTypesFromContext, cursorFromContextualData } from "../contextuallyTyped.js"; @@ -51,6 +52,7 @@ import { type FlexibleNodeSubSequence, TreeStatus, flexTreeMarker, + flexTreeSlot, } from "./flexTreeTypes.js"; import { LazyEntity, @@ -60,7 +62,7 @@ import { isFreedSymbol, tryMoveCursorToAnchorSymbol, } from "./lazyEntity.js"; -import { makeTree } from "./lazyNode.js"; +import { type LazyTreeNode, makeTree } from "./lazyNode.js"; import { unboxedUnion } from "./unboxed.js"; import { indexForAt, @@ -69,20 +71,69 @@ import { } from "./utilities.js"; import { UsageError } from "@fluidframework/telemetry-utils/internal"; +/** + * Reuse fields. + * Since field currently own cursors and register themselves for disposal when the node hit end of life, + * not reusing them results in memory leaks every time the field is accessed. + * Since the fields stay alive until the node is end of life reusing them this way is safe. + * + * This ins't a perfect solution: + * + * - This can cause leaks, like map nodes will keep all accessed field objects around. Since other things cause this same leak already, its not too bad. + * - This does not cache the root. + * - Finding the parent anchor to do the caching on has significant cost. + * + * Despite these limitations, this cache provides a large performance win in some common cases (over 10x), especially with how simple tree requests far more field objects than necessary currently. + */ +const fieldCache: WeakMap> = new WeakMap(); + export function makeField( context: Context, schema: FlexFieldSchema, cursor: ITreeSubscriptionCursor, ): FlexTreeField { const fieldAnchor = cursor.buildFieldAnchor(); + let usedAnchor = false; + + const makeFlexTreeField = (): FlexTreeField => { + usedAnchor = true; + const field = new (kindToClass.get(schema.kind) ?? fail("missing field implementation"))( + context, + schema, + cursor, + fieldAnchor, + ); + return field; + }; + + if (fieldAnchor.parent === undefined) { + return makeFlexTreeField(); + } - const field = new (kindToClass.get(schema.kind) ?? fail("missing field implementation"))( - context, - schema, - cursor, - fieldAnchor, + // For the common case (all but roots), cache field associated with its node's anchor and field key. + const anchorNode = + context.checkout.forest.anchors.locate(fieldAnchor.parent) ?? fail("missing anchor"); + + // Since anchor-set could be reused across a flex tree context getting disposed, key off the flex tree node not the anchor. + const cacheKey = anchorNode.slots.get(flexTreeSlot); + + // If there is no flex tree parent node, skip caching: this is not expected to be a hot path, but should probably be fixed at some point. + if (cacheKey === undefined) { + return makeFlexTreeField(); + } + + const innerCache = getOrCreate( + fieldCache, + cacheKey, + () => new Map(), ); - return field; + const result = getOrCreate(innerCache, fieldAnchor.fieldKey, makeFlexTreeField); + if (!usedAnchor) { + // The anchor must be disposed to avoid leaking. In the case of a cache hit, + // we are not transferring ownership to a new FlexTreeField, so it must be disposed of here to avoid the leak. + context.checkout.forest.anchors.forget(fieldAnchor.parent); + } + return result; } /**