Basic idea building

Author: rchiodo

src/client/datascience/baseJupyterSession.ts

@@ -19,7 +19,7 @@ import { Telemetry } from './constants';
 import { JupyterKernelPromiseFailedError } from './jupyter/kernels/jupyterKernelPromiseFailedError';
 import { KernelSelector } from './jupyter/kernels/kernelSelector';
 import { LiveKernelModel } from './jupyter/kernels/types';
-import { IJupyterKernelSpec, IJupyterSession, IKernelSession, KernelSocketInformation } from './types';
+import { IJupyterKernelSpec, IJupyterSession, ISessionWithSocket, KernelSocketInformation } from './types';
 
 /**
  * Exception raised when starting a Jupyter Session fails.
@@ -37,10 +37,10 @@ export class JupyterSessionStartError extends Error {
 }
 
 export abstract class BaseJupyterSession implements IJupyterSession {
- protected get session(): IKernelSession | undefined {
+ protected get session(): ISessionWithSocket | undefined {
  return this._session;
  }
- protected set session(session: IKernelSession | undefined) {
+ protected set session(session: ISessionWithSocket | undefined) {
  const oldSession = this._session;
  this._session = session;
 
@@ -89,10 +89,10 @@ export abstract class BaseJupyterSession implements IJupyterSession {
  return this.connected;
  }
  protected onStatusChangedEvent: EventEmitter<ServerStatus> = new EventEmitter<ServerStatus>();
- protected statusHandler: Slot<IKernelSession, Kernel.Status>;
+ protected statusHandler: Slot<ISessionWithSocket, Kernel.Status>;
  protected connected: boolean = false;
- protected restartSessionPromise: Promise<IKernelSession | undefined> | undefined;
- private _session: IKernelSession | undefined;
+ protected restartSessionPromise: Promise<ISessionWithSocket | undefined> | undefined;
+ private _session: ISessionWithSocket | undefined;
  private _kernelSocket = new ReplaySubject<KernelSocketInformation | undefined>();
  private _jupyterLab?: typeof import('@jupyterlab/services');
 
@@ -141,7 +141,7 @@ export abstract class BaseJupyterSession implements IJupyterSession {
  }
 
  public async changeKernel(kernel: IJupyterKernelSpec | LiveKernelModel, timeoutMS: number): Promise<void> {
- let newSession: IKernelSession | undefined;
+ let newSession: ISessionWithSocket | undefined;
 
  // If we are already using this kernel in an active session just return back
  if (this.kernelSpec?.name === kernel.name && this.session) {
@@ -319,19 +319,19 @@ export abstract class BaseJupyterSession implements IJupyterSession {
  protected abstract startRestartSession(): void;
  protected abstract async createRestartSession(
  kernelSpec: IJupyterKernelSpec | LiveKernelModel | undefined,
- session: IKernelSession,
+ session: ISessionWithSocket,
  cancelToken?: CancellationToken
- ): Promise<IKernelSession>;
+ ): Promise<ISessionWithSocket>;
 
  // Sub classes need to implement their own kernel change specific code
  protected abstract createNewKernelSession(
  kernel: IJupyterKernelSpec | LiveKernelModel,
  timeoutMS: number
- ): Promise<IKernelSession>;
+ ): Promise<ISessionWithSocket>;
 
  protected async shutdownSession(
- session: IKernelSession | undefined,
- statusHandler: Slot<IKernelSession, Kernel.Status> | undefined
+ session: ISessionWithSocket | undefined,
+ statusHandler: Slot<ISessionWithSocket, Kernel.Status> | undefined
  ): Promise<void> {
  if (session && session.kernel) {
  const kernelId = session.kernel.id;

src/client/datascience/jupyter/jupyterSession.ts

@@ -26,7 +26,7 @@ import { BaseJupyterSession, JupyterSessionStartError } from '../baseJupyterSess
 import { Identifiers, Telemetry } from '../constants';
 import { reportAction } from '../progress/decorator';
 import { ReportableAction } from '../progress/types';
-import { IJupyterConnection, IJupyterKernelSpec, IKernelSession } from '../types';
+import { IJupyterConnection, IJupyterKernelSpec, ISessionWithSocket } from '../types';
 import { JupyterInvalidKernelError } from './jupyterInvalidKernelError';
 import { JupyterWaitForIdleError } from './jupyterWaitForIdleError';
 import { JupyterWebSockets } from './jupyterWebSocket';
@@ -107,8 +107,8 @@ export class JupyterSession extends BaseJupyterSession {
  public async createNewKernelSession(
  kernel: IJupyterKernelSpec | LiveKernelModel,
  timeoutMS: number
- ): Promise<IKernelSession> {
- let newSession: IKernelSession | undefined;
+ ): Promise<ISessionWithSocket> {
+ let newSession: ISessionWithSocket | undefined;
 
  try {
  // Don't immediately assume this kernel is valid. Try creating a session with it first.
@@ -144,15 +144,15 @@ export class JupyterSession extends BaseJupyterSession {
 
  protected async createRestartSession(
  kernelSpec: IJupyterKernelSpec | LiveKernelModel | undefined,
- session: IKernelSession,
+ session: ISessionWithSocket,
  cancelToken?: CancellationToken
- ): Promise<IKernelSession> {
+ ): Promise<ISessionWithSocket> {
  // We need all of the above to create a restart session
  if (!session || !this.contentsManager || !this.sessionManager) {
  throw new Error(localize.DataScience.sessionDisposed());
  }
 
- let result: IKernelSession | undefined;
+ let result: ISessionWithSocket | undefined;
  let tryCount = 0;
  // tslint:disable-next-line: no-any
  let exception: any;
@@ -187,7 +187,7 @@ export class JupyterSession extends BaseJupyterSession {
  }
 
  @captureTelemetry(Telemetry.WaitForIdleJupyter, undefined, true)
- private async waitForIdleOnSession(session: IKernelSession | undefined, timeout: number): Promise<void> {
+ private async waitForIdleOnSession(session: ISessionWithSocket | undefined, timeout: number): Promise<void> {
  if (session && session.kernel) {
  traceInfo(`Waiting for idle on (kernel): ${session.kernel.id} -> ${session.kernel.status}`);
  // tslint:disable-next-line: no-any
@@ -209,14 +209,14 @@ export class JupyterSession extends BaseJupyterSession {
  }
  };
 
- let statusChangeHandler: Slot<IKernelSession, Kernel.Status> | undefined;
+ let statusChangeHandler: Slot<ISessionWithSocket, Kernel.Status> | undefined;
  const kernelStatusChangedPromise = new Promise((resolve, reject) => {
- statusChangeHandler = (_: IKernelSession, e: Kernel.Status) => statusHandler(resolve, reject, e);
+ statusChangeHandler = (_: ISessionWithSocket, e: Kernel.Status) => statusHandler(resolve, reject, e);
  session.statusChanged.connect(statusChangeHandler);
  });
- let kernelChangedHandler: Slot<IKernelSession, Session.IKernelChangedArgs> | undefined;
+ let kernelChangedHandler: Slot<ISessionWithSocket, Session.IKernelChangedArgs> | undefined;
  const statusChangedPromise = new Promise((resolve, reject) => {
- kernelChangedHandler = (_: IKernelSession, e: Session.IKernelChangedArgs) =>
+ kernelChangedHandler = (_: ISessionWithSocket, e: Session.IKernelChangedArgs) =>
  statusHandler(resolve, reject, e.newValue?.status);
  session.kernelChanged.connect(kernelChangedHandler);
  });
@@ -265,7 +265,7 @@ export class JupyterSession extends BaseJupyterSession {
  kernelSpec: IJupyterKernelSpec | LiveKernelModel | undefined,
  contentsManager: ContentsManager,
  cancelToken?: CancellationToken
- ): Promise<IKernelSession> {
+ ): Promise<ISessionWithSocket> {
  // Create a temporary notebook for this session.
  this.notebookFiles.push(await contentsManager.newUntitled({ type: 'notebook' }));
 

src/client/datascience/jupyter/jupyterWebSocket.ts

@@ -3,6 +3,7 @@
 'use strict';
 import * as WebSocketWS from 'ws';
 import { traceError } from '../../common/logger';
+import { KernelSocketWrapper } from '../kernelSocketWrapper';
 import { IKernelSocket } from '../types';
 
 // tslint:disable: no-any
@@ -11,12 +12,8 @@ export const JupyterWebSockets = new Map<string, WebSocketWS & IKernelSocket>();
 // We need to override the websocket that jupyter lab services uses to put in our cookie information
 // Do this as a function so that we can pass in variables the the socket will have local access to
 export function createJupyterWebSocket(cookieString?: string, allowUnauthorized?: boolean) {
- class JupyterWebSocket extends WebSocketWS {
+ class JupyterWebSocket extends KernelSocketWrapper(WebSocketWS) {
  private kernelId: string | undefined;
- private receiveHooks: ((data: WebSocketWS.Data) => Promise<void>)[];
- private sendHooks: ((data: any, cb?: (err?: Error) => void) => Promise<void>)[];
- private msgChain: Promise<any>;
- private sendChain: Promise<any>;
 
  constructor(url: string, protocols?: string | string[] | undefined) {
  let co: WebSocketWS.ClientOptions = {};
@@ -36,12 +33,8 @@ export function createJupyterWebSocket(cookieString?: string, allowUnauthorized?
 
  super(url, protocols, co);
 
- // Make sure the message chain is initialized
- this.msgChain = Promise.resolve();
- this.sendChain = Promise.resolve();
-
  // Parse the url for the kernel id
- const parsed = /.*\/kernels\/(.*)\/.*/.exec(this.url);
+ const parsed = /.*\/kernels\/(.*)\/.*/.exec(url);
  if (parsed && parsed.length > 1) {
  this.kernelId = parsed[1];
  }
