improvement(container-runtime): Pass all resubmit-related info in a single parameter when flushing a resubmitted batch (#24516)

Prerequisite refactor needed for #24510 

This lets all the flush code know whether this flush is for resubmit or
not by checking for the existence of that `BatchResubmitInfo` param.

Author: markfields

packages/runtime/container-runtime/src/containerRuntime.ts

@@ -210,6 +210,7 @@ import {
 	Outbox,
 	RemoteMessageProcessor,
 	type OutboundBatch,
+	type BatchResubmitInfo,
 } from "./opLifecycle/index.js";
 import { pkgVersion } from "./packageVersion.js";
 import {
@@ -3190,23 +3191,20 @@ export class ContainerRuntime
 	}
 
 	/**
-	 * Flush the pending ops manually.
-	 * This method is expected to be called at the end of a batch.
+	 * Flush the current batch of ops to the ordering service for sequencing
+	 * This method is not expected to be called in the middle of a batch.
 	 * @remarks - If it throws (e.g. if the batch is too large to send), the container will be closed.
 	 *
-	 * @param resubmittingBatchId - If defined, indicates this is a resubmission of a batch
-	 * with the given Batch ID, which must be preserved
-	 * @param resubmittingStagedBatch - If defined, indicates this is a resubmission of a batch that is staged,
-	 * meaning it should not be sent to the ordering service yet.
+	 * @param resubmitInfo - If defined, indicates this is a resubmission of a batch with the given Batch info needed for resubmit.
 	 */
-	private flush(resubmittingBatchId?: BatchId, resubmittingStagedBatch?: boolean): void {
+	private flush(resubmitInfo?: BatchResubmitInfo): void {
 		try {
 			assert(
 				!this.batchRunner.running,
 				0x24c /* "Cannot call `flush()` while manually accumulating a batch (e.g. under orderSequentially) */,
 			);
 
-			this.outbox.flush(resubmittingBatchId, resubmittingStagedBatch);
+			this.outbox.flush(resubmitInfo);
 			assert(this.outbox.isEmpty, 0x3cf /* reentrancy */);
 		} catch (error) {
 			const error2 = normalizeError(error, {
@@ -3320,7 +3318,7 @@ export class ContainerRuntime
 		this.outbox.flush();
 		const exitStagingMode = (discardOrCommit: () => void) => (): void => {
 			// Final flush of any last staged changes
-			this.outbox.flush(undefined, true /* staged */);
+			this.outbox.flush();
 
 			this.stageControls = undefined;
 
@@ -4558,7 +4556,7 @@ export class ContainerRuntime
 
 		// Only include Batch ID if "Offline Load" feature is enabled
 		// It's only needed to identify batches across container forks arising from misuse of offline load.
-		this.flush(this.offlineEnabled ? batchId : undefined, staged);
+		this.flush({ batchId: this.offlineEnabled ? batchId : undefined, staged });
 	}
 
 	private reSubmit(message: PendingMessageResubmitData): void {

packages/runtime/container-runtime/src/opLifecycle/index.ts

@@ -27,6 +27,7 @@ export {
 	ensureContentsDeserialized,
 } from "./opSerialization.js";
 export {
+	BatchResubmitInfo,
 	estimateSocketSize,
 	localBatchToOutboundBatch,
 	Outbox,

packages/runtime/container-runtime/src/opLifecycle/outbox.ts

@@ -70,6 +70,23 @@ export interface IOutboxParameters {
 	readonly opReentrancy: () => boolean;
 }
 
+/**
+ * Info needed to correctly resubmit a batch
+ */
+export interface BatchResubmitInfo {
+	/**
+	 * If defined, indicates the Batch ID of the batch being resubmitted.
+	 * This must be preserved on the new batch about to be submitted so they can be correlated/deduped in case both are sent.
+	 */
+	batchId?: string;
+	/**
+	 * Indicates whether or not this batch is "staged", meaning it should not be sent to the ordering service yet
+	 * This is important on resubmit because we may be in Staging Mode for new changes,
+	 * but resubmitting a non-staged change from before entering Staging Mode
+	 */
+	staged: boolean;
+}
+
 /**
  * Temporarily increase the stack limit while executing the provided action.
  * If a negative value is provided for `length`, no stack frames will be collected.
@@ -335,37 +352,33 @@ export class Outbox {
 	 * This method is expected to be called at the end of a batch.
 	 *
 	 * @throws If called from a reentrant context, or if the batch being flushed is too large.
-	 * @param resubmittingBatchId - If defined, indicates this is a resubmission of a batch
-	 * with the given Batch ID, which must be preserved
-	 * @param resubmittingStagedBatch - If defined, indicates this is a resubmission of a batch that is staged,
-	 * meaning it should not be sent to the ordering service yet.
+	 * @param resubmitInfo - Key information when flushing a resubmitted batch. Undefined means this is not resubmit.
 	 */
-	public flush(resubmittingBatchId?: BatchId, resubmittingStagedBatch?: boolean): void {
+	public flush(resubmitInfo?: BatchResubmitInfo): void {
 		assert(
 			!this.isContextReentrant(),
 			0xb7b /* Flushing must not happen while incoming changes are being processed */,
 		);
-
-		this.flushAll(resubmittingBatchId, resubmittingStagedBatch);
+		this.flushAll(resubmitInfo);
 	}
 
-	private flushAll(resubmittingBatchId?: BatchId, resubmittingStagedBatch?: boolean): void {
+	private flushAll(resubmitInfo?: BatchResubmitInfo): void {
 		const allBatchesEmpty =
 			this.idAllocationBatch.empty && this.blobAttachBatch.empty && this.mainBatch.empty;
 		if (allBatchesEmpty) {
-			// If we're resubmitting and all batches are empty, we need to flush an empty batch.
-			// Note that we currently resubmit one batch at a time, so on resubmit, 2 of the 3 batches will *always* be empty.
+			// If we're resubmitting with a batchId and all batches are empty, we need to flush an empty batch.
+			// Note that we currently resubmit one batch at a time, so on resubmit, 1 of the 2 batches will *always* be empty.
 			// It's theoretically possible that we don't *need* to resubmit this empty batch, and in those cases, it'll safely be ignored
 			// by the rest of the system, including remote clients.
 			// In some cases we *must* resubmit the empty batch (to match up with a non-empty version tracked locally by a container fork), so we do it always.
-			if (resubmittingBatchId) {
-				this.flushEmptyBatch(resubmittingBatchId, resubmittingStagedBatch === true);
+			if (resubmitInfo?.batchId !== undefined) {
+				this.flushEmptyBatch(resubmitInfo.batchId, resubmitInfo.staged);
 			}
 			return;
 		}
 
 		// Don't use resubmittingBatchId for idAllocationBatch.
-		// ID Allocation messages are not directly resubmitted so we don't want to reuse the batch ID.
+		// ID Allocation messages are not directly resubmitted so don't pass the resubmitInfo
 		this.flushInternal({
 			batchManager: this.idAllocationBatch,
 			// Note: For now, we will never stage ID Allocation messages.
@@ -374,13 +387,11 @@ export class Outbox {
 		this.flushInternal({
 			batchManager: this.blobAttachBatch,
 			disableGroupedBatching: true,
-			resubmittingBatchId,
-			resubmittingStagedBatch,
+			resubmitInfo,
 		});
 		this.flushInternal({
 			batchManager: this.mainBatch,
-			resubmittingBatchId,
-			resubmittingStagedBatch,
+			resubmitInfo,
 		});
 	}
 
@@ -416,25 +427,19 @@ export class Outbox {
 	private flushInternal(params: {
 		batchManager: BatchManager;
 		disableGroupedBatching?: boolean;
-		resubmittingBatchId?: BatchId; // undefined if not resubmitting
-		resubmittingStagedBatch?: boolean; // undefined if not resubmitting
+		resubmitInfo?: BatchResubmitInfo; // undefined if not resubmitting
 	}): void {
-		const {
-			batchManager,
-			disableGroupedBatching = false,
-			resubmittingBatchId,
-			resubmittingStagedBatch,
-		} = params;
+		const { batchManager, disableGroupedBatching = false, resubmitInfo } = params;
 		if (batchManager.empty) {
 			return;
 		}
 
-		const rawBatch = batchManager.popBatch(resubmittingBatchId);
+		const rawBatch = batchManager.popBatch(resubmitInfo?.batchId);
 
 		// When resubmitting, we respect the staged state of the original batch.
 		// In this case rawBatch.staged will match the state of inStagingMode when
 		// the resubmit occurred, which is not relevant.
-		const staged = resubmittingStagedBatch ?? rawBatch.staged === true;
+		const staged = resubmitInfo?.staged ?? rawBatch.staged === true;
 
 		const groupingEnabled =
 			!disableGroupedBatching && this.params.groupingManager.groupedBatchingEnabled();

packages/runtime/container-runtime/src/test/opLifecycle/outbox.spec.ts

@@ -405,7 +405,7 @@ describe("Outbox", () => {
 		assert.equal(state.pendingOpContents.length, 0);
 		const batchId = "batchId";
 		// ...But if batchId is provided, it's resubmit, and we need to send an empty batch with the batchId
-		outbox.flush(batchId);
+		outbox.flush({ batchId, staged: false });
 		assert.equal(state.opsSubmitted, 1);
 		assert.equal(state.batchesSubmitted.length, 1);
 		assert.equal(
@@ -430,17 +430,17 @@ describe("Outbox", () => {
 		outbox.submitIdAllocation(createMessage(ContainerMessageType.IdAllocation, "0")); // Separate batch, batch ID not used
 		outbox.submit(createMessage(ContainerMessageType.FluidDataStoreOp, "1"));
 		outbox.submit(createMessage(ContainerMessageType.FluidDataStoreOp, "2"));
-		outbox.flush("batchId-A");
+		outbox.flush({ batchId: "batchId-A", staged: false });
 
 		// Flush 2 - resubmit single-message batch
 		outbox.submit(createMessage(ContainerMessageType.FluidDataStoreOp, "3"));
-		outbox.flush("batchId-B");
+		outbox.flush({ batchId: "batchId-B", staged: false });
 
 		// Flush 3 - resubmit blob attach batch
 		outbox.submitBlobAttach(createMessage(ContainerMessageType.BlobAttach, "4"));
 		outbox.submitBlobAttach(createMessage(ContainerMessageType.BlobAttach, "5"));
 		currentSeqNumbers.referenceSequenceNumber = 0;
-		outbox.flush("batchId-C");
+		outbox.flush({ batchId: "batchId-C", staged: false });
 
 		// Flush 4 - no batch ID given
 		outbox.submit(createMessage(ContainerMessageType.FluidDataStoreOp, "6"));