Break up automodeling into batches

Author: charisk

extensions/ql-vscode/src/data-extensions-editor/auto-model-v2.ts

@@ -8,9 +8,6 @@ import { ExternalApiUsage, MethodSignature } from "./external-api-usage";
 import { ModeledMethod } from "./modeled-method";
 import { groupMethods, sortGroupNames, sortMethods } from "./shared/sorting";
 
-// Soft limit on the number of candidates to send to the model.
-// Note that the model may return fewer than this number of candidates.
-const candidateLimit = 20;
 /**
  * Return the candidates that the model should be run on. This includes limiting the number of
  * candidates to the candidate limit and filtering out anything that is already modeled and respecting
@@ -41,11 +38,6 @@ export function getCandidates(
       type: "none",
     };
 
-    // If we have reached the max number of candidates then stop
-    if (candidates.length >= candidateLimit) {
-      break;
-    }
-
     // Anything that is modeled is not a candidate
     if (modeledMethod.type !== "none") {
       continue;

extensions/ql-vscode/src/data-extensions-editor/auto-modeler.ts

@@ -16,13 +16,21 @@ import { QueryRunner } from "../query-server";
 import { DatabaseItem } from "../databases/local-databases";
 import { Mode } from "./shared/mode";
 
+// Limit the number of candidates we send to the model in each request
+// to avoid long requests.
+// Note that the model may return fewer than this number of candidates.
+const candidateBatchSize = 20;
+
 export class AutoModeler {
   constructor(
     private readonly app: App,
     private readonly cliServer: CodeQLCliServer,
     private readonly queryRunner: QueryRunner,
     private readonly queryStorageDir: string,
     private readonly databaseItem: DatabaseItem,
+    private readonly setInProgressMethods: (
+      inProgressMethods: string[],
+    ) => Promise<void>,
     private readonly addModeledMethods: (
       modeledMethods: Record<string, ModeledMethod>,
     ) => Promise<void>,
@@ -52,26 +60,52 @@ export class AutoModeler {
     await withProgress(async (progress) => {
       const maxStep = 3000;
 
-      progress({
-        step: 0,
-        maxStep,
-        message: "Retrieving usages",
-      });
-
       // Fetch the candidates to send to the model
-      const candidateMethods = getCandidates(
+      const allCandidateMethods = getCandidates(
         mode,
         externalApiUsages,
         modeledMethods,
       );
 
       // If there are no candidates, there is nothing to model and we just return
-      if (candidateMethods.length === 0) {
+      if (allCandidateMethods.length === 0) {
         void extLogger.log("No candidates to model. Stopping.");
         return;
       }
 
-      await this.modelCandidates(candidateMethods, mode, progress, maxStep);
+      // Find number of slices to make
+      const batchNumber = Math.ceil(
+        allCandidateMethods.length / candidateBatchSize,
+      );
+      try {
+        for (let i = 0; i < batchNumber; i++) {
+          const start = i * candidateBatchSize;
+          const end = start + candidateBatchSize;
+          const candidatesToProcess = allCandidateMethods.slice(start, end);
+
+          await this.setInProgressMethods(
+            candidatesToProcess.map((c) => c.signature),
+          );
+
+          progress({
+            step: 1800 + i * 100,
+            maxStep,
+            message: `Automodeling candidates, batch ${
+              i + 1
+            } of ${batchNumber}`,
+          });
+
+          await this.modelCandidates(
+            candidatesToProcess,
+            mode,
+            progress,
+            maxStep,
+          );
+        }
+      } finally {
+        // Clear out in progress methods
+        await this.setInProgressMethods([]);
+      }
     });
   }
 

extensions/ql-vscode/src/data-extensions-editor/data-extensions-editor-view.ts

@@ -82,6 +82,12 @@ export class DataExtensionsEditorView extends AbstractWebview<
       queryRunner,
       queryStorageDir,
       databaseItem,
+      async (inProgressMethods) => {
+        await this.postMessage({
+          t: "setInProgressMethods",
+          inProgressMethods,
+        });
+      },
       async (modeledMethods) => {
         await this.postMessage({ t: "addModeledMethods", modeledMethods });
       },

extensions/ql-vscode/test/unit-tests/data-extensions-editor/auto-model-v2.test.ts

@@ -165,28 +165,4 @@ describe("getCandidates", () => {
     );
     expect(candidates.length).toEqual(1);
   });
-
-  it("respects the limit", () => {
-    const externalApiUsages: ExternalApiUsage[] = [];
-    for (let i = 0; i < 30; i++) {
-      externalApiUsages.push({
-        library: "my.jar",
-        signature: `org.my.A#x${i}()`,
-        packageName: "org.my",
-        typeName: "A",
-        methodName: `x${i}`,
-        methodParameters: "()",
-        supported: false,
-        supportedType: "none",
-        usages: [],
-      });
-    }
-    const modeledMethods = {};
-    const candidates = getCandidates(
-      Mode.Application,
-      externalApiUsages,
-      modeledMethods,
-    );
-    expect(candidates.length).toEqual(20);
-  });
 });