[main > release/client/2.0]: Add snapshotWithLoadingGroupId to cacheC…

…ontentType for caching in odsp driver #21793 (#21794)

## Description

If a snapshot with loading groupid is cached when enabled, and then a
loaded from cache in environment where snapshot with loading groupid is
not supported, then we will get a partial snapshot with missing blobs,
and runtime when loading datastores, will not determine whether it needs
to fetch LoadingGroupId or not and then it will just continue, and each
individual blob will be fetched separately which can cause potential
throttling issues.

Sol: Need to cache the snapshot with different key based on the snapshot
type.

packages/drivers/odsp-driver-definitions/api-report/odsp-driver-definitions.alpha.api.md

@@ -10,7 +10,7 @@ import { IDriverErrorBase } from '@fluidframework/driver-definitions/internal';
 import { IResolvedUrl } from '@fluidframework/driver-definitions/internal';
 
 // @alpha (undocumented)
-export type CacheContentType = "snapshot" | "ops";
+export type CacheContentType = "snapshot" | "ops" | "snapshotWithLoadingGroupId";
 
 // @alpha (undocumented)
 export interface HostStoragePolicy {

packages/drivers/odsp-driver-definitions/package.json

@@ -114,6 +114,16 @@
 		"typescript": "~5.4.5"
 	},
 	"typeValidation": {
-		"broken": {}
+		"broken": {
+			"TypeAliasDeclaration_CacheContentType": {
+				"backCompat": false
+			},
+			"InterfaceDeclaration_ICacheEntry": {
+				"backCompat": false
+			},
+			"InterfaceDeclaration_IEntry": {
+				"backCompat": false
+			}
+		}
 	}
 }

packages/drivers/odsp-driver-definitions/src/index.ts

