Rate limit CLI update checks

henrymercer · henrymercer · commit 93ed82033360 · 2019-12-20T11:26:22.000Z
This helps to avoid hitting the GitHub API limits of 60 requests per
hour for unauthenticated IPs.
diff --git a/extensions/ql-vscode/src/distribution.ts b/extensions/ql-vscode/src/distribution.ts
@@ -6,7 +6,7 @@ import * as unzipper from "unzipper";
 import * as url from "url";
 import { ExtensionContext, Event } from "vscode";
 import { DistributionConfig } from "./config";
-import { ProgressUpdate, showAndLogErrorMessage } from "./helpers";
+import { InvocationRateLimiter, InvocationRateLimiterResultKind, ProgressUpdate, showAndLogErrorMessage } from "./helpers";
 import { logger } from "./logging";
 import { getCodeQlCliVersion, tryParseVersionString, Version } from "./cli-version";
 
@@ -55,6 +55,11 @@ export class DistributionManager implements DistributionProvider {
     this._config = config;
     this._extensionSpecificDistributionManager = new ExtensionSpecificDistributionManager(extensionContext, config, versionConstraint);
     this._onDidChangeDistribution = config.onDidChangeDistributionConfiguration;
+    this._updateCheckRateLimiter = new InvocationRateLimiter(
+      extensionContext,
+      "extensionSpecificDistributionUpdateCheck",
+      () => this._extensionSpecificDistributionManager.checkForUpdatesToDistribution()
+    );
     this._versionConstraint = versionConstraint;
   }
 
