Skip to content

Commit

Permalink
Fix race condition in notebook kernel association (#13364)
Browse files Browse the repository at this point in the history
  • Loading branch information
@msujew
msujew authored Feb 8, 2024
1 parent ba9cf88 commit c4fb1cd
Show file tree
Hide file tree
Showing 5 changed files with 44 additions and 34 deletions.
34 changes: 13 additions & 21 deletions packages/notebook/src/browser/service/notebook-kernel-service.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -144,41 +144,42 @@ export class SourceCommand implements Disposable {

const NOTEBOOK_KERNEL_BINDING_STORAGE_KEY = 'notebook.kernel.bindings';
@injectable()
export class NotebookKernelService implements Disposable {
export class NotebookKernelService {

@inject(NotebookService)
protected notebookService: NotebookService;

@inject(StorageService)
protected storageService: StorageService;

private readonly kernels = new Map<string, KernelInfo>();
protected readonly kernels = new Map<string, KernelInfo>();

private notebookBindings: { [key: string]: string } = {};
protected notebookBindings: Record<string, string> = {};

private readonly kernelDetectionTasks = new Map<string, string[]>();
private readonly onDidChangeKernelDetectionTasksEmitter = new Emitter<string>();
protected readonly kernelDetectionTasks = new Map<string, string[]>();
protected readonly onDidChangeKernelDetectionTasksEmitter = new Emitter<string>();
readonly onDidChangeKernelDetectionTasks = this.onDidChangeKernelDetectionTasksEmitter.event;

private readonly onDidChangeSourceActionsEmitter = new Emitter<NotebookSourceActionChangeEvent>();
private readonly kernelSourceActionProviders = new Map<string, KernelSourceActionProvider[]>();
protected readonly onDidChangeSourceActionsEmitter = new Emitter<NotebookSourceActionChangeEvent>();
protected readonly kernelSourceActionProviders = new Map<string, KernelSourceActionProvider[]>();
readonly onDidChangeSourceActions: Event<NotebookSourceActionChangeEvent> = this.onDidChangeSourceActionsEmitter.event;

private readonly onDidAddKernelEmitter = new Emitter<NotebookKernel>();
protected readonly onDidAddKernelEmitter = new Emitter<NotebookKernel>();
readonly onDidAddKernel: Event<NotebookKernel> = this.onDidAddKernelEmitter.event;

private readonly onDidRemoveKernelEmitter = new Emitter<NotebookKernel>();
protected readonly onDidRemoveKernelEmitter = new Emitter<NotebookKernel>();
readonly onDidRemoveKernel: Event<NotebookKernel> = this.onDidRemoveKernelEmitter.event;

private readonly onDidChangeSelectedNotebookKernelBindingEmitter = new Emitter<SelectedNotebookKernelChangeEvent>();
protected readonly onDidChangeSelectedNotebookKernelBindingEmitter = new Emitter<SelectedNotebookKernelChangeEvent>();
readonly onDidChangeSelectedKernel: Event<SelectedNotebookKernelChangeEvent> = this.onDidChangeSelectedNotebookKernelBindingEmitter.event;

private readonly onDidChangeNotebookAffinityEmitter = new Emitter<void>();
protected readonly onDidChangeNotebookAffinityEmitter = new Emitter<void>();
readonly onDidChangeNotebookAffinity: Event<void> = this.onDidChangeNotebookAffinityEmitter.event;

@postConstruct()
init(): void {
this.storageService.getData(NOTEBOOK_KERNEL_BINDING_STORAGE_KEY).then((value: { [key: string]: string } | undefined) => {
this.notebookService.onDidAddNotebookDocument(model => this.tryAutoBindNotebook(model));
this.storageService.getData(NOTEBOOK_KERNEL_BINDING_STORAGE_KEY).then((value: Record<string, string> | undefined) => {
if (value) {
this.notebookBindings = value;
}
Expand Down Expand Up @@ -344,13 +345,4 @@ export class NotebookKernelService implements Disposable {
const allActions = await Promise.all(promises);
return allActions.flat();
}

dispose(): void {
this.onDidChangeKernelDetectionTasksEmitter.dispose();
this.onDidChangeSourceActionsEmitter.dispose();
this.onDidAddKernelEmitter.dispose();
this.onDidRemoveKernelEmitter.dispose();
this.onDidChangeSelectedNotebookKernelBindingEmitter.dispose();
this.onDidChangeNotebookAffinityEmitter.dispose();
}
}
13 changes: 5 additions & 8 deletions packages/plugin-ext/src/main/browser/notebooks/notebook-documents-and-editors-main.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,6 @@ import { NotebookEditorsMainImpl } from './notebook-editors-main';
import { NotebookDocumentsMainImpl } from './notebook-documents-main';
import { diffMaps, diffSets } from '../../../common/collections';
import { Mutex } from 'async-mutex';
import throttle = require('@theia/core/shared/lodash.throttle');

interface NotebookAndEditorDelta {
removedDocuments: UriComponents[];
Expand Down Expand Up @@ -107,12 +106,12 @@ export class NotebooksAndEditorsMain implements NotebookDocumentsAndEditorsMain
this.notebookEditorService = container.get(NotebookEditorWidgetService);
this.WidgetManager = container.get(WidgetManager);

this.notebookService.onDidAddNotebookDocument(async () => this.throttleStateUpdate(), this, this.disposables);
this.notebookService.onDidRemoveNotebookDocument(async () => this.throttleStateUpdate(), this, this.disposables);
this.notebookService.onDidAddNotebookDocument(async () => this.updateState(), this, this.disposables);
this.notebookService.onDidRemoveNotebookDocument(async () => this.updateState(), this, this.disposables);
// this.WidgetManager.onActiveEditorChanged(() => this.updateState(), this, this.disposables);
this.notebookEditorService.onDidAddNotebookEditor(async editor => this.handleEditorAdd(editor), this, this.disposables);
this.notebookEditorService.onDidRemoveNotebookEditor(async editor => this.handleEditorRemove(editor), this, this.disposables);
this.notebookEditorService.onDidChangeFocusedEditor(async editor => this.throttleStateUpdate(editor), this, this.disposables);
this.notebookEditorService.onDidChangeFocusedEditor(async editor => this.updateState(editor), this, this.disposables);
}

dispose(): void {
Expand All @@ -130,18 +129,16 @@ export class NotebooksAndEditorsMain implements NotebookDocumentsAndEditorsMain
} else {
this.editorListeners.set(editor.id, [disposable]);
}
await this.throttleStateUpdate();
await this.updateState();
}

private handleEditorRemove(editor: NotebookEditorWidget): void {
const listeners = this.editorListeners.get(editor.id);
listeners?.forEach(listener => listener.dispose());
this.editorListeners.delete(editor.id);
this.throttleStateUpdate();
this.updateState();
}

private throttleStateUpdate = throttle((focusedEditor?: NotebookEditorWidget) => this.updateState(focusedEditor), 100);

private async updateState(focusedEditor?: NotebookEditorWidget): Promise<void> {
await this.updateMutex.runExclusive(async () => this.doUpdateState(focusedEditor));
}
Expand Down
8 changes: 4 additions & 4 deletions packages/plugin-ext/src/plugin/notebook/notebook-kernels.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -294,11 +294,11 @@ export class NotebookKernelsExtImpl implements NotebookKernelsExt {
};
}

$acceptNotebookAssociation(handle: number, uri: UriComponents, value: boolean): void {
async $acceptNotebookAssociation(handle: number, uri: UriComponents, value: boolean): Promise<void> {
const obj = this.kernelData.get(handle);
if (obj) {
// update data structure
const notebook = this.notebooks.getNotebookDocument(URI.from(uri))!;
const notebook = await this.notebooks.waitForNotebookDocument(URI.from(uri));
if (value) {
obj.associatedNotebooks.set(notebook.uri.toString(), true);
} else {
Expand All @@ -320,7 +320,7 @@ export class NotebookKernelsExtImpl implements NotebookKernelsExt {
// extension can dispose kernels in the meantime
return Promise.resolve();
}
const document = this.notebooks.getNotebookDocument(URI.from(uri));
const document = await this.notebooks.waitForNotebookDocument(URI.from(uri));
const cells: theia.NotebookCell[] = [];
for (const cellHandle of handles) {
const cell = document.getCell(cellHandle);
Expand All @@ -347,7 +347,7 @@ export class NotebookKernelsExtImpl implements NotebookKernelsExt {

// cancel or interrupt depends on the controller. When an interrupt handler is used we
// don't trigger the cancelation token of executions.N
const document = this.notebooks.getNotebookDocument(URI.from(uri));
const document = await this.notebooks.waitForNotebookDocument(URI.from(uri));
if (obj.controller.interruptHandler) {
await obj.controller.interruptHandler.call(obj.controller, document.apiNotebook);

Expand Down
20 changes: 20 additions & 0 deletions packages/plugin-ext/src/plugin/notebook/notebooks.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -312,6 +312,26 @@ export class NotebooksExtImpl implements NotebooksExt {
return result;
}

waitForNotebookDocument(uri: TheiaURI, duration = 2000): Promise<NotebookDocument> {
const existing = this.getNotebookDocument(uri, true);
if (existing) {
return Promise.resolve(existing);
}
return new Promise<NotebookDocument>((resolve, reject) => {
const listener = this.onDidOpenNotebookDocument(event => {
if (event.uri.toString() === uri.toString()) {
clearTimeout(timeout);
listener.dispose();
resolve(this.getNotebookDocument(uri));
}
});
const timeout = setTimeout(() => {
listener.dispose();
reject(new Error(`Notebook document did NOT open in ${duration}ms: ${uri}`));
}, duration);
});
}

private createExtHostEditor(document: NotebookDocument, editorId: string, data: NotebookEditorAddData): void {

if (this.editors.has(editorId)) {
Expand Down
3 changes: 2 additions & 1 deletion packages/plugin-ext/src/plugin/plugin-context.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -709,7 +709,8 @@ export function createAPIFactory(
} else {
throw new Error('Invalid arguments');
}
return notebooksExt.getNotebookDocument(uri).apiNotebook;
const result = await notebooksExt.waitForNotebookDocument(uri);
return result.apiNotebook;

},
createFileSystemWatcher: (pattern, ignoreCreate, ignoreChange, ignoreDelete): theia.FileSystemWatcher =>
Expand Down

0 comments on commit c4fb1cd

Please sign in to comment.