send controller selection event when selection knowledge is restored …

…from memento, #121904
microsoft · Apr 22, 2021 · 1df9d89 · 1df9d89
1 parent b3151c1
commit 1df9d89
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 15 deletions.
diff --git a/src/vs/workbench/api/browser/mainThreadNotebookKernels.ts b/src/vs/workbench/api/browser/mainThreadNotebookKernels.ts
@@ -4,8 +4,9 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { flatten, isNonEmptyArray } from 'vs/base/common/arrays';
+import { runWhenIdle } from 'vs/base/common/async';
 import { Emitter, Event } from 'vs/base/common/event';
-import { combinedDisposable, DisposableStore, IDisposable } from 'vs/base/common/lifecycle';
+import { combinedDisposable, DisposableStore, IDisposable, MutableDisposable } from 'vs/base/common/lifecycle';
 import { URI } from 'vs/base/common/uri';
 import { IModeService } from 'vs/editor/common/services/modeService';
 import { ExtensionIdentifier } from 'vs/platform/extensions/common/extensions';
@@ -197,15 +198,21 @@ export class MainThreadNotebookKernels implements MainThreadNotebookKernelsShape
 		}(data, this._modeService);
 		const registration = this._notebookKernelService.registerKernel(kernel);
 
+		let idleHandle = new MutableDisposable();
+
 		const listener = this._notebookKernelService.onDidChangeNotebookKernelBinding(e => {
-			if (e.oldKernel === kernel.id) {
-				this._proxy.$acceptSelection(handle, e.notebook, false);
-			} else if (e.newKernel === kernel.id) {
-				this._proxy.$acceptSelection(handle, e.notebook, true);
-			}
+
+			idleHandle.value = runWhenIdle(() => {
+
+				if (e.oldKernel === kernel.id) {
+					this._proxy.$acceptSelection(handle, e.notebook, false);
+				} else if (e.newKernel === kernel.id) {
+					this._proxy.$acceptSelection(handle, e.notebook, true);
+				}
+			}, 100);
 		});
 
-		this._kernels.set(handle, [kernel, combinedDisposable(listener, registration)]);
+		this._kernels.set(handle, [kernel, combinedDisposable(listener, registration, idleHandle)]);
 	}
 
 	$updateKernel(handle: number, data: Partial<INotebookKernelDto2>): void {

diff --git a/src/vs/workbench/contrib/notebook/browser/notebookKernelServiceImpl.ts b/src/vs/workbench/contrib/notebook/browser/notebookKernelServiceImpl.ts
@@ -4,7 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { Event, Emitter } from 'vs/base/common/event';
-import { IDisposable, toDisposable } from 'vs/base/common/lifecycle';
+import { DisposableStore, IDisposable, toDisposable } from 'vs/base/common/lifecycle';
 import { INotebookKernel, INotebookTextModel } from 'vs/workbench/contrib/notebook/common/notebookCommon';
 import { INotebookKernelBindEvent, INotebookKernelService, INotebookTextModelLike } from 'vs/workbench/contrib/notebook/common/notebookKernelService';
 import { score } from 'vs/workbench/contrib/notebook/common/notebookSelector';
@@ -14,6 +14,7 @@ import { URI } from 'vs/base/common/uri';
 import { runWhenIdle } from 'vs/base/common/async';
 import { ILogService } from 'vs/platform/log/common/log';
 import { isEqual } from 'vs/base/common/resources';
+import { INotebookService } from 'vs/workbench/contrib/notebook/common/notebookService';
 
 class KernelInfo {
 
@@ -58,6 +59,7 @@ export class NotebookKernelService implements INotebookKernelService {
 
 	private static _storageKey = 'notebook.kernelBindings';
 
+	private readonly _disposables = new DisposableStore();
 	private readonly _kernels = new Map<string, KernelInfo>();
 	private readonly _kernelBindings = new LRUCache<string, string>(1000, 0.7);
 	private readonly _scoreInfo = new ScoreInfo();
@@ -71,6 +73,7 @@ export class NotebookKernelService implements INotebookKernelService {
 	readonly onDidRemoveKernel: Event<INotebookKernel> = this._onDidRemoveKernel.event;
 
 	constructor(
+		@INotebookService private readonly _notebookService: INotebookService,
 		@IStorageService private _storageService: IStorageService,
 		@ILogService logService: ILogService,
 	) {
@@ -82,22 +85,43 @@ export class NotebookKernelService implements INotebookKernelService {
 		} catch {
 			logService.warn('FAILED to restore kernel bindings');
 		}
-	}
 
-	private _persistBindings(): void {
-		runWhenIdle(() => {
-			const raw = JSON.stringify(this._kernelBindings);
-			this._storageService.store(NotebookKernelService._storageKey, raw, StorageScope.WORKSPACE, StorageTarget.MACHINE);
-		}, 100);
+		// auto associate kernels to new notebook documents
+		this._disposables.add(_notebookService.onDidAddNotebookDocument(this._autoAssociateNotebook, this));
 	}
 
 	dispose() {
+		this._disposables.dispose();
 		this._onDidChangeNotebookKernelBinding.dispose();
 		this._onDidAddKernel.dispose();
 		this._onDidRemoveKernel.dispose();
 		this._kernels.clear();
 	}
 
+	private _persistBindingsSoon(): void {
+		runWhenIdle(() => {
+			const raw = JSON.stringify(this._kernelBindings);
+			this._storageService.store(NotebookKernelService._storageKey, raw, StorageScope.WORKSPACE, StorageTarget.MACHINE);
+		}, 100);
+	}
+
+	private _autoAssociateNotebook(notebook: INotebookTextModel, onlyThisKernel?: INotebookKernel): void {
+
+		const id = this._kernelBindings.get(notebook.uri.toString());
+		if (!id) {
+			// no kernel associated
+			return;
+		}
+		const existingKernel = this._kernels.get(id);
+		if (!existingKernel) {
+			// associated kernel not known
+			return;
+		}
+		if (!onlyThisKernel || existingKernel.kernel === onlyThisKernel) {
+			this._onDidChangeNotebookKernelBinding.fire({ notebook: notebook.uri, oldKernel: undefined, newKernel: existingKernel.kernel.id });
+		}
+	}
+
 	registerKernel(kernel: INotebookKernel): IDisposable {
 		if (this._kernels.has(kernel.id)) {
 			throw new Error(`NOTEBOOK CONTROLLER with id '${kernel.id}' already exists`);
@@ -107,6 +131,12 @@ export class NotebookKernelService implements INotebookKernelService {
 		this._scoreInfo.reset();
 		this._onDidAddKernel.fire(kernel);
 
+		// auto associate the new kernel to existing notebooks it was
+		// associated to in the past.
+		for (const notebook of this._notebookService.getNotebookTextModels()) {
+			this._autoAssociateNotebook(notebook, kernel);
+		}
+
 		return toDisposable(() => {
 			this._scoreInfo.reset();
 			if (this._kernels.delete(kernel.id)) {
@@ -170,7 +200,7 @@ export class NotebookKernelService implements INotebookKernelService {
 				this._kernelBindings.delete(key);
 			}
 			this._onDidChangeNotebookKernelBinding.fire({ notebook: notebook.uri, oldKernel, newKernel: kernel?.id });
-			this._persistBindings();
+			this._persistBindingsSoon();
 		}
 	}
 }