@@ -53,57 +46,6 @@ export function createJupyterWebSocket(cookieString?: string, allowUnauthorized?
  } else {
  traceError('KernelId not extracted from Kernel WebSocket URL');
  }
- this.receiveHooks = [];
- this.sendHooks = [];
- }
-
- public send(data: any, a2: any): void {
- if (this.sendHooks) {
- // Stick the send hooks into the send chain. We use chain
- // to ensure that:
- // a) Hooks finish before we fire the event for real
- // b) Event fires
- // c) Next message happens after this one (so the UI can handle the message before another event goes through)
- this.sendChain = this.sendChain
- .then(() => Promise.all(this.sendHooks.map((s) => s(data, a2))))
- .then(() => super.send(data, a2));
- } else {
- super.send(data, a2);
- }
- }
-
- public emit(event: string | symbol, ...args: any[]): boolean {
- if (event === 'message' && this.receiveHooks.length) {
- // Stick the receive hooks into the message chain. We use chain
- // to ensure that:
- // a) Hooks finish before we fire the event for real
- // b) Event fires
- // c) Next message happens after this one (so this side can handle the message before another event goes through)
- this.msgChain = this.msgChain
- .then(() => Promise.all(this.receiveHooks.map((p) => p(args[0]))))
- .then(() => super.emit(event, ...args));
- // True value indicates there were handlers. We definitely have 'message' handlers.
- return true;
- } else {
- return super.emit(event, ...args);
- }
- }
-
- public addReceiveHook(hook: (data: WebSocketWS.Data) => Promise<void>) {
- this.receiveHooks.push(hook);
- }
- public removeReceiveHook(hook: (data: WebSocketWS.Data) => Promise<void>) {
- this.receiveHooks = this.receiveHooks.filter((l) => l !== hook);
- }
-
- // tslint:disable-next-line: no-any
- public addSendHook(patch: (data: any, cb?: (err?: Error) => void) => Promise<void>): void {
- this.sendHooks.push(patch);
- }
-
- // tslint:disable-next-line: no-any
- public removeSendHook(patch: (data: any, cb?: (err?: Error) => void) => Promise<void>): void {
- this.sendHooks = this.sendHooks.filter((p) => p !== patch);
  }
  }
  return JupyterWebSocket;

