Merge pull request #1599 from github/elenatanasoiu/add-batching-to-download-take-two

Download variant analysis results in batches - take two

Author: elenatanasoiu

extensions/ql-vscode/package-lock.json

@@ -1230,6 +1230,7 @@
     "minimist": "~1.2.6",
     "nanoid": "^3.2.0",
     "node-fetch": "~2.6.7",
+    "p-queue": "^6.0.0",
     "path-browserify": "^1.0.1",
     "react": "^17.0.2",
     "react-dom": "^17.0.2",

extensions/ql-vscode/package.json

@@ -951,7 +951,7 @@ async function activateWithInstalledDistribution(
       variantAnalysisSummary: VariantAnalysisApiResponse,
       token: CancellationToken
     ) => {
-      await variantAnalysisManager.autoDownloadVariantAnalysisResult(scannedRepo, variantAnalysisSummary, token);
+      await variantAnalysisManager.enqueueDownload(scannedRepo, variantAnalysisSummary, token);
     })
   );
 

extensions/ql-vscode/src/extension.ts

@@ -22,6 +22,7 @@ import { VariantAnalysisResultsManager } from './variant-analysis-results-manage
 import { CodeQLCliServer } from '../cli';
 import { getControllerRepo } from './run-remote-query';
 import { processUpdatedVariantAnalysis } from './variant-analysis-processor';
+import PQueue from 'p-queue';
 
 export class VariantAnalysisManager extends DisposableObject implements VariantAnalysisViewManager<VariantAnalysisView> {
   private readonly _onVariantAnalysisAdded = this.push(new EventEmitter<VariantAnalysis>());
@@ -31,6 +32,8 @@ export class VariantAnalysisManager extends DisposableObject implements VariantA
   private readonly variantAnalysisResultsManager: VariantAnalysisResultsManager;
   private readonly variantAnalyses = new Map<number, VariantAnalysis>();
   private readonly views = new Map<number, VariantAnalysisView>();
+  private static readonly maxConcurrentDownloads = 3;
+  private readonly queue = new PQueue({ concurrency: VariantAnalysisManager.maxConcurrentDownloads });
 
   constructor(
     private readonly ctx: ExtensionContext,
@@ -161,6 +164,18 @@ export class VariantAnalysisManager extends DisposableObject implements VariantA
     await this.onRepoStateUpdated(variantAnalysisSummary.id, repoState);
   }
 
+  public async enqueueDownload(
+    scannedRepo: ApiVariantAnalysisScannedRepository,
+    variantAnalysisSummary: VariantAnalysisApiResponse,
+    token: CancellationToken
+  ): Promise<void> {
+    await this.queue.add(() => this.autoDownloadVariantAnalysisResult(scannedRepo, variantAnalysisSummary, token));
+  }
+
+  public downloadsQueueSize(): number {
+    return this.queue.pending;
+  }
+
   public async promptOpenVariantAnalysis() {
     const credentials = await Credentials.initialize(this.ctx);
     if (!credentials) { throw Error('Error authenticating with GitHub'); }

extensions/ql-vscode/src/remote-queries/variant-analysis-manager.ts

@@ -4,7 +4,8 @@ import * as ghApiClient from './gh-api/gh-api-client';
 
 import { VariantAnalysis, VariantAnalysisStatus } from './shared/variant-analysis';
 import {
-  VariantAnalysis as VariantAnalysisApiResponse
+  VariantAnalysis as VariantAnalysisApiResponse,
+  VariantAnalysisScannedRepository
 } from './gh-api/variant-analysis';
 import { VariantAnalysisMonitorResult } from './shared/variant-analysis-monitor-result';
 import { processFailureReason, processUpdatedVariantAnalysis } from './variant-analysis-processor';
@@ -71,14 +72,8 @@ export class VariantAnalysisMonitor extends DisposableObject {
 
       this._onVariantAnalysisChange.fire(variantAnalysis);
 
-      if (variantAnalysisSummary.scanned_repositories) {
-        variantAnalysisSummary.scanned_repositories.forEach(scannedRepo => {
-          if (!scannedReposDownloaded.includes(scannedRepo.repository.id) && scannedRepo.analysis_status === 'succeeded') {
-            void commands.executeCommand('codeQL.autoDownloadVariantAnalysisResult', scannedRepo, variantAnalysisSummary);
-            scannedReposDownloaded.push(scannedRepo.repository.id);
-          }
-        });
-      }
+      const downloadedRepos = this.downloadVariantAnalysisResults(variantAnalysisSummary, scannedReposDownloaded);
+      scannedReposDownloaded.push(...downloadedRepos);
 
       if (variantAnalysisSummary.status === 'completed') {
         break;
@@ -90,6 +85,46 @@ export class VariantAnalysisMonitor extends DisposableObject {
     return { status: 'CompletedSuccessfully', scannedReposDownloaded: scannedReposDownloaded };
   }
 
+  private scheduleForDownload(
+    scannedRepo: VariantAnalysisScannedRepository,
+    variantAnalysisSummary: VariantAnalysisApiResponse
+  ) {
+    void commands.executeCommand('codeQL.autoDownloadVariantAnalysisResult', scannedRepo, variantAnalysisSummary);
+  }
+
+  private shouldDownload(
+    scannedRepo: VariantAnalysisScannedRepository,
+    alreadyDownloaded: number[]
+  ): boolean {
+    return !alreadyDownloaded.includes(scannedRepo.repository.id) && scannedRepo.analysis_status === 'succeeded';
+  }
+
+  private getReposToDownload(
+    variantAnalysisSummary: VariantAnalysisApiResponse,
+    alreadyDownloaded: number[]
+  ): VariantAnalysisScannedRepository[] {
+    if (variantAnalysisSummary.scanned_repositories) {
+      return variantAnalysisSummary.scanned_repositories.filter(scannedRepo => this.shouldDownload(scannedRepo, alreadyDownloaded));
+    } else {
+      return [];
+    }
+  }
+
+  private downloadVariantAnalysisResults(
+    variantAnalysisSummary: VariantAnalysisApiResponse,
+    scannedReposDownloaded: number[]
+  ): number[] {
+    const repoResultsToDownload = this.getReposToDownload(variantAnalysisSummary, scannedReposDownloaded);
+    const downloadedRepos: number[] = [];
+
+    repoResultsToDownload.forEach(scannedRepo => {
+      downloadedRepos.push(scannedRepo.repository.id);
+      this.scheduleForDownload(scannedRepo, variantAnalysisSummary);
+    });
+
+    return downloadedRepos;
+  }
+
   private async sleep(ms: number) {
     return new Promise(resolve => setTimeout(resolve, ms));
   }

extensions/ql-vscode/src/remote-queries/variant-analysis-monitor.ts

@@ -149,6 +149,17 @@ describe('Variant Analysis Manager', async function() {
 
         expect(getVariantAnalysisRepoResultStub.calledOnce).to.be.true;
       });
+
+      it('should pop download tasks off the queue', async () => {
+        const getResultsSpy = sandbox.spy(variantAnalysisManager, 'autoDownloadVariantAnalysisResult');
+
+        await variantAnalysisManager.enqueueDownload(scannedRepos[0], variantAnalysis, cancellationTokenSource.token);
+        await variantAnalysisManager.enqueueDownload(scannedRepos[1], variantAnalysis, cancellationTokenSource.token);
+        await variantAnalysisManager.enqueueDownload(scannedRepos[2], variantAnalysis, cancellationTokenSource.token);
+
+        expect(variantAnalysisManager.downloadsQueueSize()).to.equal(0);
+        expect(getResultsSpy).to.have.been.calledThrice;
+      });
     });
   });
 });

