move NotebookCellExecutionTask into extHostNotebookKernels, fyi @roblourens

src/vs/workbench/api/common/extHostNotebookKernels.ts

@@ -4,8 +4,8 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { Emitter } from 'vs/base/common/event';
-import { DisposableStore } from 'vs/base/common/lifecycle';
-import { ExtHostNotebookKernelsShape, IMainContext, INotebookKernelDto2, MainContext, MainThreadNotebookKernelsShape } from 'vs/workbench/api/common/extHost.protocol';
+import { Disposable, DisposableStore } from 'vs/base/common/lifecycle';
+import { ExtHostNotebookKernelsShape, IMainContext, INotebookKernelDto2, MainContext, MainThreadNotebookDocumentsShape, MainThreadNotebookKernelsShape } from 'vs/workbench/api/common/extHost.protocol';
 import * as vscode from 'vscode';
 import { ExtHostNotebookController } from 'vs/workbench/api/common/extHostNotebook';
 import { ExtensionIdentifier, IExtensionDescription } from 'vs/platform/extensions/common/extensions';
@@ -14,6 +14,11 @@ import * as extHostTypeConverters from 'vs/workbench/api/common/extHostTypeConve
 import { IExtHostInitDataService } from 'vs/workbench/api/common/extHostInitDataService';
 import { asWebviewUri } from 'vs/workbench/api/common/shared/webview';
 import { ResourceMap } from 'vs/base/common/map';
