Add logic to stop a model evaluation run

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,30 +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,
-              false,
-            );
-        } 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 model evaluation",
         cancellable: false,
@@ -84,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() {
@@ -128,6 +140,60 @@ export class ModelEvaluator extends DisposableObject {
     return undefined;
   }
 
+  private async runVariantAnalysis(
+    qlPack: QlPackDetails,
+    progress: ProgressCallback,
+    token: CancellationToken,
+  ): Promise<number | void> {
+    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,
+          false,
+        ),
+        new Promise<void>((_, 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(

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,

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<ModelingStore>({
     initializeStateForDb,
     getStateForActiveDb,
+    getModelEvaluationRun,
+    updateModelEvaluationRun,
   });
 }

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<CodeQLCliServer>({});
+    getModelEvaluationRunMock = jest.fn();
+    modelingStore = createMockModelingStore({
+      getModelEvaluationRun: getModelEvaluationRunMock,
+    });
+    modelingEvents = createMockModelingEvents();
+    variantAnalysisManager = mockedObject<VariantAnalysisManager>({
+      cancelVariantAnalysis: jest.fn(),
+    });
+    dbItem = mockedObject<DatabaseItem>({});
+    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,
+      );
+    });
+  });
+});