Support abort signal set-up in coq-lsp client wrappers

Pass abort signal into coq-lsp client correctly in benchmarks

src/benchmark/framework/benchmarkingCore/singleCompletionGeneration/benchmarkSingleCompletionGeneration.ts

@@ -162,7 +162,8 @@ export async function benchmarkSingleCompletionGeneration<
             generationArgs.completionContext,
             generationArgs.sourceFileEnvironment,
             generationArgs.workspaceRoot,
-            logger
+            logger,
+            abortSignal
         );
     } catch (error) {
         if (error instanceof ProofsCheckFailedError) {

src/benchmark/framework/benchmarkingCore/singleCompletionGeneration/proofsCheckers/abstractProofsChecker.ts

@@ -29,6 +29,7 @@ export abstract class AbstractProofsChecker {
         completionContext: CompletionContext,
         sourceFileEnvironment: SourceFileEnvironment,
         workspaceRoot: WorkspaceRoot,
-        logger: BenchmarkingLogger
+        logger: BenchmarkingLogger,
+        abortSignal?: AbortSignal
     ): Promise<ProofsCheckResult>;
 }

src/benchmark/framework/benchmarkingCore/singleCompletionGeneration/proofsCheckers/implementation/checkProofs.ts

@@ -28,15 +28,19 @@ export namespace CheckProofsImpl {
 
     export async function checkProofsMeasured(
         args: Signature.Args,
-        providedLogger: ProvidedLogger
+        providedLogger: ProvidedLogger,
+        abortSignal?: AbortSignal
     ): Promise<Signature.Result> {
         const fileUri = deserializeUri(args.serializedFileUri);
         const timeMark = new TimeMark();
 
         try {
             const proofCheckResults = await withDocumentOpenedByTestCoqLsp(
                 { uri: fileUri, version: args.documentVersion },
-                args.workspaceRootPath,
+                {
+                    workspaceRootPath: args.workspaceRootPath,
+                    abortSignal: abortSignal,
+                },
                 (coqLspClient) =>
                     new CoqProofChecker(coqLspClient).checkProofs(
                         fileUri,

src/benchmark/framework/benchmarkingCore/singleCompletionGeneration/proofsCheckers/localProofsChecker.ts

@@ -23,7 +23,8 @@ export class LocalProofsChecker extends AbstractProofsChecker {
         completionContext: CompletionContext,
         sourceFileEnvironment: SourceFileEnvironment,
         workspaceRoot: WorkspaceRoot,
-        logger: BenchmarkingLogger
+        logger: BenchmarkingLogger,
+        abortSignal: AbortSignal
     ): Promise<ProofsCheckResult> {
         const args = ProofsCheckerUtils.buildArgs(
             preparedProofs,
@@ -36,7 +37,8 @@ export class LocalProofsChecker extends AbstractProofsChecker {
             {
                 logger: logger,
                 logSuccess: false,
-            }
+            },
+            abortSignal
         );
         return ProofsCheckerUtils.unpackSuccessResultOrThrow(proofsCheckResult);
     }

src/benchmark/framework/parseDataset/coqProjectParser/implementation/parseCoqProject.ts

@@ -44,8 +44,10 @@ export namespace ParseCoqProjectImpl {
     ): Promise<Signature.ResultModels.Result> {
         const parsedWorkspace: Signature.ResultModels.Result = {};
 
+        // Note: specific abort controller is not passed here, since
+        // the abort behaviour is not supported (and not needed) at the parsing stage.
         await withTestCoqLspClient(
-            args.workspaceRootPath,
+            { workspaceRootPath: args.workspaceRootPath },
             async (coqLspClient) => {
                 for (const filePath in args.workspaceTargets) {
                     parsedWorkspace[filePath] =
@@ -99,6 +101,7 @@ export namespace ParseCoqProjectImpl {
         const mockDocumentVersion = 1;
         // TODO: [@Gleb Solovev] Do not create this Abort Controller but pass
         // the one created at the top
+        // TODO + check coq-lsp creation in benchmarks
         const abortController = new AbortController();
         const sourceFileEnvironment = await createSourceFileEnvironment(
             mockDocumentVersion,

src/benchmark/framework/utils/subprocessUtils/ipc/onParentProcessCallExecutor/executeOnParentProcessCall.ts

@@ -24,7 +24,8 @@ import { OnParentProcessCallExecutorUtils } from "./utils";
 import Utils = OnParentProcessCallExecutorUtils;
 
 // TODO: document
-// TODO: design better logging through actual file
+// TODO: design better logging through actual
+// TODO: support abort behaviour
 export async function executeAsFunctionOnParentProcessCall<
     ArgsType,
     ResultType,

src/coqLsp/coqLspBuilders.ts

@@ -14,32 +14,40 @@ export async function createCoqLspClient(
     coqLspServerPath: string,
     logOutputChannel?: OutputChannel,
     eventLogger?: EventLogger,
-    abortController?: AbortController
+    abortSignal?: AbortSignal
 ): Promise<CoqLspClient> {
     return createAbstractCoqLspClient(
         CoqLspConfig.createClientConfig(coqLspServerPath),
         logOutputChannel,
         eventLogger,
-        abortController
+        abortSignal
     );
 }
 
+export interface TestCoqLspClientOptions {
+    workspaceRootPath?: string;
+    abortSignal?: AbortSignal;
+}
+
 export async function createTestCoqLspClient(
-    workspaceRootPath?: string
+    options: TestCoqLspClientOptions = {}
 ): Promise<CoqLspClient> {
     return createAbstractCoqLspClient(
         CoqLspConfig.createClientConfig(
             process.env.COQ_LSP_PATH || "coq-lsp",
-            workspaceRootPath
-        )
+            options.workspaceRootPath
+        ),
+        undefined,
+        undefined,
+        options.abortSignal
     );
 }
 
 export async function withTestCoqLspClient<T>(
-    workspaceRootPath: string | undefined,
+    options: TestCoqLspClientOptions,
     block: (coqLspClient: CoqLspClient) => Promise<T>
 ) {
-    const coqLspClient = await createTestCoqLspClient(workspaceRootPath);
+    const coqLspClient = await createTestCoqLspClient(options);
     try {
         return await block(coqLspClient);
     } finally {
@@ -49,13 +57,13 @@ export async function withTestCoqLspClient<T>(
 
 export async function withDocumentOpenedByTestCoqLsp<T>(
     documentSpec: DocumentSpec,
-    workspaceRootPath: string | undefined,
+    options: TestCoqLspClientOptions,
     block: (
         coqLspClient: CoqLspClient,
         openedDocDiagnostic: DiagnosticMessage
     ) => Promise<T>
 ): Promise<T> {
-    return withTestCoqLspClient(workspaceRootPath, (coqLspClient) =>
+    return withTestCoqLspClient(options, (coqLspClient) =>
         coqLspClient.withTextDocument(documentSpec, (openedDocDiagnostic) =>
             block(coqLspClient, openedDocDiagnostic)
         )
@@ -68,14 +76,14 @@ async function createAbstractCoqLspClient(
         "CoqPilot: coq-lsp events"
     ),
     eventLogger?: EventLogger,
-    abortController?: AbortController
+    abortSignal?: AbortSignal
 ): Promise<CoqLspClient> {
     const coqLspServerConfig = CoqLspConfig.createServerConfig();
     return await CoqLspClientImpl.create(
         coqLspServerConfig,
         coqLspClientConfig,
         logOutputChannel,
         eventLogger,
-        abortController
+        abortSignal
     );
 }

src/coqLsp/coqLspClient.ts

@@ -108,7 +108,7 @@ export class CoqLspClientImpl implements CoqLspClient {
     private constructor(
         coqLspConnector: CoqLspConnector,
         public readonly eventLogger?: EventLogger,
-        public readonly abortController?: AbortController
+        public readonly abortSignal?: AbortSignal
     ) {
         this.client = coqLspConnector;
     }
@@ -118,7 +118,7 @@ export class CoqLspClientImpl implements CoqLspClient {
         clientConfig: CoqLspClientConfig,
         logOutputChannel: OutputChannel,
         eventLogger?: EventLogger,
-        abortController?: AbortController
+        abortSignal?: AbortSignal
     ): Promise<CoqLspClientImpl> {
         const connector = new CoqLspConnector(
             serverConfig,
@@ -131,7 +131,7 @@ export class CoqLspClientImpl implements CoqLspClient {
                 clientConfig.coq_lsp_server_path
             );
         });
-        return new CoqLspClientImpl(connector, eventLogger, abortController);
+        return new CoqLspClientImpl(connector, eventLogger, abortSignal);
     }
 
     async getGoalsAtPoint(
@@ -141,7 +141,7 @@ export class CoqLspClientImpl implements CoqLspClient {
         command?: string
     ): Promise<Result<ProofGoal[], Error>> {
         return await this.mutex.runExclusive(async () => {
-            throwOnAbort(this.abortController?.signal);
+            throwOnAbort(this.abortSignal);
             return this.getGoalsAtPointUnsafe(
                 position,
                 documentUri,
@@ -158,7 +158,7 @@ export class CoqLspClientImpl implements CoqLspClient {
         command?: string
     ): Promise<ProofGoal> {
         return await this.mutex.runExclusive(async () => {
-            throwOnAbort(this.abortController?.signal);
+            throwOnAbort(this.abortSignal);
             const goals = await this.getGoalsAtPointUnsafe(
                 position,
                 documentUri,
@@ -181,14 +181,14 @@ export class CoqLspClientImpl implements CoqLspClient {
         version: number = 1
     ): Promise<DiagnosticMessage> {
         return await this.mutex.runExclusive(async () => {
-            throwOnAbort(this.abortController?.signal);
+            throwOnAbort(this.abortSignal);
             return this.openTextDocumentUnsafe(uri, version);
         });
     }
 
     async closeTextDocument(uri: Uri): Promise<void> {
         return await this.mutex.runExclusive(async () => {
-            throwOnAbort(this.abortController?.signal);
+            throwOnAbort(this.abortSignal);
             return this.closeTextDocumentUnsafe(uri);
         });
     }
@@ -212,7 +212,7 @@ export class CoqLspClientImpl implements CoqLspClient {
 
     async getFlecheDocument(uri: Uri): Promise<FlecheDocument> {
         return await this.mutex.runExclusive(async () => {
-            throwOnAbort(this.abortController?.signal);
+            throwOnAbort(this.abortSignal);
             return this.getFlecheDocumentUnsafe(uri);
         });
     }

src/extension/sessionState.ts

@@ -69,7 +69,7 @@ export class SessionState implements Disposable {
             coqLspServerPath,
             logOutputChannel,
             eventLogger,
-            abortController
+            abortController.signal
         );
 
         pluginStatusIndicator.updateStatusBar(true);
@@ -126,7 +126,7 @@ export class SessionState implements Disposable {
             coqLspServerPath,
             this.logOutputChannel,
             this.eventLogger,
-            this._abortController
+            this._abortController.signal
         );
 
         this.pluginStatusIndicator.updateStatusBar(this._isActive);

src/test/commonTestFunctions/coqFileParser.ts

@@ -18,7 +18,7 @@ export async function parseTheoremsFromCoqFile(
 
     return await withDocumentOpenedByTestCoqLsp(
         { uri: fileUri },
-        rootDir,
+        { workspaceRootPath: rootDir },
         (coqLspClient) =>
             parseCoqFile(fileUri, coqLspClient, new AbortController().signal)
     );

src/test/commonTestFunctions/prepareEnvironment.ts

@@ -36,7 +36,7 @@ export async function withPreparedEnvironment<T>(
     );
     const fileUri = Uri.fromPath(filePath);
 
-    const client = await createTestCoqLspClient(rootDir);
+    const client = await createTestCoqLspClient({ workspaceRootPath: rootDir });
     const coqProofChecker = new CoqProofChecker(client);
     try {
         const [completionContexts, sourceFileEnvironment] =

src/test/coqLsp/coqLspGetGoals.test.ts

@@ -21,7 +21,7 @@ suite("Retrieve goals from Coq file", () => {
 
         return withDocumentOpenedByTestCoqLsp(
             { uri: fileUri },
-            rootDir,
+            { workspaceRootPath: rootDir },
             (coqLspClient) =>
                 Promise.all(
                     points.map(async (point) => {

src/test/coqLsp/coqLspGetGoalsWithCommandApplied.test.ts

@@ -22,7 +22,7 @@ suite("Request goals with `command/pretac` argument", () => {
 
         return withDocumentOpenedByTestCoqLsp(
             { uri: fileUri },
-            rootDir,
+            { workspaceRootPath: rootDir },
             (coqLspClient) =>
                 coqLspClient.getGoalsAtPoint(position, fileUri, 1, command)
         );

src/test/core/coqProofChecker.test.ts

@@ -24,7 +24,7 @@ suite("`CoqProofChecker` tests", () => {
 
         return withDocumentOpenedByTestCoqLsp(
             { uri: fileUri, version: documentVersion },
-            rootDir,
+            { workspaceRootPath: rootDir },
             (coqLspClient) => {
                 const coqProofChecker = new CoqProofChecker(coqLspClient);
                 const preparedProofs = proofsToCheck.map((proofs) =>

src/test/legacyBenchmark/benchmarkingFramework.ts

@@ -82,16 +82,26 @@ export async function runTestBenchmark(
         resolvedOptions.filePath,
         resolvedOptions.additionalImports
     );
+    /**
+     * Note: so far the abort signal is never triggered;
+     * however, such behaviour can be supported:
+     * the same `AbortController` object is passed throughout the run properly.
+     */
+    const abortController = new AbortController();
 
     return withDocumentOpenedByTestCoqLsp(
         { uri: fileUri },
-        inputOptions.workspaceRootPath,
+        {
+            workspaceRootPath: inputOptions.workspaceRootPath,
+            abortSignal: abortController.signal,
+        },
         (coqLspClient) =>
             runTestBenchmarkOnPreparedFile(
                 resolvedOptions,
                 coqLspClient,
                 fileUri,
-                isNewlyCreatedFile
+                isNewlyCreatedFile,
+                abortController
             )
     );
 }
@@ -116,7 +126,8 @@ export async function runTestBenchmarkOnPreparedFile(
     options: TestBenchmarkOptions,
     coqLspClient: CoqLspClient,
     fileUri: Uri,
-    isNewlyCreatedFile: boolean
+    isNewlyCreatedFile: boolean,
+    abortController: AbortController
 ): Promise<BenchmarkReport> {
     consoleLog(`run benchmarks for file: ${options.filePath}\n`, "blue");
     const shouldCompleteHole = (_hole: ProofStep) => true;
@@ -160,6 +171,7 @@ export async function runTestBenchmarkOnPreparedFile(
             getSingleModelId(options.inputModelsParams),
             options.relativePathToFile,
             options.groupName,
+            abortController,
             eventLogger,
             options.maxPremisesNumber,
             options.reportHolder,
@@ -184,6 +196,7 @@ export async function runTestBenchmarkOnPreparedFile(
             getSingleModelId(options.inputModelsParams),
             options.relativePathToFile,
             options.groupName,
+            abortController,
             eventLogger,
             options.maxPremisesNumber,
             options.reportHolder,
@@ -260,6 +273,7 @@ export async function benchmarkTargets(
     modelId: string,
     checkedFilePath: string,
     groupName: string,
+    abortController: AbortController,
     eventLogger: EventLogger,
     maxPremisesNumber?: number,
     reportHolder?: BenchmarkReportHolder,
@@ -275,6 +289,7 @@ export async function benchmarkTargets(
             modelId,
             checkedFilePath,
             groupName,
+            abortController,
             eventLogger,
             maxPremisesNumber,
             reportHolder,
@@ -297,6 +312,7 @@ async function benchmarkCompletionGeneration(
     modelId: string,
     checkedFilePath: string,
     groupName: string,
+    abortController: AbortController,
     eventLogger: EventLogger,
     maxPremisesNumber?: number,
     reportHolder?: BenchmarkReportHolder,
@@ -331,7 +347,6 @@ async function benchmarkCompletionGeneration(
         premisesNumber: maxPremisesNumber,
     };
 
-    const abortController = new AbortController();
     const result = await generateCompletion(
         completionContext,
         sourceFileEnvironmentWithFilteredContext,