Batch mark/unmark files as viewed

src/commands.ts

@@ -1163,7 +1163,7 @@ ${contents}
 						}
 					}
 					const manager = reposManager.getManagerForFile(treeNode);
-					await manager?.activePullRequest?.markFileAsViewed(treeNode.path, true);
+					await manager?.activePullRequest?.markFiles([treeNode.path], true, 'viewed');
 					manager?.setFileViewedContext();
 				}
 			} catch (e) {
@@ -1184,7 +1184,7 @@ ${contents}
 					treeNode.unmarkFileAsViewed(false);
 				} else if (treeNode) {
 					const manager = reposManager.getManagerForFile(treeNode);
-					await manager?.activePullRequest?.unmarkFileAsViewed(treeNode.path, true);
+					await manager?.activePullRequest?.markFiles([treeNode.path], true, 'unviewed');
 					manager?.setFileViewedContext();
 				}
 			} catch (e) {

src/github/pullRequestModel.ts

@@ -6,6 +6,7 @@
 import * as buffer from 'buffer';
 import * as path from 'path';
 import equals from 'fast-deep-equal';
+import gql from 'graphql-tag';
 import * as vscode from 'vscode';
 import { Repository } from '../api/api';
 import { DiffSide, IComment, IReviewThread, SubjectType, ViewedState } from '../common/comment';
@@ -101,6 +102,8 @@ export interface FileViewedStateChangeEvent {
 
 export type FileViewedState = { [key: string]: ViewedState };
 
+const BATCH_SIZE = 100;
+
 export class PullRequestModel extends IssueModel<PullRequest> implements IPullRequestModel {
 	static ID = 'PullRequestModel';
 
@@ -1751,42 +1754,56 @@ export class PullRequestModel extends IssueModel<PullRequest> implements IPullRe
 		}
 	}
 
-	async markFileAsViewed(filePathOrSubpath: string, event: boolean): Promise<void> {
-		const { mutate, schema } = await this.githubRepository.ensure();
-		const fileName = filePathOrSubpath.startsWith(this.githubRepository.rootUri.path) ?
-			filePathOrSubpath.substring(this.githubRepository.rootUri.path.length + 1) : filePathOrSubpath;
-		await mutate<void>({
-			mutation: schema.MarkFileAsViewed,
-			variables: {
-				input: {
-					path: fileName,
-					pullRequestId: this.graphNodeId,
-				},
-			},
-		});
-
-		this.setFileViewedState(fileName, ViewedState.VIEWED, event);
-	}
-
-	async unmarkFileAsViewed(filePathOrSubpath: string, event: boolean): Promise<void> {
-		const { mutate, schema } = await this.githubRepository.ensure();
-		const fileName = filePathOrSubpath.startsWith(this.githubRepository.rootUri.path) ?
-			filePathOrSubpath.substring(this.githubRepository.rootUri.path.length + 1) : filePathOrSubpath;
-		await mutate<void>({
-			mutation: schema.UnmarkFileAsViewed,
-			variables: {
-				input: {
-					path: fileName,
-					pullRequestId: this.graphNodeId,
-				},
-			},
-		});
+	async markFiles(filePathOrSubpaths: string[], event: boolean, state: 'viewed' | 'unviewed'): Promise<void> {
+		const { mutate } = await this.githubRepository.ensure();
+		const pullRequestId = this.graphNodeId;
 
-		this.setFileViewedState(fileName, ViewedState.UNVIEWED, event);
-	}
+		const allFilenames = filePathOrSubpaths
+			.map((f) =>
+				f.startsWith(this.githubRepository.rootUri.path)
+					? f.substring(this.githubRepository.rootUri.path.length + 1)
+					: f
+			);
 
-	async unmarkAllFilesAsViewed(): Promise<void[]> {
-		return Promise.all(Array.from(this.fileChanges.keys()).map(change => this.unmarkFileAsViewed(change, true)));
+		const mutationName = state === 'viewed'
+			? 'markFileAsViewed'
+			: 'unmarkFileAsViewed';
+
+		// We only ever send 100 mutations at once. Any more than this and
+		// we risk a timeout from GitHub.
+		for (let i = 0; i < allFilenames.length; i += BATCH_SIZE) {
+			const batch = allFilenames.slice(i, i + BATCH_SIZE);
+			// See below for an example of what a mutation produced by this
+			// will look like
+			const mutation = gql`mutation Batch${mutationName}{
+				${batch.map((filename, i) =>
+				`alias${i}: ${mutationName}(
+						input: {path: "${filename}", pullRequestId: "${pullRequestId}"}
+					) { clientMutationId }
+					`
+			)}
+			}`;
+			await mutate<void>({ mutation });
+		}
+
+		// mutation BatchUnmarkFileAsViewedInline {
+		// 	alias0: unmarkFileAsViewed(
+		// 	  input: { path: "some_folder/subfolder/A.txt", pullRequestId: "PR_someid" }
+		// 	) {
+		// 	  clientMutationId
+		// 	}
+		// 	alias1: unmarkFileAsViewed(
+		// 	  input: { path: "some_folder/subfolder/B.txt", pullRequestId: "PR_someid" }
+		// 	) {
+		// 	  clientMutationId
+		// 	}
+		// }
+
+		filePathOrSubpaths.forEach(path => this.setFileViewedState(path, state === 'viewed' ? ViewedState.VIEWED : ViewedState.UNVIEWED, event));
+	}
+
+	async unmarkAllFilesAsViewed(): Promise<void> {
+		return this.markFiles(Array.from(this.fileChanges.keys()), true, 'unviewed');
 	}
 
 	private setFileViewedState(fileSubpath: string, viewedState: ViewedState, event: boolean) {

src/view/prChangesTreeDataProvider.ts

@@ -17,6 +17,7 @@ import { DescriptionNode } from './treeNodes/descriptionNode';
 import { GitFileChangeNode } from './treeNodes/fileChangeNode';
 import { RepositoryChangesNode } from './treeNodes/repositoryChangesNode';
 import { BaseTreeNode, TreeNode } from './treeNodes/treeNode';
+import { TreeUtils } from './treeNodes/treeUtils';
 
 export class PullRequestChangesTreeDataProvider extends vscode.Disposable implements vscode.TreeDataProvider<TreeNode>, BaseTreeNode {
 	private _onDidChangeTreeData = new vscode.EventEmitter<TreeNode | void>();
@@ -53,13 +54,7 @@ export class PullRequestChangesTreeDataProvider extends vscode.Disposable implem
 			}),
 		);
 
-		this._disposables.push(this._view.onDidChangeCheckboxState(checkboxUpdates => {
-			checkboxUpdates.items.forEach(checkboxUpdate => {
-				const node = checkboxUpdate[0];
-				const newState = checkboxUpdate[1];
-				node.updateFromCheckboxChanged(newState);
-			});
-		}));
+		this._disposables.push(this._view.onDidChangeCheckboxState(TreeUtils.processCheckboxUpdates));
 	}
 
 	refresh(treeNode?: TreeNode) {

src/view/prsTreeDataProvider.ts

@@ -24,6 +24,7 @@ import { CategoryTreeNode, PRCategoryActionNode, PRCategoryActionType } from './
 import { InMemFileChangeNode } from './treeNodes/fileChangeNode';
 import { BaseTreeNode, EXPANDED_QUERIES_STATE, TreeNode } from './treeNodes/treeNode';
 import { WorkspaceFolderNode } from './treeNodes/workspaceFolderNode';
+import { TreeUtils } from './treeNodes/treeUtils';
 
 export class PullRequestsTreeDataProvider implements vscode.TreeDataProvider<TreeNode>, BaseTreeNode, vscode.Disposable {
 	private _onDidChangeTreeData = new vscode.EventEmitter<TreeNode | void>();
@@ -107,13 +108,7 @@ export class PullRequestsTreeDataProvider implements vscode.TreeDataProvider<Tre
 			}),
 		);
 
