From 8f67f4d574c56cf64c2ed42121207183f290ace9 Mon Sep 17 00:00:00 2001 From: Charis Kyriakou Date: Wed, 28 Feb 2024 15:21:06 +0000 Subject: [PATCH] Add logic to stop an evaluation run --- .../src/model-editor/model-evaluator.ts | 124 +++++++++++++----- .../src/model-editor/modeling-store.ts | 6 + .../model-editor/modelingStoreMock.ts | 6 + .../model-editor/model-evaluator.test.ts | 97 ++++++++++++++ 4 files changed, 202 insertions(+), 31 deletions(-) create mode 100644 extensions/ql-vscode/test/vscode-tests/activated-extension/model-editor/model-evaluator.test.ts diff --git a/extensions/ql-vscode/src/model-editor/model-evaluator.ts b/extensions/ql-vscode/src/model-editor/model-evaluator.ts index 9ad202def1d..70a228f309e 100644 --- a/extensions/ql-vscode/src/model-editor/model-evaluator.ts +++ b/extensions/ql-vscode/src/model-editor/model-evaluator.ts @@ -9,10 +9,22 @@ import type { CodeQLCliServer } from "../codeql-cli/cli"; import type { VariantAnalysisManager } from "../variant-analysis/variant-analysis-manager"; import type { QueryLanguage } from "../common/query-language"; import { resolveCodeScanningQueryPack } from "../variant-analysis/code-scanning-pack"; -import { withProgress } from "../common/vscode/progress"; +import type { ProgressCallback } from "../common/vscode/progress"; +import { + UserCancellationException, + withProgress, +} from "../common/vscode/progress"; import type { VariantAnalysis } from "../variant-analysis/shared/variant-analysis"; +import type { CancellationToken } from "vscode"; +import { CancellationTokenSource } from "vscode"; +import type { QlPackDetails } from "../variant-analysis/ql-pack-details"; export class ModelEvaluator extends DisposableObject { + // Cancellation token source to allow cancelling of the current run + // before a variant analysis has been submitted. Once it has been + // submitted, we use the variant analysis manager's cancellation support. + private cancellationSource: CancellationTokenSource; + public constructor( private readonly logger: BaseLogger, private readonly cliServer: CodeQLCliServer, @@ -28,6 +40,8 @@ export class ModelEvaluator extends DisposableObject { super(); this.registerToModelingEvents(); + + this.cancellationSource = new CancellationTokenSource(); } public async startEvaluation() { @@ -52,29 +66,12 @@ export class ModelEvaluator extends DisposableObject { // Submit variant analysis and monitor progress return withProgress( - async (progress, token) => { - let variantAnalysisId: number | undefined = undefined; - try { - variantAnalysisId = - await this.variantAnalysisManager.runVariantAnalysis( - qlPack, - progress, - token, - ); - } catch (e) { - this.modelingStore.updateModelEvaluationRun(this.dbItem, undefined); - throw e; - } - - if (variantAnalysisId) { - this.monitorVariantAnalysis(variantAnalysisId); - } else { - this.modelingStore.updateModelEvaluationRun(this.dbItem, undefined); - throw new Error( - "Unable to trigger variant analysis for evaluation run", - ); - } - }, + (progress) => + this.runVariantAnalysis( + qlPack, + progress, + this.cancellationSource.token, + ), { title: "Run Variant Analysis", cancellable: true, @@ -83,13 +80,29 @@ export class ModelEvaluator extends DisposableObject { } public async stopEvaluation() { - // For now just update the store. - // This will be fleshed out in the near future. - const evaluationRun: ModelEvaluationRun = { - isPreparing: false, - variantAnalysisId: undefined, - }; - this.modelingStore.updateModelEvaluationRun(this.dbItem, evaluationRun); + const evaluationRun = this.modelingStore.getModelEvaluationRun(this.dbItem); + if (!evaluationRun) { + void this.logger.log("No active evaluation run to stop"); + return; + } + + this.cancellationSource.cancel(); + + if (evaluationRun.variantAnalysisId === undefined) { + // If the variant analysis has not been submitted yet, we can just + // update the store. + this.modelingStore.updateModelEvaluationRun(this.dbItem, { + ...evaluationRun, + isPreparing: false, + }); + } else { + // If the variant analysis has been submitted, we need to cancel it. We + // don't need to update the store here, as the event handler for + // onVariantAnalysisStatusUpdated will do that for us. + await this.variantAnalysisManager.cancelVariantAnalysis( + evaluationRun.variantAnalysisId, + ); + } } private registerToModelingEvents() { @@ -127,6 +140,55 @@ export class ModelEvaluator extends DisposableObject { return undefined; } + private async runVariantAnalysis( + qlPack: QlPackDetails, + progress: ProgressCallback, + token: CancellationToken, + ): Promise { + let result: number | void = undefined; + try { + // Use Promise.race to make sure to stop the variant analysis processing when the + // user has stopped the evaluation run. We can't simply rely on the cancellation token + // because we haven't fully implemented cancellation support for variant analysis. + // Using this approach we make sure that the process is stopped from a user's point + // of view (the notification goes away too). It won't necessarily stop any tasks + // that are not aware of the cancellation token. + result = await Promise.race([ + this.variantAnalysisManager.runVariantAnalysis(qlPack, progress, token), + new Promise((_, reject) => { + token.onCancellationRequested(() => + reject(new UserCancellationException(undefined, true)), + ); + }), + ]); + } catch (e) { + this.modelingStore.updateModelEvaluationRun(this.dbItem, undefined); + if (!(e instanceof UserCancellationException)) { + throw e; + } else { + return; + } + } finally { + // Renew the cancellation token source for the new evaluation run. + // This is necessary because we don't want the next evaluation run + // to start as cancelled. + this.cancellationSource = new CancellationTokenSource(); + } + + // If the result is a number, it means the variant analysis was successfully submitted, + // so we need to update the store and start up the monitor. + if (typeof result === "number") { + this.modelingStore.updateModelEvaluationRun(this.dbItem, { + isPreparing: true, + variantAnalysisId: result, + }); + this.monitorVariantAnalysis(result); + } else { + this.modelingStore.updateModelEvaluationRun(this.dbItem, undefined); + throw new Error("Unable to trigger variant analysis for evaluation run"); + } + } + private monitorVariantAnalysis(variantAnalysisId: number) { this.push( this.variantAnalysisManager.onVariantAnalysisStatusUpdated( diff --git a/extensions/ql-vscode/src/model-editor/modeling-store.ts b/extensions/ql-vscode/src/model-editor/modeling-store.ts index 1d170891919..097148863d0 100644 --- a/extensions/ql-vscode/src/model-editor/modeling-store.ts +++ b/extensions/ql-vscode/src/model-editor/modeling-store.ts @@ -423,6 +423,12 @@ export class ModelingStore extends DisposableObject { return this.state.get(databaseItem.databaseUri.toString())!; } + public getModelEvaluationRun( + dbItem: DatabaseItem, + ): ModelEvaluationRun | undefined { + return this.getState(dbItem).modelEvaluationRun; + } + private changeMethods( dbItem: DatabaseItem, updateState: (state: InternalDbModelingState) => void, diff --git a/extensions/ql-vscode/test/__mocks__/model-editor/modelingStoreMock.ts b/extensions/ql-vscode/test/__mocks__/model-editor/modelingStoreMock.ts index 4d8c6799156..7cbad68c62c 100644 --- a/extensions/ql-vscode/test/__mocks__/model-editor/modelingStoreMock.ts +++ b/extensions/ql-vscode/test/__mocks__/model-editor/modelingStoreMock.ts @@ -4,12 +4,18 @@ import type { ModelingStore } from "../../../src/model-editor/modeling-store"; export function createMockModelingStore({ initializeStateForDb = jest.fn(), getStateForActiveDb = jest.fn(), + getModelEvaluationRun = jest.fn(), + updateModelEvaluationRun = jest.fn(), }: { initializeStateForDb?: ModelingStore["initializeStateForDb"]; getStateForActiveDb?: ModelingStore["getStateForActiveDb"]; + getModelEvaluationRun?: ModelingStore["getModelEvaluationRun"]; + updateModelEvaluationRun?: ModelingStore["updateModelEvaluationRun"]; } = {}): ModelingStore { return mockedObject({ initializeStateForDb, getStateForActiveDb, + getModelEvaluationRun, + updateModelEvaluationRun, }); } diff --git a/extensions/ql-vscode/test/vscode-tests/activated-extension/model-editor/model-evaluator.test.ts b/extensions/ql-vscode/test/vscode-tests/activated-extension/model-editor/model-evaluator.test.ts new file mode 100644 index 00000000000..e5a15d0500b --- /dev/null +++ b/extensions/ql-vscode/test/vscode-tests/activated-extension/model-editor/model-evaluator.test.ts @@ -0,0 +1,97 @@ +import type { CodeQLCliServer } from "../../../../src/codeql-cli/cli"; +import type { BaseLogger } from "../../../../src/common/logging"; +import { QueryLanguage } from "../../../../src/common/query-language"; +import type { DatabaseItem } from "../../../../src/databases/local-databases"; +import type { ModelEvaluationRun } from "../../../../src/model-editor/model-evaluation-run"; +import { ModelEvaluator } from "../../../../src/model-editor/model-evaluator"; +import type { ModelingEvents } from "../../../../src/model-editor/modeling-events"; +import type { ModelingStore } from "../../../../src/model-editor/modeling-store"; +import type { VariantAnalysisManager } from "../../../../src/variant-analysis/variant-analysis-manager"; +import { createMockLogger } from "../../../__mocks__/loggerMock"; +import { createMockModelingEvents } from "../../../__mocks__/model-editor/modelingEventsMock"; +import { createMockModelingStore } from "../../../__mocks__/model-editor/modelingStoreMock"; +import { mockedObject } from "../../../mocked-object"; + +describe("Model Evaluator", () => { + let modelEvaluator: ModelEvaluator; + let logger: BaseLogger; + let cliServer: CodeQLCliServer; + let modelingStore: ModelingStore; + let modelingEvents: ModelingEvents; + let variantAnalysisManager: VariantAnalysisManager; + let dbItem: DatabaseItem; + let language: QueryLanguage; + let updateView: jest.Mock; + let getModelEvaluationRunMock = jest.fn(); + + beforeEach(() => { + logger = createMockLogger(); + cliServer = mockedObject({}); + getModelEvaluationRunMock = jest.fn(); + modelingStore = createMockModelingStore({ + getModelEvaluationRun: getModelEvaluationRunMock, + }); + modelingEvents = createMockModelingEvents(); + variantAnalysisManager = mockedObject({ + cancelVariantAnalysis: jest.fn(), + }); + dbItem = mockedObject({}); + language = QueryLanguage.Java; + updateView = jest.fn(); + + modelEvaluator = new ModelEvaluator( + logger, + cliServer, + modelingStore, + modelingEvents, + variantAnalysisManager, + dbItem, + language, + updateView, + ); + }); + + describe("stopping evaluation", () => { + it("should just log a message if it never started", async () => { + getModelEvaluationRunMock.mockReturnValue(undefined); + + await modelEvaluator.stopEvaluation(); + + expect(logger.log).toHaveBeenCalledWith( + "No active evaluation run to stop", + ); + }); + + it("should update the store if evaluation run exists", async () => { + getModelEvaluationRunMock.mockReturnValue({ + isPreparing: true, + variantAnalysisId: undefined, + }); + + await modelEvaluator.stopEvaluation(); + + expect(modelingStore.updateModelEvaluationRun).toHaveBeenCalledWith( + dbItem, + { + isPreparing: false, + varianAnalysis: undefined, + }, + ); + }); + + it("should cancel the variant analysis if one has been started", async () => { + const evaluationRun: ModelEvaluationRun = { + isPreparing: false, + variantAnalysisId: 123, + }; + getModelEvaluationRunMock.mockReturnValue(evaluationRun); + + await modelEvaluator.stopEvaluation(); + + expect(modelingStore.updateModelEvaluationRun).not.toHaveBeenCalled(); + expect(variantAnalysisManager.cancelVariantAnalysis).toHaveBeenCalledWith( + evaluationRun.variantAnalysisId, + ); + }); + }); +});