@@ -128,14 +133,21 @@ export class DistributionManager implements DistributionProvider {
    * 
    * Returns a failed promise if an unexpected error occurs during installation.
    */
-  public async checkForUpdatesToExtensionManagedDistribution(): Promise<DistributionUpdateCheckResult> {
+  public async checkForUpdatesToExtensionManagedDistribution(
+    minSecondsSinceLastUpdateCheck: number): Promise<DistributionUpdateCheckResult> {
     const codeQlPath = await this.getCodeQlPathWithoutVersionCheck();
     const extensionManagedCodeQlPath = await this._extensionSpecificDistributionManager.getCodeQlPathWithoutVersionCheck();
     if (codeQlPath !== undefined && codeQlPath !== extensionManagedCodeQlPath) {
       // A distribution is present but it isn't managed by the extension.
-      return createInvalidDistributionLocationResult();
+      return createInvalidLocationResult();
+    }
+    const updateCheckResult = await this._updateCheckRateLimiter.invokeFunctionIfIntervalElapsed(minSecondsSinceLastUpdateCheck);
+    switch (updateCheckResult.kind) {
+      case InvocationRateLimiterResultKind.Invoked:
+        return updateCheckResult.result;
+      case InvocationRateLimiterResultKind.RateLimited:
+        return createAlreadyCheckedRecentlyResult();
     }
-    return this._extensionSpecificDistributionManager.checkForUpdatesToDistribution();
   }
 
   /**
@@ -154,6 +166,7 @@ export class DistributionManager implements DistributionProvider {
 
   private readonly _config: DistributionConfig;
   private readonly _extensionSpecificDistributionManager: ExtensionSpecificDistributionManager;
+  private readonly _updateCheckRateLimiter: InvocationRateLimiter<DistributionUpdateCheckResult>;
   private readonly _onDidChangeDistribution: Event<void> | undefined;
   private readonly _versionConstraint: VersionConstraint;
 }
@@ -196,7 +209,7 @@ class ExtensionSpecificDistributionManager {
     const latestRelease = await this.getLatestRelease();
 
     if (extensionSpecificRelease !== undefined && codeQlPath !== undefined && latestRelease.id === extensionSpecificRelease.id) {
-      return createDistributionAlreadyUpToDateResult();
+      return createAlreadyUpToDateResult();
     }
     return createUpdateAvailableResult(latestRelease);
   }
@@ -234,7 +247,7 @@ class ExtensionSpecificDistributionManager {
 
       if (progressCallback && contentLength !== null) {
         const totalNumBytes = parseInt(contentLength, 10);
-        const bytesToDisplayMB = (numBytes: number) => `${(numBytes/(1024*1024)).toFixed(1)} MB`;
+        const bytesToDisplayMB = (numBytes: number) => `${(numBytes / (1024 * 1024)).toFixed(1)} MB`;
         const updateProgress = () => {
           progressCallback({
             step: numBytesDownloaded,
@@ -258,7 +271,7 @@ class ExtensionSpecificDistributionManager {
           .on("error", reject)
       );
 
-      this.bumpDistributionFolderIndex();
+      await this.bumpDistributionFolderIndex();
 
       logger.log(`Extracting CodeQL CLI to ${this.getDistributionStoragePath()}`);
       await extractZipArchive(archivePath, this.getDistributionStoragePath());
@@ -293,10 +306,10 @@ class ExtensionSpecificDistributionManager {
     return new ReleasesApiConsumer(ownerName, repositoryName, this._config.personalAccessToken);
   }
 
-  private bumpDistributionFolderIndex(): void {
+  private async bumpDistributionFolderIndex(): Promise<void> {
     const index = this._extensionContext.globalState.get(
       ExtensionSpecificDistributionManager._currentDistributionFolderIndexStateKey, 0);
-    this._extensionContext.globalState.update(
+    await this._extensionContext.globalState.update(
       ExtensionSpecificDistributionManager._currentDistributionFolderIndexStateKey, index + 1);
   }
 
@@ -317,8 +330,8 @@ class ExtensionSpecificDistributionManager {
     return this._extensionContext.globalState.get(ExtensionSpecificDistributionManager._installedReleaseStateKey);
   }
 
-  private storeInstalledRelease(release: Release | undefined): void {
-    this._extensionContext.globalState.update(ExtensionSpecificDistributionManager._installedReleaseStateKey, release);
+  private async storeInstalledRelease(release: Release | undefined): Promise<void> {
+    await this._extensionContext.globalState.update(ExtensionSpecificDistributionManager._installedReleaseStateKey, release);
   }
 
   private readonly _config: DistributionConfig;
@@ -532,39 +545,50 @@ interface NoDistributionResult {
 }
 
 export enum DistributionUpdateCheckResultKind {
+  AlreadyCheckedRecentlyResult,
   AlreadyUpToDate,
-  InvalidDistributionLocation,
+  InvalidLocation,
   UpdateAvailable
 }
 
-type DistributionUpdateCheckResult = DistributionAlreadyUpToDateResult | InvalidDistributionLocationResult |
+type DistributionUpdateCheckResult = AlreadyCheckedRecentlyResult | AlreadyUpToDateResult | InvalidLocationResult |
   UpdateAvailableResult;
 
-export interface DistributionAlreadyUpToDateResult {
+export interface AlreadyCheckedRecentlyResult {
+  kind: DistributionUpdateCheckResultKind.AlreadyCheckedRecentlyResult
+}
+
+export interface AlreadyUpToDateResult {
   kind: DistributionUpdateCheckResultKind.AlreadyUpToDate;
 }
 
 /**
  * The distribution could not be installed or updated because it is not managed by the extension.
  */
-export interface InvalidDistributionLocationResult {
-  kind: DistributionUpdateCheckResultKind.InvalidDistributionLocation;
+export interface InvalidLocationResult {
+  kind: DistributionUpdateCheckResultKind.InvalidLocation;
 }
 
 export interface UpdateAvailableResult {
   kind: DistributionUpdateCheckResultKind.UpdateAvailable;
   updatedRelease: Release;
 }
 
-function createDistributionAlreadyUpToDateResult(): DistributionAlreadyUpToDateResult {
+function createAlreadyCheckedRecentlyResult(): AlreadyCheckedRecentlyResult {
+  return {
+    kind: DistributionUpdateCheckResultKind.AlreadyCheckedRecentlyResult
+  };
+}
+
+function createAlreadyUpToDateResult(): AlreadyUpToDateResult {
   return {
     kind: DistributionUpdateCheckResultKind.AlreadyUpToDate
   };
 }
 
-function createInvalidDistributionLocationResult(): InvalidDistributionLocationResult {
+function createInvalidLocationResult(): InvalidLocationResult {
   return {
-    kind: DistributionUpdateCheckResultKind.InvalidDistributionLocation
+    kind: DistributionUpdateCheckResultKind.InvalidLocation
   };
 }
 
diff --git a/extensions/ql-vscode/src/extension.ts b/extensions/ql-vscode/src/extension.ts
@@ -4,8 +4,10 @@ import * as archiveFilesystemProvider from './archive-filesystem-provider';
 import { DistributionConfigListener, QueryServerConfigListener, QueryHistoryConfigListener } from './config';
 import { DatabaseManager } from './databases';
 import { DatabaseUI } from './databases-ui';
-import { DistributionUpdateCheckResultKind, DistributionManager, FindDistributionResult, FindDistributionResultKind, GithubApiError,
-  DEFAULT_DISTRIBUTION_VERSION_CONSTRAINT, GithubRateLimitedError } from './distribution';
+import {
+  DistributionUpdateCheckResultKind, DistributionManager, FindDistributionResult, FindDistributionResultKind, GithubApiError,
+  DEFAULT_DISTRIBUTION_VERSION_CONSTRAINT, GithubRateLimitedError
+} from './distribution';
 import * as helpers from './helpers';
 import { spawnIdeServer } from './ide-server';
 import { InterfaceManager, WebviewReveal } from './interface';
@@ -83,29 +85,32 @@ export async function activate(ctx: ExtensionContext): Promise<void> {
     helpers.showAndLogErrorMessage(`Can't execute ${command}: waiting to finish loading CodeQL CLI.`);
   });
 
-  interface ReportingConfig {
+  interface DistributionUpdateConfig {
+    isUserInitiated: boolean;
     shouldDisplayMessageWhenNoUpdates: boolean;
-    shouldErrorIfUpdateFails: boolean;
   }
 
-  async function installOrUpdateDistributionWithProgressTitle(progressTitle: string, reportingConfig: ReportingConfig): Promise<void> {
-    const result = await distributionManager.checkForUpdatesToExtensionManagedDistribution();
+  async function installOrUpdateDistributionWithProgressTitle(progressTitle: string, config: DistributionUpdateConfig): Promise<void> {
+    const minSecondsSinceLastUpdateCheck = config.isUserInitiated ? 0 : 86400;
+    const noUpdatesLoggingFunc = config.shouldDisplayMessageWhenNoUpdates ?
+      helpers.showAndLogInformationMessage : async (message: string) => logger.log(message);
+    const result = await distributionManager.checkForUpdatesToExtensionManagedDistribution(minSecondsSinceLastUpdateCheck);
     switch (result.kind) {
+      case DistributionUpdateCheckResultKind.AlreadyCheckedRecentlyResult:
+        logger.log("Didn't perform CodeQL CLI update check since a check was already performed within the previous " +
+          `${minSecondsSinceLastUpdateCheck} seconds.`);
+        break;
       case DistributionUpdateCheckResultKind.AlreadyUpToDate:
-        if (reportingConfig.shouldDisplayMessageWhenNoUpdates) {
-          helpers.showAndLogInformationMessage("CodeQL CLI already up to date.");
-        }
+        await noUpdatesLoggingFunc("CodeQL CLI already up to date.");
         break;
-      case DistributionUpdateCheckResultKind.InvalidDistributionLocation:
-        if (reportingConfig.shouldDisplayMessageWhenNoUpdates) {
-          helpers.showAndLogErrorMessage("CodeQL CLI is installed externally so could not be updated.");
-        }
+      case DistributionUpdateCheckResultKind.InvalidLocation:
+        await noUpdatesLoggingFunc("CodeQL CLI is installed externally so could not be updated.");
         break;
       case DistributionUpdateCheckResultKind.UpdateAvailable:
         if (beganMainExtensionActivation) {
           const updateAvailableMessage = `Version "${result.updatedRelease.name}" of the CodeQL CLI is now available. ` +
             "The update will be installed after Visual Studio Code restarts. Restart now to upgrade?";
-          ctx.globalState.update(shouldUpdateOnNextActivationKey, true);
+          await ctx.globalState.update(shouldUpdateOnNextActivationKey, true);
           if (await helpers.showInformationMessageWithAction(updateAvailableMessage, "Restart and Upgrade")) {
             await commands.executeCommand("workbench.action.reloadWindow");
           }
@@ -118,7 +123,7 @@ export async function activate(ctx: ExtensionContext): Promise<void> {
           await helpers.withProgress(progressOptions, progress =>
             distributionManager.installExtensionManagedDistributionRelease(result.updatedRelease, progress));
 
-          ctx.globalState.update(shouldUpdateOnNextActivationKey, false);
+          await ctx.globalState.update(shouldUpdateOnNextActivationKey, false);
           helpers.showAndLogInformationMessage(`CodeQL CLI updated to version "${result.updatedRelease.name}".`);
         }
         break;
@@ -127,7 +132,7 @@ export async function activate(ctx: ExtensionContext): Promise<void> {
     }
   }
 
-  async function installOrUpdateDistribution(reportingConfig: ReportingConfig): Promise<void> {
+  async function installOrUpdateDistribution(config: DistributionUpdateConfig): Promise<void> {
     if (isInstallingOrUpdatingDistribution) {
       throw new Error("Already installing or updating CodeQL CLI");
     }
@@ -137,11 +142,11 @@ export async function activate(ctx: ExtensionContext): Promise<void> {
     const messageText = willUpdateCodeQl ? "Updating CodeQL CLI" :
       codeQlInstalled ? "Checking for updates to CodeQL CLI" : "Installing CodeQL CLI";
     try {
-      await installOrUpdateDistributionWithProgressTitle(messageText, reportingConfig);
+      await installOrUpdateDistributionWithProgressTitle(messageText, config);
     } catch (e) {
       // Don't rethrow the exception, because if the config is changed, we want to be able to retry installing
       // or updating the distribution.
-      const alertFunction = (codeQlInstalled && !reportingConfig.shouldErrorIfUpdateFails) ?
+      const alertFunction = (codeQlInstalled && !config.isUserInitiated) ?
         helpers.showAndLogWarningMessage : helpers.showAndLogErrorMessage;
       const taskDescription = (willUpdateCodeQl ? "update" :
         codeQlInstalled ? "check for updates to" : "install") + " CodeQL CLI";
@@ -180,8 +185,8 @@ export async function activate(ctx: ExtensionContext): Promise<void> {
     return result;
   }
 
-  async function installOrUpdateThenTryActivate(reportingConfig: ReportingConfig): Promise<void> {
-    await installOrUpdateDistribution(reportingConfig);
+  async function installOrUpdateThenTryActivate(config: DistributionUpdateConfig): Promise<void> {
+    await installOrUpdateDistribution(config);
 
     // Display the warnings even if the extension has already activated.
     const distributionResult = await getDistributionDisplayingDistributionWarnings();
@@ -194,26 +199,26 @@ export async function activate(ctx: ExtensionContext): Promise<void> {
         const chosenAction = await helpers.showAndLogErrorMessage(`Can't execute ${command}: missing CodeQL CLI.`, installActionName);
         if (chosenAction === installActionName) {
           installOrUpdateThenTryActivate({
-            shouldDisplayMessageWhenNoUpdates: false,
-            shouldErrorIfUpdateFails: true
+            isUserInitiated: true,
+            shouldDisplayMessageWhenNoUpdates: false
           });
         }
       });
     }
   }
 
   ctx.subscriptions.push(distributionConfigListener.onDidChangeDistributionConfiguration(() => installOrUpdateThenTryActivate({
-    shouldDisplayMessageWhenNoUpdates: false,
-    shouldErrorIfUpdateFails: true
+    isUserInitiated: true,
+    shouldDisplayMessageWhenNoUpdates: false
   })));
   ctx.subscriptions.push(commands.registerCommand(checkForUpdatesCommand, () => installOrUpdateThenTryActivate({
-    shouldDisplayMessageWhenNoUpdates: true,
-    shouldErrorIfUpdateFails: true
+    isUserInitiated: true,
+    shouldDisplayMessageWhenNoUpdates: true
   })));
 
   await installOrUpdateThenTryActivate({
-    shouldDisplayMessageWhenNoUpdates: false,
-    shouldErrorIfUpdateFails: !!ctx.globalState.get(shouldUpdateOnNextActivationKey)
+    isUserInitiated: !!ctx.globalState.get(shouldUpdateOnNextActivationKey),
+    shouldDisplayMessageWhenNoUpdates: false
   });
 }
 
