@fluidframework/container-runtime 0.59.4002 → 1.0.2

package/src/dataStoreContext.ts

@@ -58,6 +58,7 @@ import {
     ISummarizeResult,
     ISummarizerNodeWithGC,
     SummarizeInternalFn,
+    ITelemetryContext,
 } from "@fluidframework/runtime-definitions";
 import { addBlobToSummary, convertSummaryTreeToITree } from "@fluidframework/runtime-utils";
 import {
@@ -289,7 +290,8 @@ export abstract class FluidDataStoreContext extends TypedEventEmitter<IFluidData
         };
 
         const thisSummarizeInternal =
-            async (fullTree: boolean, trackState: boolean) => this.summarizeInternal(fullTree, trackState);
+            async (fullTree: boolean, trackState: boolean, telemetryContext?: ITelemetryContext) =>
+            this.summarizeInternal(fullTree, trackState, telemetryContext);
 
         this.summarizerNode = props.createSummarizerNodeFn(
             thisSummarizeInternal,
@@ -445,16 +447,25 @@ export abstract class FluidDataStoreContext extends TypedEventEmitter<IFluidData
      * Returns a summary at the current sequence number.
      * @param fullTree - true to bypass optimizations and force a full summary tree
      * @param trackState - This tells whether we should track state from this summary.
+     * @param telemetryContext - summary data passed through the layers for telemetry purposes
      */
-    public async summarize(fullTree: boolean = false, trackState: boolean = true): Promise<ISummarizeResult> {
-        return this.summarizerNode.summarize(fullTree, trackState);
+    public async summarize(
+        fullTree: boolean = false,
+        trackState: boolean = true,
+        telemetryContext?: ITelemetryContext,
+    ): Promise<ISummarizeResult> {
+        return this.summarizerNode.summarize(fullTree, trackState, telemetryContext);
     }
 
-    private async summarizeInternal(fullTree: boolean, trackState: boolean): Promise<ISummarizeInternalResult> {
+    private async summarizeInternal(
+        fullTree: boolean,
+        trackState: boolean,
+        telemetryContext?: ITelemetryContext,
+    ): Promise<ISummarizeInternalResult> {
         await this.realize();
 
         // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-        const summarizeResult = await this.channel!.summarize(fullTree, trackState);
+        const summarizeResult = await this.channel!.summarize(fullTree, trackState, telemetryContext);
         let pathPartsForChildren: string[] | undefined;
 
         if (!this.disableIsolatedChannels) {
@@ -717,6 +728,17 @@ export abstract class FluidDataStoreContext extends TypedEventEmitter<IFluidData
         this.channel.reSubmit(innerContents.type, innerContents.content, localOpMetadata);
     }
 
+    public rollback(contents: any, localOpMetadata: unknown) {
+        if (!this.channel) {
+            throw new Error("Channel must exist when rolling back ops");
+        }
+        if (!this.channel.rollback) {
+            throw new Error("Channel doesn't support rollback");
+        }
+        const innerContents = contents as FluidDataStoreMessage;
+        this.channel.rollback(innerContents.type, innerContents.content, localOpMetadata);
+    }
+
     public async applyStashedOp(contents: any): Promise<unknown> {
         if (!this.channel) {
             await this.realize();

package/src/dataStores.ts

@@ -25,6 +25,7 @@ import {
     InboundAttachMessage,
     ISummarizeResult,
     ISummaryTreeWithStats,
+    ITelemetryContext,
 } from "@fluidframework/runtime-definitions";
 import {
      convertSnapshotTreeToSummaryTree,
@@ -380,6 +381,13 @@ export class DataStores implements IDisposable {
         context.reSubmit(envelope.contents, localOpMetadata);
     }
 
+    public rollbackDataStoreOp(content: any, localOpMetadata: unknown) {
+        const envelope = content as IEnvelope;
+        const context = this.contexts.get(envelope.address);
+        assert(!!context, 0x2e8 /* "There should be a store context for the op" */);
+        context.rollback(envelope.contents, localOpMetadata);
+    }
+
     public async applyStashedOp(content: any): Promise<unknown> {
         const envelope = content as IEnvelope;
         const context = this.contexts.get(envelope.address);
@@ -473,7 +481,11 @@ export class DataStores implements IDisposable {
         return this.contexts.size;
     }
 
-    public async summarize(fullTree: boolean, trackState: boolean): Promise<ISummaryTreeWithStats> {
+    public async summarize(
+        fullTree: boolean,
+        trackState: boolean,
+        telemetryContext?: ITelemetryContext,
+    ): Promise<ISummaryTreeWithStats> {
         const summaryBuilder = new SummaryTreeBuilder();
 
         // Iterate over each store and ask it to snapshot
@@ -484,14 +496,14 @@ export class DataStores implements IDisposable {
                     0x165 /* "Summarizer cannot work if client has local changes" */);
                 return context.attachState === AttachState.Attached;
             }).map(async ([contextId, context]) => {
-                const contextSummary = await context.summarize(fullTree, trackState);
+                const contextSummary = await context.summarize(fullTree, trackState, telemetryContext);
                 summaryBuilder.addWithStats(contextId, contextSummary);
             }));
 
         return summaryBuilder.getSummaryTree();
     }
 
-    public createSummary(): ISummaryTreeWithStats {
+    public createSummary(telemetryContext?: ITelemetryContext): ISummaryTreeWithStats {
         const builder = new SummaryTreeBuilder();
         // Attaching graph of some stores can cause other stores to get bound too.
         // So keep taking summary until no new stores get bound.

package/src/garbageCollection.ts

@@ -22,8 +22,9 @@ import {
     IGarbageCollectionData,
     IGarbageCollectionState,
     IGarbageCollectionDetailsBase,
-    IGarbageCollectionNodeData,
     ISummarizeResult,
+    ITelemetryContext,
+    IGarbageCollectionNodeData,
 } from "@fluidframework/runtime-definitions";
 import {
     mergeStats,
@@ -162,7 +163,11 @@ export interface IGarbageCollector {
         options: { logger?: ITelemetryLogger; runGC?: boolean; runSweep?: boolean; fullGC?: boolean; },
     ): Promise<IGCStats>;
     /** Summarizes the GC data and returns it as a summary tree. */
-    summarize(fullTree: boolean, trackState: boolean): ISummarizeResult | undefined;
+    summarize(
+        fullTree: boolean,
+        trackState: boolean,
+        telemetryContext?: ITelemetryContext,
+    ): ISummarizeResult | undefined;
     /** Returns the garbage collector specific metadata to be written into the summary. */
     getMetadata(): IGCMetadata;
     /** Returns a map of each node id to its base GC details in the base summary. */
@@ -330,7 +335,7 @@ export class GarbageCollector implements IGarbageCollector {
     private _writeDataAtRoot: boolean = false;
     public get writeDataAtRoot(): boolean {
         return this._writeDataAtRoot;
-     }
+    }
 
     /**
      * Tells whether the initial GC state needs to be reset. This can happen under 2 conditions:
@@ -420,12 +425,14 @@ export class GarbageCollector implements IGarbageCollector {
         } else {
             // Sweep should not be enabled without enabling GC mark phase. We could silently disable sweep in this
             // scenario but explicitly failing makes it clearer and promotes correct usage.
-            if (gcOptions.sweepAllowed && !gcOptions.gcAllowed) {
+            if (gcOptions.sweepAllowed && gcOptions.gcAllowed === false) {
                 throw new UsageError("GC sweep phase cannot be enabled without enabling GC mark phase");
             }
 
-            // For new documents, GC has to be explicitly enabled via the flags in GC options.
-            this.gcEnabled = gcOptions.gcAllowed === true;
+            // For new documents, GC is enabled by default. It can be explicitly disabled by setting the gcAllowed
+            // flag in GC options to false.
+            this.gcEnabled = gcOptions.gcAllowed !== false;
+            // The sweep phase has to be explicitly enabled by setting the sweepAllowed flag in GC options to true.
             this.sweepEnabled = gcOptions.sweepAllowed === true;
 
             // Set the Session Expiry only if the flag is enabled or the test option is set.
@@ -588,7 +595,7 @@ export class GarbageCollector implements IGarbageCollector {
                 return;
             }
 
-            const gcNodes: { [ id: string ]: string[]; } = {};
+            const gcNodes: { [id: string]: string[]; } = {};
             for (const [nodeId, nodeData] of Object.entries(baseState.gcNodes)) {
                 if (nodeData.unreferencedTimestampMs !== undefined) {
                     this.unreferencedNodesState.set(
@@ -613,7 +620,7 @@ export class GarbageCollector implements IGarbageCollector {
                 return new Map();
             }
 
-            const gcNodes: { [ id: string ]: string[]; } = {};
+            const gcNodes: { [id: string]: string[]; } = {};
             for (const [nodeId, nodeData] of Object.entries(baseState.gcNodes)) {
                 gcNodes[nodeId] = Array.from(nodeData.outboundRoutes);
             }
@@ -737,8 +744,7 @@ export class GarbageCollector implements IGarbageCollector {
             this.completedRuns++;
 
             return gcStats;
-        },
-        { end: true, cancel: "error" });
+        }, { end: true, cancel: "error" });
     }
 
     /**
@@ -749,6 +755,7 @@ export class GarbageCollector implements IGarbageCollector {
     public summarize(
         fullTree: boolean,
         trackState: boolean,
+        telemetryContext?: ITelemetryContext,
     ): ISummarizeResult | undefined {
         if (!this.shouldRunGC || this.previousGCDataFromLastRun === undefined) {
             return;
@@ -1338,16 +1345,16 @@ function meetsMinimumVersionRequirement(currentVersion: string, minimumVersion:
  */
 export function semverCompare(currentVersion: string, minimumVersion: string): number {
     const minimumValues = minimumVersion.split(".").map((value): number => {
-        assert(isNaN(+value) === false, 0x2fa /* Expected real numbers in minimum version! */);
+        assert(isNaN(+value) === false, 0x2f8 /* Expected real numbers in minimum version! */);
         return Number.parseInt(value, 10);
     });
-    assert(minimumValues.length === 3, 0x2fb /* Expected minimumVersion to be [major].[minor].[patch] */);
+    assert(minimumValues.length === 3, 0x2f9 /* Expected minimumVersion to be [major].[minor].[patch] */);
     const [minMajor, minMinor, minPatch] = minimumValues;
 
     const currentValuesString = currentVersion.split(/\W/);
-    assert(currentValuesString.length >= 3, 0x2fc /* Expected version to match semver rules! */);
+    assert(currentValuesString.length >= 3, 0x2fa /* Expected version to match semver rules! */);
     const currentValues = currentValuesString.slice(0, 3).map((value) => {
-        assert(isNaN(+value) === false, 0x2fd /* Expected real numbers in minimum version! */);
+        assert(isNaN(+value) === false, 0x2fb /* Expected real numbers in minimum version! */);
         return Number.parseInt(value, 10);
     });
     const [cMajor, cMinor, cPatch] = currentValues;

package/src/index.ts

@@ -9,7 +9,12 @@ export {
     ContainerRuntimeMessage,
     IGCRuntimeOptions,
     ISummaryRuntimeOptions,
+    ISummaryBaseConfiguration,
+    ISummaryConfigurationHeuristics,
+    ISummaryConfigurationDisableSummarizer,
+    ISummaryConfigurationDisableHeuristics,
     IContainerRuntimeOptions,
+    IPendingRuntimeState,
     IRootSummaryTreeWithStats,
     isRuntimeMessage,
     RuntimeMessage,
@@ -18,6 +23,8 @@ export {
     agentSchedulerId,
     ContainerRuntime,
     RuntimeHeaders,
+    ISummaryConfiguration,
+    DefaultSummaryConfiguration,
 } from "./containerRuntime";
 export { DeltaScheduler } from "./deltaScheduler";
 export { FluidDataStoreRegistry } from "./dataStoreRegistry";
@@ -55,7 +62,6 @@ export {
     ISummarizer,
     ISummarizerEvents,
     ISummarizerInternalsProvider,
-    ISummarizerOptions,
     ISummarizerRuntime,
     ISummarizingWarning,
     ISummaryCancellationToken,

package/src/orderedClientElection.ts

@@ -2,7 +2,7 @@
  * Copyright (c) Microsoft Corporation and contributors. All rights reserved.
  * Licensed under the MIT License.
  */
-
+/* eslint-disable @rushstack/no-new-null */
 import { IEvent, IEventProvider, ITelemetryLogger } from "@fluidframework/common-definitions";
 import { assert, TypedEventEmitter } from "@fluidframework/common-utils";
 import { IDeltaManager } from "@fluidframework/container-definitions";

package/src/packageVersion.ts

@@ -6,4 +6,4 @@
  */
 
 export const pkgName = "@fluidframework/container-runtime";
-export const pkgVersion = "0.59.4002";
+export const pkgVersion = "1.0.2";

package/src/pendingStateManager.ts

@@ -5,13 +5,15 @@
 
 import { IDisposable } from "@fluidframework/common-definitions";
 import { assert, Lazy } from "@fluidframework/common-utils";
+import { ICriticalContainerError } from "@fluidframework/container-definitions";
 import { DataProcessingError } from "@fluidframework/container-utils";
 import {
     ISequencedDocumentMessage,
 } from "@fluidframework/protocol-definitions";
 import { FlushMode } from "@fluidframework/runtime-definitions";
+import { wrapError } from "@fluidframework/telemetry-utils";
 import Deque from "double-ended-queue";
-import { ContainerRuntime, ContainerMessageType, isRuntimeMessage } from "./containerRuntime";
+import { ContainerMessageType } from "./containerRuntime";
 
 /**
  * This represents a message that has been submitted and is added to the pending queue when `submit` is called on the
@@ -47,16 +49,31 @@ export interface IPendingFlush {
 export type IPendingState = IPendingMessage | IPendingFlushMode | IPendingFlush;
 
 export interface IPendingLocalState {
-    /**
-     * client ID we most recently connected with, or undefined if we never connected
-     */
-    clientId?: string;
     /**
      * list of pending states, including ops and batch information
      */
     pendingStates: IPendingState[];
 }
 
+export interface IRuntimeStateHandler{
+    connected(): boolean;
+    clientId(): string | undefined;
+    flushMode(): FlushMode;
+    setFlushMode(mode: FlushMode): void;
+    close(error?: ICriticalContainerError): void;
+    applyStashedOp: (type: ContainerMessageType, content: ISequencedDocumentMessage) => Promise<unknown>;
+    flush(): void;
+    reSubmit(
+        type: ContainerMessageType,
+        content: any,
+        localOpMetadata: unknown,
+        opMetadata: Record<string, unknown> | undefined): void;
+    rollback(
+        type: ContainerMessageType,
+        content: any,
+        localOpMetadata: unknown): void;
+}
+
 /**
  * PendingStateManager is responsible for maintaining the messages that have not been sent or have not yet been
  * acknowledged by the server. It also maintains the batch information for both automatically and manually flushed
@@ -69,15 +86,16 @@ export interface IPendingLocalState {
 export class PendingStateManager implements IDisposable {
     private readonly pendingStates = new Deque<IPendingState>();
     private readonly initialStates: Deque<IPendingState>;
-    private readonly previousClientIds = new Set<string>();
-    private readonly firstStashedCSN: number = -1;
     private readonly disposeOnce = new Lazy<void>(() => {
         this.initialStates.clear();
         this.pendingStates.clear();
     });
 
     // Maintains the count of messages that are currently unacked.
-    private pendingMessagesCount: number = 0;
+    private _pendingMessagesCount: number = 0;
+    public get pendingMessagesCount(): number {
+        return this._pendingMessagesCount;
+    }
 
     // Indicates whether we are processing a batch.
     private isProcessingBatch: boolean = false;
@@ -95,22 +113,18 @@ export class PendingStateManager implements IDisposable {
 
     private clientId: string | undefined;
 
-    private get connected(): boolean {
-        return this.containerRuntime.connected;
-    }
-
     /**
      * Called to check if there are any pending messages in the pending state queue.
      * @returns A boolean indicating whether there are messages or not.
      */
     public hasPendingMessages(): boolean {
-        return this.pendingMessagesCount !== 0;
+        return this._pendingMessagesCount !== 0 || !this.initialStates.isEmpty();
     }
 
     public getLocalState(): IPendingLocalState | undefined {
+        assert(this.initialStates.isEmpty(), 0x2e9 /* "Must call getLocalState() after applying initial states" */);
         if (this.hasPendingMessages()) {
             return {
-                clientId: this.clientId,
                 pendingStates: this.pendingStates.toArray().map(
                     // delete localOpMetadata since it may not be serializable
                     // and will be regenerated by applyStashedOp()
@@ -120,23 +134,12 @@ export class PendingStateManager implements IDisposable {
     }
 
     constructor(
-        private readonly containerRuntime: ContainerRuntime,
-        private readonly applyStashedOp: (type, content) => Promise<unknown>,
+        private readonly stateHandler: IRuntimeStateHandler,
         initialFlushMode: FlushMode,
         initialLocalState: IPendingLocalState | undefined,
     ) {
         this.initialStates = new Deque<IPendingState>(initialLocalState?.pendingStates ?? []);
 
-        if (initialLocalState) {
-            if (initialLocalState?.clientId) {
-                this.previousClientIds.add(initialLocalState.clientId);
-            }
-            // get stashed op count and client sequence number of first op
-            const messages = initialLocalState.pendingStates
-                .filter((state) => state.type === "message") as IPendingMessage[];
-            this.firstStashedCSN = messages[0].clientSequenceNumber;
-        }
-
         this.flushModeForNextMessage = initialFlushMode;
         this.onFlushModeUpdated(initialFlushMode);
     }
@@ -172,7 +175,7 @@ export class PendingStateManager implements IDisposable {
 
         this.pendingStates.push(pendingMessage);
 
-        this.pendingMessagesCount++;
+        this._pendingMessagesCount++;
     }
 
     /**
@@ -189,7 +192,7 @@ export class PendingStateManager implements IDisposable {
     public onFlush() {
         // If the FlushMode is Immediate, we don't need to track an explicit flush call because every message is
         // automatically flushed. So, flush is a no-op.
-        if (this.containerRuntime.flushMode === FlushMode.Immediate) {
+        if (this.stateHandler.flushMode() === FlushMode.Immediate) {
             return;
         }
 
@@ -205,21 +208,25 @@ export class PendingStateManager implements IDisposable {
 
     /**
      * Applies stashed ops at their reference sequence number so they are ready to be ACKed or resubmitted
+     * @param seqNum - Sequence number at which to apply ops. Will apply all ops if seqNum is undefined.
      */
-    public async applyStashedOpsAt(seqNum: number) {
+    public async applyStashedOpsAt(seqNum?: number) {
         // apply stashed ops at sequence number
         while (!this.initialStates.isEmpty()) {
             // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
             const nextState = this.initialStates.peekFront()!;
             if (nextState.type === "message") {
-                if (nextState.referenceSequenceNumber > seqNum) {
-                    break; // nothing left to do at this sequence number
-                } else if (nextState.referenceSequenceNumber > 0 && nextState.referenceSequenceNumber < seqNum) {
-                    throw new Error("loaded from snapshot too recent to apply stashed ops");
+                if (seqNum !== undefined) {
+                    if (nextState.referenceSequenceNumber > seqNum) {
+                        break; // nothing left to do at this sequence number
+                    } else if (nextState.referenceSequenceNumber < seqNum) {
+                        throw new Error("loaded from snapshot too recent to apply stashed ops");
+                    }
                 }
 
                 // applyStashedOp will cause the DDS to behave as if it has sent the op but not actually send it
-                const localOpMetadata = await this.applyStashedOp(nextState.messageType, nextState.content);
+                const localOpMetadata =
+                    await this.stateHandler.applyStashedOp(nextState.messageType, nextState.content);
                 nextState.localOpMetadata = localOpMetadata;
             }
 
@@ -229,88 +236,12 @@ export class PendingStateManager implements IDisposable {
         }
     }
 
-    /**
-     * Processes a local message once it's ack'd by the server to verify that there was no data corruption and that
-     * the batch information was preserved for batch messages. Also process remote messages that might have been
-     * sent from a previous container.
-     * @param message - The message that got ack'd and needs to be processed.
-     */
-    public processMessage(message: ISequencedDocumentMessage, local: boolean) {
-        // Do not process chunked ops until all pieces are available.
-        if (message.type === ContainerMessageType.ChunkedOp) {
-            return { localAck: false, localOpMetadata: undefined };
-        }
-
-        if (local) {
-            return { localAck: false, localOpMetadata: this.processPendingLocalMessage(message) };
-        } else {
-            return this.processRemoteMessage(message);
-        }
-    }
-
-    /**
-     * Listens for ACKs of stashed ops
-     */
-    private processRemoteMessage(message: ISequencedDocumentMessage) {
-        if (!isRuntimeMessage(message)) {
-            return { localAck: false, localOpMetadata: undefined };
-        }
-
-        // this message was a pending op that was actually sent successfully
-        const isOriginalClientId = message.clientId === Array.from(this.previousClientIds)[0] &&
-            message.clientSequenceNumber >= this.firstStashedCSN;
-        // this message is a pending or stashed op that was resubmitted
-        const isNewClientId = Array.from(this.previousClientIds).indexOf(message.clientId) > 0;
-
-        // if this is an ack for a stashed op, dequeue one message.
-        // we should have seen its ref seq num by now and the DDS should be ready for it to be ACKed
-        if (isOriginalClientId || isNewClientId) {
-            assert(this.clientId === undefined, 0x28b /* "multiple clients connected with stashed ops" */);
-            while (!this.pendingStates.isEmpty()) {
-                // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-                const nextState = this.pendingStates.shift()!;
-                // if it's not a message just drop it and keep looking
-                if (nextState.type === "message") {
-                    this.assertOpMatch(nextState, message, isOriginalClientId);
-                    return { localAck: true, localOpMetadata: nextState.localOpMetadata };
-                }
-            }
-        }
-
-        if (message.type === ContainerMessageType.Rejoin && this.previousClientIds.has(message.contents?.clientId)) {
-            this.previousClientIds.add(message.clientId);
-        }
-
-        return { localAck: false, localOpMetadata: undefined };
-    }
-
-    private assertOpMatch(state: IPendingMessage, message: ISequencedDocumentMessage, isOriginalClientId: boolean) {
-        assert(message.type === state.messageType, 0x28c /* "different message type" */);
-        assert(message.clientSequenceNumber === state.clientSequenceNumber || !isOriginalClientId,
-            0x28d /* "client sequence number doesn't match" */);
-        switch (message.type) {
-            case ContainerMessageType.Attach:
-                assert(message.contents.id === state.content.id, 0x28e /* "datastore ID doesn't match" */);
-                break;
-            case ContainerMessageType.FluidDataStoreOp:
-                assert(message.contents.address === state.content.address, 0x28f /* "address doesn't match" */);
-                break;
-            case ContainerMessageType.BlobAttach:
-                // todo: assert we have blob storage, assert blob IDs match, remove blob from blob storage since it made
-                // it through successfully
-                break;
-            case ContainerMessageType.Rejoin:
-            default:
-                throw new Error(`${message.type} not expected`);
-        }
-    }
-
     /**
      * Processes a local message once its ack'd by the server. It verifies that there was no data corruption and that
      * the batch information was preserved for batch messages.
      * @param message - The message that got ack'd and needs to be processed.
      */
-    private processPendingLocalMessage(message: ISequencedDocumentMessage): unknown {
+    public processPendingLocalMessage(message: ISequencedDocumentMessage): unknown {
         // Pre-processing part - This may be the start of a batch.
         this.maybeProcessBatchBegin(message);
 
@@ -330,11 +261,11 @@ export class PendingStateManager implements IDisposable {
                 { expectedClientSequenceNumber: pendingState.clientSequenceNumber },
             );
 
-            this.containerRuntime.closeFn(error);
+            this.stateHandler.close(error);
             return;
         }
 
-        this.pendingMessagesCount--;
+        this._pendingMessagesCount--;
 
         // Post-processing part - If we are processing a batch then this could be the last message in the batch.
         this.maybeProcessBatchEnd(message);
@@ -446,6 +377,31 @@ export class PendingStateManager implements IDisposable {
         this.isProcessingBatch = false;
     }
 
+    /**
+     * Capture the pending state at this point
+     */
+    public checkpoint() {
+        const checkpointHead = this.pendingStates.peekBack();
+        return {
+            rollback: () => {
+                try {
+                    while (this.pendingStates.peekBack() !== checkpointHead) {
+                        this.rollbackNextPendingState();
+                    }
+                } catch (err) {
+                    const error = wrapError(err, (message) => {
+                        return DataProcessingError.create(
+                            `RollbackError: ${message}`,
+                            "checkpointRollback",
+                            undefined) as DataProcessingError;
+                    });
+                    this.stateHandler.close(error);
+                    throw error;
+                }
+            },
+        };
+    }
+
     /**
      * Returns the next pending state from the pending state queue.
      */
@@ -455,17 +411,42 @@ export class PendingStateManager implements IDisposable {
         return nextPendingState;
     }
 
+    /**
+     * Undo the last pending state
+     */
+    private rollbackNextPendingState() {
+        const pendingStatesCount = this.pendingStates.length;
+        if (pendingStatesCount === 0) {
+            return;
+        }
+
+        this._pendingMessagesCount--;
+
+        // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+        const pendingState = this.pendingStates.pop()!;
+        switch (pendingState.type) {
+            case "message":
+                this.stateHandler.rollback(
+                    pendingState.messageType,
+                    pendingState.content,
+                    pendingState.localOpMetadata);
+                break;
+            default:
+                throw new Error(`Can't rollback state ${pendingState.type}`);
+        }
+    }
+
     /**
      * Called when the Container's connection state changes. If the Container gets connected, it replays all the pending
      * states in its queue. This includes setting the FlushMode and triggering resubmission of unacked ops.
      */
     public replayPendingStates() {
-        assert(this.connected, 0x172 /* "The connection state is not consistent with the runtime" */);
+        assert(this.stateHandler.connected(), 0x172 /* "The connection state is not consistent with the runtime" */);
 
         // This assert suggests we are about to send same ops twice, which will result in data loss.
-        assert(this.clientId !== this.containerRuntime.clientId,
+        assert(this.clientId !== this.stateHandler.clientId(),
             0x173 /* "replayPendingStates called twice for same clientId!" */);
-        this.clientId = this.containerRuntime.clientId;
+        this.clientId = this.stateHandler.clientId();
 
         assert(this.initialStates.isEmpty(), 0x174 /* "initial states should be empty before replaying pending" */);
 
@@ -475,14 +456,14 @@ export class PendingStateManager implements IDisposable {
         }
 
         // Reset the pending message count because all these messages will be removed from the queue.
-        this.pendingMessagesCount = 0;
+        this._pendingMessagesCount = 0;
 
         // Save the current FlushMode so that we can revert it back after replaying the states.
-        const savedFlushMode = this.containerRuntime.flushMode;
+        const savedFlushMode = this.stateHandler.flushMode();
 
         // Set the flush mode for the next message. This step is important because the flush mode may have been changed
         // after the next pending message was sent.
-        this.containerRuntime.setFlushMode(this.flushModeForNextMessage);
+        this.stateHandler.setFlushMode(this.flushModeForNextMessage);
 
         // Process exactly `pendingStatesCount` items in the queue as it represents the number of states that were
         // pending when we connected. This is important because the `reSubmitFn` might add more items in the queue
@@ -492,17 +473,17 @@ export class PendingStateManager implements IDisposable {
             const pendingState = this.pendingStates.shift()!;
             switch (pendingState.type) {
                 case "message":
-                    this.containerRuntime.reSubmitFn(
+                    this.stateHandler.reSubmit(
                         pendingState.messageType,
                         pendingState.content,
                         pendingState.localOpMetadata,
                         pendingState.opMetadata);
                     break;
                 case "flushMode":
-                    this.containerRuntime.setFlushMode(pendingState.flushMode);
+                    this.stateHandler.setFlushMode(pendingState.flushMode);
                     break;
                 case "flush":
-                    this.containerRuntime.flush();
+                    this.stateHandler.flush();
                     break;
                 default:
                     break;
@@ -511,6 +492,6 @@ export class PendingStateManager implements IDisposable {
         }
 
         // Revert the FlushMode.
-        this.containerRuntime.setFlushMode(savedFlushMode);
+        this.stateHandler.setFlushMode(savedFlushMode);
     }
 }