Skip to content

Commit

Permalink
store associated notebook by uri _and_ type, send unselect event on k…
Browse files Browse the repository at this point in the history 
…ernel- or notebook-remove but keep the memento untouched, #121904
  • Loading branch information
@jrieken
jrieken committed Apr 26, 2021
1 parent 594b7b0 commit b8c9fbe
Show file tree
Hide file tree
Showing 3 changed files with 70 additions and 33 deletions.
97 changes: 67 additions & 30 deletions src/vs/workbench/contrib/notebook/browser/notebookKernelServiceImpl.ts
View file
Expand Up @@ -11,6 +11,7 @@ import { LRUCache, ResourceMap } from 'vs/base/common/map';
import { IStorageService, StorageScope, StorageTarget } from 'vs/platform/storage/common/storage';
import { URI } from 'vs/base/common/uri';
import { INotebookService } from 'vs/workbench/contrib/notebook/common/notebookService';
import { revive } from 'vs/base/common/marshalling';

class KernelInfo {

Expand All @@ -29,36 +30,55 @@ class KernelInfo {
}
}

class LRUMemento<K, V> {
class LRUMemento<K = string> {

readonly map: LRUCache<K, V>;
private readonly _map: LRUCache<string, string>;

constructor(
limit: number,
private readonly _keys: { asString(key: K): string, asKey(s: string): K },
private readonly _key: string,
private readonly _scope: StorageScope,
private readonly _target: StorageTarget,
@IStorageService private readonly _storageService: IStorageService,
) {
this.map = new LRUCache(limit, 0.7);
this._map = new LRUCache(limit, 0.7);
this.restore();
}

restore(): this {
try {
const value = this._storageService.get(this._key, this._scope, '[]');
const data = JSON.parse(value);
this.map.fromJSON(data);
this._map.fromJSON(data);
} catch {
// ignore
}
return this;
}

store(): this {
this._storageService.store(this._key, JSON.stringify(this.map), this._scope, this._target);
this._storageService.store(this._key, JSON.stringify(this._map), this._scope, this._target);
return this;
}

set(key: K, value: string): void {
this._map.set(this._keys.asString(key), value);
}

get(key: K): string | undefined {
return this._map.get(this._keys.asString(key));
}

delete(key: K): boolean {
return this._map.delete(this._keys.asString(key));
}

*[Symbol.iterator](): IterableIterator<[K, string]> {
for (const [key, value] of this._map) {
yield [this._keys.asKey(key), value];
}
}
}

export class NotebookKernelService implements INotebookKernelService {
Expand All @@ -68,8 +88,8 @@ export class NotebookKernelService implements INotebookKernelService {
private readonly _disposables = new DisposableStore();
private readonly _kernels = new Map<string, KernelInfo>();

private readonly _notebookInstanceBindings: LRUMemento<string, string>;
private readonly _notebookTypeBindings: LRUMemento<string, string>;
private readonly _notebookInstanceBindings: LRUMemento<INotebookTextModelLike>;
private readonly _notebookTypeBindings: LRUMemento<string>;

private readonly _onDidChangeNotebookKernelBinding = new Emitter<INotebookKernelBindEvent>();
private readonly _onDidAddKernel = new Emitter<INotebookKernel>();
Expand All @@ -86,11 +106,18 @@ export class NotebookKernelService implements INotebookKernelService {
@IStorageService storageService: IStorageService,
) {

this._notebookInstanceBindings = new LRUMemento(1000, 'notebook.kernelBinding', StorageScope.WORKSPACE, StorageTarget.MACHINE, storageService);
this._notebookTypeBindings = new LRUMemento(100, 'notebook.typeBinding', StorageScope.GLOBAL, StorageTarget.USER, storageService);
this._notebookInstanceBindings = new LRUMemento(1000, { asString: notebook => JSON.stringify({ viewType: notebook.viewType, uri: notebook.viewType }), asKey: s => revive(JSON.parse(s)) }, 'notebook.controllerInstanceBinding', StorageScope.WORKSPACE, StorageTarget.MACHINE, storageService);
this._notebookTypeBindings = new LRUMemento(100, { asString: s => s, asKey: s => s }, 'notebook.controllerTypeBinding', StorageScope.GLOBAL, StorageTarget.USER, storageService);

// auto associate kernels to new notebook documents
this._disposables.add(_notebookService.onDidAddNotebookDocument(this._autoAssociateNotebook, this));
// auto associate kernels to new notebook documents, also emit event when
// a notebook has been closed (but don't update the memento)
this._disposables.add(_notebookService.onDidAddNotebookDocument(this._tryAutoBindNotebook, this));
this._disposables.add(_notebookService.onDidRemoveNotebookDocument(e => {
const kernelId = this._notebookInstanceBindings.get(e);
if (kernelId) {
this._onDidChangeNotebookKernelBinding.fire({ notebook: e.uri, oldKernel: kernelId, newKernel: undefined });
}
}));
}

dispose() {
Expand All @@ -101,16 +128,16 @@ export class NotebookKernelService implements INotebookKernelService {
this._kernels.clear();
}

private _autoAssociateNotebook(notebook: INotebookTextModel, onlyThisKernel?: INotebookKernel): void {
private _tryAutoBindNotebook(notebook: INotebookTextModel, onlyThisKernel?: INotebookKernel): void {

const id = this._notebookInstanceBindings.map.get(notebook.uri.toString());
const id = this._notebookInstanceBindings.get(notebook);
if (!id) {
// no kernel associated
return;
}
const existingKernel = this._kernels.get(id);
if (!existingKernel) {
// associated kernel not known
if (!existingKernel || !NotebookKernelService._score(existingKernel.kernel, notebook)) {
// associated kernel not known, not matching
return;
}
if (!onlyThisKernel || existingKernel.kernel === onlyThisKernel) {
Expand All @@ -129,17 +156,16 @@ export class NotebookKernelService implements INotebookKernelService {
// 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);
this._tryAutoBindNotebook(notebook, kernel);
}

return toDisposable(() => {
if (this._kernels.delete(kernel.id)) {
this._onDidRemoveKernel.fire(kernel);
}
for (let [uri, candidate] of Array.from(this._notebookInstanceBindings.map)) {
for (const [key, candidate] of Array.from(this._notebookInstanceBindings)) {
if (candidate === kernel.id) {
this._notebookInstanceBindings.map.delete(uri);
this._onDidChangeNotebookKernelBinding.fire({ notebook: URI.parse(uri), oldKernel: kernel.id, newKernel: undefined });
this._onDidChangeNotebookKernelBinding.fire({ notebook: key.uri, oldKernel: kernel.id, newKernel: undefined });
}
}
});
Expand All @@ -148,33 +174,45 @@ export class NotebookKernelService implements INotebookKernelService {
getNotebookKernels(notebook: INotebookTextModelLike): { bound: INotebookKernel | undefined, all: INotebookKernel[] } {

// all applicable kernels
const kernels: { kernel: INotebookKernel, instanceAffinity: number, typeAffinity: number }[] = [];
const kernels: { kernel: INotebookKernel, instanceAffinity: number, typeAffinity: number, score: number }[] = [];
for (const info of this._kernels.values()) {
if (info.kernel.viewType === notebook.viewType || info.kernel.viewType === '*') {
const score = NotebookKernelService._score(info.kernel, notebook);
if (score) {
kernels.push({
score,
kernel: info.kernel,
instanceAffinity: info.notebookPriorities.get(notebook.uri) ?? 1 /* vscode.NotebookControllerPriority.Default */,
typeAffinity: this._notebookTypeBindings.map.get(info.kernel.viewType) === info.kernel.id ? 1 : 0
typeAffinity: this._notebookTypeBindings.get(info.kernel.viewType) === info.kernel.id ? 1 : 0
});
}
}

const all = kernels
.sort((a, b) => b.instanceAffinity - a.instanceAffinity || b.typeAffinity - a.typeAffinity || a.kernel.label.localeCompare(b.kernel.label))
.sort((a, b) => b.instanceAffinity - a.instanceAffinity || b.typeAffinity - a.typeAffinity || a.score - b.score || a.kernel.label.localeCompare(b.kernel.label))
.map(obj => obj.kernel);

// bound kernel
const boundId = this._notebookInstanceBindings.map.get(notebook.uri.toString());
const boundId = this._notebookInstanceBindings.get(notebook);
const bound = boundId ? this._kernels.get(boundId)?.kernel : undefined;

return { all, bound };
}

private static _score(kernel: INotebookKernel, notebook: INotebookTextModelLike): number {
if (kernel.viewType === '*') {
return 5;
} else if (kernel.viewType === notebook.viewType) {
return 10;
} else {
return 0;
}
}

// default kernel for notebookType
updateNotebookTypeKernelBinding(typeId: string, kernel: INotebookKernel): void {
const existing = this._notebookTypeBindings.map.get(typeId);
const existing = this._notebookTypeBindings.get(typeId);
if (existing !== kernel.id) {
this._notebookTypeBindings.map.set(typeId, kernel.id);
this._notebookTypeBindings.set(typeId, kernel.id);
this._notebookTypeBindings.store();
this._onDidChangeNotebookAffinity.fire();
}
Expand All @@ -183,13 +221,12 @@ export class NotebookKernelService implements INotebookKernelService {
// a notebook has one kernel, a kernel has N notebooks
// notebook <-1----N-> kernel
updateNotebookInstanceKernelBinding(notebook: INotebookTextModel, kernel: INotebookKernel | undefined): void {
const key = notebook.uri.toString();
const oldKernel = this._notebookInstanceBindings.map.get(key);
const oldKernel = this._notebookInstanceBindings.get(notebook);
if (oldKernel !== kernel?.id) {
if (kernel) {
this._notebookInstanceBindings.map.set(key, kernel.id);
this._notebookInstanceBindings.set(notebook, kernel.id);
} else {
this._notebookInstanceBindings.map.delete(key);
this._notebookInstanceBindings.delete(notebook);
}
this._onDidChangeNotebookKernelBinding.fire({ notebook: notebook.uri, oldKernel, newKernel: kernel?.id });
this._notebookInstanceBindings.store();
Expand Down
4 changes: 2 additions & 2 deletions src/vs/workbench/contrib/notebook/browser/notebookServiceImpl.ts
View file
Expand Up @@ -277,7 +277,7 @@ export class NotebookService extends Disposable implements INotebookService, IEd

private readonly _onDidCreateNotebookDocument = this._register(new Emitter<NotebookTextModel>());
private readonly _onDidAddNotebookDocument = this._register(new Emitter<NotebookTextModel>());
private readonly _onDidRemoveNotebookDocument = this._register(new Emitter<URI>());
private readonly _onDidRemoveNotebookDocument = this._register(new Emitter<NotebookTextModel>());

readonly onDidCreateNotebookDocument = this._onDidCreateNotebookDocument.event;
readonly onDidAddNotebookDocument = this._onDidAddNotebookDocument.event;
Expand Down Expand Up @@ -552,7 +552,7 @@ export class NotebookService extends Disposable implements INotebookService, IEd
if (modelData) {
this._models.delete(model.uri);
modelData.dispose();
this._onDidRemoveNotebookDocument.fire(modelData.model.uri);
this._onDidRemoveNotebookDocument.fire(modelData.model);
}
}

Expand Down
2 changes: 1 addition & 1 deletion src/vs/workbench/contrib/notebook/common/notebookService.ts
View file
Expand Up @@ -62,7 +62,7 @@ export interface INotebookService {

onDidCreateNotebookDocument: Event<NotebookTextModel>;
onDidAddNotebookDocument: Event<NotebookTextModel>;
onDidRemoveNotebookDocument: Event<URI>;
onDidRemoveNotebookDocument: Event<NotebookTextModel>;

registerNotebookController(viewType: string, extensionData: NotebookExtensionDescription, controller: INotebookContentProvider): IDisposable;
registerNotebookSerializer(viewType: string, extensionData: NotebookExtensionDescription, serializer: INotebookSerializer): IDisposable;
Expand Down

0 comments on commit b8c9fbe

Please sign in to comment.