src/client/datascience/kernelSocketWrapper.ts

@@ -0,0 +1,110 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+import * as WebSocketWS from 'ws';
+import { ClassType } from '../ioc/types';
+import { IKernelSocket } from './types';
+
+// tslint:disable: no-any prefer-method-signature
+export type IWebSocketLike = {
+ onopen: (event: { target: any }) => void;
+ onerror: (event: { error: any; message: string; type: string; target: any }) => void;
+ onclose: (event: { wasClean: boolean; code: number; reason: string; target: any }) => void;
+ onmessage: (event: { data: WebSocketWS.Data; type: string; target: any }) => void;
+ emit(event: string | symbol, ...args: any[]): boolean;
+ send(data: any, a2: any): void;
+};
+
+/**
+ * This is called a mixin class in TypeScript.
+ * Allows us to have different base classes but inherit behavior (workaround for not allowing multiple inheritance).
+ * Essentially it sticks a temp class in between the base class and the class you're writing.
+ * Something like this:
+ *
+ * class Base {
+ * doStuff() {
+ *
+ * }
+ * }
+ *
+ * function Mixin = (SuperClass) {
+ * return class extends SuperClass {
+ * doExtraStuff() {
+ * super.doStuff();
+ * }
+ * }
+ * }
+ *
+ * function SubClass extends Mixin(Base) {
+ * doBar() : {
+ * super.doExtraStuff();
+ * }
+ * }
+ *
+ */
+
+export function KernelSocketWrapper<T extends ClassType<IWebSocketLike>>(SuperClass: T) {
+ return class BaseKernelSocket extends SuperClass implements IKernelSocket {
+ private receiveHooks: ((data: WebSocketWS.Data) => Promise<void>)[];
+ private sendHooks: ((data: any, cb?: (err?: Error) => void) => Promise<void>)[];
+ private msgChain: Promise<any>;
+ private sendChain: Promise<any>;
+
+ constructor(...rest: any[]) {
+ super(...rest);
+ // Make sure the message chain is initialized
+ this.msgChain = Promise.resolve();
+ this.sendChain = Promise.resolve();
+ this.receiveHooks = [];
+ this.sendHooks = [];
+ }
+
+ public send(data: any, a2: any): void {
+ if (this.sendHooks) {
+ // Stick the send hooks into the send chain. We use chain
+ // to ensure that:
+ // a) Hooks finish before we fire the event for real
+ // b) Event fires
+ // c) Next message happens after this one (so the UI can handle the message before another event goes through)
+ this.sendChain = this.sendChain
+ .then(() => Promise.all(this.sendHooks.map((s) => s(data, a2))))
+ .then(() => super.send(data, a2));
+ } else {
+ super.send(data, a2);
+ }
+ }
+
+ public emit(event: string | symbol, ...args: any[]): boolean {
+ if (event === 'message' && this.receiveHooks.length) {
+ // Stick the receive hooks into the message chain. We use chain
+ // to ensure that:
+ // a) Hooks finish before we fire the event for real
+ // b) Event fires
+ // c) Next message happens after this one (so this side can handle the message before another event goes through)
+ this.msgChain = this.msgChain
+ .then(() => Promise.all(this.receiveHooks.map((p) => p(args[0]))))
+ .then(() => super.emit(event, ...args));
+ // True value indicates there were handlers. We definitely have 'message' handlers.
+ return true;
+ } else {
+ return super.emit(event, ...args);
+ }
+ }
+
+ public addReceiveHook(hook: (data: WebSocketWS.Data) => Promise<void>) {
+ this.receiveHooks.push(hook);
+ }
+ public removeReceiveHook(hook: (data: WebSocketWS.Data) => Promise<void>) {
+ this.receiveHooks = this.receiveHooks.filter((l) => l !== hook);
+ }
+
+ // tslint:disable-next-line: no-any
+ public addSendHook(patch: (data: any, cb?: (err?: Error) => void) => Promise<void>): void {
+ this.sendHooks.push(patch);
+ }
+
+ // tslint:disable-next-line: no-any
+ public removeSendHook(patch: (data: any, cb?: (err?: Error) => void) => Promise<void>): void {
+ this.sendHooks = this.sendHooks.filter((p) => p !== patch);
+ }
+ };
+}