fix kernel cache and add unit test for kernel reg.

This commit is contained in:
rebornix 2021-02-22 16:56:19 -08:00
parent eff172b7dd
commit ba3bd11ac7
2 changed files with 126 additions and 16 deletions

View file

@ -73,8 +73,8 @@ class ExtHostWebviewCommWrapper extends Disposable {
}
export class ExtHostNotebookKernelProviderAdapter extends Disposable {
private _kernelToFriendlyId = new Map<vscode.NotebookKernel, string>();
private _friendlyIdToKernel = new Map<string, vscode.NotebookKernel>();
private _kernelToFriendlyId = new ResourceMap<Map<vscode.NotebookKernel, string>>();
private _friendlyIdToKernel = new ResourceMap<Map<string, vscode.NotebookKernel>>();
constructor(
private readonly _proxy: MainThreadNotebookShape,
private readonly _handle: number,
@ -98,16 +98,16 @@ export class ExtHostNotebookKernelProviderAdapter extends Disposable {
let kernel_unique_pool = 0;
const kernelFriendlyIdCache = new Set<string>();
const kernelToFriendlyId = this._kernelToFriendlyId.get(document.uri);
const transformedData: INotebookKernelInfoDto2[] = data.map(kernel => {
let friendlyId = this._kernelToFriendlyId.get(kernel);
let friendlyId = kernelToFriendlyId?.get(kernel);
if (friendlyId === undefined) {
if (kernel.id && kernelFriendlyIdCache.has(kernel.id)) {
friendlyId = `${this._extension.identifier.value}_${kernel.id}_${kernel_unique_pool++}`;
} else {
friendlyId = `${this._extension.identifier.value}_${kernel.id || UUID.generateUuid()}`;
}
this._kernelToFriendlyId.set(kernel, friendlyId);
}
newMap.set(kernel, friendlyId);
@ -127,22 +127,22 @@ export class ExtHostNotebookKernelProviderAdapter extends Disposable {
};
});
this._kernelToFriendlyId = newMap;
this._friendlyIdToKernel.clear();
this._kernelToFriendlyId.forEach((value, key) => {
this._friendlyIdToKernel.set(value, key);
this._kernelToFriendlyId.set(document.uri, newMap);
const friendlyIdToKernel = new Map<string, vscode.NotebookKernel>();
newMap.forEach((value, key) => {
friendlyIdToKernel.set(value, key);
});
this._friendlyIdToKernel.set(document.uri, friendlyIdToKernel);
return transformedData;
}
getKernelByFriendlyId(kernelId: string) {
return this._friendlyIdToKernel.get(kernelId);
getKernelByFriendlyId(uri: URI, kernelId: string) {
return this._friendlyIdToKernel.get(uri)?.get(kernelId);
}
async resolveNotebook(kernelId: string, document: ExtHostNotebookDocument, webview: vscode.NotebookCommunication, token: CancellationToken) {
const kernel = this._friendlyIdToKernel.get(kernelId);
const kernel = this._friendlyIdToKernel.get(document.uri)?.get(kernelId);
if (kernel && this._provider.resolveKernel) {
return this._provider.resolveKernel(kernel, document.notebookDocument, webview, token);
@ -150,7 +150,7 @@ export class ExtHostNotebookKernelProviderAdapter extends Disposable {
}
async executeNotebook(kernelId: string, document: ExtHostNotebookDocument, cell: ExtHostCell | undefined) {
const kernel = this._friendlyIdToKernel.get(kernelId);
const kernel = this._friendlyIdToKernel.get(document.uri)?.get(kernelId);
if (!kernel) {
return;
@ -164,7 +164,7 @@ export class ExtHostNotebookKernelProviderAdapter extends Disposable {
}
async cancelNotebook(kernelId: string, document: ExtHostNotebookDocument, cell: ExtHostCell | undefined) {
const kernel = this._friendlyIdToKernel.get(kernelId);
const kernel = this._friendlyIdToKernel.get(document.uri)?.get(kernelId);
if (!kernel) {
return;
@ -536,7 +536,7 @@ export class ExtHostNotebookController implements ExtHostNotebookShape {
$acceptNotebookActiveKernelChange(event: { uri: UriComponents, providerHandle: number | undefined, kernelFriendlyId: string | undefined; }) {
if (event.providerHandle !== undefined) {
this._withAdapter(event.providerHandle, event.uri, async (adapter, document) => {
const kernel = event.kernelFriendlyId ? adapter.getKernelByFriendlyId(event.kernelFriendlyId) : undefined;
const kernel = event.kernelFriendlyId ? adapter.getKernelByFriendlyId(URI.revive(event.uri), event.kernelFriendlyId) : undefined;
this._editors.forEach(editor => {
if (editor.editor.notebookData === document) {
editor.editor._acceptKernel(kernel);

View file

@ -0,0 +1,110 @@
/*---------------------------------------------------------------------------------------------
* Copyright (c) Microsoft Corporation. All rights reserved.
* Licensed under the MIT License. See License.txt in the project root for license information.
*--------------------------------------------------------------------------------------------*/
import * as assert from 'assert';
import { TestRPCProtocol } from 'vs/workbench/test/browser/api/testRPCProtocol';
import { ExtHostDocuments } from 'vs/workbench/api/common/extHostDocuments';
import { ExtHostDocumentsAndEditors } from 'vs/workbench/api/common/extHostDocumentsAndEditors';
import { NullLogService } from 'vs/platform/log/common/log';
import { ExtHostNotebookController } from 'vs/workbench/api/common/extHostNotebook';
import { ExtHostNotebookDocument } from 'vs/workbench/api/common/extHostNotebookDocument';
import { URI } from 'vs/base/common/uri';
import { CellKind, CellUri } from 'vs/workbench/contrib/notebook/common/notebookCommon';
import { ExtHostCommands } from 'vs/workbench/api/common/extHostCommands';
import { nullExtensionDescription } from 'vs/workbench/services/extensions/common/extensions';
import * as vscode from 'vscode';
import { mock } from 'vs/workbench/test/common/workbenchTestServices';
import { MainContext, MainThreadCommandsShape, MainThreadNotebookShape } from 'vs/workbench/api/common/extHost.protocol';
import { DisposableStore } from 'vs/base/common/lifecycle';
import { IExtensionStoragePaths } from 'vs/workbench/api/common/extHostStoragePaths';
import { generateUuid } from 'vs/base/common/uuid';
import { CancellationToken } from 'vs/base/common/cancellation';
suite('NotebookKernel', function () {
let rpcProtocol: TestRPCProtocol;
let notebook: ExtHostNotebookDocument;
let extHostDocumentsAndEditors: ExtHostDocumentsAndEditors;
let extHostDocuments: ExtHostDocuments;
let extHostNotebooks: ExtHostNotebookController;
const notebookUri = URI.parse('test:///notebook.file');
const disposables = new DisposableStore();
setup(async function () {
disposables.clear();
rpcProtocol = new TestRPCProtocol();
rpcProtocol.set(MainContext.MainThreadCommands, new class extends mock<MainThreadCommandsShape>() {
$registerCommand() { }
});
rpcProtocol.set(MainContext.MainThreadNotebook, new class extends mock<MainThreadNotebookShape>() {
async $registerNotebookProvider() { }
async $unregisterNotebookProvider() { }
async $registerNotebookKernelProvider() { }
async $unregisterNotebookKernelProvider() { }
});
extHostDocumentsAndEditors = new ExtHostDocumentsAndEditors(rpcProtocol, new NullLogService());
extHostDocuments = new ExtHostDocuments(rpcProtocol, extHostDocumentsAndEditors);
const extHostStoragePaths = new class extends mock<IExtensionStoragePaths>() {
workspaceValue() {
return URI.from({ scheme: 'test', path: generateUuid() });
}
};
extHostNotebooks = new ExtHostNotebookController(rpcProtocol, new ExtHostCommands(rpcProtocol, new NullLogService()), extHostDocumentsAndEditors, { isExtensionDevelopmentDebug: false, webviewCspSource: '', webviewResourceRoot: '' }, new NullLogService(), extHostStoragePaths);
let reg = extHostNotebooks.registerNotebookContentProvider(nullExtensionDescription, 'test', new class extends mock<vscode.NotebookContentProvider>() {
// async openNotebook() { }
});
const kernels = [new class extends mock<vscode.NotebookKernel>() {
id = 'first';
}, new class extends mock<vscode.NotebookKernel>() {
id = 'second';
}];
let kernelReg = extHostNotebooks.registerNotebookKernelProvider(nullExtensionDescription, { viewType: 'test' }, new class extends mock<vscode.NotebookKernelProvider>() {
async provideKernels() { return kernels; }
});
// init
extHostNotebooks.$acceptDocumentAndEditorsDelta({
addedDocuments: [{
uri: notebookUri,
viewType: 'test',
cells: [{
handle: 0,
uri: CellUri.generate(notebookUri, 0),
source: ['console.log'],
eol: '\n',
language: 'javascript',
cellKind: CellKind.Code,
outputs: [],
}],
contentOptions: { transientOutputs: false, transientMetadata: {} },
versionId: 0
}],
addedEditors: [
{
documentUri: notebookUri,
id: '_notebook_editor_0',
selections: [{ start: 0, end: 1 }],
visibleRanges: []
}
]
});
extHostNotebooks.$acceptDocumentAndEditorsDelta({ newActiveEditor: '_notebook_editor_0' });
notebook = extHostNotebooks.notebookDocuments[0]!;
disposables.add(reg);
disposables.add(kernelReg);
disposables.add(notebook);
disposables.add(extHostDocuments);
});
test('provide kernels', async function () {
const dto = await extHostNotebooks.$provideNotebookKernels(0, notebook.uri, CancellationToken.None);
assert.deepStrictEqual(dto.map(kernel => kernel.id), ['first', 'second']);
});
});