fix(tree): commit enrichment after transaction abort (microsoft#25978)

## Description

Aborting a transaction used to put the tree in a state that would
trigger an assert when sending some undo/redo edits to peers.
This would prevent some undo/redo edits from being sent and would put
the tree in a broken state that prevented any further edits.
This issue could not have caused document corruption, so reopening the
document was a possible remedy.

After this PR, aborting a transaction no longer puts the tree in such a
state, so it is safe to perform undo/redo edits after that.

This bug was due to the fact that `TreeCheckout` overwrote its
`DetachedFieldIndex` member with a new instance when performing a
transaction rollback but had no way of informing the
`SharedTreeReadonlyChangeEnricher` that it should now be reading from
the new instance during change enrichment.

This PR fixes the bug by ensuring that a given `TreeCheckout` instance
always uses the same `DetachedFieldIndex` instance. This requires
introducing a way of snapshotting and (on transaction rollback)
restoring the `DetachedFieldIndex` state that does not require replacing
the instance.

This PR also reduces the `DetachedFieldIndex` API surface made available
to `SharedTreeReadonlyChangeEnricher` by introducing the
`ReadOnlyDetachedFieldIndex` interface. This is not required as part of
this change and could be done in a separate PR.

## Breaking Changes

None

## Reviewer Guidance

Let me know if you think it's worth spinning up a separate PR for the
introduction of `ReadOnlyDetachedFieldIndex`.

Author: yann-achard-MS

.changeset/spicy-grapes-leave.md

@@ -0,0 +1,11 @@
+---
+"@fluidframework/tree": minor
+"fluid-framework": minor
+"__section": tree
+---
+Fixed bug in sending of revert edits after an aborted transaction
+
+Aborting a transaction used to put the tree in a state that would trigger an assert when sending some undo/redo edits to peers.
+This would prevent some undo/redo edits from being sent and would put the tree in a broken state that prevented any further edits.
+This issue could not have caused document corruption, so reopening the document was a possible remedy.
+Aborting a transaction no longer puts the tree in such a state, so it is safe to perform undo/redo edits after that.

packages/dds/tree/src/core/index.ts

@@ -83,6 +83,8 @@ export {
 	CursorMarker,
 	isCursor,
 	DetachedFieldIndex,
+	type ReadOnlyDetachedFieldIndex,
+	type DetachedFieldIndexCheckpoint,
 	type ForestRootId,
 	getDetachedFieldContainingPath,
 	aboveRootPlaceholder,

packages/dds/tree/src/core/tree/detachedFieldIndex.ts

@@ -37,10 +37,44 @@ import type {
 } from "./detachedFieldIndexTypes.js";
 import { makeDetachedFieldIndexCodec } from "./detachedFieldIndexCodecs.js";
 
+/**
+ * Readonly interface for {@link DetachedFieldIndex}.
+ */
+export interface ReadOnlyDetachedFieldIndex {
+	/**
+	 * Creates a deep clone of this `DetachedFieldIndex`.
+	 */
+	clone(): DetachedFieldIndex;
+
+	/**
+	 * Returns a field key for the given ID.
+	 * This does not save the field key on the index. To do so, call {@link createEntry}.
+	 */
+	toFieldKey(id: ForestRootId): FieldKey;
+
+	/**
+	 * Returns the `ForestRootId` associated with the given id.
+	 * Returns undefined if no such id is known to the index.
+	 */
+	tryGetEntry(id: Delta.DetachedNodeId): ForestRootId | undefined;
+
+	/**
+	 * Returns the `ForestRootId` associated with the given id.
+	 * Fails if no such id is known to the index.
+	 */
+	getEntry(id: Delta.DetachedNodeId): ForestRootId;
+}
+
+/**
+ * Restores the originating DetachedFieldIndex to the state it was in when the checkpoint was created.
+ * Can be invoked multiple times.
+ */
+export type DetachedFieldIndexCheckpoint = () => void;
+
 /**
  * The tree index records detached field IDs and associates them with a change atom ID.
  */
-export class DetachedFieldIndex {
+export class DetachedFieldIndex implements ReadOnlyDetachedFieldIndex {
 	/**
 	 * A mapping from detached node ids to detached fields.
 	 */
@@ -106,6 +140,25 @@ export class DetachedFieldIndex {
 		return clone;
 	}
 
+	/**
+	 * Creates a restorable checkpoint of the current state of the DetachedFieldIndex.
+	 */
+	public createCheckpoint(): DetachedFieldIndexCheckpoint {
+		const clone = this.clone();
+		return () => {
+			this.purge();
+			populateNestedMap(clone.detachedNodeToField, this.detachedNodeToField, true);
+			populateNestedMap(
+				clone.latestRelevantRevisionToFields,
+				this.latestRelevantRevisionToFields,
+				true,
+			);
+			this.rootIdAllocator = idAllocatorFromMaxId(
+				clone.rootIdAllocator.getMaxId(),
+			) as IdAllocator<ForestRootId>;
+		};
+	}
+
 	public *entries(): Generator<{
 		root: ForestRootId;
 		latestRelevantRevision?: RevisionTag;
@@ -201,26 +254,14 @@ export class DetachedFieldIndex {
 		}
 	}
 
-	/**
-	 * Returns a field key for the given ID.
-	 * This does not save the field key on the index. To do so, call {@link createEntry}.
-	 */
 	public toFieldKey(id: ForestRootId): FieldKey {
 		return brand(`${this.name}-${id}`);
 	}
 
-	/**
-	 * Returns the FieldKey associated with the given id.
-	 * Returns undefined if no such id is known to the index.
-	 */
 	public tryGetEntry(id: Delta.DetachedNodeId): ForestRootId | undefined {
 		return tryGetFromNestedMap(this.detachedNodeToField, id.major, id.minor)?.root;
 	}
 
-	/**
-	 * Returns the FieldKey associated with the given id.
-	 * Fails if no such id is known to the index.
-	 */
 	public getEntry(id: Delta.DetachedNodeId): ForestRootId {
 		const key = this.tryGetEntry(id);
 		assert(key !== undefined, 0x7aa /* Unknown removed node ID */);

packages/dds/tree/src/core/tree/index.ts

@@ -117,7 +117,11 @@ export {
 	type ChunkedCursor,
 } from "./chunk.js";
 
-export { DetachedFieldIndex } from "./detachedFieldIndex.js";
+export {
+	DetachedFieldIndex,
+	type DetachedFieldIndexCheckpoint,
+	type ReadOnlyDetachedFieldIndex,
+} from "./detachedFieldIndex.js";
 
 export { getCodecTreeForDetachedFieldIndexFormat } from "./detachedFieldIndexCodecs.js";
 export { DetachedFieldIndexFormatVersion } from "./detachedFieldIndexFormatCommon.js";

packages/dds/tree/src/shared-tree/sharedTreeChangeEnricher.ts

@@ -9,6 +9,8 @@ import {
 	type DeltaDetachedNodeId,
 	type DetachedFieldIndex,
 	type IEditableForest,
+	type IForestSubscription,
+	type ReadOnlyDetachedFieldIndex,
 	type RevisionTag,
 	type TreeStoredSchemaRepository,
 	tagChange,
@@ -35,24 +37,25 @@ export class SharedTreeReadonlyChangeEnricher
 	implements ChangeEnricherReadonlyCheckout<SharedTreeChange>
 {
 	/**
-	 * @param forest - The state based on which to enrich changes.
-	 * Exclusively owned by the constructed instance.
+	 * @param borrowedForest - The state based on which to enrich changes.
+	 * Not owned by the constructed instance.
 	 * @param schema - The schema that corresponds to the forest.
-	 * @param removedRoots - The set of removed roots based on which to enrich changes.
-	 * Exclusively owned by the constructed instance.
+	 * @param borrowedRemovedRoots - The set of removed roots based on which to enrich changes.
+	 * Not owned by the constructed instance.
+	 * @param idCompressor - The id compressor to use when chunking trees.
 	 */
 	public constructor(
-		protected readonly forest: IEditableForest,
+		protected readonly borrowedForest: IForestSubscription,
 		private readonly schema: TreeStoredSchemaRepository,
-		protected readonly removedRoots: DetachedFieldIndex,
+		protected readonly borrowedRemovedRoots: ReadOnlyDetachedFieldIndex,
 		private readonly idCompressor?: IIdCompressor,
 	) {}
 
 	public fork(): ChangeEnricherMutableCheckout<SharedTreeChange> {
 		return new SharedTreeMutableChangeEnricher(
-			this.forest.clone(this.schema, new AnchorSet()),
+			this.borrowedForest.clone(this.schema, new AnchorSet()),
 			this.schema,
-			this.removedRoots.clone(),
+			this.borrowedRemovedRoots.clone(),
 		);
 	}
 
@@ -66,10 +69,10 @@ export class SharedTreeReadonlyChangeEnricher
 	}
 
 	private readonly getDetachedRoot = (id: DeltaDetachedNodeId): TreeChunk | undefined => {
-		const root = this.removedRoots.tryGetEntry(id);
+		const root = this.borrowedRemovedRoots.tryGetEntry(id);
 		if (root !== undefined) {
-			const cursor = this.forest.getCursorAboveDetachedFields();
-			const parentField = this.removedRoots.toFieldKey(root);
+			const cursor = this.borrowedForest.getCursorAboveDetachedFields();
+			const parentField = this.borrowedRemovedRoots.toFieldKey(root);
 			cursor.enterField(parentField);
 			cursor.enterNode(0);
 			return chunkTree(cursor, {
@@ -85,6 +88,23 @@ export class SharedTreeMutableChangeEnricher
 	extends SharedTreeReadonlyChangeEnricher
 	implements ChangeEnricherMutableCheckout<SharedTreeChange>
 {
+	/**
+	 * @param forest - The state based on which to enrich changes.
+	 * Owned by the constructed instance.
+	 * @param schema - The schema that corresponds to the forest.
+	 * @param removedRoots - The set of removed roots based on which to enrich changes.
+	 * Owned by the constructed instance.
+	 * @param idCompressor - The id compressor to use when chunking trees.
+	 */
+	public constructor(
+		private readonly forest: IEditableForest,
+		schema: TreeStoredSchemaRepository,
+		private readonly removedRoots: DetachedFieldIndex,
+		idCompressor?: IIdCompressor,
+	) {
+		super(forest, schema, removedRoots, idCompressor);
+	}
+
 	public applyTipChange(change: SharedTreeChange, revision?: RevisionTag): void {
 		for (const dataOrSchemaChange of change.changes) {
 			const type = dataOrSchemaChange.type;

packages/dds/tree/src/shared-tree/treeCheckout.ts

@@ -48,6 +48,7 @@ import {
 	type TreeNodeStoredSchema,
 	LeafNodeStoredSchema,
 	diffHistories,
+	type ReadOnlyDetachedFieldIndex,
 } from "../core/index.js";
 import {
 	type FieldBatchCodec,
@@ -406,7 +407,7 @@ export class TreeCheckout implements ITreeCheckoutFork {
 		private readonly mintRevisionTag: () => RevisionTag,
 		private readonly revisionTagCodec: RevisionTagCodec,
 		private readonly idCompressor: IIdCompressor,
-		public removedRoots: DetachedFieldIndex = makeDetachedFieldIndex(
+		private readonly _removedRoots: DetachedFieldIndex = makeDetachedFieldIndex(
 			"repair",
 			revisionTagCodec,
 			idCompressor,
@@ -421,6 +422,10 @@ export class TreeCheckout implements ITreeCheckoutFork {
 		this.registerForBranchEvents();
 	}
 
+	public get removedRoots(): ReadOnlyDetachedFieldIndex {
+		return this._removedRoots;
+	}
+
 	private registerForBranchEvents(): void {
 		this.#transaction.branch.events.on("afterChange", this.onAfterBranchChange);
 		this.#transaction.activeBranchEvents.on("afterChange", this.onAfterChange);
@@ -441,7 +446,7 @@ export class TreeCheckout implements ITreeCheckoutFork {
 			(commits) => {
 				const revision = this.mintRevisionTag();
 				for (const transactionStep of commits) {
-					this.removedRoots.updateMajor(transactionStep.revision, revision);
+					this._removedRoots.updateMajor(transactionStep.revision, revision);
 				}
 
 				const squashedChange = this.changeFamily.rebaser.compose(commits);
@@ -452,12 +457,12 @@ export class TreeCheckout implements ITreeCheckoutFork {
 				const disposeForks = this.disposeForksAfterTransaction
 					? trackForksForDisposal(this)
 					: undefined;
-				// When each transaction is started, take a snapshot of the current state of removed roots
-				const removedRootsSnapshot = this.removedRoots.clone();
+				// When each transaction is started, make a restorable checkpoint of the current state of removed roots
+				const restoreRemovedRoots = this._removedRoots.createCheckpoint();
 				return (result) => {
 					switch (result) {
 						case TransactionResult.Abort:
-							this.removedRoots = removedRootsSnapshot;
+							restoreRemovedRoots();
 							break;
 						case TransactionResult.Commit:
 							if (!this.transaction.isInProgress()) {
@@ -566,7 +571,7 @@ export class TreeCheckout implements ITreeCheckoutFork {
 			if (innerChange.type === "data") {
 				const delta = intoDelta(tagChange(innerChange.innerChange, revision));
 				this.withCombinedVisitor((visitor) => {
-					visitDelta(delta, visitor, this.removedRoots, revision);
+					visitDelta(delta, visitor, this._removedRoots, revision);
 				});
 			} else if (innerChange.type === "schema") {
 				// Schema changes from a current to a new schema are expected to be backwards compatible.
@@ -597,14 +602,14 @@ export class TreeCheckout implements ITreeCheckoutFork {
 		this.withCombinedVisitor((visitor) => {
 			revisions.forEach((revision) => {
 				// get all the roots last created or used by the revision
-				const roots = this.removedRoots.getRootsLastTouchedByRevision(revision);
+				const roots = this._removedRoots.getRootsLastTouchedByRevision(revision);
 
 				// get the detached field for the root and delete it from the removed roots
 				for (const root of roots) {
-					visitor.destroy(this.removedRoots.toFieldKey(root), 1);
+					visitor.destroy(this._removedRoots.toFieldKey(root), 1);
 				}
 
-				this.removedRoots.deleteRootsLastTouchedByRevision(revision);
+				this._removedRoots.deleteRootsLastTouchedByRevision(revision);
 			});
 		});
 	};
@@ -774,7 +779,7 @@ export class TreeCheckout implements ITreeCheckoutFork {
 			this.mintRevisionTag,
 			this.revisionTagCodec,
 			this.idCompressor,
-			this.removedRoots.clone(),
+			this._removedRoots.clone(),
 			this.logger,
 			this.breaker,
 			this.disposeForksAfterTransaction,
@@ -894,8 +899,8 @@ export class TreeCheckout implements ITreeCheckoutFork {
 		this.assertNoUntrackedRoots();
 		const trees: [string | number | undefined, number, JsonableTree][] = [];
 		const cursor = this.forest.allocateCursor("getRemovedRoots");
-		for (const { id, root } of this.removedRoots.entries()) {
-			const parentField = this.removedRoots.toFieldKey(root);
+		for (const { id, root } of this._removedRoots.entries()) {
+			const parentField = this._removedRoots.toFieldKey(root);
 			this.forest.moveCursorToPath({ parent: undefined, parentField, parentIndex: 0 }, cursor);
 			const tree = jsonableTreeFromCursor(cursor);
 			// This method is used for tree consistency comparison.
@@ -915,7 +920,7 @@ export class TreeCheckout implements ITreeCheckoutFork {
 	public load(): void {
 		// Set the tip revision as the latest relevant revision for any removed roots that are loaded from a summary - this allows them to be garbage collected later.
 		// When a load happens, the head of the trunk and the head of the local/main branch must be the same (this is enforced by SharedTree).
-		this.removedRoots.setRevisionsForLoadedData(this.#transaction.branch.getHead().revision);
+		this._removedRoots.setRevisionsForLoadedData(this.#transaction.branch.getHead().revision);
 		// The content of the checkout (e.g. the forest) has (maybe) changed, so fire an afterBatch event.
 		this.#events.emit("afterBatch");
 	}
@@ -987,8 +992,8 @@ export class TreeCheckout implements ITreeCheckoutFork {
 	private assertNoUntrackedRoots(): void {
 		const cursor = this.forest.getCursorAboveDetachedFields();
 		const rootFields = new Set([rootFieldKey]);
-		for (const { root } of this.removedRoots.entries()) {
-			rootFields.add(this.removedRoots.toFieldKey(root));
+		for (const { root } of this._removedRoots.entries()) {
+			rootFields.add(this._removedRoots.toFieldKey(root));
 		}
 
 		if (!cursor.firstField()) {