Matrix: Enable rollback and local server stress

packages/dds/matrix/src/matrix.ts

@@ -213,6 +213,22 @@ type FirstWriterWinsPolicy =
 			cellLastWriteTracker: SparseArray2D<CellLastWriteTrackerItem>;
 	  };
 
+/**
+ * Tracks pending local changes for a cell.
+ */
+interface PendingCellChanges<T> {
+	/**
+	 * The local changes including the local seq, and the value set at that local seq.
+	 */
+	local: { localSeq: number; value: MatrixItem<T> }[];
+	/**
+	 * The latest consensus value across all clients.
+	 * this will either be a remote value or ack'd local
+	 * value.
+	 */
+	consensus?: MatrixItem<T>;
+}
+
 /**
  * A SharedMatrix holds a rectangular 2D array of values.  Supported operations
  * include setting values and inserting/removing rows and columns.
@@ -252,7 +268,7 @@ export class SharedMatrix<T = any>
 	private readonly cols: PermutationVector; // Map logical col to storage handle (if any)
 
 	private cells = new SparseArray2D<MatrixItem<T>>(); // Stores cell values.
-	private readonly pending = new SparseArray2D<number>(); // Tracks pending writes.
+	private readonly pending = new SparseArray2D<PendingCellChanges<T>>(); // Tracks pending writes.
 
 	private fwwPolicy: FirstWriterWinsPolicy = {
 		state: "off",
@@ -418,21 +434,22 @@ export class SharedMatrix<T = any>
 		value: MatrixItem<T>,
 		rowHandle = this.rows.getAllocatedHandle(row),
 		colHandle = this.cols.getAllocatedHandle(col),
+		rollback?: boolean,
 	): void {
 		this.protectAgainstReentrancy(() => {
-			if (this.undo !== undefined) {
-				let oldValue = this.cells.getCell(rowHandle, colHandle);
-				if (oldValue === null) {
-					oldValue = undefined;
-				}
+			const oldValue = this.cells.getCell(rowHandle, colHandle) ?? undefined;
 
+			if (this.undo !== undefined) {
 				this.undo.cellSet(rowHandle, colHandle, oldValue);
 			}
 
 			this.cells.setCell(rowHandle, colHandle, value);
 
-			if (this.isAttached()) {
-				this.sendSetCellOp(row, col, value, rowHandle, colHandle);
+			if (this.isAttached() && rollback !== true) {
+				const pending = this.sendSetCellOp(row, col, value, rowHandle, colHandle);
+				if (pending.local.length === 1) {
+					pending.consensus ??= oldValue;
+				}
 			}
 
 			// Avoid reentrancy by raising change notifications after the op is queued.
@@ -467,7 +484,7 @@ export class SharedMatrix<T = any>
 		rowHandle: Handle,
 		colHandle: Handle,
 		localSeq = this.nextLocalSeq(),
-	): void {
+	): PendingCellChanges<T> {
 		assert(
 			this.isAttached(),
 			0x1e2 /* "Caller must ensure 'isAttached()' before calling 'sendSetCellOp'." */,
@@ -493,7 +510,12 @@ export class SharedMatrix<T = any>
 		};
 
 		this.submitLocalMessage(op, metadata);
