Fix race condition in notebook kernel association

packages/notebook/src/browser/service/notebook-kernel-service.ts

@@ -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;
             }
@@ -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();
-    }
 }

packages/plugin-ext/src/main/browser/notebooks/notebook-documents-and-editors-main.ts

@@ -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[];
@@ -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 {
@@ -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));
     }

packages/plugin-ext/src/plugin/notebook/notebook-kernels.ts

@@ -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 {
@@ -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);
@@ -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);
 

packages/plugin-ext/src/plugin/notebook/notebooks.ts

@@ -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)) {

packages/plugin-ext/src/plugin/plugin-context.ts

@@ -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 =>