@fluidframework/container-runtime 2.0.0-internal.1.4.4 → 2.0.0-internal.2.0.1

package/src/dataStores.ts

@@ -97,7 +97,6 @@ export class DataStores implements IDisposable {
         private readonly gcNodeUpdated: (
             nodePath: string, timestampMs: number, packagePath?: readonly string[]) => void,
         private readonly aliasMap: Map<string, string>,
-        private readonly writeGCDataAtRoot: boolean,
         private readonly contexts: DataStoreContexts = new DataStoreContexts(baseLogger),
     ) {
         this.logger = ChildLogger.create(baseLogger);
@@ -142,7 +141,6 @@ export class DataStores implements IDisposable {
                         key,
                         { type: CreateSummarizerNodeSource.FromSummary },
                     ),
-                    writeGCDataAtRoot: this.writeGCDataAtRoot,
                 });
             } else {
                 if (typeof value !== "object") {
@@ -162,7 +160,6 @@ export class DataStores implements IDisposable {
                     makeLocallyVisibleFn: () => this.makeDataStoreLocallyVisible(key),
                     snapshotTree,
                     isRootDataStore: undefined,
-                    writeGCDataAtRoot: this.writeGCDataAtRoot,
                 });
             }
             this.contexts.addBoundOrRemoted(dataStoreContext);
@@ -247,7 +244,6 @@ export class DataStores implements IDisposable {
                     },
                 },
             ),
-            writeGCDataAtRoot: this.writeGCDataAtRoot,
             pkg,
         });
 
@@ -351,7 +347,6 @@ export class DataStores implements IDisposable {
             makeLocallyVisibleFn: () => this.makeDataStoreLocallyVisible(id),
             snapshotTree: undefined,
             isRootDataStore: isRoot,
-            writeGCDataAtRoot: this.writeGCDataAtRoot,
         });
         this.contexts.addUnbound(context);
         return context;
@@ -372,7 +367,6 @@ export class DataStores implements IDisposable {
             makeLocallyVisibleFn: () => this.makeDataStoreLocallyVisible(id),
             snapshotTree: undefined,
             isRootDataStore: false,
-            writeGCDataAtRoot: this.writeGCDataAtRoot,
             createProps: props,
         });
         this.contexts.addUnbound(context);
