Retry review comments with smaller pages

src/github/githubRepository.ts

@@ -351,7 +351,7 @@ export class GitHubRepository extends Disposable {
 			Logger.error(`Error querying GraphQL API (${logInfo}): ${e.message}${gqlErrors ? `. ${gqlErrors.map(error => error.extensions?.code).join(',')}` : ''}`, this.id);
 			if (legacyFallback) {
 				query.query = legacyFallback.query;
-				query.variables = legacyFallback.variables;
+				query.variables = legacyFallback.variables ?? query.variables;
 				return this.query(query, ignoreSamlErrors);
 			}
 

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,44 @@ 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) {
+				let data: PullRequestCommentsResponse | null;
+				try {
+					({ data } = await query<PullRequestCommentsResponse>({
+						query: schema.PullRequestComments,
+						variables: {
+							owner: remote.owner,
+							name: remote.repositoryName,
+							number: this.number,
+							first: pageSize,
+							after
+						},
+					}, false, { query: schema.LegacyPullRequestComments }));
+				} 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,37 @@ describe('GitHubRepository', function () {
 		sinon.restore();
 	});
 
+	describe('query', function () {
+		for (const replacement of [undefined, { owner: 'other', name: 'repo', number: 2, first: 20 }]) {
+			it(`uses ${replacement ? 'replacement' : 'original'} variables for a legacy fallback`, 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', number: 1, first: 20, after: 'cursor' };
+				const response = { data: {}, loading: false, stale: false, networkStatus: NetworkStatus.ready };
+				graphql.query.onFirstCall().rejects(new Error('Bad Gateway'));
+				graphql.query.onSecondCall().resolves(response);
+
+				try {
+					const result = await repo.query({ query: repo.schema.PullRequestComments, variables }, false, {
+						query: repo.schema.LegacyPullRequestComments,
+						variables: replacement,
+					});
+
+					assert.strictEqual(result, response);
+					assert.strictEqual(graphql.query.callCount, 2);
+					assert.strictEqual(graphql.query.secondCall.args[0].query, repo.schema.LegacyPullRequestComments);
+					assert.deepStrictEqual(graphql.query.secondCall.args[0].variables, replacement ?? 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,7 @@ import { NetworkStatus } from 'apollo-client';
 import { MockExtensionContext } from '../mocks/mockExtensionContext';
 import { GitHubServerType } from '../../common/authentication';
 import { mergeQuerySchemaWithShared } from '../../github/common';
+import Logger from '../../common/logger';
 const queries = mergeQuerySchemaWithShared(require('../../github/queries.gql'), require('../../github/queriesShared.gql')) as any;
 
 const telemetry = new MockTelemetry();
@@ -96,6 +97,73 @@ 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('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(