Patch (#21570) tree: Reuse flex tree fields in most cases (#21577)

This is as cherry-pick of #21570 

Currently flex-tree fields tend to be leaked until the nodes they are
part of are disposed (if ever) or the context is disposed.

Due to this, allocating tons of separate objects for the same field is
really bad. This change reuses flex tree fields in most cases,
mitigating this issue.

This was observed to make bubble bench >10x faster, however most of that
win was due to avoiding the O(number of leaked copies of the field)
creation cost in the event subscription logic which is being removed
independently.

Co-authored-by: Craig Macomber (Microsoft) <[email protected]>

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<LazyTreeNode, Map<FieldKey, FlexTreeField>> = 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<FieldKey, FlexTreeField>(),
 	);
-	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;
 }
 
 /**