-		this.pending.setCell(rowHandle, colHandle, localSeq);
+		const pendingCell: PendingCellChanges<T> = this.pending.getCell(rowHandle, colHandle) ?? {
+			local: [],
+		};
+		pendingCell.local.push({ localSeq, value });
+		this.pending.setCell(rowHandle, colHandle, pendingCell);
+		return pendingCell;
 	}
 
 	/**
@@ -679,7 +701,16 @@ export class SharedMatrix<T = any>
 			| undefined
 			| number
 			| ReturnType<SparseArray2D<MatrixItem<T> | number>["snapshot"]>
-		)[] = [this.cells.snapshot(), this.pending.snapshot()];
+		)[] = [
+			this.cells.snapshot(),
+			/**
+			 * we used to write this.pending.snapshot(). this should have never been done, as pending is only for local
+			 * changes, and there should never be local changes in the summarizer. This was also never used on load
+			 * as there is no way to understand a previous clients pending changes. so we just set this to a constant
+			 * which matches an empty this.pending.snapshot() for back-compat in terms of the array length
+			 */
+			[undefined],
+		];
 
 		// Only need to store it in the snapshot if we have switched the policy already.
 		if (this.fwwPolicy.state === "on") {
@@ -806,22 +837,30 @@ export class SharedMatrix<T = any>
 			const col = this.rebasePosition(this.cols, colsRef, localSeq);
 			this.rows.removeLocalReferencePosition(rowsRef);
 			this.cols.removeLocalReferencePosition(colsRef);
-			if (row !== undefined && col !== undefined && row >= 0 && col >= 0) {
-				// If the mode is LWW, then send the op.
+
+			const pendingCell = this.pending.getCell(rowHandle, colHandle);
+			assert(pendingCell !== undefined, "local operation must have a pending array");
+			const { local } = pendingCell;
+			assert(local !== undefined, "local operation must have a pending array");
+			const localSeqIndex = local.findIndex((p) => p.localSeq === localSeq);
+			assert(localSeqIndex >= 0, "local operation must have a pending entry");
+			const [change] = local.splice(localSeqIndex, 1);
+			assert(change.localSeq === localSeq, "must match");
+
+			if (
+				row !== undefined &&
+				col !== undefined &&
+				row >= 0 &&
+				col >= 0 && // If the mode is LWW, then send the op.
 				// Otherwise if the current mode is FWW and if we generated this op, after seeing the
 				// last set op, or it is the first set op for the cell, then regenerate the op,
 				// otherwise raise conflict. We want to check the current mode here and not that
 				// whether op was made in FWW or not.
-				if (
-					this.fwwPolicy.state !== "on" ||
+				(this.fwwPolicy.state !== "on" ||
 					referenceSeqNumber >=
-						(this.fwwPolicy.cellLastWriteTracker.getCell(rowHandle, colHandle)?.seqNum ?? 0)
-				) {
-					this.sendSetCellOp(row, col, setOp.value, rowHandle, colHandle, localSeq);
-				} else if (this.pending.getCell(rowHandle, colHandle) !== undefined) {
-					// Clear the pending changes if any as we are not sending the op.
-					this.pending.setCell(rowHandle, colHandle, undefined);
-				}
+						(this.fwwPolicy.cellLastWriteTracker.getCell(rowHandle, colHandle)?.seqNum ?? 0))
+			) {
+				this.sendSetCellOp(row, col, setOp.value, rowHandle, colHandle, localSeq);
 			}
 		} else {
 			switch (content.target) {
@@ -840,6 +879,47 @@ export class SharedMatrix<T = any>
 		}
 	}
 
+	protected rollback(content: unknown, localOpMetadata: unknown): void {
+		const contents = content as MatrixSetOrVectorOp<T>;
+		const target = contents.target;
+
+		switch (target) {
+			case SnapshotPath.cols: {
+				this.cols.rollback(content, localOpMetadata);
+				break;
+			}
+			case SnapshotPath.rows: {
+				this.rows.rollback(content, localOpMetadata);
+				break;
+			}
+			case undefined: {
+				assert(contents.type === MatrixOp.set, "only sets supported");
+				const setMetadata = localOpMetadata as ISetOpMetadata;
+
+				const pendingCell = this.pending.getCell(setMetadata.rowHandle, setMetadata.colHandle);
+				assert(pendingCell !== undefined, "must have pending");
+
+				const change = pendingCell.local.pop();
+				assert(change?.localSeq === setMetadata.localSeq, "must have change");
+
+				const previous =
+					pendingCell.local.length > 0
+						? pendingCell.local[pendingCell.local.length - 1].value
+						: pendingCell.consensus;
+
+				this.setCellCore(
+					contents.row,
+					contents.col,
+					previous,
+					setMetadata.rowHandle,
+					setMetadata.colHandle,
+					true,
+				);
+			}
+			default:
+		}
+	}
+
 	protected onDisconnect(): void {}
 
 	/**
@@ -968,73 +1048,86 @@ export class SharedMatrix<T = any>
 					// We are receiving the ACK for a local pending set operation.
 					const { rowHandle, colHandle, localSeq, rowsRef, colsRef } =
 						localOpMetadata as ISetOpMetadata;
-					const isLatestPendingOp = this.isLatestPendingWrite(rowHandle, colHandle, localSeq);
 					this.rows.removeLocalReferencePosition(rowsRef);
 					this.cols.removeLocalReferencePosition(colsRef);
+
+					const pendingCell = this.pending.getCell(rowHandle, colHandle);
+					const ackedChange = pendingCell?.local.shift();
+					assert(ackedChange?.localSeq === localSeq, "must match");
+					if (pendingCell?.local.length === 0) {
+						this.pending.setCell(rowHandle, colHandle, undefined);
+					}
+
 					// If policy is switched and cell should be modified too based on policy, then update the tracker.
 					// If policy is not switched, then also update the tracker in case it is the latest.
 					if (
 						this.fwwPolicy.state === "on" &&
 						this.shouldSetCellBasedOnFWW(rowHandle, colHandle, msg)
 					) {
+						// eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+						pendingCell!.consensus = ackedChange.value;
 						this.fwwPolicy.cellLastWriteTracker.setCell(rowHandle, colHandle, {
 							seqNum: msg.sequenceNumber,
 							clientId: msg.clientId,
 						});
 					}
-
-					if (isLatestPendingOp) {
-						this.pending.setCell(rowHandle, colHandle, undefined);
-					}
 				} else {
 					const adjustedRow = this.rows.adjustPosition(row, msg);
-					if (adjustedRow !== undefined) {
-						const adjustedCol = this.cols.adjustPosition(col, msg);
-
-						if (adjustedCol !== undefined) {
-							const rowHandle = this.rows.getAllocatedHandle(adjustedRow);
-							const colHandle = this.cols.getAllocatedHandle(adjustedCol);
-
-							assert(
-								isHandleValid(rowHandle) && isHandleValid(colHandle),
-								0x022 /* "SharedMatrix row and/or col handles are invalid!" */,
-							);
-							if (this.fwwPolicy.state === "on") {
-								// If someone tried to Overwrite the cell value or first write on this cell or
-								// same client tried to modify the cell or if the previous mode was LWW, then we need to still
-								// overwrite the cell and raise conflict if we have pending changes as our change is going to be lost.
-								if (this.shouldSetCellBasedOnFWW(rowHandle, colHandle, msg)) {
-									const previousValue = this.cells.getCell(rowHandle, colHandle);
-									this.cells.setCell(rowHandle, colHandle, value);
-									this.fwwPolicy.cellLastWriteTracker.setCell(rowHandle, colHandle, {
-										seqNum: msg.sequenceNumber,
-										clientId: msg.clientId,
-									});
-									for (const consumer of this.consumers.values()) {
-										consumer.cellsChanged(adjustedRow, adjustedCol, 1, 1, this);
-									}
-									// Check is there are any pending changes, which will be rejected. If so raise conflict.
-									if (this.pending.getCell(rowHandle, colHandle) !== undefined) {
-										// Don't reset the pending value yet, as there maybe more fww op from same client, so we want
-										// to raise conflict event for that op also.
-										this.emit(
-											"conflict",
-											row,
-											col,
-											value, // Current value
-											previousValue, // Ignored local value
-											this,
-										);
-									}
+					const adjustedCol = this.cols.adjustPosition(col, msg);
+
+					const rowHandle = adjustedRow.handle;
+					const colHandle = adjustedCol.handle;
+
+					assert(
+						isHandleValid(rowHandle) && isHandleValid(colHandle),
+						0x022 /* "SharedMatrix row and/or col handles are invalid!" */,
+					);
+					const pendingCell = this.pending.getCell(rowHandle, colHandle);
+					if (this.fwwPolicy.state === "on") {
+						// If someone tried to Overwrite the cell value or first write on this cell or
+						// same client tried to modify the cell or if the previous mode was LWW, then we need to still
+						// overwrite the cell and raise conflict if we have pending changes as our change is going to be lost.
+						if (this.shouldSetCellBasedOnFWW(rowHandle, colHandle, msg)) {
+							const previousValue = this.cells.getCell(rowHandle, colHandle);
+							this.cells.setCell(rowHandle, colHandle, value);
+							this.fwwPolicy.cellLastWriteTracker.setCell(rowHandle, colHandle, {
+								seqNum: msg.sequenceNumber,
+								clientId: msg.clientId,
+							});
+							if (pendingCell !== undefined) {
+								pendingCell.consensus = value;
+							}
+							if (adjustedRow.pos !== undefined && adjustedCol.pos !== undefined) {
+								for (const consumer of this.consumers.values()) {
+									consumer.cellsChanged(adjustedRow.pos, adjustedCol.pos, 1, 1, this);
+								}
+								// Check is there are any pending changes, which will be rejected. If so raise conflict.
+								if (pendingCell !== undefined && pendingCell.local.length > 0) {
+									// Don't reset the pending value yet, as there maybe more fww op from same client, so we want
+									// to raise conflict event for that op also.
+									this.emit(
+										"conflict",
+										row,
+										col,
+										value, // Current value
+										previousValue, // Ignored local value
+										this,
+									);
 								}
-							} else if (this.pending.getCell(rowHandle, colHandle) === undefined) {
-								// If there is a pending (unACKed) local write to the same cell, skip the current op
-								// since it "happened before" the pending write.
-								this.cells.setCell(rowHandle, colHandle, value);
+							}
+						}
+					} else {
+						if (pendingCell === undefined || pendingCell.local.length === 0) {
+							// If there is a pending (unACKed) local write to the same cell, skip the current op
+							// since it "happened before" the pending write.
+							this.cells.setCell(rowHandle, colHandle, value);
+							if (adjustedRow.pos !== undefined && adjustedCol.pos !== undefined) {
 								for (const consumer of this.consumers.values()) {
-									consumer.cellsChanged(adjustedRow, adjustedCol, 1, 1, this);
+									consumer.cellsChanged(adjustedRow.pos, adjustedCol.pos, 1, 1, this);
 								}
 							}
+						} else {
+							pendingCell.consensus = value;
 						}
 					}
 				}
@@ -1106,35 +1199,6 @@ export class SharedMatrix<T = any>
 		}
 	}
 