@@ -19,6 +19,7 @@ export {
 	IFileEntry,
 	IPersistedCache,
 	snapshotKey,
+	snapshotWithLoadingGroupIdKey,
 } from "./odspCache.js";
 export {
 	IOdspResolvedUrl,

packages/drivers/odsp-driver-definitions/src/odspCache.ts

@@ -20,10 +20,17 @@ export const maximumCacheDurationMs: FiveDaysMs = 432_000_000; // 5 days in ms
  * @internal
  */
 export const snapshotKey = "snapshot";
+
+/**
+ * Describes key for partial snapshot with loading GroupId in cache entry.
+ * @internal
+ */
+export const snapshotWithLoadingGroupIdKey = "snapshotWithLoadingGroupId";
+
 /**
  * @alpha
  */
-export type CacheContentType = "snapshot" | "ops";
+export type CacheContentType = "snapshot" | "ops" | "snapshotWithLoadingGroupId";
 
 /*
  * File / container identifier.

packages/drivers/odsp-driver-definitions/src/test/types/validateOdspDriverDefinitionsPrevious.generated.ts

@@ -41,6 +41,7 @@ declare function get_current_TypeAliasDeclaration_CacheContentType():
 declare function use_old_TypeAliasDeclaration_CacheContentType(
     use: TypeOnly<old.CacheContentType>): void;
 use_old_TypeAliasDeclaration_CacheContentType(
+    // @ts-expect-error compatibility expected to be broken
     get_current_TypeAliasDeclaration_CacheContentType());
 
 /*
@@ -97,6 +98,7 @@ declare function get_current_InterfaceDeclaration_ICacheEntry():
 declare function use_old_InterfaceDeclaration_ICacheEntry(
     use: TypeOnly<old.ICacheEntry>): void;
 use_old_InterfaceDeclaration_ICacheEntry(
+    // @ts-expect-error compatibility expected to be broken
     get_current_InterfaceDeclaration_ICacheEntry());
 
 /*
@@ -153,6 +155,7 @@ declare function get_current_InterfaceDeclaration_IEntry():
 declare function use_old_InterfaceDeclaration_IEntry(
     use: TypeOnly<old.IEntry>): void;
 use_old_InterfaceDeclaration_IEntry(
+    // @ts-expect-error compatibility expected to be broken
     get_current_InterfaceDeclaration_IEntry());
 
 /*

packages/drivers/odsp-driver/src/createFile.ts

@@ -16,6 +16,7 @@ import {
 } from "@fluidframework/odsp-driver-definitions/internal";
 import {
 	ITelemetryLoggerExt,
+	loggerToMonitoringContext,
 	PerformanceEvent,
 } from "@fluidframework/telemetry-utils/internal";
 
@@ -36,6 +37,7 @@ import {
 	buildOdspShareLinkReqParams,
 	createCacheSnapshotKey,
 	getWithRetryForTokenRefresh,
+	snapshotWithLoadingGroupIdSupported,
 } from "./odspUtils.js";
 import { pkgVersion as driverVersion } from "./packageVersion.js";
 import { runWithRetry } from "./retryUtils.js";
@@ -110,7 +112,13 @@ export async function createNewFluidFile(
 			summaryHandle,
 		);
 		// caching the converted summary
-		await epochTracker.put(createCacheSnapshotKey(odspResolvedUrl), snapshot);
+		await epochTracker.put(
+			createCacheSnapshotKey(
+				odspResolvedUrl,
+				snapshotWithLoadingGroupIdSupported(loggerToMonitoringContext(logger).config),
+			),
+			snapshot,
+		);
 	}
 	return odspResolvedUrl;
 }

packages/drivers/odsp-driver/src/createNewContainerOnExistingFile.ts

@@ -11,7 +11,10 @@ import {
 	IOdspResolvedUrl,
 	InstrumentedStorageTokenFetcher,
 } from "@fluidframework/odsp-driver-definitions/internal";
-import { ITelemetryLoggerExt } from "@fluidframework/telemetry-utils/internal";
+import {
+	ITelemetryLoggerExt,
+	loggerToMonitoringContext,
+} from "@fluidframework/telemetry-utils/internal";
 
 import { IWriteSummaryResponse } from "./contracts.js";
 import { ClpCompliantAppHeader } from "./contractsPublic.js";
@@ -24,7 +27,11 @@ import { createOdspUrl } from "./createOdspUrl.js";
 import { EpochTracker } from "./epochTracker.js";
 import { OdspDriverUrlResolver } from "./odspDriverUrlResolver.js";
 import { getApiRoot } from "./odspUrlHelper.js";
-import { IExistingFileInfo, createCacheSnapshotKey } from "./odspUtils.js";
+import {
+	IExistingFileInfo,
+	createCacheSnapshotKey,
+	snapshotWithLoadingGroupIdSupported,
+} from "./odspUtils.js";
 
 /**
  * Creates a new Fluid container on an existing file.
@@ -87,7 +94,13 @@ export async function createNewContainerOnExistingFile(
 			summaryHandle,
 		);
 		// caching the converted summary
-		await epochTracker.put(createCacheSnapshotKey(odspResolvedUrl), snapshot);
+		await epochTracker.put(
+			createCacheSnapshotKey(
+				odspResolvedUrl,
+				snapshotWithLoadingGroupIdSupported(loggerToMonitoringContext(logger).config),
+			),
+			snapshot,
+		);
 	}
 
 	return odspResolvedUrl;

packages/drivers/odsp-driver/src/odspDocumentStorageManager.ts

@@ -32,6 +32,7 @@ import {
 	loggerToMonitoringContext,
 	normalizeError,
 	overwriteStack,
+	type IConfigProvider,
 } from "@fluidframework/telemetry-utils/internal";
 
 import {
@@ -62,6 +63,7 @@ import {
 	getWithRetryForTokenRefresh,
 	isInstanceOfISnapshot,
 	isSnapshotFetchForLoadingGroup,
+	snapshotWithLoadingGroupIdSupported,
 	useLegacyFlowWithoutGroupsForSnapshotFetch,
 	type TokenFetchOptionsEx,
 } from "./odspUtils.js";
@@ -85,6 +87,7 @@ export class OdspDocumentStorageService extends OdspDocumentStorageServiceBase {
 	private readonly snapshotUrl: string | undefined;
 	private readonly attachmentPOSTUrl: string | undefined;
 	private readonly attachmentGETUrl: string | undefined;
+	private readonly config: IConfigProvider;
 	// Driver specified limits for snapshot size and time.
 	/**
 	 * NOTE: While commit cfff6e3 added restrictions to prevent large payloads, snapshot failures will continue to
@@ -115,6 +118,7 @@ export class OdspDocumentStorageService extends OdspDocumentStorageServiceBase {
 		this.snapshotUrl = this.odspResolvedUrl.endpoints.snapshotStorageUrl;
 		this.attachmentPOSTUrl = this.odspResolvedUrl.endpoints.attachmentPOSTStorageUrl;
 		this.attachmentGETUrl = this.odspResolvedUrl.endpoints.attachmentGETStorageUrl;
+		this.config = loggerToMonitoringContext(logger).config;
 	}
 
 	public get isFirstSnapshotFromNetwork(): boolean | undefined {
@@ -278,7 +282,12 @@ export class OdspDocumentStorageService extends OdspDocumentStorageServiceBase {
 					// Here's the logic to grab the persistent cache snapshot implemented by the host
 					// Epoch tracker is responsible for communicating with the persistent cache, handling epochs and cache versions
 					const cachedSnapshotP: Promise<ISnapshot | undefined> = this.epochTracker
-						.get(createCacheSnapshotKey(this.odspResolvedUrl))
+						.get(
+							createCacheSnapshotKey(
+								this.odspResolvedUrl,
+								snapshotWithLoadingGroupIdSupported(this.config),
+							),
+						)
 						.then(
 							async (
 								// eslint-disable-next-line import/no-deprecated
@@ -564,7 +573,10 @@ export class OdspDocumentStorageService extends OdspDocumentStorageServiceBase {
 		// for initial snapshot, don't consult the prefetch cache.
 		if (!this.hostPolicy.avoidPrefetchSnapshotCache && this.firstSnapshotFetchCall) {
 			const prefetchCacheKey = getKeyForCacheEntry(
-				createCacheSnapshotKey(this.odspResolvedUrl),
+				createCacheSnapshotKey(
+					this.odspResolvedUrl,
+					snapshotWithLoadingGroupIdSupported(this.config),
+				),
 			);
 			const result = await this.cache.snapshotPrefetchResultCache
 				?.get(prefetchCacheKey)
@@ -627,7 +639,10 @@ export class OdspDocumentStorageService extends OdspDocumentStorageServiceBase {
 		};
 		const putInCache = async (valueWithEpoch: IVersionedValueWithEpoch): Promise<void> => {
 			return this.cache.persistedCache.put(
-				createCacheSnapshotKey(this.odspResolvedUrl),
+				createCacheSnapshotKey(
+					this.odspResolvedUrl,
+					snapshotWithLoadingGroupIdSupported(this.config),
+				),
 				// Epoch tracker will add the epoch and version to the value here. So just send value to cache.
 				valueWithEpoch.value,
 			);

packages/drivers/odsp-driver/src/odspUtils.ts

@@ -38,8 +38,10 @@ import {
 	isTokenFromCache,
 	snapshotKey,
 	tokenFromResponse,
+	snapshotWithLoadingGroupIdKey,
 } from "@fluidframework/odsp-driver-definitions/internal";
 import {
+	type IConfigProvider,
 	type IFluidErrorBase,
 	ITelemetryLoggerExt,
 	PerformanceEvent,
@@ -452,9 +454,12 @@ export function toInstrumentedOdspTokenFetcher(
 	};
 }
 
-export function createCacheSnapshotKey(odspResolvedUrl: IOdspResolvedUrl): ICacheEntry {
+export function createCacheSnapshotKey(
+	odspResolvedUrl: IOdspResolvedUrl,
+	snapshotWithLoadingGroupId: boolean | undefined,
+): ICacheEntry {
 	const cacheEntry: ICacheEntry = {
-		type: snapshotKey,
+		type: snapshotWithLoadingGroupId ? snapshotWithLoadingGroupIdKey : snapshotKey,
 		key: odspResolvedUrl.fileVersion ?? "",
 		file: {
 			resolvedUrl: odspResolvedUrl,
@@ -464,6 +469,12 @@ export function createCacheSnapshotKey(odspResolvedUrl: IOdspResolvedUrl): ICach
 	return cacheEntry;
 }
 
+export function snapshotWithLoadingGroupIdSupported(
+	config: IConfigProvider,
+): boolean | undefined {
+	return config.getBoolean("Fluid.Container.UseLoadingGroupIdForSnapshotFetch2");
+}
+
 // 80KB is the max body size that we can put in ump post body for server to be able to accept it.
 // Keeping it 78KB to be a little cautious. As per the telemetry 99p is less than 78KB.
 export const maxUmpPostBodySize = 79872;

packages/drivers/odsp-driver/src/prefetchLatestSnapshot.ts

@@ -35,6 +35,7 @@ import {
 	createCacheSnapshotKey,
 	createOdspLogger,
 	getOdspResolvedUrl,
+	snapshotWithLoadingGroupIdSupported,
 	toInstrumentedOdspStorageTokenFetcher,
 	type TokenFetchOptionsEx,
 } from "./odspUtils.js";
@@ -75,9 +76,7 @@ export async function prefetchLatestSnapshot(
 ): Promise<boolean> {
 	const mc = createChildMonitoringContext({ logger, namespace: "PrefetchSnapshot" });
 	const odspLogger = createOdspLogger(mc.logger);
-	const useGroupIdsForSnapshotFetch = mc.config.getBoolean(
-		"Fluid.Container.UseLoadingGroupIdForSnapshotFetch2",
-	);
+	const useGroupIdsForSnapshotFetch = snapshotWithLoadingGroupIdSupported(mc.config);
 	// For prefetch, we just want to fetch the ungrouped data and want to use the new API if the
 	// feature gate is set, so provide an empty array.
 	const loadingGroupIds = useGroupIdsForSnapshotFetch ? [] : undefined;
@@ -112,7 +111,7 @@ export async function prefetchLatestSnapshot(
 			controller,
 		);
 	};
-	const snapshotKey = createCacheSnapshotKey(odspResolvedUrl);
+	const snapshotKey = createCacheSnapshotKey(odspResolvedUrl, useGroupIdsForSnapshotFetch);
 	let cacheP: Promise<void> | undefined;
 	let snapshotEpoch: string | undefined;
 	const putInCache = async (valueWithEpoch: IVersionedValueWithEpoch): Promise<void> => {

packages/drivers/odsp-driver/src/test/createNewUtilsTests.spec.ts

@@ -165,7 +165,12 @@ describe("Create New Utils Tests", () => {
 			{ "x-fluid-epoch": "epoch1" },
 		);
 		// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
-		const snapshot = await epochTracker.get(createCacheSnapshotKey(odspResolvedUrl));
+		const snapshot = await epochTracker.get(createCacheSnapshotKey(odspResolvedUrl, false));
+		// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
+		const snapshotWithLoadingGroupId = await epochTracker.get(
+			createCacheSnapshotKey(odspResolvedUrl, true),
+		);
+		assert(snapshotWithLoadingGroupId === undefined, "snapshot should not exist");
 		// eslint-disable-next-line @typescript-eslint/no-unsafe-argument
 		test(snapshot);
 		await epochTracker.removeEntries().catch(() => {});
@@ -194,7 +199,12 @@ describe("Create New Utils Tests", () => {
 			{ "x-fluid-epoch": "epoch1" },
 		);
 		// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
-		const snapshot = await epochTracker.get(createCacheSnapshotKey(odspResolvedUrl));
+		const snapshot = await epochTracker.get(createCacheSnapshotKey(odspResolvedUrl, false));
+		// eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
+		const snapshotWithLoadingGroupId = await epochTracker.get(
+			createCacheSnapshotKey(odspResolvedUrl, true),
+		);
+		assert(snapshotWithLoadingGroupId === undefined, "snapshot should not exist");
 		// eslint-disable-next-line @typescript-eslint/no-unsafe-argument
 		test(snapshot);
 		await epochTracker.removeEntries().catch(() => {});

packages/drivers/odsp-driver/src/test/fetchSnapshot.spec.ts

@@ -318,8 +318,12 @@ describe("Tests1 for snapshot fetch", () => {
 		}
 		assert(ungroupedData, "should have asked for ungroupedData");
 		const cachedValue = (await epochTracker.get(
-			createCacheSnapshotKey(resolved),
+			createCacheSnapshotKey(resolved, false),
 		)) as ISnapshot;
+		const cachedValueWithLoadingGroupId = (await epochTracker.get(
+			createCacheSnapshotKey(resolved, true),
+		)) as ISnapshot;
+		assert(cachedValueWithLoadingGroupId === undefined, "snapshot should not exist");
 		assert(cachedValue.snapshotTree.id === "SnapshotId", "snapshot should have been cached");
 		assert(service["blobCache"].value.size > 0, "blobs should be cached locally");
 		assert(service["commitCache"].size > 0, "no trees should be cached");
@@ -439,7 +443,7 @@ describe("Tests1 for snapshot fetch", () => {
 			assert.fail("the getSnapshot request should succeed");
 		}
 		const cachedValue = (await epochTracker.get(
-			createCacheSnapshotKey(resolved),
+			createCacheSnapshotKey(resolved, false),
 		)) as ISnapshot;
 		assert(cachedValue.snapshotTree.id === "SnapshotId", "snapshot should have been cached");
 		assert(service["blobCache"].value.size > 0, "blobs should still be cached locally");

packages/drivers/odsp-driver/src/test/prefetchSnapshotTests.spec.ts

@@ -167,7 +167,7 @@ describe("Tests for prefetching snapshot", () => {
 				localCache,
 				GetHostStoragePolicyInternal(),
 			);
-			snapshotPrefetchCacheKey = getKeyForCacheEntry(createCacheSnapshotKey(resolved));
+			snapshotPrefetchCacheKey = getKeyForCacheEntry(createCacheSnapshotKey(resolved, false));
 			const documentservice = await odspDocumentServiceFactory.createDocumentService(
 				resolved,
 				mockLogger,
@@ -498,7 +498,7 @@ describe("Tests for prefetching snapshot", () => {
 				localCache,
 				hostPolicy,
 			);
-			snapshotPrefetchCacheKey = getKeyForCacheEntry(createCacheSnapshotKey(resolved));
+			snapshotPrefetchCacheKey = getKeyForCacheEntry(createCacheSnapshotKey(resolved, true));
 			const documentservice = await odspDocumentServiceFactory.createDocumentService(
 				resolved,
 				mockLogger,
@@ -643,7 +643,7 @@ describe("Tests for prefetching snapshot", () => {
 				localCache,
 				GetHostStoragePolicyInternal(),
 			);
-			snapshotPrefetchCacheKey = getKeyForCacheEntry(createCacheSnapshotKey(resolved));
+			snapshotPrefetchCacheKey = getKeyForCacheEntry(createCacheSnapshotKey(resolved, false));
 			const documentservice = await odspDocumentServiceFactory.createDocumentService(
 				resolved,
 				mockLogger,
@@ -885,7 +885,7 @@ describe("Tests for prefetching snapshot", () => {
 				localCache,
 				hostPolicy,
 			);
-			snapshotPrefetchCacheKey = getKeyForCacheEntry(createCacheSnapshotKey(resolved));
+			snapshotPrefetchCacheKey = getKeyForCacheEntry(createCacheSnapshotKey(resolved, false));
 			const documentservice = await odspDocumentServiceFactory.createDocumentService(
 				resolved,
 				mockLogger,