Adds cancellation support to Graph refresh

Author: eamodio

src/webviews/plus/graph/graphWebview.ts

@@ -113,7 +113,6 @@ import { isDarkTheme, isLightTheme } from '../../../system/-webview/vscode';
 import { openUrl } from '../../../system/-webview/vscode/uris';
 import type { OpenWorkspaceLocation } from '../../../system/-webview/vscode/workspaces';
 import { openWorkspace } from '../../../system/-webview/vscode/workspaces';
-import { gate } from '../../../system/decorators/-webview/gate';
 import { debug, log } from '../../../system/decorators/log';
 import { disposableInterval } from '../../../system/function';
 import type { Deferrable } from '../../../system/function/debounce';
@@ -247,7 +246,7 @@ const compactGraphColumnsSettings: GraphColumnsSettings = {
 	sha: { width: 130, isHidden: false, order: 6 },
 };
 
-type CancellableOperations = 'hover' | 'computeIncludedRefs' | 'search' | 'state';
+type CancellableOperations = 'branchState' | 'hover' | 'computeIncludedRefs' | 'search' | 'state';
 
 export class GraphWebviewProvider implements WebviewProvider<State, State, GraphWebviewShowingArgs> {
 	private _repository?: Repository;
@@ -1473,7 +1472,6 @@ export class GraphWebviewProvider implements WebviewProvider<State, State, Graph
 		this.updateRefsMetadata();
 	}
 
-	@gate()
 	@debug()
 	private async onGetMoreRows(e: GetMoreRowsParams, sendSelectedRows: boolean = false) {
 		if (this._graph?.paging == null) return;
@@ -1483,13 +1481,7 @@ export class GraphWebviewProvider implements WebviewProvider<State, State, Graph
 			return;
 		}
 
-		using sw = new Stopwatch(`GraohWebviewProvider.onGetMoreRows(${this.host.id})`);
 		await this.updateGraphWithMoreRows(this._graph, e.id, this._search);
-		this.container.telemetry.sendEvent('graph/rows/loaded', {
-			...this.getTelemetryContext(),
-			duration: sw.elapsed(),
-			rows: this._graph.rows.length ?? 0,
-		});
 		void this.notifyDidChangeRows(sendSelectedRows);
 	}
 
@@ -2201,14 +2193,18 @@ export class GraphWebviewProvider implements WebviewProvider<State, State, Graph
 				const cancellation = this.createCancellation('computeIncludedRefs');
 
 				const [baseResult, defaultResult, targetResult] = await Promise.allSettled([
-					this.container.git.branches(current.repoPath).getBaseBranchName?.(current.name),
-					getDefaultBranchName(this.container, current.repoPath, current.getRemoteName()),
+					this.container.git.branches(current.repoPath).getBaseBranchName?.(current.name, cancellation.token),
+					getDefaultBranchName(this.container, current.repoPath, current.getRemoteName(), {
+						cancellation: cancellation.token,
+					}),
 					getTargetBranchName(this.container, current, {
 						cancellation: cancellation.token,
 						timeout: options?.timeout,
 					}),
 				]);
 
+				if (cancellation.token.isCancellationRequested) return { refs: {} };
+
 				const baseBranchName = getSettledValue(baseResult);
 				const defaultBranchName = getSettledValue(defaultResult);
 				const targetMaybeResult = getSettledValue(targetResult);
@@ -2467,13 +2463,13 @@ export class GraphWebviewProvider implements WebviewProvider<State, State, Graph
 		return item;
 	}
 
