manage scheduling of getting children to prevent too many requests (#202904)

Author: amunger

src/vs/workbench/contrib/notebook/browser/contrib/notebookVariables/notebookVariablesDataSource.ts

@@ -4,7 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { IAsyncDataSource } from 'vs/base/browser/ui/tree/tree';
-import { CancellationToken } from 'vs/base/common/cancellation';
+import { CancellationTokenSource } from 'vs/base/common/cancellation';
 import { NotebookTextModel } from 'vs/workbench/contrib/notebook/common/model/notebookTextModel';
 import { INotebookKernel, INotebookKernelService, VariablesResult, variablePageSize } from 'vs/workbench/contrib/notebook/common/notebookKernelService';
 
@@ -28,12 +28,22 @@ export interface INotebookVariableElement {
 
 export class NotebookVariableDataSource implements IAsyncDataSource<INotebookScope, INotebookVariableElement> {
 
-	constructor(private readonly notebookKernelService: INotebookKernelService) { }
+	private cancellationTokenSource: CancellationTokenSource;
+
+	constructor(private readonly notebookKernelService: INotebookKernelService) {
+		this.cancellationTokenSource = new CancellationTokenSource();
+	}
 
 	hasChildren(element: INotebookScope | INotebookVariableElement): boolean {
 		return element.kind === 'root' || element.hasNamedChildren || element.indexedChildrenCount > 0;
 	}
 
+	public cancel(): void {
+		this.cancellationTokenSource.cancel();
+		this.cancellationTokenSource.dispose();
+		this.cancellationTokenSource = new CancellationTokenSource();
+	}
+
 	async getChildren(element: INotebookScope | INotebookVariableElement): Promise<Array<INotebookVariableElement>> {
 		if (element.kind === 'root') {
 			return this.getRootVariables(element.notebook);
@@ -48,7 +58,7 @@ export class NotebookVariableDataSource implements IAsyncDataSource<INotebookSco
 
 			let children: INotebookVariableElement[] = [];
 			if (parent.hasNamedChildren) {
-				const variables = selectedKernel.provideVariables(parent.notebook.uri, parent.extHostId, 'named', 0, CancellationToken.None);
+				const variables = selectedKernel.provideVariables(parent.notebook.uri, parent.extHostId, 'named', 0, this.cancellationTokenSource.token);
 				const childNodes = await variables
 					.map(variable => { return this.createVariableElement(variable, parent.notebook); })
 					.toPromise();
@@ -88,7 +98,7 @@ export class NotebookVariableDataSource implements IAsyncDataSource<INotebookSco
 			}
 		}
 		else if (parent.indexedChildrenCount > 0) {
-			const variables = kernel.provideVariables(parent.notebook.uri, parent.extHostId, 'indexed', parent.indexStart ?? 0, CancellationToken.None);
+			const variables = kernel.provideVariables(parent.notebook.uri, parent.extHostId, 'indexed', parent.indexStart ?? 0, this.cancellationTokenSource.token);
 
 			for await (const variable of variables) {
 				childNodes.push(this.createVariableElement(variable, parent.notebook));
@@ -104,7 +114,7 @@ export class NotebookVariableDataSource implements IAsyncDataSource<INotebookSco
 	async getRootVariables(notebook: NotebookTextModel): Promise<INotebookVariableElement[]> {
 		const selectedKernel = this.notebookKernelService.getMatchingKernel(notebook).selected;
 		if (selectedKernel && selectedKernel.hasVariableProvider) {
-			const variables = selectedKernel.provideVariables(notebook.uri, undefined, 'named', 0, CancellationToken.None);
+			const variables = selectedKernel.provideVariables(notebook.uri, undefined, 'named', 0, this.cancellationTokenSource.token);
 			return await variables
 				.map(variable => { return this.createVariableElement(variable, notebook); })
 				.toPromise();

src/vs/workbench/contrib/notebook/browser/contrib/notebookVariables/notebookVariablesView.ts

@@ -3,6 +3,7 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
+import { RunOnceScheduler } from 'vs/base/common/async';
 import { URI } from 'vs/base/common/uri';
 import * as nls from 'vs/nls';
 import { ILocalizedString } from 'vs/platform/action/common/action';
@@ -34,6 +35,9 @@ export class NotebookVariablesView extends ViewPane {
 
 	private tree: WorkbenchAsyncDataTree<INotebookScope, INotebookVariableElement> | undefined;
 	private activeNotebook: NotebookTextModel | undefined;
+	private readonly dataSource: NotebookVariableDataSource;
+
+	private updateScheduler: RunOnceScheduler;
 
 	constructor(
 		options: IViewPaneOptions,
@@ -59,6 +63,9 @@ export class NotebookVariablesView extends ViewPane {
 		this._register(this.notebookExecutionStateService.onDidChangeExecution(this.handleExecutionStateChange.bind(this)));
 
 		this.setActiveNotebook();
+
+		this.dataSource = new NotebookVariableDataSource(this.notebookKernelService);
+		this.updateScheduler = new RunOnceScheduler(() => this.tree?.updateChildren(), 100);
 	}
 
 	protected override renderBody(container: HTMLElement): void {
@@ -70,7 +77,7 @@ export class NotebookVariablesView extends ViewPane {
 			container,
 			new NotebookVariablesDelegate(),
 			[new NotebookVariableRenderer()],
-			new NotebookVariableDataSource(this.notebookKernelService),
+			this.dataSource,
 			{
 				accessibilityProvider: new NotebookVariableAccessibilityProvider(),
 				identityProvider: { getId: (e: INotebookVariableElement) => e.id },
@@ -87,7 +94,7 @@ export class NotebookVariablesView extends ViewPane {
 		this.tree?.layout(height, width);
 	}
 
-	setActiveNotebook() {
+	private setActiveNotebook() {
 		const current = this.activeNotebook;
 		const activeEditorPane = this.editorService.activeEditorPane;
 		if (activeEditorPane && activeEditorPane.getId() === 'workbench.editor.notebook') {
@@ -101,23 +108,31 @@ export class NotebookVariablesView extends ViewPane {
 	private handleActiveEditorChange() {
 		if (this.setActiveNotebook() && this.activeNotebook) {
 			this.tree?.setInput({ kind: 'root', notebook: this.activeNotebook });
-			this.tree?.updateChildren();
+			this.updateScheduler.schedule();
 		}
 	}
 
 	private handleExecutionStateChange(event: ICellExecutionStateChangedEvent | IExecutionStateChangedEvent) {
 		if (this.activeNotebook) {
-			// changed === undefined -> excecution ended
-			if (event.changed === undefined && event.affectsNotebook(this.activeNotebook?.uri)) {
-				this.tree?.updateChildren();
+			if (event.affectsNotebook(this.activeNotebook.uri)) {
+				// new execution state means either new variables or the kernel is busy so we shouldn't ask
+				this.dataSource.cancel();
+
+				// changed === undefined -> excecution ended
+				if (event.changed === undefined) {
+					this.updateScheduler.schedule();
+				}
+				else {
+					this.updateScheduler.cancel();
+				}
 			}
 		}
 	}
 
 	private handleVariablesChanged(notebookUri: URI) {
 		if (this.activeNotebook && notebookUri.toString() === this.activeNotebook.uri.toString()) {
 			this.tree?.setInput({ kind: 'root', notebook: this.activeNotebook });
-			this.tree?.updateChildren();
+			this.updateScheduler.schedule();
 		}
 	}
 }