Add an internal rate limit (#4610)

* Add an internal rate limit
Fixes #4546

* Fix some error/info messages

package.json

@@ -2362,6 +2362,7 @@
 		"dayjs": "1.10.4",
 		"events": "3.2.0",
 		"fast-deep-equal": "^3.1.3",
+		"limiter": "^2.1.0",
 		"lru-cache": "6.0.0",
 		"marked": "^4.0.10",
 		"react": "^16.12.0",

src/@types/vscode.proposed.contribEditorLineNumberMenu.d.ts

@@ -0,0 +1,8 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+// empty placeholder declaration for the `editor/lineNumber/context` menu contribution point
+
+// https://github.com/microsoft/vscode/issues/175945 @joyceerhl

src/@types/vscode.proposed.contribShareMenu.d.ts

@@ -4,3 +4,4 @@
  *--------------------------------------------------------------------------------------------*/
 
 // empty placeholder declaration for the `file/share`-submenu contribution point
+// https://github.com/microsoft/vscode/issues/176316

src/@types/vscode.proposed.quickDiffProvider.d.ts

@@ -10,4 +10,8 @@ declare module 'vscode' {
 	export namespace window {
 		export function registerQuickDiffProvider(selector: DocumentSelector, quickDiffProvider: QuickDiffProvider, label: string, rootUri?: Uri): Disposable;
 	}
+
+	interface QuickDiffProvider {
+		label?: string;
+	}
 }

src/github/credentials.ts

@@ -347,7 +347,7 @@ export class CredentialStore implements vscode.Disposable {
 			},
 		});
 