extensions/ql-vscode/src/vscode-tests/cli-integration/remote-queries/variant-analysis-manager.test.ts

@@ -1,31 +1,35 @@
 import * as sinon from 'sinon';
 import { expect } from 'chai';
-import { CancellationTokenSource, extensions } from 'vscode';
+import { CancellationTokenSource, commands, extensions } from 'vscode';
 import { CodeQLExtensionInterface } from '../../../extension';
 import * as config from '../../../config';
 
 import * as ghApiClient from '../../../remote-queries/gh-api/gh-api-client';
 import { VariantAnalysisMonitor } from '../../../remote-queries/variant-analysis-monitor';
 import {
   VariantAnalysis as VariantAnalysisApiResponse,
+  VariantAnalysisFailureReason,
   VariantAnalysisScannedRepository as ApiVariantAnalysisScannedRepository,
-  VariantAnalysisFailureReason
 } from '../../../remote-queries/gh-api/variant-analysis';
 import { createFailedMockApiResponse, createMockApiResponse } from '../../factories/remote-queries/gh-api/variant-analysis-api-response';
 import { VariantAnalysis, VariantAnalysisStatus } from '../../../remote-queries/shared/variant-analysis';
 import { createMockScannedRepos } from '../../factories/remote-queries/gh-api/scanned-repositories';
 import { processFailureReason } from '../../../remote-queries/variant-analysis-processor';
 import { Credentials } from '../../../authentication';
 import { createMockVariantAnalysis } from '../../factories/remote-queries/shared/variant-analysis';
