LiveShare support for raw kernels (microsoft#11560)

Author: IanMatthewHuff

news/1 Enhancements/10988.md

@@ -0,0 +1 @@
+Add basic liveshare support for raw kernels.

src/client/datascience/constants.ts

@@ -450,4 +450,6 @@ export namespace LiveShareCommands {
     export const guestCheck = 'guestCheck';
     export const createNotebook = 'createNotebook';
     export const inspect = 'inspect';
+    export const rawKernelSupported = 'rawKernelSupported';
+    export const createRawNotebook = 'createRawNotebook';
 }

src/client/datascience/interactive-common/notebookProvider.ts

@@ -6,13 +6,10 @@
 import { inject, injectable } from 'inversify';
 import { EventEmitter, Uri } from 'vscode';
 import { IWorkspaceService } from '../../common/application/types';
-import { LocalZMQKernel } from '../../common/experimentGroups';
-import { traceError, traceInfo } from '../../common/logger';
 import { IFileSystem } from '../../common/platform/types';
-import { IConfigurationService, IDisposableRegistry, IExperimentsManager, Resource } from '../../common/types';
+import { IDisposableRegistry, Resource } from '../../common/types';
 import { noop } from '../../common/utils/misc';
-import { sendTelemetryEvent } from '../../telemetry';
-import { Identifiers, Settings, Telemetry } from '../constants';
+import { Identifiers } from '../constants';
 import {
     ConnectNotebookProviderOptions,
     GetNotebookOptions,
@@ -30,7 +27,6 @@ import {
 export class NotebookProvider implements INotebookProvider {
     private readonly notebooks = new Map<string, Promise<INotebook>>();
     private _notebookCreated = new EventEmitter<{ identity: Uri; notebook: INotebook }>();
-    private _zmqSupported: boolean | undefined;
     public get activeNotebooks() {
         return [...this.notebooks.values()];
     }
@@ -41,8 +37,6 @@ export class NotebookProvider implements INotebookProvider {
         @inject(IDisposableRegistry) disposables: IDisposableRegistry,
         @inject(IRawNotebookProvider) private readonly rawNotebookProvider: IRawNotebookProvider,
         @inject(IJupyterNotebookProvider) private readonly jupyterNotebookProvider: IJupyterNotebookProvider,
-        @inject(IConfigurationService) private readonly configuration: IConfigurationService,
-        @inject(IExperimentsManager) private readonly experimentsManager: IExperimentsManager,
         @inject(IWorkspaceService) private readonly workspaceService: IWorkspaceService
     ) {
         disposables.push(editorProvider.onDidCloseNotebookEditor(this.onDidCloseNotebookEditor, this));
@@ -57,23 +51,23 @@ export class NotebookProvider implements INotebookProvider {
     // Disconnect from the specified provider
     public async disconnect(options: ConnectNotebookProviderOptions): Promise<void> {
         // Only need to disconnect from actual jupyter servers
-        if (!(await this.rawKernelSupported())) {
+        if (!(await this.rawNotebookProvider.supported())) {
             return this.jupyterNotebookProvider.disconnect(options);
         }
     }
 
     // Attempt to connect to our server provider, and if we do, return the connection info
     public async connect(options: ConnectNotebookProviderOptions): Promise<INotebookProviderConnection | undefined> {
         // Connect to either a jupyter server or a stubbed out raw notebook "connection"
-        if (await this.rawKernelSupported()) {
+        if (await this.rawNotebookProvider.supported()) {
             return this.rawNotebookProvider.connect();
         } else {
             return this.jupyterNotebookProvider.connect(options);
         }
     }
 
     public async getOrCreateNotebook(options: GetNotebookOptions): Promise<INotebook | undefined> {
-        const rawKernel = await this.rawKernelSupported();
+        const rawKernel = await this.rawNotebookProvider.supported();
 
         // Check to see if our provider already has this notebook
         const notebook = rawKernel
@@ -117,45 +111,6 @@ export class NotebookProvider implements INotebookProvider {
         return promise;
     }
 
-    // Check to see if we have all that we need for supporting raw kernel launch
-    private async rawKernelSupported(): Promise<boolean> {
-        const zmqOk = await this.zmqSupported();
-
-        return zmqOk && this.localLaunch() && this.experimentsManager.inExperiment(LocalZMQKernel.experiment)
-            ? true
-            : false;
-    }
-
-    private localLaunch(): boolean {
-        const settings = this.configuration.getSettings(undefined);
-        const serverURI: string | undefined = settings.datascience.jupyterServerURI;
-
-        if (!serverURI || serverURI.toLowerCase() === Settings.JupyterServerLocalLaunch) {
-            return true;
-        }
-
-        return false;
-    }
-
-    // Check to see if this machine supports our local ZMQ launching
-    private async zmqSupported(): Promise<boolean> {
-        if (this._zmqSupported !== undefined) {
-            return this._zmqSupported;
-        }
-
-        try {
-            await import('zeromq');
-            traceInfo(`ZMQ install verified.`);
-            this._zmqSupported = true;
-        } catch (e) {
-            traceError(`Exception while attempting zmq :`, e);
-            sendTelemetryEvent(Telemetry.ZMQNotSupported);
-            this._zmqSupported = false;
-        }
-
-        return this._zmqSupported;
-    }
-
     // Cache the promise that will return a notebook
     private cacheNotebookPromise(identity: Uri, promise: Promise<INotebook>) {
         this.notebooks.set(identity.fsPath, promise);

src/client/datascience/raw-kernel/liveshare/guestRawNotebookProvider.ts

@@ -8,23 +8,32 @@ import * as vsls from 'vsls/vscode';
 import { IApplicationShell, ILiveShareApi, IWorkspaceService } from '../../../common/application/types';
 import { IFileSystem } from '../../../common/platform/types';
 import { IAsyncDisposableRegistry, IConfigurationService, IDisposableRegistry, Resource } from '../../../common/types';
+import { createDeferred } from '../../../common/utils/async';
+import * as localize from '../../../common/utils/localize';
 import { IServiceContainer } from '../../../ioc/types';
-import { LiveShare } from '../../constants';
+import { LiveShare, LiveShareCommands } from '../../constants';
+import { GuestJupyterNotebook } from '../../jupyter/liveshare/guestJupyterNotebook';
 import {
     LiveShareParticipantDefault,
     LiveShareParticipantGuest
 } from '../../jupyter/liveshare/liveShareParticipantMixin';
 import { ILiveShareParticipant } from '../../jupyter/liveshare/types';
-import { INotebook, IRawConnection, IRawNotebookProvider } from '../../types';
+import { IDataScience, INotebook, IRawConnection, IRawNotebookProvider } from '../../types';
+import { RawConnection } from '../rawNotebookProvider';
 
 export class GuestRawNotebookProvider
     extends LiveShareParticipantGuest(LiveShareParticipantDefault, LiveShare.RawNotebookProviderService)
     implements IRawNotebookProvider, ILiveShareParticipant {
+    // Keep track of guest notebooks on this side
+    private notebooks = new Map<string, Promise<INotebook>>();
+    private rawConnection = new RawConnection();
+
     constructor(
-        liveShare: ILiveShareApi,
-        _disposableRegistry: IDisposableRegistry,
+        private readonly liveShare: ILiveShareApi,
+        private readonly dataScience: IDataScience,
+        private readonly disposableRegistry: IDisposableRegistry,
         _asyncRegistry: IAsyncDisposableRegistry,
-        _configService: IConfigurationService,
+        private readonly configService: IConfigurationService,
         _workspaceService: IWorkspaceService,
         _appShell: IApplicationShell,
         _fs: IFileSystem,
@@ -34,41 +43,97 @@ export class GuestRawNotebookProvider
     }
 
     public async supported(): Promise<boolean> {
-        // For now just false, but when implemented will message the host
-        return false;
+        // Query the host to see if liveshare is supported
+        const service = await this.waitForService();
+        let result = false;
+        if (service) {
+            result = await service.request(LiveShareCommands.rawKernelSupported, []);
+        }
+
+        return result;
     }
 
     public async createNotebook(
-        _identity: Uri,
-        _resource: Resource,
+        identity: Uri,
+        resource: Resource,
         _disableUI: boolean,
-        _notebookMetadata: nbformat.INotebookMetadata,
+        notebookMetadata: nbformat.INotebookMetadata,
         _cancelToken: CancellationToken
     ): Promise<INotebook> {
-        throw new Error('Not implemented');
-    }
+        // Remember we can have multiple native editors opened against the same ipynb file.
+        if (this.notebooks.get(identity.toString())) {
+            return this.notebooks.get(identity.toString())!;
+        }
+
+        const deferred = createDeferred<INotebook>();
+        this.notebooks.set(identity.toString(), deferred.promise);
+        // Tell the host side to generate a notebook for this uri
+        const service = await this.waitForService();
+        if (service) {
+            const resourceString = resource ? resource.toString() : undefined;
+            const identityString = identity.toString();
+            const notebookMetadataString = JSON.stringify(notebookMetadata);
+            await service.request(LiveShareCommands.createRawNotebook, [
+                resourceString,
+                identityString,
+                notebookMetadataString
+            ]);
+        }
 
-    public connect(): Promise<IRawConnection> {
-        throw new Error('Not implemented');
+        // Return a new notebook to listen to
+        const result = new GuestJupyterNotebook(
+            this.liveShare,
+            this.disposableRegistry,
+            this.configService,
+            resource,
+            identity,
+            undefined,
+            this.dataScience.activationStartTime
+        );
+        deferred.resolve(result);
+        const oldDispose = result.dispose.bind(result);
+        result.dispose = () => {
+            this.notebooks.delete(identity.toString());
+            return oldDispose();
+        };
+
+        return result;
     }
 
-    public async onSessionChange(_api: vsls.LiveShare | null): Promise<void> {
-        // Not implemented yet
+    public async connect(): Promise<IRawConnection> {
+        return Promise.resolve(this.rawConnection);
     }
 
-    public async getNotebook(_resource: Uri): Promise<INotebook | undefined> {
-        throw new Error('Not implemented');
+    public async onSessionChange(api: vsls.LiveShare | null): Promise<void> {
+        await super.onSessionChange(api);
+
+        this.notebooks.forEach(async (notebook) => {
+            const guestNotebook = (await notebook) as GuestJupyterNotebook;
+            if (guestNotebook) {
+                await guestNotebook.onSessionChange(api);
+            }
+        });
     }
 
-    public async shutdown(): Promise<void> {
-        throw new Error('Not implemented');
+    public async getNotebook(resource: Uri): Promise<INotebook | undefined> {
+        return this.notebooks.get(resource.toString());
     }
 
-    public dispose(): Promise<void> {
-        throw new Error('Not implemented');
+    public async onAttach(api: vsls.LiveShare | null): Promise<void> {
+        await super.onAttach(api);
+
+        if (api) {
+            const service = await this.waitForService();
+
+            // Wait for sync up
+            const synced = service ? await service.request(LiveShareCommands.syncRequest, []) : undefined;
+            if (!synced && api.session && api.session.role !== vsls.Role.None) {
+                throw new Error(localize.DataScience.liveShareSyncFailure());
+            }
+        }
     }
 
-    public async onAttach(_api: vsls.LiveShare | null): Promise<void> {
-        // Not implemented yet
+    public async waitForServiceName(): Promise<string> {
+        return LiveShare.RawNotebookProviderService;
     }
 }

src/client/datascience/raw-kernel/liveshare/hostRawNotebookProvider.ts

@@ -15,20 +15,22 @@ import {
     IAsyncDisposableRegistry,
     IConfigurationService,
     IDisposableRegistry,
+    IExperimentsManager,
     IOutputChannel,
     Resource
 } from '../../../common/types';
 import { createDeferred } from '../../../common/utils/async';
 import * as localize from '../../../common/utils/localize';
 import { IServiceContainer } from '../../../ioc/types';
-import { Identifiers, LiveShare, Settings } from '../../constants';
+import { Identifiers, LiveShare, LiveShareCommands, Settings } from '../../constants';
 import { KernelSelector } from '../../jupyter/kernels/kernelSelector';
 import { HostJupyterNotebook } from '../../jupyter/liveshare/hostJupyterNotebook';
 import { LiveShareParticipantHost } from '../../jupyter/liveshare/liveShareParticipantMixin';
 import { IRoleBasedObject } from '../../jupyter/liveshare/roleBasedFactory';
 import { IKernelLauncher } from '../../kernel-launcher/types';
 import { ProgressReporter } from '../../progress/progressReporter';
 import {
+    IDataScience,
     IJupyterKernelSpec,
     INotebook,
     INotebookExecutionInfo,
@@ -48,6 +50,7 @@ export class HostRawNotebookProvider
     private disposed = false;
     constructor(
         private liveShare: ILiveShareApi,
+        _dataScience: IDataScience,
         private disposableRegistry: IDisposableRegistry,
         asyncRegistry: IAsyncDisposableRegistry,
         private configService: IConfigurationService,
@@ -58,9 +61,10 @@ export class HostRawNotebookProvider
         private kernelLauncher: IKernelLauncher,
         private kernelSelector: KernelSelector,
         private progressReporter: ProgressReporter,
-        private outputChannel: IOutputChannel
+        private outputChannel: IOutputChannel,
+        experimentsManager: IExperimentsManager
     ) {
-        super(liveShare, asyncRegistry);
+        super(liveShare, asyncRegistry, configService, experimentsManager);
     }
 
     public async dispose(): Promise<void> {
@@ -70,20 +74,57 @@ export class HostRawNotebookProvider
         }
     }
 
-    public async onAttach(_api: vsls.LiveShare | null): Promise<void> {
-        // Not implemented yet
+    public async onAttach(api: vsls.LiveShare | null): Promise<void> {
+        await super.onAttach(api);
+        if (api && !this.disposed) {
+            const service = await this.waitForService();
+            // Attach event handlers to different requests
+            if (service) {
+                service.onRequest(LiveShareCommands.syncRequest, (_args: any[], _cancellation: CancellationToken) =>
+                    this.onSync()
+                );
+                service.onRequest(
+                    LiveShareCommands.rawKernelSupported,
+                    (_args: any[], _cancellation: CancellationToken) => this.supported()
+                );
+                service.onRequest(
+                    LiveShareCommands.createRawNotebook,
+                    async (args: any[], _cancellation: CancellationToken) => {
+                        const resource = this.parseUri(args[0]);
+                        const identity = this.parseUri(args[1]);
+                        const notebookMetadata = JSON.parse(args[2]) as nbformat.INotebookMetadata;
+                        // Don't return the notebook. We don't want it to be serialized. We just want its live share server to be started.
+                        const notebook = (await this.createNotebook(
+                            identity!,
+                            resource,
+                            true, // Disable UI for this creation
+                            notebookMetadata,
+                            undefined
+                        )) as HostJupyterNotebook;
+                        await notebook.onAttach(api);
+                    }
+                );
+            }
+        }
     }
 
-    public async onSessionChange(_api: vsls.LiveShare | null): Promise<void> {
-        // Not implemented yet
+    public async onSessionChange(api: vsls.LiveShare | null): Promise<void> {
+        await super.onSessionChange(api);
+
+        this.getNotebooks().forEach(async (notebook) => {
+            const hostNotebook = (await notebook) as HostJupyterNotebook;
+            if (hostNotebook) {
+                await hostNotebook.onSessionChange(api);
+            }
+        });
     }
 
-    public async onDetach(_api: vsls.LiveShare | null): Promise<void> {
-        // Not implemented yet
+    public async onDetach(api: vsls.LiveShare | null): Promise<void> {
+        await super.onDetach(api);
     }
 
     public async waitForServiceName(): Promise<string> {
-        return 'Not implemented';
+        return LiveShare.RawNotebookProviderService;
     }
 
     protected async createNotebookInstance(
@@ -186,4 +227,18 @@ export class HostRawNotebookProvider
             purpose: Identifiers.RawPurpose
         };
     }
+
+    private parseUri(uri: string | undefined): Resource {
+        const parsed = uri ? vscode.Uri.parse(uri) : undefined;
+        return parsed &&
+            parsed.scheme &&
+            parsed.scheme !== Identifiers.InteractiveWindowIdentityScheme &&
+            parsed.scheme === 'vsls'
+            ? this.finishedApi!.convertSharedUriToLocal(parsed)
+            : parsed;
+    }
+
+    private onSync(): Promise<any> {
+        return Promise.resolve(true);
+    }
 }