-		this._disposables.push(this._view.onDidChangeCheckboxState(checkboxUpdates => {
-			checkboxUpdates.items.forEach(checkboxUpdate => {
-				const node = checkboxUpdate[0];
-				const newState = checkboxUpdate[1];
-				node.updateFromCheckboxChanged(newState);
-			});
-		}));
+		this._disposables.push(this._view.onDidChangeCheckboxState(TreeUtils.processCheckboxUpdates));
 
 		this._disposables.push(this._view.onDidExpandElement(expanded => {
 			this._updateExpandedQueries(expanded.element, true);

src/view/treeNodes/fileChangeNode.ts

@@ -167,24 +167,18 @@ export class FileChangeNode extends TreeNode implements vscode.TreeItem {
 	}
 
 	public async markFileAsViewed(fromCheckboxChanged: boolean = true) {
-		await this.pullRequest.markFileAsViewed(this.fileName, !fromCheckboxChanged);
+		await this.pullRequest.markFiles([this.fileName], !fromCheckboxChanged, 'viewed');
 		this.pullRequestManager.setFileViewedContext();
 	}
 
 	public async unmarkFileAsViewed(fromCheckboxChanged: boolean = true) {
-		await this.pullRequest.unmarkFileAsViewed(this.fileName, !fromCheckboxChanged);
+		await this.pullRequest.markFiles([this.fileName], !fromCheckboxChanged, 'unviewed');
 		this.pullRequestManager.setFileViewedContext();
 	}
 
 	updateFromCheckboxChanged(newState: vscode.TreeItemCheckboxState) {
 		const viewed = newState === vscode.TreeItemCheckboxState.Checked ? ViewedState.VIEWED : ViewedState.UNVIEWED;
 		this.updateViewed(viewed);
-
-		if (newState === vscode.TreeItemCheckboxState.Checked) {
-			this.markFileAsViewed();
-		} else {
-			this.unmarkFileAsViewed();
-		}
 	}
 
 	updateShowOptions() {

src/view/treeNodes/treeNode.ts

@@ -66,7 +66,6 @@ export abstract class TreeNode implements vscode.Disposable {
 
 	updateFromCheckboxChanged(_newState: vscode.TreeItemCheckboxState): void { }
 
-
 	dispose(): void {
 		if (this.childrenDisposables) {
 			dispose(this.childrenDisposables);

src/view/treeNodes/treeUtils.ts

@@ -0,0 +1,41 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as vscode from 'vscode';
+import { FileChangeNode } from './fileChangeNode';
+import { TreeNode } from './treeNode';
+
+export namespace TreeUtils {
+	export function processCheckboxUpdates(checkboxUpdates: vscode.TreeCheckboxChangeEvent<TreeNode>) {
+		const checkedNodes: FileChangeNode[] = [];
+		const uncheckedNodes: FileChangeNode[] = [];
+
+		checkboxUpdates.items.forEach(checkboxUpdate => {
+			const node = checkboxUpdate[0];
+			const newState = checkboxUpdate[1];
+
+			if (node instanceof FileChangeNode) {
+				if (newState == vscode.TreeItemCheckboxState.Checked) {
+					checkedNodes.push(node);
+				} else {
+					uncheckedNodes.push(node);
+				}
+			}
+
+			node.updateFromCheckboxChanged(newState);
+		});
+
+		if (checkedNodes.length > 0) {
+			const prModel = checkedNodes[0].pullRequest;
+			const filenames = checkedNodes.map(n => n.fileName);
+			prModel.markFiles(filenames, true, 'viewed');
+		}
+		if (uncheckedNodes.length > 0) {
+			const prModel = uncheckedNodes[0].pullRequest;
+			const filenames = uncheckedNodes.map(n => n.fileName);
+			prModel.markFiles(filenames, true, 'unviewed');
+		}
+	}
+}