@fluidframework/container-loader 2.63.0-359461 → 2.63.0-359962

package/src/serializedStateManager.ts

@@ -10,6 +10,7 @@ import type {
 	IEvent,
 	ITelemetryBaseLogger,
 } from "@fluidframework/core-interfaces";
+import type { IDisposable } from "@fluidframework/core-interfaces/internal";
 import { Timer, assert } from "@fluidframework/core-utils/internal";
 import {
 	FetchSource,
@@ -21,7 +22,7 @@ import {
 	type IVersion,
 	type ISequencedDocumentMessage,
 } from "@fluidframework/driver-definitions/internal";
-import { getSnapshotTree } from "@fluidframework/driver-utils/internal";
+import { getSnapshotTree, isInstanceOfISnapshot } from "@fluidframework/driver-utils/internal";
 import {
 	type MonitoringContext,
 	PerformanceEvent,
@@ -30,10 +31,14 @@ import {
 } from "@fluidframework/telemetry-utils/internal";
 
 import {
-	type ISerializableBlobContents,
 	getBlobContentsFromTree,
+	type ISerializableBlobContents,
 } from "./containerStorageAdapter.js";
-import { convertSnapshotToSnapshotInfo, getDocumentAttributes } from "./utils.js";
+import {
+	convertISnapshotToSnapshotWithBlobs,
+	convertSnapshotToSnapshotInfo,
+	getDocumentAttributes,
+} from "./utils.js";
 
 /**
  * This is very similar to {@link @fluidframework/protocol-definitions/internal#ISnapshot}, but the difference is
@@ -75,7 +80,7 @@ export interface IPendingContainerState extends SnapshotWithBlobs {
 	/**
 	 * Any group snapshots (aka delay-loaded) we've downloaded from the service for this container
 	 */
-	loadedGroupIdSnapshots?: Record<string, ISnapshotInfo>;
+	loadedGroupIdSnapshots?: Record<string, SerializedSnapshotInfo>;
 	/**
 	 * All ops since base snapshot sequence number up to the latest op
 	 * seen when the container was closed. Used to apply stashed (saved pending)
@@ -117,9 +122,14 @@ export interface IPendingDetachedContainerState extends SnapshotWithBlobs {
 	pendingRuntimeState?: unknown;
 }
 
-export interface ISnapshotInfo extends SnapshotWithBlobs {
+export interface SerializedSnapshotInfo extends SnapshotWithBlobs {
+	snapshotSequenceNumber: number;
+}
+
+interface ISnapshotInfo {
 	snapshotSequenceNumber: number;
 	snapshotFetchedTime?: number | undefined;
+	snapshot: ISnapshot | ISnapshotTree;
 }
 
 export type ISerializedStateManagerDocumentStorageService = Pick<
@@ -161,10 +171,10 @@ class RefreshPromiseTracker {
  * as well as the snapshot to be used for serialization.
  * It also keeps track of container dirty state and which local ops have been processed
  */
-export class SerializedStateManager {
+export class SerializedStateManager implements IDisposable {
 	private readonly processedOps: ISequencedDocumentMessage[] = [];
 	private readonly mc: MonitoringContext;
-	private snapshot: ISnapshotInfo | undefined;
+	private snapshotInfo: ISnapshotInfo | undefined;
 	private latestSnapshot: ISnapshotInfo | undefined;
 	private readonly refreshTracker = new RefreshPromiseTracker(
 		// eslint-disable-next-line unicorn/consistent-function-scoping
@@ -176,12 +186,13 @@ export class SerializedStateManager {
 				error,
 			),
 	);
-	private readonly lastSavedOpSequenceNumber: number = 0;
-	private readonly refreshTimer: Timer;
+	private lastSavedOpSequenceNumber: number = 0;
+	private readonly refreshTimer: Timer | undefined;
 	private readonly snapshotRefreshTimeoutMs: number = 60 * 60 * 24 * 1000;
+	readonly #snapshotRefreshEnabled: boolean;
+	#disposed: boolean = false;
 
 	/**
-	 * @param pendingLocalState - The pendingLocalState being rehydrated, if any (undefined when loading directly from storage)
 	 * @param subLogger - Container's logger to use as parent for our logger
 	 * @param storageAdapter - Storage adapter for fetching snapshots
 	 * @param _offlineLoadEnabled - Is serializing/rehydrating containers allowed?
@@ -189,7 +200,6 @@ export class SerializedStateManager {
 	 * @param containerDirty - Is the container "dirty"? That's the opposite of "saved" - there is pending state that may not have been received yet by the service.
 	 */
 	constructor(
-		private readonly pendingLocalState: IPendingContainerState | undefined,
 		subLogger: ITelemetryBaseLogger,
 		private readonly storageAdapter: ISerializedStateManagerDocumentStorageService,
 		private readonly _offlineLoadEnabled: boolean,
@@ -204,19 +214,30 @@ export class SerializedStateManager {
 		});
 
 		this.snapshotRefreshTimeoutMs = snapshotRefreshTimeoutMs ?? this.snapshotRefreshTimeoutMs;
-		this.refreshTimer = new Timer(this.snapshotRefreshTimeoutMs, () =>
-			this.tryRefreshSnapshot(),
-		);
-		// special case handle. Obtaining the last saved op seq num to avoid
-		// refreshing the snapshot before we have processed it. It could cause
-		// a subsequent stashing to have a newer snapshot than allowed.
-		if (pendingLocalState && pendingLocalState.savedOps.length > 0) {
-			const savedOpsSize = pendingLocalState.savedOps.length;
-			this.lastSavedOpSequenceNumber =
-				pendingLocalState.savedOps[savedOpsSize - 1].sequenceNumber;
-		}
+
+		this.#snapshotRefreshEnabled =
+			_offlineLoadEnabled &&
+			(this.mc.config.getBoolean("Fluid.Container.enableOfflineSnapshotRefresh") ??
+				this.mc.config.getBoolean("Fluid.Container.enableOfflineFull")) === true;
+
+		this.refreshTimer = this.#snapshotRefreshEnabled
+			? new Timer(this.snapshotRefreshTimeoutMs, () => this.tryRefreshSnapshot())
+			: undefined;
 		containerEvent.on("saved", () => this.updateSnapshotAndProcessedOpsMaybe());
 	}
+	public get disposed(): boolean {
+		return this.#disposed;
+	}
+	dispose(): void {
+		this.#disposed = true;
+		this.refreshTimer?.clear();
+	}
+
+	private verifyNotDisposed(): void {
+		if (this.#disposed) {
+			throw new Error("SerializedStateManager used after dispose.");
+		}
+	}
 
 	public get offlineLoadEnabled(): boolean {
 		return this._offlineLoadEnabled;
@@ -248,69 +269,74 @@ export class SerializedStateManager {
 	 * Otherwise, fetch it from storage (according to specifiedVersion if provided).
 	 *
 	 * @param specifiedVersion - If a version is specified and we don't have pendingLocalState, fetch this version from storage.
-	 * @param supportGetSnapshotApi - a boolean indicating whether to use the fetchISnapshot or fetchISnapshotTree.
+	 * @param pendingLocalState - The pendingLocalState being rehydrated, if any (undefined when loading directly from storage)
 	 * @returns The snapshot to boot the container from
 	 */
-	public async fetchSnapshot(specifiedVersion: string | undefined): Promise<{
-		baseSnapshot: ISnapshot | ISnapshotTree;
+	public async fetchSnapshot(
+		specifiedVersion: string | undefined,
+		pendingLocalState: IPendingContainerState | undefined,
+	): Promise<{
+		snapshot: ISnapshot | ISnapshotTree;
 		version: IVersion | undefined;
 		attributes: IDocumentAttributes;
 	}> {
-		if (this.pendingLocalState === undefined) {
-			const { baseSnapshot, version } = await getSnapshot(
+		this.verifyNotDisposed();
+		if (pendingLocalState === undefined) {
+			const { snapshot, version } = await getSnapshot(
 				this.mc,
 				this.storageAdapter,
 				this.supportGetSnapshotApi(),
 				specifiedVersion,
 			);
-			const baseSnapshotTree: ISnapshotTree | undefined = getSnapshotTree(baseSnapshot);
+			const baseSnapshotTree: ISnapshotTree | undefined = getSnapshotTree(snapshot);
 			const attributes = await getDocumentAttributes(this.storageAdapter, baseSnapshotTree);
-			// non-interactive clients will not have any pending state we want to save
 			if (this.offlineLoadEnabled) {
-				// we defer getting the blobs to not impact the container load flow
-				// only getPendingState depends on the resolution of this promise
-				this.refreshTracker.setPromise(
-					getBlobContentsFromTree(baseSnapshot, this.storageAdapter).then((snapshotBlobs) => {
-						this.snapshot = {
-							baseSnapshot: baseSnapshotTree,
-							snapshotBlobs,
-							snapshotSequenceNumber: attributes.sequenceNumber,
-						};
-						this.refreshTimer.start();
-						return attributes.sequenceNumber;
-					}),
-				);
+				this.refreshTimer?.start();
+				this.snapshotInfo = {
+					snapshot,
+					snapshotSequenceNumber: attributes.sequenceNumber,
+				};
 			}
-			return { baseSnapshot, version, attributes };
+			return { snapshot, version, attributes };
 		} else {
-			const { baseSnapshot, snapshotBlobs } = this.pendingLocalState;
+			const { baseSnapshot, snapshotBlobs, savedOps } = pendingLocalState;
 			const attributes = await getDocumentAttributes(this.storageAdapter, baseSnapshot);
-			this.snapshot = {
-				baseSnapshot,
-				snapshotBlobs,
-				snapshotSequenceNumber: attributes.sequenceNumber,
-			};
-			this.tryRefreshSnapshot();
 			const blobContents = new Map<string, ArrayBuffer>();
 			for (const [id, value] of Object.entries(snapshotBlobs)) {
 				blobContents.set(id, stringToBuffer(value, "utf8"));
 			}
-			const iSnapshot: ISnapshot = {
-				sequenceNumber: this.snapshot.snapshotSequenceNumber,
+			const snapshot: ISnapshot = {
+				sequenceNumber: attributes.sequenceNumber,
 				snapshotTree: baseSnapshot,
 				blobContents,
 				latestSequenceNumber: undefined,
 				ops: [],
 				snapshotFormatV: 1,
 			};
-			return { baseSnapshot: iSnapshot, version: undefined, attributes };
+
+			if (this.offlineLoadEnabled) {
+				// special case handle. Obtaining the last saved op seq num to avoid
+				// refreshing the snapshot before we have processed it. It could cause
+				// a subsequent stashing to have a newer snapshot than allowed.
+				if (savedOps.length > 0) {
+					const savedOpsSize = savedOps.length;
+					this.lastSavedOpSequenceNumber = savedOps[savedOpsSize - 1].sequenceNumber;
+				}
+
+				this.snapshotInfo = {
+					snapshot,
+					snapshotSequenceNumber: attributes.sequenceNumber,
+				};
+				this.tryRefreshSnapshot();
+			}
+			return { snapshot, version: undefined, attributes };
 		}
 	}
 
 	private tryRefreshSnapshot(): void {
 		if (
-			(this.mc.config.getBoolean("Fluid.Container.enableOfflineSnapshotRefresh") ??
-				this.mc.config.getBoolean("Fluid.Container.enableOfflineFull")) === true &&
+			this.#snapshotRefreshEnabled &&
+			!this.#disposed &&
 			!this.refreshTracker.hasPromise &&
 			this.latestSnapshot === undefined
 		) {
@@ -332,6 +358,10 @@ export class SerializedStateManager {
 			supportGetSnapshotApi,
 		);
 
+		if (this.#disposed) {
+			return -1;
+		}
+
 		// These are loading groupIds that the containerRuntime has requested over its lifetime.
 		// We will fetch the latest snapshot for the groupIds, which will update storageAdapter.loadedGroupIdSnapshots's cache
 		const downloadedGroupIds = Object.keys(this.storageAdapter.loadedGroupIdSnapshots);
@@ -360,6 +390,7 @@ export class SerializedStateManager {
 	private updateSnapshotAndProcessedOpsMaybe(): number {
 		const snapshotSequenceNumber = this.latestSnapshot?.snapshotSequenceNumber;
 		if (
+			this.#disposed ||
 			snapshotSequenceNumber === undefined ||
 			this.processedOps.length === 0 ||
 			this.processedOps[this.processedOps.length - 1].sequenceNumber <
@@ -382,17 +413,17 @@ export class SerializedStateManager {
 				snapshotSequenceNumber,
 				firstProcessedOpSequenceNumber,
 				lastProcessedOpSequenceNumber,
-				stashedSnapshotSequenceNumber: this.snapshot?.snapshotSequenceNumber,
+				stashedSnapshotSequenceNumber: this.snapshotInfo?.snapshotSequenceNumber,
 			});
 			this.latestSnapshot = undefined;
-			this.refreshTimer.restart();
+			this.refreshTimer?.restart();
 		} else if (snapshotSequenceNumber <= lastProcessedOpSequenceNumber) {
 			// Snapshot seq num is between the first and last processed op.
 			// Remove the ops that are already part of the snapshot
 			this.processedOps.splice(0, snapshotSequenceNumber - firstProcessedOpSequenceNumber + 1);
-			this.snapshot = this.latestSnapshot;
+			this.snapshotInfo = this.latestSnapshot;
 			this.latestSnapshot = undefined;
-			this.refreshTimer.restart();
+			this.refreshTimer?.restart();
 			this.mc.logger.sendTelemetryEvent({
 				eventName: "SnapshotRefreshed",
 				snapshotSequenceNumber,
@@ -410,31 +441,15 @@ export class SerializedStateManager {
 	 * base snapshot when attaching.
 	 * @param snapshot - snapshot and blobs collected while attaching (a form of the attach summary)
 	 */
-	public setInitialSnapshot(snapshot: SnapshotWithBlobs | undefined): void {
+	public setInitialSnapshot(snapshot: ISnapshot): void {
+		this.verifyNotDisposed();
 		if (this.offlineLoadEnabled) {
-			assert(
-				this.snapshot === undefined,
-				0x937 /* inital snapshot should only be defined once */,
-			);
-			assert(snapshot !== undefined, 0x938 /* attachment snapshot should be defined */);
-			const { baseSnapshot, snapshotBlobs } = snapshot;
-			const attributesHash =
-				".protocol" in baseSnapshot.trees
-					? baseSnapshot.trees[".protocol"].blobs.attributes
-					: baseSnapshot.blobs[".attributes"];
-			// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
-			const attributes = JSON.parse(snapshotBlobs[attributesHash]);
-			assert(
-				// eslint-disable-next-line @typescript-eslint/no-unsafe-member-access
-				attributes.sequenceNumber === 0,
-				0x939 /* trying to set a non attachment snapshot */,
-			);
-			this.snapshot = {
-				...snapshot,
-				// eslint-disable-next-line @typescript-eslint/no-unsafe-member-access
-				snapshotSequenceNumber: attributes.sequenceNumber as number,
+			this.snapshotInfo = {
+				snapshot,
+				snapshotSequenceNumber: snapshot.sequenceNumber ?? 0,
+				snapshotFetchedTime: Date.now(),
 			};
-			this.refreshTimer.start();
+			this.refreshTimer?.start();
 		}
 	}
 
@@ -447,6 +462,11 @@ export class SerializedStateManager {
 		runtime: Pick<IRuntime, "getPendingLocalState">,
 		resolvedUrl: IResolvedUrl,
 	): Promise<string> {
+		this.verifyNotDisposed();
+		if (!this.offlineLoadEnabled) {
+			throw new UsageError("Can't get pending local state unless offline load is enabled");
+		}
+
 		return PerformanceEvent.timedExecAsync(
 			this.mc.logger,
 			{
@@ -460,25 +480,14 @@ export class SerializedStateManager {
 				clientId,
 			},
 			async () => {
-				if (!this.offlineLoadEnabled) {
-					throw new UsageError("Can't get pending local state unless offline load is enabled");
-				}
-				if (this.snapshot === undefined && this.refreshTracker.hasPromise) {
-					// we deferred the initial download of the snapshot to not block
-					// the container load flow, so if it is not resolved
-					// and we don't have a snapshot, we will wait for the download
-					// to finish.
-					await this.refreshTracker.Promise;
-				}
-
-				assert(this.snapshot !== undefined, 0x8e5 /* no base data */);
+				assert(this.snapshotInfo !== undefined, 0x8e5 /* no base data */);
 				const pendingRuntimeState = await runtime.getPendingLocalState({
 					notifyImminentClosure: false,
-					snapshotSequenceNumber: this.snapshot.snapshotSequenceNumber,
-					sessionExpiryTimerStarted: this.snapshot.snapshotFetchedTime,
+					snapshotSequenceNumber: this.snapshotInfo.snapshotSequenceNumber,
+					sessionExpiryTimerStarted: this.snapshotInfo.snapshotFetchedTime,
 				});
 				// This conversion is required because ArrayBufferLike doesn't survive JSON.stringify
-				const loadedGroupIdSnapshots = {};
+				const loadedGroupIdSnapshots: Record<string, SerializedSnapshotInfo> = {};
 				let hasGroupIdSnapshots = false;
 				const groupIdSnapshots = Object.entries(this.storageAdapter.loadedGroupIdSnapshots);
 				if (groupIdSnapshots.length > 0) {
@@ -487,11 +496,20 @@ export class SerializedStateManager {
 						loadedGroupIdSnapshots[groupId] = convertSnapshotToSnapshotInfo(snapshot);
 					}
 				}
+
+				const snapshotWithBlobs: SnapshotWithBlobs = isInstanceOfISnapshot(
+					this.snapshotInfo.snapshot,
+				)
+					? convertISnapshotToSnapshotWithBlobs(this.snapshotInfo.snapshot)
+					: await convertSnapshotTreeToSnapshotWithBlobs(
+							this.snapshotInfo.snapshot,
+							this.storageAdapter,
+						);
+
 				const pendingState: IPendingContainerState = {
 					attached: true,
 					pendingRuntimeState,
-					baseSnapshot: this.snapshot.baseSnapshot,
-					snapshotBlobs: this.snapshot.snapshotBlobs,
+					...snapshotWithBlobs,
 					loadedGroupIdSnapshots: hasGroupIdSnapshots ? loadedGroupIdSnapshots : undefined,
 					savedOps: this.processedOps,
 					url: resolvedUrl.url,
@@ -504,6 +522,17 @@ export class SerializedStateManager {
 	}
 }
 
+async function convertSnapshotTreeToSnapshotWithBlobs(
+	snapshot: ISnapshotTree,
+	storageAdapter: ISerializedStateManagerDocumentStorageService,
+): Promise<SnapshotWithBlobs> {
+	const snapshotBlobs = await getBlobContentsFromTree(snapshot, storageAdapter);
+	return {
+		baseSnapshot: snapshot,
+		snapshotBlobs,
+	};
+}
+
 /**
  * Retrieves the most recent snapshot and returns its info.
  *
@@ -517,41 +546,46 @@ export async function getLatestSnapshotInfo(
 	storageAdapter: ISerializedStateManagerDocumentStorageService,
 	supportGetSnapshotApi: boolean,
 ): Promise<ISnapshotInfo | undefined> {
-	return PerformanceEvent.timedExecAsync(
+	return PerformanceEvent.timedExecAsync<ISnapshotInfo | undefined>(
 		mc.logger,
 		{ eventName: "GetLatestSnapshotInfo" },
-		async () => {
-			// get the latest non cached snapshot version
-			const specifiedVersion: IVersion[] = await storageAdapter.getVersions(
-				// eslint-disable-next-line unicorn/no-null
-				null,
-				1,
-				"getLatestSnapshotInfo",
-				FetchSource.noCache,
-			);
-			const { baseSnapshot } = await getSnapshot(
-				mc,
-				storageAdapter,
-				supportGetSnapshotApi,
-				specifiedVersion[0]?.id,
-			);
+		async (event) => {
+			try {
+				// get the latest non cached snapshot version
+				const specifiedVersion: IVersion[] = await storageAdapter.getVersions(
+					// eslint-disable-next-line unicorn/no-null
+					null,
+					1,
+					"getLatestSnapshotInfo",
+					FetchSource.noCache,
+				);
+				const { snapshot: baseSnapshot } = await getSnapshot(
+					mc,
+					storageAdapter,
+					supportGetSnapshotApi,
+					specifiedVersion[0]?.id,
+				);
 
-			const baseSnapshotTree: ISnapshotTree | undefined = getSnapshotTree(baseSnapshot);
-			const snapshotFetchedTime = Date.now();
-			const snapshotBlobs = await getBlobContentsFromTree(baseSnapshot, storageAdapter);
-			const attributes: IDocumentAttributes = await getDocumentAttributes(
-				storageAdapter,
-				baseSnapshotTree,
-			);
-			const snapshotSequenceNumber = attributes.sequenceNumber;
-			return {
-				baseSnapshot: baseSnapshotTree,
-				snapshotBlobs,
-				snapshotSequenceNumber,
-				snapshotFetchedTime,
-			};
+				const { sequenceNumber, snapshotTree } = isInstanceOfISnapshot(baseSnapshot)
+					? baseSnapshot
+					: { snapshotTree: baseSnapshot, sequenceNumber: undefined };
+
+				const snapshotSequenceNumber: number =
+					sequenceNumber ??
+					(await getDocumentAttributes(storageAdapter, snapshotTree).then(
+						(a) => a.sequenceNumber,
+					));
+				return {
+					snapshot: baseSnapshot,
+					snapshotSequenceNumber,
+					snapshotFetchedTime: Date.now(),
+				};
+			} catch (error) {
+				event.cancel(undefined, error);
+			}
+			return undefined;
 		},
-	).catch(() => undefined);
+	);
 }
 
 /**
@@ -571,12 +605,12 @@ async function getSnapshot(
 	>,
 	supportGetSnapshotApi: boolean,
 	specifiedVersion: string | undefined,
-): Promise<{ baseSnapshot: ISnapshot | ISnapshotTree; version?: IVersion }> {
+): Promise<{ snapshot: ISnapshot | ISnapshotTree; version?: IVersion }> {
 	const { snapshot, version } = supportGetSnapshotApi
 		? await fetchISnapshot(mc, storageAdapter, specifiedVersion)
 		: await fetchISnapshotTree(mc, storageAdapter, specifiedVersion);
 	assert(snapshot !== undefined, 0x8e4 /* Snapshot should exist */);
-	return { baseSnapshot: snapshot, version };
+	return { snapshot, version };
 }
 
 /**