Tree: Local Rollback (#24954)

This change implements local rollback in shared tree and adds some basic
unit test for the scenario. This is modeled after the existing code for
aborting a transaction, which behaves similarly.

This is a precursor change to expanded fuzz coverage for rollback for
all dds via the dds fuzz and incorporating tree into local server
stress. There is a prototype of that here:
#24930

Author: anthony-murphy

packages/dds/tree/src/shared-tree-core/defaultResubmitMachine.ts

@@ -59,6 +59,19 @@ export class DefaultResubmitMachine<TChange> implements ResubmitMachine<TChange>
 		this.inFlightQueue.push(commit);
 	}
 
+	public onCommitRollback(commit: GraphCommit<TChange>): void {
+		assert(
+			this.inFlightQueue.length > 0 &&
+				commit.revision === this.inFlightQueue[this.inFlightQueue.length - 1]?.revision,
+			"must rollback latest commit in the in flight queue",
+		);
+
+		if (this.latestInFlightCommitWithStaleEnrichments === this.inFlightQueue.length - 1) {
+			this.latestInFlightCommitWithStaleEnrichments -= 1;
+		}
+		this.inFlightQueue.pop();
+	}
+
 	public prepareForResubmit(toResubmit: readonly GraphCommit<TChange>[]): void {
 		assert(
 			!this.isInResubmitPhase,

packages/dds/tree/src/shared-tree-core/resubmitMachine.ts

@@ -37,6 +37,13 @@ export interface ResubmitMachine<TChange> {
 	 */
 	onCommitSubmitted(commit: GraphCommit<TChange>): void;
 
+	/**
+	 * Must be called on a commit after rollback, so it can be removed
+	 * as it will never be (re)submitted.
+	 * @param commit - The commit that was rolled back
+	 */
+	onCommitRollback(commit: GraphCommit<TChange>): void;
+
 	/**
 	 * Must be called after a sequenced commit is applied.
 	 * Note that this may be called multiples times in a row after a number of sequenced commits have been applied

packages/dds/tree/src/shared-tree-core/sharedTreeCore.ts

@@ -418,6 +418,20 @@ export class SharedTreeCore<TEditor extends ChangeFamilyEditor, TChange>
 		const enrichedCommit = this.resubmitMachine.peekNextCommit();
 		this.submitCommit(enrichedCommit, localOpMetadata, true);
 	}
+	public rollback(content: JsonCompatibleReadOnly, localOpMetadata: unknown): void {
+		// Empty context object is passed in, as our decode function is schema-agnostic.
+		const {
+			commit: { revision },
+		} = this.messageCodec.decode(this.serializer.decode(content), {
+			idCompressor: this.idCompressor,
+		});
+		const [commit] = this.editManager.findLocalCommit(revision);
+		const { parent } = commit;
+		assert(parent !== undefined, "must have parent");
+		const [precedingCommit] = this.editManager.findLocalCommit(parent.revision);
+		this.editManager.localBranch.removeAfter(precedingCommit);
+		this.resubmitMachine.onCommitRollback(commit);
+	}
 
 	public applyStashedOp(content: JsonCompatibleReadOnly): void {
 		// Empty context object is passed in, as our decode function is schema-agnostic.

packages/dds/tree/src/test/shared-tree-core/sharedTreeCore.rollback.spec.ts

@@ -0,0 +1,198 @@
+/*!
+ * Copyright (c) Microsoft Corporation and contributors. All rights reserved.
+ * Licensed under the MIT License.
+ */
+
+import { strict as assert } from "node:assert";
+
+import { createIdCompressor } from "@fluidframework/id-compressor/internal";
+
+import {
+	MockContainerRuntimeFactory,
+	MockFluidDataStoreRuntime,
+	MockStorage,
+} from "@fluidframework/test-runtime-utils/internal";
+
+import { SharedTreeTestFactory } from "../utils.js";
+
+import { SchemaFactory, TreeViewConfiguration } from "../../simple-tree/index.js";
+import { FlushMode } from "@fluidframework/runtime-definitions/internal";
+
+const sf = new SchemaFactory("Test");
+class TestNode extends sf.objectRecursive("test node", {
+	child: sf.optionalRecursive([sf.number]),
+}) {}
+
+function setupTree() {
+	const containerRuntimeFactory = new MockContainerRuntimeFactory({
+		flushMode: FlushMode.TurnBased,
+	});
+	const dataStoreRuntime1 = new MockFluidDataStoreRuntime({
+		idCompressor: createIdCompressor(),
+	});
+
+	const factory = new SharedTreeTestFactory(() => {});
+
+	const containerRuntime = containerRuntimeFactory.createContainerRuntime(dataStoreRuntime1);
+	const tree1 = factory.create(dataStoreRuntime1, "A");
+	tree1.connect({
+		deltaConnection: dataStoreRuntime1.createDeltaConnection(),
+		objectStorage: new MockStorage(),
+	});
+
+	const view = tree1.viewWith(new TreeViewConfiguration({ schema: TestNode }));
+	view.initialize(new TestNode({}));
+	containerRuntime.flush();
+	return { view, containerRuntime, containerRuntimeFactory };
+}
+
+describe("SharedTreeCore rollback", () => {
+	it("should rollback a local insert operation", async () => {
+		const { view, containerRuntime, containerRuntimeFactory } = setupTree();
+		view.root = new TestNode({ child: 0 });
+		assert.deepEqual(view.root.child, 0, "after local insert");
+
+		// Rollback local change
+		containerRuntime.rollback?.();
+		assert.deepEqual(view.root.child, undefined, "after rollback of insert");
+
+		// Process messages to ensure no-op
+		containerRuntime.flush();
+		containerRuntimeFactory.processAllMessages();
+		assert.deepEqual(view.root.child, undefined, "after processAllMessages post-rollback");
+	});
+
+	it("should rollback a local update operation", async () => {
+		const { view, containerRuntime, containerRuntimeFactory } = setupTree();
+		view.root = new TestNode({ child: 1 });
+		containerRuntime.flush();
+		containerRuntimeFactory.processAllMessages();
+		assert.deepEqual(view.root.child, 1, "after initial insert");
+
+		// Local update
+		view.root.child = 2;
+		assert.deepEqual(view.root.child, 2, "after local update");
+
+		containerRuntime.rollback?.();
+		assert.deepEqual(view.root.child, 1, "after rollback of update");
+	});
+
+	it("should rollback a local delete operation", async () => {
+		const { view, containerRuntime, containerRuntimeFactory } = setupTree();
+		view.root = new TestNode({ child: 5 });
+		containerRuntime.flush();
+		containerRuntimeFactory.processAllMessages();
+		assert.deepEqual(view.root.child, 5, "after initial insert");
+
+		// Local delete
+		view.root.child = undefined;
+		assert.deepEqual(view.root.child, undefined, "after local delete");
+
+		containerRuntime.rollback?.();
+		assert.deepEqual(view.root.child, 5, "after rollback of delete");
+	});
+
+	it("should rollback multiple local operations in sequence", async () => {
+		const { view, containerRuntime, containerRuntimeFactory } = setupTree();
+		view.root = new TestNode({ child: 10 });
+		containerRuntime.flush();
+		containerRuntimeFactory.processAllMessages();
+		assert.deepEqual(view.root.child, 10, "after initial insert");
+
+		// Multiple local changes
+		view.root.child = 20;
+		view.root.child = undefined;
+		view.root.child = 30;
+		assert.deepEqual(view.root.child, 30, "after multiple local ops");
+
+		containerRuntime.rollback?.();
+		assert.deepEqual(view.root.child, 10, "after rollback of multiple ops");
+	});
+
+	it("should not rollback already flushed (acked) operations", async () => {
+		const { view, containerRuntime, containerRuntimeFactory } = setupTree();
+		view.root = new TestNode({ child: 100 });
+		containerRuntime.flush();
+		containerRuntimeFactory.processAllMessages();
+		assert.deepEqual(view.root.child, 100, "after flush and process");
+
+		// Local change and flush again
+		view.root.child = 200;
+		containerRuntime.flush();
+		containerRuntimeFactory.processAllMessages();
+		assert.deepEqual(view.root.child, 200, "after second flush");
+
+		// Rollback should not affect acked changes
+		containerRuntime.rollback?.();
+		assert.deepEqual(view.root.child, 200, "rollback after flush (no effect)");
+	});
+
+	it("should be a no-op if rollback is called with no pending changes", async () => {
+		const { view, containerRuntime, containerRuntimeFactory } = setupTree();
+		view.root = new TestNode({ child: 7 });
+		containerRuntime.flush();
+		containerRuntimeFactory.processAllMessages();
+		assert.deepEqual(view.root.child, 7, "after flush");
+
+		containerRuntime.rollback?.();
+		assert.deepEqual(view.root.child, 7, "rollback with no pending changes");
+	});
+
+	it("should rollback local changes in presence of remote changes from another client", async () => {
+		const containerRuntimeFactory = new MockContainerRuntimeFactory({
+			flushMode: FlushMode.TurnBased,
+		});
+		const factory = new SharedTreeTestFactory(() => {});
+
+		// Client 1
+		const dataStoreRuntime1 = new MockFluidDataStoreRuntime({
+			idCompressor: createIdCompressor(),
+			clientId: "1",
+		});
+		const containerRuntime1 =
+			containerRuntimeFactory.createContainerRuntime(dataStoreRuntime1);
+		const tree1 = factory.create(dataStoreRuntime1, "A");
+		tree1.connect({
+			deltaConnection: dataStoreRuntime1.createDeltaConnection(),
+			objectStorage: new MockStorage(),
+		});
+		const view1 = tree1.viewWith(new TreeViewConfiguration({ schema: TestNode }));
+		view1.initialize(new TestNode({}));
+		containerRuntime1.flush();
+
+		// Client 2
+		const dataStoreRuntime2 = new MockFluidDataStoreRuntime({
+			idCompressor: createIdCompressor(),
+			clientId: "2",
+		});
+		const containerRuntime2 =
+			containerRuntimeFactory.createContainerRuntime(dataStoreRuntime2);
+		const tree2 = factory.create(dataStoreRuntime2, "A");
+		tree2.connect({
+			deltaConnection: dataStoreRuntime2.createDeltaConnection(),
+			objectStorage: new MockStorage(),
+		});
+		const view2 = tree2.viewWith(new TreeViewConfiguration({ schema: TestNode }));
+		view2.initialize(new TestNode({}));
+		containerRuntime2.flush();
+
+		containerRuntimeFactory.processAllMessages();
+
+		// Client 1 makes a local change (not flushed)
+		view1.root.child = 1;
+		assert.deepEqual(view1.root.child, 1, "client 1 local change");
+
+		// Client 2 makes a local change and flushes
+		view2.root.child = 2;
+		containerRuntime2.flush();
+		containerRuntimeFactory.processAllMessages();
+
+		// Rollback local change in client 1
+		containerRuntime1.rollback?.();
+		containerRuntime1.flush();
+		containerRuntimeFactory.processAllMessages();
+
+		// Should reflect remote change from client 2
+		assert.deepEqual(view1.root.child, 2, "client 1 after rollback and remote change");
+	});
+});

packages/dds/tree/src/test/shared-tree-core/sharedTreeCore.spec.ts

@@ -410,6 +410,9 @@ describe("SharedTreeCore", () => {
 			public onSequencedCommitApplied(isLocal: boolean): void {
 				this.sequencingLog.push(isLocal);
 			}
+			public onCommitRollback(): void {
+				throw new Error("not implemented");
+			}
 		}
 
 		interface Enrichment<T extends object> {