-	private async getWorkingTreeStats(): Promise<GraphWorkingTreeStats | undefined> {
+	private async getWorkingTreeStats(cancellation?: CancellationToken): Promise<GraphWorkingTreeStats | undefined> {
 		if (this.repository == null || this.container.git.repositoryCount === 0) return undefined;
 
-		const statusProivder = this.container.git.status(this.repository.path);
-		const status = await statusProivder.getStatus();
+		const statusProvider = this.container.git.status(this.repository.path);
+		const status = await statusProvider.getStatus(cancellation);
 		const workingTreeStatus = status?.getDiffStatus();
-		const pausedOpStatus = await statusProivder.getPausedOperationStatus?.();
+		const pausedOpStatus = await statusProvider.getPausedOperationStatus?.(cancellation);
 
 		return {
 			added: workingTreeStatus?.added ?? 0,
@@ -2492,6 +2488,9 @@ export class GraphWebviewProvider implements WebviewProvider<State, State, Graph
 	}
 
 	private async getState(deferRows?: boolean): Promise<State> {
+		this.cancelOperation('branchState');
+		this.cancelOperation('state');
+
 		if (this.container.git.repositoryCount === 0) {
 			return { ...this.host.baseWebviewState, allowed: true, repositories: [] };
 		}
@@ -2503,6 +2502,8 @@ export class GraphWebviewProvider implements WebviewProvider<State, State, Graph
 			}
 		}
 
+		const cancellation = this.createCancellation('state');
+
 		this._etagRepository = this.repository?.etag;
 		this.host.title = `${this.host.originalTitle}: ${this.repository.formattedName}`;
 
@@ -2517,35 +2518,42 @@ export class GraphWebviewProvider implements WebviewProvider<State, State, Graph
 		const columns = this.getColumns();
 		const columnSettings = this.getColumnSettings(columns);
 
-		const dataPromise = this.repository.git.graph().getGraph(rev, uri => this.host.asWebviewUri(uri), {
-			include: {
-				stats:
-					(configuration.get('graph.minimap.enabled') &&
-						configuration.get('graph.minimap.dataType') === 'lines') ||
-					!columnSettings.changes.isHidden,
+		const dataPromise = this.repository.git.graph().getGraph(
+			rev,
+			uri => this.host.asWebviewUri(uri),
+			{
+				include: {
+					stats:
+						(configuration.get('graph.minimap.enabled') &&
+							configuration.get('graph.minimap.dataType') === 'lines') ||
+						!columnSettings.changes.isHidden,
+				},
+				limit: limit,
 			},
-			limit: limit,
-		});
+			cancellation.token,
+		);
 
 		// Check for access and working tree stats
 		const promises = Promise.allSettled([
 			this.getGraphAccess(),
-			this.getWorkingTreeStats(),
-			this.repository.git.branches().getBranch(),
+			this.getWorkingTreeStats(cancellation.token),
+			this.repository.git.branches().getBranch(undefined, cancellation.token),
 			this.repository.getLastFetched(),
 		]);
 
 		let data;
 		if (deferRows) {
 			queueMicrotask(async () => {
-				const data = await dataPromise;
-				this.setGraph(data);
-				if (selectedId !== uncommitted) {
-					this.setSelectedRows(data.id);
-				}
+				try {
+					const data = await dataPromise;
+					this.setGraph(data);
+					if (selectedId !== uncommitted) {
+						this.setSelectedRows(data.id);
+					}
 
-				void this.notifyDidChangeRefsVisibility();
-				void this.notifyDidChangeRows(true);
+					void this.notifyDidChangeRefsVisibility();
+					void this.notifyDidChangeRows(true);
+				} catch {}
 			});
 		} else {
 			data = await dataPromise;
@@ -2556,6 +2564,8 @@ export class GraphWebviewProvider implements WebviewProvider<State, State, Graph
 		}
 
 		const [accessResult, workingStatsResult, branchResult, lastFetchedResult] = await promises;
+		if (cancellation.token.isCancellationRequested) throw new CancellationError();
+
 		const [access, visibility] = getSettledValue(accessResult) ?? [];
 
 		let branchState: BranchState | undefined;
@@ -2564,17 +2574,18 @@ export class GraphWebviewProvider implements WebviewProvider<State, State, Graph
 		if (branch != null) {
 			branchState = { ...(branch.upstream?.state ?? { ahead: 0, behind: 0 }) };
 
-			const worktreesByBranch = data?.worktreesByBranch ?? (await getWorktreesByBranch(this.repository));
+			const worktreesByBranch =
+				data?.worktreesByBranch ?? (await getWorktreesByBranch(this.repository, undefined, cancellation.token));
 			branchState.worktree = worktreesByBranch?.has(branch.id) ?? false;
 
 			if (branch.upstream != null) {
 				branchState.upstream = branch.upstream.name;
 
-				const cancellation = this.createCancellation('state');
+				const branchStateCancellation = this.createCancellation('branchState');
 
 				const [remoteResult, prResult] = await Promise.allSettled([
 					branch.getRemote(),
-					pauseOnCancelOrTimeout(branch.getAssociatedPullRequest(), cancellation.token, 100),
+					pauseOnCancelOrTimeout(branch.getAssociatedPullRequest(), branchStateCancellation.token, 100),
 				]);
 
 				const remote = getSettledValue(remoteResult);
@@ -2590,7 +2601,7 @@ export class GraphWebviewProvider implements WebviewProvider<State, State, Graph
 				if (maybePr?.paused) {
 					const updatedBranchState = { ...branchState };
 					void maybePr.value.then(pr => {
-						if (cancellation?.token.isCancellationRequested) return;
+						if (branchStateCancellation?.token.isCancellationRequested) return;
 
 						if (pr != null) {
 							updatedBranchState.pr = serializePullRequest(pr);
@@ -3044,9 +3055,68 @@ export class GraphWebviewProvider implements WebviewProvider<State, State, Graph
 		}
 	}
 
+	private _pendingRowsQuery:
+		| {
+				promise: Promise<void>;
+				cancellable: CancellationTokenSource;
+				id?: string | undefined;
+				search?: GitGraphSearch;
+		  }
+		| undefined;
 	private async updateGraphWithMoreRows(graph: GitGraph, id: string | undefined, search?: GitGraphSearch) {
+		if (this._pendingRowsQuery != null) {
+			const { id: pendingId, search: pendingSearch } = this._pendingRowsQuery;
+			if (pendingSearch === search && (pendingId === id || (pendingId != null && id == null))) {
+				return this._pendingRowsQuery.promise;
+			}
+
+			this._pendingRowsQuery.cancellable.cancel();
+			this._pendingRowsQuery.cancellable.dispose();
+			this._pendingRowsQuery = undefined;
+		}
+
+		const sw = new Stopwatch(undefined);
+
+		const cancellable = new CancellationTokenSource();
+		const cancellation = cancellable.token;
+
+		this._pendingRowsQuery = {
+			promise: this.updateGraphWithMoreRowsCore(graph, id, search, cancellation).catch((ex: unknown) => {
+				if (cancellation.isCancellationRequested) return;
+
+				throw ex;
+			}),
+			cancellable: cancellable,
+			id: id,
+			search: search,
+		};
+
+		void this._pendingRowsQuery.promise.finally(() => {
+			if (cancellation.isCancellationRequested) return;
+
+			this.container.telemetry.sendEvent('graph/rows/loaded', {
+				...this.getTelemetryContext(),
+				duration: sw.elapsed(),
+				rows: graph.rows.length ?? 0,
+			});
+			sw.stop();
+
+			this._pendingRowsQuery = undefined;
+		});
+
+		return this._pendingRowsQuery.promise;
+	}
+
+	private async updateGraphWithMoreRowsCore(
+		graph: GitGraph,
+		id: string | undefined,
+		search?: GitGraphSearch,
+		cancellation?: CancellationToken,
+	) {
+		console.warn('##### updateGraphWithMoreRows', id, search);
+
 		const { defaultItemLimit, pageItemLimit } = configuration.get('graph');
-		const updatedGraph = await graph.more?.(pageItemLimit ?? defaultItemLimit, id);
+		const updatedGraph = await graph.more?.(pageItemLimit ?? defaultItemLimit, id ?? undefined, cancellation);
 		if (updatedGraph != null) {
 			this.setGraph(updatedGraph);
 

src/webviews/webviewController.ts

@@ -5,6 +5,7 @@ import type { WebviewCommands, WebviewViewCommands } from '../constants.commands
 import type { WebviewTelemetryContext } from '../constants.telemetry';
 import type { CustomEditorTypes, WebviewIds, WebviewTypes, WebviewViewIds, WebviewViewTypes } from '../constants.views';
 import type { Container } from '../container';
+import { isCancellationError } from '../errors';
 import { executeCommand, executeCoreCommand } from '../system/-webview/command';
 import { setContext } from '../system/-webview/context';
 import { getScopedCounter } from '../system/counter';
@@ -363,7 +364,16 @@ export class WebviewController<
 
 		if (loading) {
 			this.cancellation ??= new CancellationTokenSource();
-			this.webview.html = await this.getHtml(this.webview);
+			try {
+				this.webview.html = await this.getHtml(this.webview);
+			} catch (ex) {
+				if (isCancellationError(ex)) {
+					this.cancellation.cancel();
+					return;
+				}
+
+				throw ex;
+			}
 		}
 
 		if (this.isHost('editor')) {
@@ -420,7 +430,18 @@ export class WebviewController<
 		const wasReady = this._ready;
 		this._ready = false;
 
-		const html = await this.getHtml(this.webview);
+		let html;
+		try {
+			html = await this.getHtml(this.webview);
+		} catch (ex) {
+			if (isCancellationError(ex)) {
+				this.cancellation.cancel();
+				return;
+			}
+
+			throw ex;
+		}
+
 		if (force) {
 			// Reset the html to get the webview to reload
 			this.webview.html = '';