+import { timeout } from 'vs/base/common/async';
+import { ExtHostCell, ExtHostNotebookDocument } from 'vs/workbench/api/common/extHostNotebookDocument';
+import { CellEditType, IImmediateCellEditOperation, NullablePartialNotebookCellMetadata } from 'vs/workbench/contrib/notebook/common/notebookCommon';
+import { CancellationTokenSource } from 'vs/base/common/cancellation';
+import { NotebookCellExecutionState } from 'vs/workbench/api/common/extHostTypes';
 
 interface IKernelData {
 	extensionId: ExtensionIdentifier,
@@ -26,16 +31,17 @@ interface IKernelData {
 export class ExtHostNotebookKernels implements ExtHostNotebookKernelsShape {
 
 	private readonly _proxy: MainThreadNotebookKernelsShape;
+	private readonly _activeExecutions = new ResourceMap<NotebookCellExecutionTask>();
 
 	private readonly _kernelData = new Map<number, IKernelData>();
 	private _handlePool: number = 0;
 
 	constructor(
-		mainContext: IMainContext,
+		private readonly _mainContext: IMainContext,
 		private readonly _initData: IExtHostInitDataService,
 		private readonly _extHostNotebook: ExtHostNotebookController
 	) {
-		this._proxy = mainContext.getProxy(MainContext.MainThreadNotebookKernels);
+		this._proxy = _mainContext.getProxy(MainContext.MainThreadNotebookKernels);
 	}
 
 	createNotebookController(extension: IExtensionDescription, id: string, viewType: string, label: string, handler?: vscode.NotebookExecuteHandler, preloads?: vscode.NotebookKernelPreload[]): vscode.NotebookController {
@@ -153,21 +159,13 @@ export class ExtHostNotebookKernels implements ExtHostNotebookKernelsShape {
 				_update();
 			},
 			createNotebookCellExecutionTask(cell) {
-				if (cell.index < 0) {
-					throw new Error('CANNOT execute cell that has been REMOVED from notebook');
-				}
 				if (isDisposed) {
 					throw new Error('notebook controller is DISPOSED');
 				}
 				if (!associatedNotebooks.has(cell.notebook.uri)) {
 					throw new Error(`notebook controller is NOT associated to notebook: ${cell.notebook.uri.toString()}`);
 				}
-				const notebook = that._extHostNotebook.getNotebookDocument(cell.notebook.uri);
-				const cellObj = notebook.getCellFromApiCell(cell);
-				if (!cellObj) {
-					throw new Error('invalid cell');
-				}
-				return that._extHostNotebook.createNotebookCellExecution(cellObj)!;
+				return that._createNotebookCellExecution(cell);
 			},
 			dispose: () => {
 				if (!isDisposed) {
@@ -259,7 +257,7 @@ export class ExtHostNotebookKernels implements ExtHostNotebookKernelsShape {
 		for (let cellHandle of handles) {
 			const cell = document.getCell(cellHandle);
 			if (cell) {
-				this._extHostNotebook.cancelOneNotebookCellExecution(cell);
+				this._activeExecutions.get(cell.uri)?.cancel();
 			}
 		}
 	}
@@ -278,4 +276,212 @@ export class ExtHostNotebookKernels implements ExtHostNotebookKernelsShape {
 
 		obj.onDidReceiveMessage.fire(Object.freeze({ editor: editor.apiEditor, message }));
 	}
+
+	// ---
+
+	_createNotebookCellExecution(cell: vscode.NotebookCell): vscode.NotebookCellExecutionTask {
+		if (cell.index < 0) {
+			throw new Error('CANNOT execute cell that has been REMOVED from notebook');
+		}
+		const notebook = this._extHostNotebook.getNotebookDocument(cell.notebook.uri);
+		const cellObj = notebook.getCellFromApiCell(cell);
+		if (!cellObj) {
+			throw new Error('invalid cell');
+		}
+		if (this._activeExecutions.has(cellObj.uri)) {
+			throw new Error(`duplicate execution for ${cellObj.uri}`);
+		}
+		const execution = new NotebookCellExecutionTask(cellObj.notebook, cellObj, this._mainContext.getProxy(MainContext.MainThreadNotebookDocuments));
+		this._activeExecutions.set(cellObj.uri, execution);
+		const listener = execution.onDidChangeState(() => {
+			if (execution.state === NotebookCellExecutionTaskState.Resolved) {
+				execution.dispose();
+				listener.dispose();
+				this._activeExecutions.delete(cellObj.uri);
+			}
+		});
+		return execution.asApiObject();
+	}
+}
+
+
+enum NotebookCellExecutionTaskState {
+	Init,
+	Started,
+	Resolved
+}
+
+class NotebookCellExecutionTask extends Disposable {
+	private _onDidChangeState = new Emitter<void>();
+	readonly onDidChangeState = this._onDidChangeState.event;
+
+	private _state = NotebookCellExecutionTaskState.Init;
+	get state(): NotebookCellExecutionTaskState { return this._state; }
+
+	private readonly _tokenSource = this._register(new CancellationTokenSource());
+
+	private readonly _collector: TimeoutBasedCollector<IImmediateCellEditOperation>;
+
+	private _executionOrder: number | undefined;
+
+	constructor(
+		private readonly _document: ExtHostNotebookDocument,
+		private readonly _cell: ExtHostCell,
+		private readonly _proxy: MainThreadNotebookDocumentsShape
+	) {
+		super();
+
+		this._collector = new TimeoutBasedCollector(10, edits => this.applyEdits(edits));
+
+		this._executionOrder = _cell.internalMetadata.executionOrder;
+		this.mixinMetadata({
+			runState: NotebookCellExecutionState.Pending,
+			executionOrder: null
+		});
+	}
+
+	cancel(): void {
+		this._tokenSource.cancel();
+	}
+
+	private async applyEditSoon(edit: IImmediateCellEditOperation): Promise<void> {
+		await this._collector.addItem(edit);
+	}
+
+	private async applyEdits(edits: IImmediateCellEditOperation[]): Promise<void> {
+		return this._proxy.$applyEdits(this._document.uri, edits, false);
+	}
+
+	private verifyStateForOutput() {
+		if (this._state === NotebookCellExecutionTaskState.Init) {
+			throw new Error('Must call start before modifying cell output');
+		}
+
+		if (this._state === NotebookCellExecutionTaskState.Resolved) {
+			throw new Error('Cannot modify cell output after calling resolve');
+		}
+	}
+
+	private mixinMetadata(mixinMetadata: NullablePartialNotebookCellMetadata) {
+		const edit: IImmediateCellEditOperation = { editType: CellEditType.PartialMetadata, handle: this._cell.handle, metadata: mixinMetadata };
+		this.applyEdits([edit]);
+	}
+
+	private cellIndexToHandle(cellIndex: number | undefined): number | undefined {
+		const cell = typeof cellIndex === 'number' ? this._document.getCellFromIndex(cellIndex) : this._cell;
+		if (!cell) {
+			return;
+		}
+
+		return cell.handle;
+	}
+
+	asApiObject(): vscode.NotebookCellExecutionTask {
+		const that = this;
+		return Object.freeze(<vscode.NotebookCellExecutionTask>{
+			get document() { return that._document.apiNotebook; },
+			get cell() { return that._cell.apiCell; },
+
+			get executionOrder() { return that._executionOrder; },
+			set executionOrder(v: number | undefined) {
+				that._executionOrder = v;
+				that.mixinMetadata({
+					executionOrder: v
+				});
+			},
+
+			start(context?: vscode.NotebookCellExecuteStartContext): void {
+				if (that._state === NotebookCellExecutionTaskState.Resolved || that._state === NotebookCellExecutionTaskState.Started) {
+					throw new Error('Cannot call start again');
+				}
+
+				that._state = NotebookCellExecutionTaskState.Started;
+				that._onDidChangeState.fire();
+
+				that.mixinMetadata({
+					runState: NotebookCellExecutionState.Executing,
+					runStartTime: context?.startTime ?? null
+				});
+			},
+
+			end(result?: vscode.NotebookCellExecuteEndContext): void {
+				if (that._state === NotebookCellExecutionTaskState.Resolved) {
+					throw new Error('Cannot call resolve twice');
+				}
+
+				that._state = NotebookCellExecutionTaskState.Resolved;
+				that._onDidChangeState.fire();
+
+				that.mixinMetadata({
+					runState: NotebookCellExecutionState.Idle,
+					lastRunSuccess: result?.success ?? null,
+					runEndTime: result?.endTime ?? null,
+				});
+			},
+
+			clearOutput(cellIndex?: number): Thenable<void> {
+				that.verifyStateForOutput();
+				return this.replaceOutput([], cellIndex);
+			},
+
+			async appendOutput(outputs: vscode.NotebookCellOutput | vscode.NotebookCellOutput[], cellIndex?: number): Promise<void> {
+				that.verifyStateForOutput();
+				const handle = that.cellIndexToHandle(cellIndex);
+				if (typeof handle !== 'number') {
+					return;
+				}
+
+				outputs = Array.isArray(outputs) ? outputs : [outputs];
+				return that.applyEditSoon({ editType: CellEditType.Output, handle, append: true, outputs: outputs.map(extHostTypeConverters.NotebookCellOutput.from) });
+			},
+
+			async replaceOutput(outputs: vscode.NotebookCellOutput | vscode.NotebookCellOutput[], cellIndex?: number): Promise<void> {
+				that.verifyStateForOutput();
+				const handle = that.cellIndexToHandle(cellIndex);
+				if (typeof handle !== 'number') {
+					return;
+				}
+
+				outputs = Array.isArray(outputs) ? outputs : [outputs];
+				return that.applyEditSoon({ editType: CellEditType.Output, handle, outputs: outputs.map(extHostTypeConverters.NotebookCellOutput.from) });
+			},
+
+			async appendOutputItems(items: vscode.NotebookCellOutputItem | vscode.NotebookCellOutputItem[], outputId: string): Promise<void> {
+				that.verifyStateForOutput();
+				items = Array.isArray(items) ? items : [items];
+				return that.applyEditSoon({ editType: CellEditType.OutputItems, append: true, items: items.map(extHostTypeConverters.NotebookCellOutputItem.from), outputId });
+			},
+
+			async replaceOutputItems(items: vscode.NotebookCellOutputItem | vscode.NotebookCellOutputItem[], outputId: string): Promise<void> {
+				that.verifyStateForOutput();
+				items = Array.isArray(items) ? items : [items];
+				return that.applyEditSoon({ editType: CellEditType.OutputItems, items: items.map(extHostTypeConverters.NotebookCellOutputItem.from), outputId });
+			},
+
+			token: that._tokenSource.token
+		});
+	}
+}
+
+class TimeoutBasedCollector<T> {
+	private batch: T[] = [];
+	private waitPromise: Promise<void> | undefined;
+
+	constructor(
+		private readonly delay: number,
+		private readonly callback: (items: T[]) => Promise<void>) { }
+
+	addItem(item: T): Promise<void> {
+		this.batch.push(item);
+		if (!this.waitPromise) {
+			this.waitPromise = timeout(this.delay).then(() => {
+				this.waitPromise = undefined;
+				const batch = this.batch;
+				this.batch = [];
+				return this.callback(batch);
+			});
+		}
+
+		return this.waitPromise;
+	}
 }