@@ -595,10 +589,8 @@ export class DataStores implements IDisposable {
     /**
      * After GC has run, called to notify this Container's data stores of routes that are used in it.
      * @param usedRoutes - The routes that are used in all data stores in this Container.
-     * @param gcTimestamp - The time when GC was run that generated these used routes. If any node node becomes
-     * unreferenced as part of this GC run, this should be used to update the time when it happens.
      */
-    public updateUsedRoutes(usedRoutes: string[], gcTimestamp?: number) {
+    public updateUsedRoutes(usedRoutes: string[]) {
         // Get a map of data store ids to routes used in it.
         const usedDataStoreRoutes = unpackChildNodesUsedRoutes(usedRoutes);
 
@@ -609,7 +601,7 @@ export class DataStores implements IDisposable {
 
         // Update the used routes in each data store. Used routes is empty for unused data stores.
         for (const [contextId, context] of this.contexts) {
-            context.updateUsedRoutes(usedDataStoreRoutes.get(contextId) ?? [], gcTimestamp);
+            context.updateUsedRoutes(usedDataStoreRoutes.get(contextId) ?? []);
         }
     }
 

package/src/garbageCollection.ts

@@ -67,12 +67,8 @@ export const runGCKey = "Fluid.GarbageCollection.RunGC";
 export const runSweepKey = "Fluid.GarbageCollection.RunSweep";
 // Feature gate key to turn GC test mode on / off.
 export const gcTestModeKey = "Fluid.GarbageCollection.GCTestMode";
-// Feature gate key to write GC data at the root of the summary tree.
-const writeAtRootKey = "Fluid.GarbageCollection.WriteDataAtRoot";
 // Feature gate key to expire a session after a set period of time.
 export const runSessionExpiryKey = "Fluid.GarbageCollection.RunSessionExpiry";
-// Feature gate key to disable expiring session after a set period of time, even if expiry value is present.
-export const disableSessionExpiryKey = "Fluid.GarbageCollection.DisableSessionExpiry";
 // Feature gate key to write the gc blob as a handle if the data is the same.
 export const trackGCStateKey = "Fluid.GarbageCollection.TrackGCState";
 // Feature gate key to turn GC sweep log off.
@@ -126,7 +122,7 @@ export interface IGarbageCollectionRuntime {
     /** Returns the garbage collection data of the runtime. */
     getGCData(fullGC?: boolean): Promise<IGarbageCollectionData>;
     /** After GC has run, called to notify the runtime of routes that are used in it. */
-    updateUsedRoutes(usedRoutes: string[], gcTimestamp?: number): void;
+    updateUsedRoutes(usedRoutes: string[]): void;
     /** After GC has run, called to delete objects in the runtime whose routes are unused. */
     deleteUnusedRoutes(unusedRoutes: string[]): void;
     /** Returns a referenced timestamp to be used to track unreferenced nodes. */
@@ -143,8 +139,6 @@ export interface IGarbageCollector {
     readonly shouldRunGC: boolean;
     /** Tells whether the GC state in summary needs to be reset in the next summary. */
     readonly summaryStateNeedsReset: boolean;
-    /** Tells whether GC data should be written to the root of the summary tree. */
-    readonly writeDataAtRoot: boolean;
     readonly trackGCState: boolean;
     /** Run garbage collection and update the reference / used state of the system. */
     collectGarbage(
@@ -190,7 +184,6 @@ export interface IGarbageCollectorCreateParams {
     readonly readAndParseBlob: ReadAndParseBlob;
     readonly activeConnection: () => boolean;
     readonly getContainerDiagnosticId: () => string;
-    readonly snapshotCacheExpiryMs?: number;
 }
 
 /** The state of node that is unreferenced. */
@@ -378,14 +371,6 @@ export class GarbageCollector implements IGarbageCollector {
     private readonly testMode: boolean;
     private readonly mc: MonitoringContext;
 
-    /**
-     * Tells whether the GC data should be written to the root of the summary tree.
-     */
-    private _writeDataAtRoot: boolean = true;
-    public get writeDataAtRoot(): boolean {
-        return this._writeDataAtRoot;
-    }
-
     /**
      * Tells whether the initial GC state needs to be reset. This can happen under 2 conditions:
      *
@@ -462,9 +447,9 @@ export class GarbageCollector implements IGarbageCollector {
             sweepEnabled: this.sweepEnabled,
             runGC: this.shouldRunGC,
             runSweep: this.shouldRunSweep,
-            writeAtRoot: this._writeDataAtRoot,
             testMode: this.testMode,
             sessionExpiry: this.sessionExpiryTimeoutMs,
+            sweepTimeout: this.sweepTimeoutMs,
             inactiveTimeout: this.inactiveTimeoutMs,
             trackGCState: this.trackGCState,
             ...this.gcOptions,
@@ -498,6 +483,21 @@ export class GarbageCollector implements IGarbageCollector {
 
         let prevSummaryGCVersion: number | undefined;
 
+        /**
+         * Sweep timeout is the time after which unreferenced content can be swept.
+         * Sweep timeout = session expiry timeout + snapshot cache expiry timeout + one day buffer.
+         *
+         * The snapshot cache expiry timeout cannot be known precisely but the upper bound is 5 days.
+         * The buffer is added to account for any clock skew or other edge cases.
+         * We use server timestamps throughout so the skew should be minimal but make it 1 day to be safe.
+         */
+        function computeSweepTimeout(sessionExpiryTimeoutMs: number | undefined) {
+            const maxSnapshotCacheExpiryMs = 5 * oneDayMs;
+            const bufferMs = oneDayMs;
+            return sessionExpiryTimeoutMs &&
+                (sessionExpiryTimeoutMs + maxSnapshotCacheExpiryMs + bufferMs);
+        }
+
         /**
          * The following GC state is enabled during container creation and cannot be changed throughout its lifetime:
          * 1. Whether running GC mark phase is allowed or not.
@@ -512,6 +512,9 @@ export class GarbageCollector implements IGarbageCollector {
             this.gcEnabled = prevSummaryGCVersion > 0;
             this.sweepEnabled = metadata?.sweepEnabled ?? false;
             this.sessionExpiryTimeoutMs = metadata?.sessionExpiryTimeoutMs;
+            this.sweepTimeoutMs =
+                metadata?.sweepTimeoutMs
+                ?? computeSweepTimeout(this.sessionExpiryTimeoutMs); // Backfill old documents that didn't persist this
         } 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.
@@ -519,20 +522,28 @@ export class GarbageCollector implements IGarbageCollector {
                 throw new UsageError("GC sweep phase cannot be enabled without enabling GC mark phase");
             }
 
+            // This Test Override only applies for new containers
+            const testOverrideSweepTimeoutMs =
+                this.mc.config.getNumber("Fluid.GarbageCollection.TestOverride.SweepTimeoutMs");
+
             // 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 = this.gcOptions.gcAllowed !== false;
             // The sweep phase has to be explicitly enabled by setting the sweepAllowed flag in GC options to true.
-            this.sweepEnabled = this.gcOptions.sweepAllowed === true;
+            // ...unless we're using the TestOverride
+            this.sweepEnabled = this.gcOptions.sweepAllowed === true || testOverrideSweepTimeoutMs !== undefined;
 
             // Set the Session Expiry only if the flag is enabled and GC is enabled.
             if (this.mc.config.getBoolean(runSessionExpiryKey) && this.gcEnabled) {
                 this.sessionExpiryTimeoutMs = this.gcOptions.sessionExpiryTimeoutMs ?? defaultSessionExpiryDurationMs;
             }
+            this.sweepTimeoutMs =
+                testOverrideSweepTimeoutMs
+                ?? computeSweepTimeout(this.sessionExpiryTimeoutMs);
         }
 
         // If session expiry is enabled, we need to close the container when the session expiry timeout expires.
-        if (this.sessionExpiryTimeoutMs !== undefined && this.mc.config.getBoolean(disableSessionExpiryKey) !== true) {
+        if (this.sessionExpiryTimeoutMs !== undefined) {
             // If Test Override config is set, override Session Expiry timeout.
             const overrideSessionExpiryTimeoutMs =
                 this.mc.config.getNumber("Fluid.GarbageCollection.TestOverride.SessionExpiryMs");
@@ -543,21 +554,6 @@ export class GarbageCollector implements IGarbageCollector {
                 () => { this.runtime.closeFn(new ClientSessionExpiredError(`Client session expired.`, timeoutMs)); },
             );
             this.sessionExpiryTimer.start();
-
-            // TEMPORARY: Hardcode a default of 5 days which will be >= the policy's value in the ODSP driver.
-            // This unblocks the Sweep Log (see logSweepEvents function).
-            // This will be removed before sweep is fully implemented.
-            const snapshotCacheExpiryMs = createParams.snapshotCacheExpiryMs ?? 5 * 24 * 60 * 60 * 1000;
-
-            /**
-             * Sweep timeout is the time after which unreferenced content can be swept.
-             * Sweep timeout = session expiry timeout + snapshot cache expiry timeout + one day buffer. The buffer is
-             * added to account for any clock skew. We use server timestamps throughout so the skew should be minimal
-             * but make it one day to be safe.
-             */
-            if (snapshotCacheExpiryMs !== undefined) {
-                this.sweepTimeoutMs = this.sessionExpiryTimeoutMs + snapshotCacheExpiryMs + oneDayMs;
-            }
         }
 
         // For existing document, the latest summary is the one that we loaded from. So, use its GC version as the
@@ -584,16 +580,16 @@ export class GarbageCollector implements IGarbageCollector {
          * Whether sweep should run or not. The following conditions have to be met to run sweep:
          *
          * 1. Overall GC or mark phase must be enabled (this.shouldRunGC).
-         *
          * 2. Sweep timeout should be available. Without this, we wouldn't know when an object should be deleted.
-         *
-         * 3. Sweep should be enabled for this container (this.sweepEnabled). This can be overridden via runSweep
+         * 3. The driver must implement the policy limiting the age of snapshots used for loading. Otherwise
+         * the Sweep Timeout calculation is not valid. We use the persisted value to ensure consistency over time.
+         * 4. Sweep should be enabled for this container (this.sweepEnabled). This can be overridden via runSweep
          * feature flag.
          */
-        this.shouldRunSweep = false; // disable while TEMPORARY measure hardcoding snapshotCacheExpiryMs is here
-            // this.shouldRunGC
-            // && this.sweepTimeoutMs !== undefined
-            // && (this.mc.config.getBoolean(runSweepKey) ?? this.sweepEnabled);
+        this.shouldRunSweep =
+            this.shouldRunGC
+            && this.sweepTimeoutMs !== undefined
+            && (this.mc.config.getBoolean(runSweepKey) ?? this.sweepEnabled);
 
         this.trackGCState = this.mc.config.getBoolean(trackGCStateKey) === true;
 
@@ -610,15 +606,10 @@ export class GarbageCollector implements IGarbageCollector {
         // Whether we are running in test mode. In this mode, unreferenced nodes are immediately deleted.
         this.testMode = this.mc.config.getBoolean(gcTestModeKey) ?? this.gcOptions.runGCInTestMode === true;
 
-        // GC state is written into root of the summary tree by default. Can be overridden via feature flag for now.
-        this._writeDataAtRoot = this.mc.config.getBoolean(writeAtRootKey) ?? true;
-
-        if (this._writeDataAtRoot) {
-            // The GC state needs to be reset if the base snapshot contains GC tree and GC is disabled or it doesn't
-            // contain GC tree and GC is enabled.
-            const gcTreePresent = baseSnapshot?.trees[gcTreeKey] !== undefined;
-            this.initialStateNeedsReset = gcTreePresent !== this.shouldRunGC;
-        }
+        // The GC state needs to be reset if the base snapshot contains GC tree and GC is disabled or it doesn't
+        // contain GC tree and GC is enabled.
+        const gcTreePresent = baseSnapshot?.trees[gcTreeKey] !== undefined;
+        this.initialStateNeedsReset = gcTreePresent !== this.shouldRunGC;
 
         // Get the GC state from the GC blob in the base snapshot. Use LazyPromise because we only want to do
         // this once since it involves fetching blobs from storage which is expensive.
@@ -631,8 +622,6 @@ export class GarbageCollector implements IGarbageCollector {
                 // For newer documents, GC data should be present in the GC tree in the root of the snapshot.
                 const gcSnapshotTree = baseSnapshot.trees[gcTreeKey];
                 if (gcSnapshotTree !== undefined) {
-                    // If the GC tree is written at root, we should also do the same.
-                    this._writeDataAtRoot = true;
                     const baseGCState = await getGCStateFromSnapshot(
                         gcSnapshotTree,
                         readAndParseBlob,
@@ -855,7 +844,7 @@ export class GarbageCollector implements IGarbageCollector {
             const gcResult = runGarbageCollection(gcData.gcNodes, ["/"]);
 
             const gcStats = await this.runPostGCSteps(gcData, gcResult, logger, currentReferenceTimestampMs);
-            event.end({ ...gcStats });
+            event.end({ ...gcStats, timestamp: currentReferenceTimestampMs });
             this.completedRuns++;
             return gcStats;
         }, { end: true, cancel: "error" });
@@ -884,7 +873,7 @@ export class GarbageCollector implements IGarbageCollector {
 
         // Update the current state and update the runtime of all routes or ids that used as per the GC run.
         this.updateCurrentState(gcData, gcResult, currentReferenceTimestampMs);
-        this.runtime.updateUsedRoutes(gcResult.referencedNodeIds, currentReferenceTimestampMs);
+        this.runtime.updateUsedRoutes(gcResult.referencedNodeIds);
 
         // Log events for objects that are ready to be deleted by sweep. When we have sweep enabled, we will
         // delete these objects here instead.
@@ -967,6 +956,7 @@ export class GarbageCollector implements IGarbageCollector {
             gcFeature: this.gcEnabled ? this.currentGCVersion : 0,
             sessionExpiryTimeoutMs: this.sessionExpiryTimeoutMs,
             sweepEnabled: this.sweepEnabled,
+            sweepTimeoutMs: this.sweepTimeoutMs,
         };
     }
 
@@ -1155,7 +1145,7 @@ export class GarbageCollector implements IGarbageCollector {
             this.newReferencesSinceLastRun,
         );
 
-        if (this.writeDataAtRoot && missingExplicitReferences.length > 0) {
+        if (missingExplicitReferences.length > 0) {
             missingExplicitReferences.forEach((missingExplicitReference) => {
                 const event: ITelemetryPerformanceEvent = {
                     eventName: "gcUnknownOutboundReferences",

package/src/gcSweepReadyUsageDetection.ts

@@ -74,16 +74,8 @@ export class SweepReadyUsageDetectionHandler {
         localStorageOverride?: Pick<Storage, "getItem" | "setItem">,
     ) {
         const noopStorage = { getItem: () => null, setItem: () => {} };
-        if (localStorageOverride !== undefined) {
-            this.localStorage = localStorageOverride;
-        } else {
-            try {
-                // localStorage is not defined in Node environment so this throws
-                this.localStorage = localStorage ?? noopStorage;
-            } catch (error) {
-                this.localStorage = noopStorage;
-            }
-        }
+        // localStorage is not defined in Node environment, so fall back to noopStorage if needed.
+        this.localStorage = localStorageOverride ?? globalThis.localStorage ?? noopStorage;
 
         if (this.localStorage === noopStorage) {
             // This means the Skip Closure Period logic will not work.

package/src/index.ts

@@ -23,8 +23,8 @@ export {
     RuntimeHeaders,
     ISummaryConfiguration,
     DefaultSummaryConfiguration,
+    ICompressionRuntimeOptions,
 } from "./containerRuntime";
-export { DeltaScheduler } from "./deltaScheduler";
 export { FluidDataStoreRegistry } from "./dataStoreRegistry";
 export {
     gcBlobPrefix,
@@ -34,12 +34,10 @@ export {
 } from "./garbageCollection";
 export {
     IPendingFlush,
-    IPendingFlushMode,
     IPendingLocalState,
     IPendingMessage,
     IPendingState,
 } from "./pendingStateManager";
-export { ScheduleManager } from "./scheduleManager";
 export { Summarizer } from "./summarizer";
 export {
     EnqueueSummarizeResult,
@@ -54,6 +52,7 @@ export {
     INackSummaryResult,
     IOnDemandSummarizeOptions,
     IProvideSummarizer,
+    IRefreshSummaryAckOptions,
     ISubmitSummaryOpResult,
     ISubmitSummaryOptions,
     ISummarizeOptions,

package/src/packageVersion.ts

@@ -6,4 +6,4 @@
  */
 
 export const pkgName = "@fluidframework/container-runtime";
-export const pkgVersion = "2.0.0-internal.1.4.4";
+export const pkgVersion = "2.0.0-internal.2.0.1";

package/src/pendingStateManager.ts

@@ -10,7 +10,6 @@ import { DataProcessingError } from "@fluidframework/container-utils";
 import {
     ISequencedDocumentMessage,
 } from "@fluidframework/protocol-definitions";
-import { FlushMode } from "@fluidframework/runtime-definitions";
 import Deque from "double-ended-queue";
 import { ContainerMessageType } from "./containerRuntime";
 import { pkgVersion } from "./packageVersion";
@@ -29,15 +28,6 @@ export interface IPendingMessage {
     opMetadata: Record<string, unknown> | undefined;
 }
 
-/**
- * This represents a FlushMode update and is added to the pending queue when `setFlushMode` is called on the
- * ContainerRuntime and the FlushMode changes.
- */
-export interface IPendingFlushMode {
-    type: "flushMode";
-    flushMode: FlushMode;
-}
-
 /**
  * This represents an explicit flush call and is added to the pending queue when flush is called on the ContainerRuntime
  * to flush pending messages.
@@ -46,7 +36,7 @@ export interface IPendingFlush {
     type: "flush";
 }
 
-export type IPendingState = IPendingMessage | IPendingFlushMode | IPendingFlush;
+export type IPendingState = IPendingMessage | IPendingFlush;
 
 export interface IPendingLocalState {
     /**
@@ -58,8 +48,6 @@ export interface IPendingLocalState {
 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;
@@ -68,13 +56,18 @@ export interface IRuntimeStateHandler{
         content: any,
         localOpMetadata: unknown,
         opMetadata: Record<string, unknown> | undefined): void;
+    rollback(
+        type: ContainerMessageType,
+        content: any,
+        localOpMetadata: unknown): void;
+    orderSequentially(callback: () => void): 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
  * batches along with the messages.
- * When the Container reconnects, it replays the pending states, which includes setting the FlushMode, manual flushing
+ * When the Container reconnects, it replays the pending states, which includes manual flushing
  * of messages and triggering resubmission of unacked ops.
  *
  * It verifies that all the ops are acked, are received in the right order and batch information is correct.
@@ -100,13 +93,6 @@ export class PendingStateManager implements IDisposable {
     // the correct batch metadata.
     private pendingBatchBeginMessage: ISequencedDocumentMessage | undefined;
 
-    /**
-     * This tracks the flush mode for the next message in the pending state queue. When replaying messages, we need to
-     * first set the flush mode to this value and then send ops. It is important to do this info because the flush
-     * mode could have been updated.
-     */
-    private flushModeForNextMessage: FlushMode;
-
     private clientId: string | undefined;
 
     /**
@@ -131,13 +117,9 @@ export class PendingStateManager implements IDisposable {
 
     constructor(
         private readonly stateHandler: IRuntimeStateHandler,
-        initialFlushMode: FlushMode,
         initialLocalState: IPendingLocalState | undefined,
     ) {
         this.initialStates = new Deque<IPendingState>(initialLocalState?.pendingStates ?? []);
-
-        this.flushModeForNextMessage = initialFlushMode;
-        this.onFlushModeUpdated(initialFlushMode);
     }
 
     public get disposed() { return this.disposeOnce.evaluated; }
@@ -174,24 +156,10 @@ export class PendingStateManager implements IDisposable {
         this._pendingMessagesCount++;
     }
 
-    /**
-     * Called when the FlushMode is updated. Adds the FlushMode to the pending state queue.
-     * @param flushMode - The flushMode that was updated.
-     */
-    public onFlushModeUpdated(flushMode: FlushMode) {
-        this.pendingStates.push({ type: "flushMode", flushMode });
-    }
-
     /**
      * Called when flush() is called on the ContainerRuntime to manually flush messages.
      */
     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.stateHandler.flushMode() === FlushMode.Immediate) {
-            return;
-        }
-
         // If the previous state is not a message, flush is a no-op.
         const previousState = this.pendingStates.peekBack();
         if (previousState?.type !== "message") {
@@ -279,11 +247,6 @@ export class PendingStateManager implements IDisposable {
      * @param message - The message that is being processed.
      */
     private maybeProcessBatchBegin(message: ISequencedDocumentMessage) {
-        // Tracks the last FlushMode that was set before this message was sent.
-        let pendingFlushMode: FlushMode | undefined;
-        // Tracks whether a flush was called before this message was sent.
-        let pendingFlush: boolean = false;
-
         /**
          * We are checking if the next message is the start of a batch. It can happen in the following scenarios:
          *
@@ -295,34 +258,12 @@ export class PendingStateManager implements IDisposable {
          * Keep reading pending states from the queue until we encounter a message. It's possible that the FlushMode was
          * updated a bunch of times without sending any messages.
          */
-        let nextPendingState = this.peekNextPendingState();
-        while (nextPendingState.type !== "message") {
-            if (nextPendingState.type === "flushMode") {
-                pendingFlushMode = nextPendingState.flushMode;
-            }
-            if (nextPendingState.type === "flush") {
-                pendingFlush = true;
-            }
+        while (this.peekNextPendingState().type !== "message") {
             this.pendingStates.shift();
-            nextPendingState = this.peekNextPendingState();
         }
 
-        if (pendingFlushMode !== undefined) {
-            this.flushModeForNextMessage = pendingFlushMode;
-        }
-
-        // If the FlushMode was set to Immediate before this message was sent, this message won't be a batch message
-        // because in Immediate mode, every message is flushed individually.
-        if (pendingFlushMode === FlushMode.Immediate) {
-            return;
-        }
-
-        /**
-         * This message is the first in a batch if before it was sent either the FlushMode was set to TurnBased or there
-         * was an explicit flush call. Note that a flush call is tracked only in TurnBased mode and it indicates the end
-         * of one batch and beginning of another.
-         */
-        if (pendingFlushMode === FlushMode.TurnBased || pendingFlush) {
+        // This message is the first in a batch if the "batch" property on the metadata is set to true
+        if (message.metadata?.batch) {
             // We should not already be processing a batch and there should be no pending batch begin message.
             assert(!this.isProcessingBatch && this.pendingBatchBeginMessage === undefined,
                 0x16b /* "The pending batch state indicates we are already processing a batch" */);
@@ -347,16 +288,6 @@ export class PendingStateManager implements IDisposable {
             return;
         }
 
-        /**
-         * We are in the middle of processing a batch. The batch ends when we see an explicit flush. We should never see
-         * a FlushMode before flush. This is true because we track batches only when FlushMode is TurnBased and in this
-         * mode, a batch ends either by calling flush or by changing the mode to Immediate which also triggers a flush.
-         */
-        assert(
-            nextPendingState.type !== "flushMode",
-            0x2bd /* "We should not see a pending FlushMode until we see a flush when processing a batch" */,
-        );
-
         // There should be a pending batch begin message.
         assert(this.pendingBatchBeginMessage !== undefined, 0x16d /* "There is no pending batch begin message" */);
 
@@ -406,7 +337,7 @@ export class PendingStateManager implements IDisposable {
 
     /**
      * 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.
+     * states in its queue. This includes triggering resubmission of unacked ops.
      */
     public replayPendingStates() {
         assert(this.stateHandler.connected(), 0x172 /* "The connection state is not consistent with the runtime" */);
@@ -426,12 +357,7 @@ export class PendingStateManager implements IDisposable {
         // Reset the pending message count because all these messages will be removed from the queue.
         this._pendingMessagesCount = 0;
 
-        // Save the current FlushMode so that we can revert it back after replaying the states.
-        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.stateHandler.setFlushMode(this.flushModeForNextMessage);
+        const messageBatchQueue = new Deque<IPendingMessage>();
 
         // 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
@@ -441,16 +367,43 @@ export class PendingStateManager implements IDisposable {
             const pendingState = this.pendingStates.shift()!;
             switch (pendingState.type) {
                 case "message":
-                    this.stateHandler.reSubmit(
-                        pendingState.messageType,
-                        pendingState.content,
-                        pendingState.localOpMetadata,
-                        pendingState.opMetadata);
-                    break;
-                case "flushMode":
-                    this.stateHandler.setFlushMode(pendingState.flushMode);
+                    assert(pendingState.opMetadata?.batch !== false || messageBatchQueue.length > 0,
+                        0x41b /* We cannot process batches in chunks */);
+                    /**
+                     * We want to ensure grouped messages get processed in a batch.
+                     * Note: It is not possible for the PendingStateManager to receive a partially acked batch. It will
+                     * either receive the whole batch ack or nothing at all.
+                     */
+                    if (messageBatchQueue.length > 0 || pendingState.opMetadata?.batch) {
+                        messageBatchQueue.enqueue(pendingState);
+                    } else {
+                        this.stateHandler.reSubmit(
+                            pendingState.messageType,
+                            pendingState.content,
+                            pendingState.localOpMetadata,
+                            pendingState.opMetadata);
+                    }
                     break;
                 case "flush":
+                    /**
+                     * A "flush" call can indicate the end of a batch.
+                     * We can't rely on the "batch" property in the message metadata as it gets
+                     * updated elsewhere and it is not the same object instance that gets updated.
+                     */
+                    if (messageBatchQueue.length > 0) {
+                        this.stateHandler.orderSequentially(() => {
+                            while (messageBatchQueue.length > 0) {
+                                // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+                                const message = messageBatchQueue.dequeue()!;
+                                this.stateHandler.reSubmit(
+                                    message.messageType,
+                                    message.content,
+                                    message.localOpMetadata,
+                                    message.opMetadata);
+                            }
+                        });
+                    }
+                    assert(messageBatchQueue.length === 0, 0x41c /* cannot flush in the middle of a batch */);
                     this.stateHandler.flush();
                     break;
                 default:
@@ -459,7 +412,15 @@ export class PendingStateManager implements IDisposable {
             pendingStatesCount--;
         }
 
-        // Revert the FlushMode.
-        this.stateHandler.setFlushMode(savedFlushMode);
+        // There are some cases where ops are stashed but not flushed. We need to ensure they are resubmitted
+        while (messageBatchQueue.length > 0) {
+            // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+            const message = messageBatchQueue.dequeue()!;
+            this.stateHandler.reSubmit(
+                message.messageType,
+                message.content,
+                message.localOpMetadata,
+                message.opMetadata);
+        }
     }
 }

package/src/scheduleManager.ts

@@ -221,6 +221,7 @@ class ScheduleManagerCore {
                 length: endBatch - startBatch,
             });
         }
+
         this.deltaManager.inbound.resume();
     }