-		const rateLogger = new RateLogger(this._context, this._telemetry);
+		const rateLogger = new RateLogger(this._telemetry);
 		const github: GitHub = {
 			octokit: new LoggingOctokit(octokit, rateLogger),
 			graphql: new LoggingApolloClient(graphql, rateLogger),

src/github/loggingOctokit.ts

@@ -5,14 +5,12 @@
 
 import { Octokit } from '@octokit/rest';
 import { ApolloClient, ApolloQueryResult, FetchResult, MutationOptions, NormalizedCacheObject, OperationVariables, QueryOptions } from 'apollo-boost';
+import { RateLimiter } from 'limiter';
 import * as vscode from 'vscode';
 import Logger from '../common/logger';
 import { ITelemetry } from '../common/telemetry';
 import { RateLimit } from './graphql';
 
-const RATE_COUNTER_LAST_WINDOW = 'rateCounterLastWindow';
-const RATE_COUNTER_COUNT = 'rateCounterCount';
-
 interface RestResponse {
 	headers: {
 		'x-ratelimit-limit': string;
@@ -21,44 +19,30 @@ interface RestResponse {
 }
 
 export class RateLogger {
-	private lastWindow: number;
-	private count: number = 0;
+	private limiter: RateLimiter;
 	private static ID = 'RateLimit';
 	private hasLoggedLowRateLimit: boolean = false;
 
-	constructor(private readonly context: vscode.ExtensionContext, private readonly telemetry: ITelemetry) {
-		// We assume the common case for this logging: only one user.
-		// We also make up our own window. This will not line up exactly with GitHub's rate limit reset time,
-		// but it will give us a nice idea of how many API calls we're making. We use an hour, just like GitHub.
-		this.lastWindow = this.context.globalState.get(RATE_COUNTER_LAST_WINDOW, 0);
-		// It looks like there might be separate rate limits for the REST and GraphQL api.
-		// We'll just count total API calls as a lower bound.
-		this.count = this.context.globalState.get(RATE_COUNTER_COUNT, 0);
-		this.tryUpdateWindow();
+	constructor(private readonly telemetry: ITelemetry) {
+		this.limiter = new RateLimiter({ tokensPerInterval: 100, interval: 'second' });
 	}
 
-	private tryUpdateWindow() {
-		const now = new Date().getTime();
-		if ((now - this.lastWindow) > (60 * 60 * 1000) /* 1 hour */) {
-			this.lastWindow = now;
-			this.context.globalState.update(RATE_COUNTER_LAST_WINDOW, this.lastWindow);
-			this.count = 0;
+	public logAndLimit(): boolean {
+		if (!this.limiter.tryRemoveTokens(1)) {
+			Logger.error('API call count has exceeded 100 calls in 1 second.', RateLogger.ID);
+			// We have hit 100 requests in 1 second. This likely indicates a bug in the extension.
+			/* __GDPR__
+				"pr.highApiCallRate" : {}
+			*/
+			this.telemetry.sendTelemetryErrorEvent('pr.highApiCallRate');
+			vscode.window.showErrorMessage(vscode.l10n.t('The GitHub Pull Requests extension is making too many requests to GitHub. This indicates a bug in the extension. Please file an issue on GitHub and include the output from "GitHub Pull Request".'));
+			return false;
 		}
+		Logger.debug(`Extension rate limit remaining: ${this.limiter.getTokensRemaining()}`, RateLogger.ID);
+		return true;
 	}
 
-	public log(info: string | undefined) {
-		this.tryUpdateWindow();
-		this.count++;
-		this.context.globalState.update(RATE_COUNTER_COUNT, this.count);
-		const countMessage = `API call count: ${this.count}${info ? ` (${info})` : ''}`;
-		if (this.count > 4000) {
-			Logger.appendLine(countMessage, RateLogger.ID);
-		} else {
-			Logger.debug(countMessage, RateLogger.ID);
-		}
-	}
-
-	public async logRateLimit(result: Promise<{ data: { rateLimit: RateLimit | undefined } | undefined } | undefined>, isRest: boolean = false) {
+	public async logRateLimit(info: string | undefined, result: Promise<{ data: { rateLimit: RateLimit | undefined } | undefined } | undefined>, isRest: boolean = false) {
 		let rateLimitInfo;
 		try {
 			rateLimitInfo = (await result)?.data?.rateLimit;
@@ -69,21 +53,22 @@ export class RateLogger {
 		if ((rateLimitInfo?.limit ?? 5000) < 5000) {
 			Logger.appendLine(`Unexpectedly low rate limit: ${rateLimitInfo?.limit}`, RateLogger.ID);
 		}
-		const remaining = `${isRest ? 'REST' : 'GraphQL'} Rate limit remaining: ${rateLimitInfo?.remaining}`;
+		const remaining = `${isRest ? 'REST' : 'GraphQL'} Rate limit remaining: ${rateLimitInfo?.remaining}, ${info}`;
 		if ((rateLimitInfo?.remaining ?? 1000) < 1000) {
-			Logger.appendLine(remaining, RateLogger.ID);
-		} else {
 			if (!this.hasLoggedLowRateLimit) {
 				/* __GDPR__
 					"pr.lowRateLimitRemaining" : {}
 				*/
 				this.telemetry.sendTelemetryErrorEvent('pr.lowRateLimitRemaining');
+				this.hasLoggedLowRateLimit = true;
 			}
+			Logger.warn(remaining, RateLogger.ID);
+		} else {
 			Logger.debug(remaining, RateLogger.ID);
 		}
 	}
 
-	public async logRestRateLimit(restResponse: Promise<RestResponse>) {
+	public async logRestRateLimit(info: string | undefined, restResponse: Promise<RestResponse>) {
 		let result;
 		try {
 			result = await restResponse;
@@ -97,35 +82,44 @@ export class RateLogger {
 			remaining: Number(result.headers['x-ratelimit-remaining']),
 			resetAt: ''
 		};
-		this.logRateLimit(Promise.resolve({ data: { rateLimit } }), true);
+		this.logRateLimit(info, Promise.resolve({ data: { rateLimit } }), true);
 	}
 }
 
 export class LoggingApolloClient {
 	constructor(private readonly _graphql: ApolloClient<NormalizedCacheObject>, private _rateLogger: RateLogger) { };
 
 	query<T = any, TVariables = OperationVariables>(options: QueryOptions<TVariables>): Promise<ApolloQueryResult<T>> {
-		this._rateLogger.log((options.query.definitions[0] as { name: { value: string } | undefined }).name?.value);
-		const result = this._graphql.query(options);
-		this._rateLogger.logRateLimit(result as any);
-		return result;
+		if (this._rateLogger.logAndLimit()) {
+			const result = this._graphql.query(options);
+			this._rateLogger.logRateLimit((options.query.definitions[0] as { name: { value: string } | undefined }).name?.value, result as any);
+			return result;
+		} else {
+			throw new Error('API call count has exceeded a rate limit.');
+		}
 	}
 
 	mutate<T = any, TVariables = OperationVariables>(options: MutationOptions<T, TVariables>): Promise<FetchResult<T>> {
-		this._rateLogger.log(options.context);
-		const result = this._graphql.mutate(options);
-		this._rateLogger.logRateLimit(result as any);
-		return result;
+		if (this._rateLogger.logAndLimit()) {
+			const result = this._graphql.mutate(options);
+			this._rateLogger.logRateLimit(options.context, result as any);
+			return result;
+		} else {
+			throw new Error('API call count has exceeded a rate limit.');
+		}
 	}
 }
 
 export class LoggingOctokit {
 	constructor(public readonly api: Octokit, private _rateLogger: RateLogger) { };
 
 	async call<T, U>(api: (T) => Promise<U>, args: T): Promise<U> {
-		this._rateLogger.log((api as unknown as { endpoint: { DEFAULTS: { url: string } | undefined } | undefined }).endpoint?.DEFAULTS?.url);
-		const result = api(args);
-		this._rateLogger.logRestRateLimit(result as Promise<unknown> as Promise<RestResponse>);
-		return result;
+		if (this._rateLogger.logAndLimit()) {
+			const result = api(args);
+			this._rateLogger.logRestRateLimit((api as unknown as { endpoint: { DEFAULTS: { url: string } | undefined } | undefined }).endpoint?.DEFAULTS?.url, result as Promise<unknown> as Promise<RestResponse>);
+			return result;
+		} else {
+			throw new Error('API call count has exceeded a rate limit.');
+		}
 	}
 }

src/test/mocks/mockGitHubRepository.ts

@@ -20,7 +20,6 @@ import {
 import { MockTelemetry } from './mockTelemetry';
 import { Uri } from 'vscode';
 import { LoggingOctokit, RateLogger } from '../../github/loggingOctokit';
-import { MockExtensionContext } from './mockExtensionContext';
 const queries = require('../../github/queries.gql');
 
 export class MockGitHubRepository extends GitHubRepository {
@@ -32,7 +31,7 @@ export class MockGitHubRepository extends GitHubRepository {
 		this.queryProvider = new QueryProvider(sinon);
 
 		this._hub = {
-			octokit: new LoggingOctokit(this.queryProvider.octokit, new RateLogger(new MockExtensionContext(), new MockTelemetry())),
+			octokit: new LoggingOctokit(this.queryProvider.octokit, new RateLogger(new MockTelemetry())),
 			graphql: null,
 		};
 

src/test/view/prsTree.test.ts

@@ -54,7 +54,7 @@ describe('GitHub Pull Requests view', function () {
 					baseUrl: 'https://github.com',
 					userAgent: 'GitHub VSCode Pull Requests',
 					previews: ['shadow-cat-preview'],
-				}), new RateLogger(context, telemetry)),
+				}), new RateLogger(telemetry)),
 				graphql: null,
 			};
 

yarn.lock

@@ -3625,6 +3625,11 @@ just-extend@^4.0.2:
   resolved "https://registry.yarnpkg.com/just-extend/-/just-extend-4.1.1.tgz#158f1fdb01f128c411dc8b286a7b4837b3545282"
   integrity sha512-aWgeGFW67BP3e5181Ep1Fv2v8z//iBJfrvyTnq8wG86vEESwmonn1zPBJ0VfmT9CJq2FIT0VsETtrNFm2a+SHA==
 
+just-performance@4.3.0:
+  version "4.3.0"
+  resolved "https://registry.yarnpkg.com/just-performance/-/just-performance-4.3.0.tgz#cc2bc8c9227f09e97b6b1df4cd0de2df7ae16db1"
+  integrity sha512-L7RjvtJsL0QO8xFs5wEoDDzzJwoiowRw6Rn/GnvldlchS2JQr9wFYPiwZcDfrbbujEKqKN0tvENdbjXdYhDp5Q==
+
 keygrip@~1.1.0:
   version "1.1.0"
   resolved "https://registry.yarnpkg.com/keygrip/-/keygrip-1.1.0.tgz#871b1681d5e159c62a445b0c74b615e0917e7226"
@@ -3741,6 +3746,13 @@ levn@~0.3.0:
     prelude-ls "~1.1.2"
     type-check "~0.3.2"
 
+limiter@^2.1.0:
+  version "2.1.0"
+  resolved "https://registry.yarnpkg.com/limiter/-/limiter-2.1.0.tgz#d38d7c5b63729bb84fb0c4d8594b7e955a5182a2"
+  integrity sha512-361TYz6iay6n+9KvUUImqdLuFigK+K79qrUtBsXhJTLdH4rIt/r1y8r1iozwh8KbZNpujbFTSh74mJ7bwbAMOw==
+  dependencies:
+    just-performance "4.3.0"
+
 lines-and-columns@^1.1.6:
   version "1.1.6"
   resolved "https://registry.yarnpkg.com/lines-and-columns/-/lines-and-columns-1.1.6.tgz#1c00c743b433cd0a4e80758f7b64a57440d9ff00"