+import { VariantAnalysisManager } from '../../../remote-queries/variant-analysis-manager';
 
 describe('Variant Analysis Monitor', async function() {
   this.timeout(60000);
 
   let sandbox: sinon.SinonSandbox;
+  let extension: CodeQLExtensionInterface | Record<string, never>;
   let mockGetVariantAnalysis: sinon.SinonStub;
   let cancellationTokenSource: CancellationTokenSource;
   let variantAnalysisMonitor: VariantAnalysisMonitor;
   let variantAnalysis: VariantAnalysis;
+  let variantAnalysisManager: VariantAnalysisManager;
+  let mockGetDownloadResult: sinon.SinonStub;
 
   beforeEach(async () => {
     sandbox = sinon.createSandbox();
@@ -36,12 +40,15 @@ describe('Variant Analysis Monitor', async function() {
     variantAnalysis = createMockVariantAnalysis();
 
     try {
-      const extension = await extensions.getExtension<CodeQLExtensionInterface | Record<string, never>>('GitHub.vscode-codeql')!.activate();
+      extension = await extensions.getExtension<CodeQLExtensionInterface | Record<string, never>>('GitHub.vscode-codeql')!.activate();
       variantAnalysisMonitor = new VariantAnalysisMonitor(extension.ctx);
     } catch (e) {
       fail(e as Error);
     }
 
+    variantAnalysisManager = extension.variantAnalysisManager;
+    mockGetDownloadResult = sandbox.stub(variantAnalysisManager, 'autoDownloadVariantAnalysisResult');
+
     limitNumberOfAttemptsToMonitor();
   });
 
@@ -110,20 +117,47 @@ describe('Variant Analysis Monitor', async function() {
     describe('when the variant analysis is in progress', async () => {
       let mockApiResponse: VariantAnalysisApiResponse;
       let scannedRepos: ApiVariantAnalysisScannedRepository[];
+      let succeededRepos: ApiVariantAnalysisScannedRepository[];
 
       describe('when there are successfully scanned repos', async () => {
         beforeEach(async function() {
-          scannedRepos = createMockScannedRepos(['pending', 'in_progress', 'succeeded']);
+          scannedRepos = createMockScannedRepos(['pending', 'pending', 'in_progress', 'in_progress', 'succeeded', 'succeeded', 'succeeded']);
           mockApiResponse = createMockApiResponse('completed', scannedRepos);
           mockGetVariantAnalysis = sandbox.stub(ghApiClient, 'getVariantAnalysis').resolves(mockApiResponse);
+          succeededRepos = scannedRepos.filter(r => r.analysis_status === 'succeeded');
         });
 
         it('should succeed and return a list of scanned repo ids', async () => {
           const result = await variantAnalysisMonitor.monitorVariantAnalysis(variantAnalysis, cancellationTokenSource.token);
-          const scannedRepoIds = scannedRepos.filter(r => r.analysis_status == 'succeeded').map(r => r.repository.id);
 
           expect(result.status).to.equal('CompletedSuccessfully');
-          expect(result.scannedReposDownloaded).to.eql(scannedRepoIds);
+          expect(result.scannedReposDownloaded).to.eql(succeededRepos.map(r => r.repository.id));
+        });
+
+        it('should trigger a download extension command for each repo', async () => {
+          const succeededRepos = scannedRepos.filter(r => r.analysis_status === 'succeeded');
+          const commandSpy = sandbox.spy(commands, 'executeCommand');
+
+          await variantAnalysisMonitor.monitorVariantAnalysis(variantAnalysis, cancellationTokenSource.token);
+
+          expect(commandSpy).to.have.callCount(succeededRepos.length);
+
+          succeededRepos.forEach((succeededRepo, index) => {
+            expect(commandSpy.getCall(index).args[0]).to.eq('codeQL.autoDownloadVariantAnalysisResult');
+            expect(commandSpy.getCall(index).args[1]).to.eq(succeededRepo);
+            expect(commandSpy.getCall(index).args[2]).to.eq(mockApiResponse);
+          });
+        });
+
+        it('should download all available results', async () => {
+          await variantAnalysisMonitor.monitorVariantAnalysis(variantAnalysis, cancellationTokenSource.token);
+
+          expect(mockGetDownloadResult).to.have.callCount(succeededRepos.length);
+
+          succeededRepos.forEach((succeededRepo, index) => {
+            expect(mockGetDownloadResult.getCall(index).args[0]).to.eq(succeededRepo);
+            expect(mockGetDownloadResult.getCall(index).args[1]).to.eq(mockApiResponse);
+          });
         });
       });
 
@@ -142,6 +176,12 @@ describe('Variant Analysis Monitor', async function() {
           expect(result.status).to.equal('CompletedSuccessfully');
           expect(result.scannedReposDownloaded).to.eql([]);
         });
+
+        it('should not try to download any repos', async () => {
+          await variantAnalysisMonitor.monitorVariantAnalysis(variantAnalysis, cancellationTokenSource.token);
+
+          expect(mockGetDownloadResult).to.not.have.been.called;
+        });
       });
 
       describe('when there are no repos to scan', async () => {
@@ -157,6 +197,12 @@ describe('Variant Analysis Monitor', async function() {
           expect(result.status).to.equal('CompletedSuccessfully');
           expect(result.scannedReposDownloaded).to.eql([]);
         });
+
+        it('should not try to download any repos', async () => {
+          await variantAnalysisMonitor.monitorVariantAnalysis(variantAnalysis, cancellationTokenSource.token);
+
+          expect(mockGetDownloadResult).to.not.have.been.called;
+        });
       });
     });
   });