From b646f9f4cb1afc637b2503e774e3b0da60309979 Mon Sep 17 00:00:00 2001 From: FAIZ CHACHIYA Date: Fri, 18 Aug 2023 15:40:35 +0530 Subject: [PATCH] Wmt priority execution feature (#1546) * Execute the queries with high/low priority * improvement made to the Cosmos client and created separate plugin class for setting priority header * added test cases for the priority execution utility * removed unwanted code * fix compile time issues * fix compile time issues * fixed lint and stylinkg issues * fixed lint and styling issues * skip the lint check for src/Utils/PriorityBasedExecutionUtils.ts * incorporating review comments, added the default priority level changes * changed the priority to default instead of low * removed the unwanted if condition --------- Co-authored-by: Faiz Chachiya --- .eslintignore | 3 +- src/Common/Constants.ts | 6 ++ src/Common/CosmosClient.ts | 10 ++++ src/Explorer/Explorer.tsx | 2 + .../Panes/SettingsPane/SettingsPane.tsx | 45 +++++++++++++- src/Platform/Hosted/extractFeatures.ts | 2 + src/Shared/ExplorerSettings.ts | 7 +++ src/Shared/StorageUtility.ts | 1 + src/Utils/PriorityBasedExecutionUtils.test.ts | 59 +++++++++++++++++++ src/Utils/PriorityBasedExecutionUtils.ts | 29 +++++++++ 10 files changed, 162 insertions(+), 2 deletions(-) create mode 100644 src/Utils/PriorityBasedExecutionUtils.test.ts create mode 100644 src/Utils/PriorityBasedExecutionUtils.ts diff --git a/.eslintignore b/.eslintignore index c1e8c4948..c188701b9 100644 --- a/.eslintignore +++ b/.eslintignore @@ -144,4 +144,5 @@ src/Explorer/Notebook/NotebookRenderer/decorators/kbd-shortcuts/index.tsx src/Explorer/Notebook/temp/inputs/connected-editors/codemirror.tsx src/Explorer/Tree/ResourceTreeAdapter.tsx __mocks__/monaco-editor.ts -src/Explorer/Tree/ResourceTree.tsx \ No newline at end of file +src/Explorer/Tree/ResourceTree.tsx +src/Utils/PriorityBasedExecutionUtils.ts \ No newline at end of file diff --git a/src/Common/Constants.ts b/src/Common/Constants.ts index 352378a70..884209fed 100644 --- a/src/Common/Constants.ts +++ b/src/Common/Constants.ts @@ -600,3 +600,9 @@ export const QueryCopilotSampleContainerSchema = { }, }, }; + +export class PriorityLevel { + public static readonly High = "high"; + public static readonly Low = "low"; + public static readonly Default = "low"; +} diff --git a/src/Common/CosmosClient.ts b/src/Common/CosmosClient.ts index f973df855..793213fa8 100644 --- a/src/Common/CosmosClient.ts +++ b/src/Common/CosmosClient.ts @@ -4,6 +4,9 @@ import { userContext } from "../UserContext"; import { logConsoleError } from "../Utils/NotificationConsoleUtils"; import { EmulatorMasterKey, HttpHeaders } from "./Constants"; import { getErrorMessage } from "./ErrorHandlingUtils"; +import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility"; +import { PriorityLevel } from "../Common/Constants"; +import * as PriorityBasedExecutionUtils from "../Utils/PriorityBasedExecutionUtils"; const _global = typeof self === "undefined" ? window : self; @@ -105,6 +108,13 @@ export function client(): Cosmos.CosmosClient { if (configContext.PROXY_PATH !== undefined) { (options as any).plugins = [{ on: "request", plugin: requestPlugin }]; } + + if (userContext.features.enablePriorityBasedThrottling && userContext.apiType === "SQL") { + const plugins = (options as any).plugins || []; + plugins.push({ on: "request", plugin: PriorityBasedExecutionUtils.requestPlugin }); + (options as any).plugins = plugins; + } + _client = new Cosmos.CosmosClient(options); return _client; } diff --git a/src/Explorer/Explorer.tsx b/src/Explorer/Explorer.tsx index 59d8847ef..19811a437 100644 --- a/src/Explorer/Explorer.tsx +++ b/src/Explorer/Explorer.tsx @@ -650,6 +650,8 @@ export default class Explorer { private _initSettings() { if (!ExplorerSettings.hasSettingsDefined()) { ExplorerSettings.createDefaultSettings(); + } else { + ExplorerSettings.ensurePriorityLevel(); } } diff --git a/src/Explorer/Panes/SettingsPane/SettingsPane.tsx b/src/Explorer/Panes/SettingsPane/SettingsPane.tsx index 5120d595c..64023f4c0 100644 --- a/src/Explorer/Panes/SettingsPane/SettingsPane.tsx +++ b/src/Explorer/Panes/SettingsPane/SettingsPane.tsx @@ -41,12 +41,18 @@ export const SettingsPane: FunctionComponent = () => { ? LocalStorageUtility.getEntryNumber(StorageKey.MaxDegreeOfParellism) : Constants.Queries.DefaultMaxDegreeOfParallelism ); + const [priorityLevel, setPriorityLevel] = useState( + LocalStorageUtility.hasItem(StorageKey.PriorityLevel) + ? LocalStorageUtility.getEntryString(StorageKey.PriorityLevel) + : Constants.PriorityLevel.Default + ); const explorerVersion = configContext.gitSha; const shouldShowQueryPageOptions = userContext.apiType === "SQL"; const shouldShowGraphAutoVizOption = userContext.apiType === "Gremlin"; const shouldShowCrossPartitionOption = userContext.apiType !== "Gremlin"; const shouldShowParallelismOption = userContext.apiType !== "Gremlin"; - + const shouldShowPriorityLevelOption = + userContext.features.enablePriorityBasedThrottling && userContext.apiType === "SQL"; const handlerOnSubmit = (e: MouseEvent) => { setIsExecuting(true); @@ -58,6 +64,7 @@ export const SettingsPane: FunctionComponent = () => { LocalStorageUtility.setEntryString(StorageKey.ContainerPaginationEnabled, containerPaginationEnabled.toString()); LocalStorageUtility.setEntryString(StorageKey.IsCrossPartitionQueryEnabled, crossPartitionQueryEnabled.toString()); LocalStorageUtility.setEntryNumber(StorageKey.MaxDegreeOfParellism, maxDegreeOfParallelism); + LocalStorageUtility.setEntryString(StorageKey.PriorityLevel, priorityLevel.toString()); if (shouldShowGraphAutoVizOption) { LocalStorageUtility.setEntryBoolean( @@ -76,6 +83,7 @@ export const SettingsPane: FunctionComponent = () => { StorageKey.MaxDegreeOfParellism )}` ); + logConsoleInfo(`Updated priority level setting to ${LocalStorageUtility.getEntryString(StorageKey.PriorityLevel)}`); if (shouldShowGraphAutoVizOption) { logConsoleInfo( @@ -116,6 +124,18 @@ export const SettingsPane: FunctionComponent = () => { { key: "true", text: "JSON" }, ]; + const priorityLevelOptionList: IChoiceGroupOption[] = [ + { key: Constants.PriorityLevel.Low, text: "Low" }, + { key: Constants.PriorityLevel.High, text: "High" }, + ]; + + const handleOnPriorityLevelOptionChange = ( + ev: React.FormEvent, + option: IChoiceGroupOption + ): void => { + setPriorityLevel(option.key); + }; + const handleOnPageOptionChange = (ev: React.FormEvent, option: IChoiceGroupOption): void => { setPageOption(option.key); }; @@ -260,6 +280,29 @@ export const SettingsPane: FunctionComponent = () => { )} + {shouldShowPriorityLevelOption && ( +
+
+
+ + Priority Level + + + Sets the priority level for data-plane requests from Data Explorer when using Priority-Based + Execution. If "None" is selected, Data Explorer will not specify priority level, and the + server-side default priority level will be used. + + +
+
+
+ )} {shouldShowGraphAutoVizOption && (
diff --git a/src/Platform/Hosted/extractFeatures.ts b/src/Platform/Hosted/extractFeatures.ts index df8136f42..f064cbfba 100644 --- a/src/Platform/Hosted/extractFeatures.ts +++ b/src/Platform/Hosted/extractFeatures.ts @@ -36,6 +36,7 @@ export type Features = { readonly enableLegacyMongoShellV2Debug: boolean; readonly loadLegacyMongoShellFromBE: boolean; readonly enableCopilot: boolean; + readonly enablePriorityBasedThrottling: boolean; readonly enableNPSSurvey: boolean; readonly copilotVersion?: string; @@ -106,6 +107,7 @@ export function extractFeatures(given = new URLSearchParams(window.location.sear enableLegacyMongoShellV2Debug: "true" === get("enablelegacymongoshellv2debug"), loadLegacyMongoShellFromBE: "true" === get("loadlegacymongoshellfrombe"), enableCopilot: "true" === get("enablecopilot"), + enablePriorityBasedThrottling: "true" === get("enableprioritybasedthrottling"), enableNPSSurvey: "true" === get("enablenpssurvey"), copilotVersion: get("copilotVersion") ? get("copilotVersion") : "v1.0", }; diff --git a/src/Shared/ExplorerSettings.ts b/src/Shared/ExplorerSettings.ts index 3c6f9e3d2..2c5691b8c 100644 --- a/src/Shared/ExplorerSettings.ts +++ b/src/Shared/ExplorerSettings.ts @@ -6,6 +6,7 @@ export const createDefaultSettings = () => { LocalStorageUtility.setEntryNumber(StorageKey.CustomItemPerPage, Constants.Queries.itemsPerPage); LocalStorageUtility.setEntryString(StorageKey.IsCrossPartitionQueryEnabled, "true"); LocalStorageUtility.setEntryNumber(StorageKey.MaxDegreeOfParellism, Constants.Queries.DefaultMaxDegreeOfParallelism); + LocalStorageUtility.setEntryString(StorageKey.PriorityLevel, Constants.PriorityLevel.Default); }; export const hasSettingsDefined = (): boolean => { @@ -15,3 +16,9 @@ export const hasSettingsDefined = (): boolean => { LocalStorageUtility.hasItem(StorageKey.MaxDegreeOfParellism) ); }; + +export const ensurePriorityLevel = () => { + if (!LocalStorageUtility.hasItem(StorageKey.PriorityLevel)) { + LocalStorageUtility.setEntryString(StorageKey.PriorityLevel, Constants.PriorityLevel.Default); + } +}; diff --git a/src/Shared/StorageUtility.ts b/src/Shared/StorageUtility.ts index ad8a932e1..d142c07e4 100644 --- a/src/Shared/StorageUtility.ts +++ b/src/Shared/StorageUtility.ts @@ -16,4 +16,5 @@ export enum StorageKey { SetPartitionKeyUndefined, GalleryCalloutDismissed, VisitedAccounts, + PriorityLevel, } diff --git a/src/Utils/PriorityBasedExecutionUtils.test.ts b/src/Utils/PriorityBasedExecutionUtils.test.ts new file mode 100644 index 000000000..40c9881fc --- /dev/null +++ b/src/Utils/PriorityBasedExecutionUtils.test.ts @@ -0,0 +1,59 @@ +import * as PriorityBasedExecutionUtils from "./PriorityBasedExecutionUtils"; +import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility"; +import { PriorityLevel } from "../Common/Constants"; +import * as Cosmos from "@azure/cosmos"; + +describe("Priority execution utility", () => { + it("check default priority level is Low", () => { + expect(PriorityBasedExecutionUtils.getPriorityLevel()).toEqual(PriorityLevel.Low); + }); + + it("check the priority level is returned as present in local storage", () => { + LocalStorageUtility.setEntryString(StorageKey.PriorityLevel, PriorityLevel.High); + expect(PriorityBasedExecutionUtils.getPriorityLevel()).toEqual(PriorityLevel.High); + }); + + it("check relevant request based on different resource types", () => { + const requestContext1: Cosmos.RequestContext = { + resourceType: Cosmos.ResourceType.item, + globalEndpointManager: null, + connectionPolicy: null, + requestAgent: null, + method: null, + options: null, + plugins: null, + }; + const requestContext2: Cosmos.RequestContext = { + resourceType: Cosmos.ResourceType.conflicts, + globalEndpointManager: null, + connectionPolicy: null, + requestAgent: null, + method: null, + options: null, + plugins: null, + }; + const requestContext3: Cosmos.RequestContext = { + resourceType: Cosmos.ResourceType.sproc, + operationType: Cosmos.OperationType.Execute, + globalEndpointManager: null, + connectionPolicy: null, + requestAgent: null, + method: null, + options: null, + plugins: null, + }; + const requestContext4: Cosmos.RequestContext = { + resourceType: Cosmos.ResourceType.database, + globalEndpointManager: null, + connectionPolicy: null, + requestAgent: null, + method: null, + options: null, + plugins: null, + }; + expect(PriorityBasedExecutionUtils.isRelevantRequest(requestContext1)).toEqual(true); + expect(PriorityBasedExecutionUtils.isRelevantRequest(requestContext2)).toEqual(true); + expect(PriorityBasedExecutionUtils.isRelevantRequest(requestContext3)).toEqual(true); + expect(PriorityBasedExecutionUtils.isRelevantRequest(requestContext4)).toEqual(false); + }); +}); diff --git a/src/Utils/PriorityBasedExecutionUtils.ts b/src/Utils/PriorityBasedExecutionUtils.ts new file mode 100644 index 000000000..08c272574 --- /dev/null +++ b/src/Utils/PriorityBasedExecutionUtils.ts @@ -0,0 +1,29 @@ +import * as Cosmos from "@azure/cosmos"; +import { LocalStorageUtility, StorageKey } from "Shared/StorageUtility"; +import { PriorityLevel } from "../Common/Constants"; + +export function isRelevantRequest(requestContext: Cosmos.RequestContext): boolean { + return ( + requestContext.resourceType === Cosmos.ResourceType.item || + requestContext.resourceType === Cosmos.ResourceType.conflicts || + (requestContext.resourceType === Cosmos.ResourceType.sproc && + requestContext.operationType === Cosmos.OperationType.Execute) + ); +} + +export function getPriorityLevel(): PriorityLevel { + const priorityLevel = LocalStorageUtility.getEntryString(StorageKey.PriorityLevel); + if (priorityLevel && Object.values(PriorityLevel).includes(priorityLevel)) { + return priorityLevel as PriorityLevel; + } else { + return PriorityLevel.Default; + } +} + +export const requestPlugin: Cosmos.Plugin = async (requestContext, next) => { + if (isRelevantRequest(requestContext)) { + const priorityLevel: PriorityLevel = getPriorityLevel(); + requestContext.headers["x-ms-cosmos-priority-level"] = priorityLevel as string; + } + return next(requestContext); +};