-	/**
-	 * Returns true if the latest pending write to the cell indicated by the given row/col handles
-	 * matches the given 'localSeq'.
-	 *
-	 * A return value of `true` indicates that there are no later local operations queued that will
-	 * clobber the write op at the given 'localSeq'.  This includes later ops that overwrite the cell
-	 * with a different value as well as row/col removals that might recycled the given row/col handles.
-	 */
-	private isLatestPendingWrite(
-		rowHandle: Handle,
-		colHandle: Handle,
-		localSeq: number,
-	): boolean {
-		// eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-		const pendingLocalSeq = this.pending.getCell(rowHandle, colHandle)!;
-
-		// Note while we're awaiting the ACK for a local set, it's possible for the row/col to be
-		// locally removed and the row/col handles recycled.  If this happens, the pendingLocalSeq will
-		// be 'undefined' or > 'localSeq'.
-		assert(
-			!(pendingLocalSeq < localSeq),
-			0x023 /* "The 'localSeq' of pending write (if any) must be <= the localSeq of the currently processed op." */,
-		);
-
-		// If this is the most recent write to the cell by the local client, the stored localSeq
-		// will be an exact match for the given 'localSeq'.
-		return pendingLocalSeq === localSeq;
-	}
-
 	public toString(): string {
 		let s = `client:${
 			this.runtime.clientId