Retry review comments with smaller pages

src/github/githubRepository.ts

@@ -328,7 +328,7 @@ export class GitHubRepository extends Disposable {
 		}
 	}
 
-	query = async <T>(query: QueryOptions, ignoreSamlErrors: boolean = false, legacyFallback?: { query: DocumentNode, variables?: OperationVariables }): Promise<ApolloQueryResult<T>> => {
+	query = async <T>(query: QueryOptions, ignoreSamlErrors: boolean = false, legacyFallback?: { query: DocumentNode, variables: OperationVariables }): Promise<ApolloQueryResult<T>> => {
 		const gql = this.authMatchesServer && this.hub && this.hub.graphql;
 		if (!gql) {
 			const logValue = (query.query.definitions[0] as { name: { value: string } | undefined }).name?.value;

src/github/pullRequestModel.ts

@@ -67,7 +67,7 @@ import {
 	ReviewEventEnum,
 } from './interface';
 import { IssueChangeEvent, IssueModel } from './issueModel';
-import { compareCommits, GraphQLError, GraphQLErrorType } from './loggingOctokit';
+import { compareCommits, getErrorCode, GraphQLError, GraphQLErrorType } from './loggingOctokit';
 import {
 	convertRESTPullRequestToRawPullRequest,
 	convertRESTReviewEvent,
@@ -1493,25 +1493,45 @@ export class PullRequestModel extends IssueModel<PullRequest> implements IPullRe
 
 		const { remote, query, schema } = await this.githubRepository.ensure();
 		let after: string | null = null;
-		let hasNextPage = false;
+		let pageSize = 20;
 		const reviewThreads: ReviewThread[] = [];
 		try {
-			do {
-				const { data } = await query<PullRequestCommentsResponse>({
-					query: schema.PullRequestComments,
-					variables: {
-						owner: remote.owner,
-						name: remote.repositoryName,
-						number: this.number,
-						after
-					},
-				}, false, { query: schema.LegacyPullRequestComments });
+			while (reviewThreads.length < 1000) {
+				const variables = {
+					owner: remote.owner,
+					name: remote.repositoryName,
+					number: this.number,
+					first: pageSize,
+					after,
+				};
+				let data: PullRequestCommentsResponse | null;
+				try {
+					({ data } = await query<PullRequestCommentsResponse>({
+						query: schema.PullRequestComments,
+						variables,
+					}, false, { query: schema.LegacyPullRequestComments, variables }));
+				} catch (e) {
+					if (getErrorCode(e) !== '502' || pageSize === 1) {
+						throw e;
+					}
+					// Large review-thread queries can fail with HTTP 502.
+					// Retry the same cursor with 5, then 1 thread, and keep that size.
+					pageSize = Math.max(1, Math.floor(pageSize / 4));
+					Logger.warn(`Retrying review comments for PR #${this.number} with ${pageSize} threads per page after HTTP 502.`, PullRequestModel.ID);
+					continue;
+				}
 
-				reviewThreads.push(...data.repository.pullRequest.reviewThreads.nodes);
+				if (!data?.repository) {
+					throw new Error('Review comments response did not include a repository.');
+				}
+				const page = data.repository.pullRequest.reviewThreads;
+				reviewThreads.push(...page.nodes);
 
-				hasNextPage = data.repository.pullRequest.reviewThreads.pageInfo.hasNextPage;
-				after = data.repository.pullRequest.reviewThreads.pageInfo.endCursor;
-			} while (hasNextPage && reviewThreads.length < 1000);
+				if (!page.pageInfo.hasNextPage) {
+					break;
+				}
+				after = page.pageInfo.endCursor;
+			}
 			Logger.debug(`Fetching review comments for PR #${this.number} - exit`, PullRequestModel.ID);
 
 			return reviewThreads;

src/github/queriesShared.gql

@@ -566,10 +566,10 @@ query GetPendingReviewId($pullRequestId: ID!, $author: String!) {
 	}
 }
 
-query PullRequestComments($owner: String!, $name: String!, $number: Int!, $after: String) {
+query PullRequestComments($owner: String!, $name: String!, $number: Int!, $first: Int!, $after: String) {
 	repository(owner: $owner, name: $name) {
 		pullRequest(number: $number) {
-			reviewThreads(first: 20, after: $after) {
+			reviewThreads(first: $first, after: $after) {
 				nodes {
 					id
 					isResolved
@@ -608,10 +608,10 @@ query PullRequestComments($owner: String!, $name: String!, $number: Int!, $after
 	}
 }
 
-query LegacyPullRequestComments($owner: String!, $name: String!, $number: Int!, $after: String) {
+query LegacyPullRequestComments($owner: String!, $name: String!, $number: Int!, $first: Int!, $after: String) {
 	repository(owner: $owner, name: $name) {
 		pullRequest(number: $number) {
-			reviewThreads(first: 20, after: $after) {
+			reviewThreads(first: $first, after: $after) {
 				nodes {
 					id
 					isResolved

src/test/github/githubRepository.test.ts

@@ -4,6 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { default as assert } from 'assert';
+import { NetworkStatus } from 'apollo-boost';
 import { SinonSandbox, createSandbox } from 'sinon';
 import { CredentialStore } from '../../github/credentials';
 import { MockCommandRegistry } from '../mocks/mockCommandRegistry';
@@ -18,6 +19,7 @@ import { GitHubServerType } from '../../common/authentication';
 import { CheckState, PullRequestCheckStatus } from '../../github/interface';
 import { PullRequestBuilder as GraphQLPullRequestBuilder } from '../builders/graphql/pullRequestBuilder';
 import Logger from '../../common/logger';
+import { LoggingApolloClient, LoggingOctokit } from '../../github/loggingOctokit';
 
 describe('GitHubRepository', function () {
 	let sinon: SinonSandbox;
@@ -38,6 +40,36 @@ describe('GitHubRepository', function () {
 		sinon.restore();
 	});
 
+	describe('query', function () {
+		it('replaces variables for a legacy query with different arguments', async function () {
+			const url = 'https://github.com/some/repo';
+			const remote = new GitHubRemote('origin', url, new Protocol(url), GitHubServerType.GitHubDotCom);
+			const repo = new GitHubRepository(1, remote, Uri.file('/workspaces/repo'), credentialStore, telemetry, true);
+			const graphql = sinon.createStubInstance(LoggingApolloClient);
+			sinon.stub(credentialStore, 'isAuthenticated').returns(true);
+			sinon.stub(repo, 'hub').get(() => ({ graphql, octokit: sinon.createStubInstance(LoggingOctokit) }));
+			const variables = { owner: 'some', name: 'repo', first: 100, after: 'cursor' };
+			const response = { data: {}, loading: false, stale: false, networkStatus: NetworkStatus.ready };
+			graphql.query.onFirstCall().rejects(new Error('Unsupported query'));
+			graphql.query.onSecondCall().resolves(response);
+
+			try {
+				const result = await repo.query({
+					query: repo.schema.GetSuggestedActors,
+					variables: { ...variables, capabilities: ['CAN_BE_ASSIGNED'] },
+				}, false, { query: repo.schema.GetAssignableUsers, variables });
+
+				assert.strictEqual(result, response);
+				assert.strictEqual(graphql.query.callCount, 2);
+				const [fallback] = graphql.query.secondCall.args;
+				assert.strictEqual(fallback.query, repo.schema.GetAssignableUsers);
+				assert.deepStrictEqual(fallback.variables, variables);
+			} finally {
+				repo.dispose();
+			}
+		});
+	});
+
 	describe('isGitHubDotCom', function () {
 		it('detects when the remote is pointing to github.com', function () {
 			const url = 'https://github.com/some/repo';

src/test/github/pullRequestModel.test.ts

@@ -19,6 +19,9 @@ import { NetworkStatus } from 'apollo-client';
 import { MockExtensionContext } from '../mocks/mockExtensionContext';
 import { GitHubServerType } from '../../common/authentication';
 import { mergeQuerySchemaWithShared } from '../../github/common';
+import { GitHubRepository } from '../../github/githubRepository';
+import { LoggingApolloClient, LoggingOctokit } from '../../github/loggingOctokit';
+import Logger from '../../common/logger';
 const queries = mergeQuerySchemaWithShared(require('../../github/queries.gql'), require('../../github/queriesShared.gql')) as any;
 
 const telemetry = new MockTelemetry();
@@ -96,6 +99,106 @@ describe('PullRequestModel', function () {
 	});
 
 	describe('reviewThreadCache', function () {
+		function page(id: string, endCursor: string | null) {
+			return {
+				data: {
+					repository: {
+						pullRequest: {
+							reviewThreads: {
+								nodes: [{ ...reviewThreadResponse, id }],
+								pageInfo: { hasNextPage: endCursor !== null, endCursor },
+							},
+						},
+					},
+				},
+				loading: false,
+				stale: false,
+				networkStatus: NetworkStatus.ready,
+			};
+		}
+
+		it('passes review comment variables to every legacy page', async function () {
+			const repository = new GitHubRepository(1, remote, repo.rootUri, credentials, telemetry, true);
+			const graphql = sinon.createStubInstance(LoggingApolloClient);
+			sinon.stub(credentials, 'isAuthenticated').returns(true);
+			sinon.stub(repository, 'hub').get(() => ({ graphql, octokit: sinon.createStubInstance(LoggingOctokit) }));
+			sinon.stub(repository, 'ensure').resolves(repository);
+			graphql.query.onCall(0).rejects(new Error('Unsupported query'));
+			graphql.query.onCall(1).resolves(page('1', 'first'));
+			const gatewayError = Object.assign(new Error('Bad Gateway'), { networkError: { statusCode: 502 } });
+			graphql.query.onCall(2).rejects(gatewayError);
+			graphql.query.onCall(3).rejects(gatewayError);
+			graphql.query.onCall(4).rejects(new Error('Unsupported query'));
+			graphql.query.onCall(5).resolves(page('2', null));
+
+			try {
+				const pr = new PullRequestBuilder().build();
+				const model = new PullRequestModel(credentials, telemetry, repository, remote, convertRESTPullRequestToRawPullRequest(pr, repository));
+				const threads = await model.getReviewThreads();
+
+				assert.deepStrictEqual(threads.map(thread => thread.id), ['1', '2']);
+				assert.strictEqual(graphql.query.callCount, 6);
+				for (const [call, after, first] of [[graphql.query.secondCall, null, 20], [graphql.query.lastCall, 'first', 5]] as const) {
+					const [fallback] = call.args;
+					assert.strictEqual(fallback.query, repository.schema.LegacyPullRequestComments);
+					assert.deepStrictEqual(fallback.variables, {
+						owner: remote.owner, name: remote.repositoryName, number: pr.number, first, after,
+					});
+				}
+			} finally {
+				repository.dispose();
+			}
+		});
+
+		it('retries gateway failures with smaller pages without losing the cursor', async function () {
+			const pr = new PullRequestBuilder().build();
+			const model = new PullRequestModel(credentials, telemetry, repo, remote, convertRESTPullRequestToRawPullRequest(pr, repo));
+			const gatewayError = Object.assign(new Error('Bad Gateway'), { networkError: { statusCode: 502 } });
+			const query = sinon.stub(repo, 'query');
+			query.onCall(0).resolves(page('1', 'first'));
+			query.onCall(1).rejects(gatewayError);
+			query.onCall(2).rejects(gatewayError);
+			query.onCall(3).resolves(page('2', 'second'));
+			query.onCall(4).resolves(page('3', null));
+
+			const threads = await model.getReviewThreads();
+
+			assert.deepStrictEqual(threads.map(thread => thread.id), ['1', '2', '3']);
+			assert.deepStrictEqual(query.getCalls().map(call => call.args[0].variables), [
+				{ owner: remote.owner, name: remote.repositoryName, number: pr.number, first: 20, after: null },
+				{ owner: remote.owner, name: remote.repositoryName, number: pr.number, first: 20, after: 'first' },
+				{ owner: remote.owner, name: remote.repositoryName, number: pr.number, first: 5, after: 'first' },
+				{ owner: remote.owner, name: remote.repositoryName, number: pr.number, first: 1, after: 'first' },
+				{ owner: remote.owner, name: remote.repositoryName, number: pr.number, first: 1, after: 'second' },
+			]);
+		});
+
+		for (const [statusCode, pageSizes] of [[502, [20, 5, 1]], [403, [20]]] as const) {
+			it(`stops retrying review comments after HTTP ${statusCode}`, async function () {
+				const pr = new PullRequestBuilder().build();
+				const model = new PullRequestModel(credentials, telemetry, repo, remote, convertRESTPullRequestToRawPullRequest(pr, repo));
+				const query = sinon.stub(repo, 'query').rejects(Object.assign(new Error('Request failed'), {
+					networkError: { statusCode },
+				}));
+
+				assert.deepStrictEqual(await model.getReviewThreads(), []);
+				assert.deepStrictEqual(query.getCalls().map(call => call.args[0].variables?.first), [...pageSizes]);
+			});
+		}
+
+		it('reports missing review data without retrying', async function () {
+			const pr = new PullRequestBuilder().build();
+			const model = new PullRequestModel(credentials, telemetry, repo, remote, convertRESTPullRequestToRawPullRequest(pr, repo));
+			const query = sinon.stub(repo, 'query').resolves({
+				data: null, loading: false, stale: false, networkStatus: NetworkStatus.error,
+			});
+			const error = sinon.stub(Logger, 'error');
+
+			assert.deepStrictEqual(await model.getReviewThreads(), []);
+			assert.strictEqual(query.callCount, 1);
+			assert.strictEqual(error.lastCall.args[0], 'Failed to get pull request review comments: Error: Review comments response did not include a repository.');
+		});
+
 		it('should update the cache when then cache is initialized', async function () {
 			const pr = new PullRequestBuilder().build();
 			const model = new PullRequestModel(