Skip to content

Commit

Permalink
Stability Kernel Status banner websocket (jupyter-server#232)
Browse files Browse the repository at this point in the history
* when websocket closes, reopen a new one

* stabilize kernel blocked websocket too

* handle websockets closed on purpose

* add new telemetry listener to consolidate some logic

* create separate websockets for each telemetry event

* turn telemetry listener into singleton and single websocket

* add a test for telemetry listerner

* add test for reconnecting when a window is not hidden

* skip test that came from cookiecutter
  • Loading branch information
Zsailer authored and GitHub Enterprise committed Nov 2, 2021
1 parent 5eccf9e commit e7a73fa
Show file tree
Hide file tree
Showing 5 changed files with 191 additions and 40 deletions.
6 changes: 5 additions & 1 deletion jest.config.js
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,11 @@ config.transform = {
'\\.[j]sx?$': 'babel-jest'
};
config.transformIgnorePatterns = [];
config.coveragePathIgnorePatterns = ['./node_modules/', './test'];
config.coveragePathIgnorePatterns = [
'./node_modules/',
'./test',
'./src/handler.ts'
];
config.coverageThreshold = {
global: {
functions: 90,
Expand Down
23 changes: 10 additions & 13 deletions src/kernelblocked.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,8 @@ import {
JupyterFrontEndPlugin
} from '@jupyterlab/application';

import { PageConfig, URLExt } from '@jupyterlab/coreutils';
import { showErrorMessage } from '@jupyterlab/apputils';
import { TelemetryListener } from './telemetrylistener';

/**
* A plugin to capture Kernel action blocking events and surface to the user.
Expand All @@ -16,18 +16,13 @@ export const KernelBlockedPlugin: JupyterFrontEndPlugin<void> = {
activate: async (app: JupyterFrontEnd) => {
console.log('JupyterLab extension "Kernel Blocked Dialog" is activated!');

let url =
PageConfig.getOption('studioSubscribeURL') ||
URLExt.join(PageConfig.getWsUrl(), 'subscribe');
// Create a listener for kernel-blocked events.
const listener = TelemetryListener.getInstance();

const ws = new WebSocket(url);

ws.addEventListener('message', function (event) {
const data = JSON.parse(event.data);
console.log('Kernel action blocked:', data);

// Use session messages to update the path map and also log them.
if (data.__schema__ == 'event.datastudio.jupyter.com/kernel-blocked') {
listener.addCallback(
'event.datastudio.jupyter.com/kernel-blocked',
function (data: any) {
console.log('Kernel action blocked:', data);
const msg =
'Cannot ' +
data.action +
Expand All @@ -38,6 +33,8 @@ export const KernelBlockedPlugin: JupyterFrontEndPlugin<void> = {

showErrorMessage('409: Kernel Action Blocked', msg);
}
});
);

listener.connect();
}
};
60 changes: 34 additions & 26 deletions src/status.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,6 @@ import {

import { Dialog } from '@jupyterlab/apputils';

import { PageConfig, URLExt } from '@jupyterlab/coreutils';
import { DocumentRegistry } from '@jupyterlab/docregistry';
import { ILoggerRegistry } from '@jupyterlab/logconsole';
import {
Expand All @@ -29,6 +28,8 @@ import theme from '@tidbits/react-tidbits/theme';

import { ISessionContext, ReactWidget } from '@jupyterlab/apputils';

import { TelemetryListener } from './telemetrylistener';

// This is managed by tbump config in pyproject.toml
const VERSION = '0.14.5';

Expand Down Expand Up @@ -129,11 +130,6 @@ export const kernelStatusPlugin: JupyterFrontEndPlugin<void> = {
notebookTracker.forEach(panel => handleAdded);

const pathCache = new LRUCache(50);
let url =
PageConfig.getOption('studioSubscribeURL') ||
URLExt.join(PageConfig.getWsUrl(), 'subscribe');

const ws = new WebSocket(url);

const getPath = (kernel_id: string) => {
const model = find(
Expand All @@ -155,23 +151,29 @@ export const kernelStatusPlugin: JupyterFrontEndPlugin<void> = {
new NotebookBannerExtension()
);

ws.addEventListener('message', function (event) {
const data = JSON.parse(event.data);
let path = '';
console.log('got raw data', data);
// Create a listener for kernel-blocked events.
const listener = TelemetryListener.getInstance();

// Use session messages to update the path map and also log them.
if (data.__schema__ == 'event.datastudio.jupyter.com/session-message') {
listener.addCallback(
'event.datastudio.jupyter.com/session-message',
function (data: any) {
let path = '';
path = data.kernel_path;
pathCache.set(data.kernel_id, data.kernel_path);
console.log(
`setting path ${data.kernel_path} for kernel ${data.kernel_id}`
);
// Log messages to the logger registry
const logger = loggerRegistry.getLogger(path);
logger.log({ type: 'text', data: data.message, level: logger.level });
logger.log({ type: 'text', data: data, level: 'debug' });
}
);

// Log kernel messages to their mapped path if available.
} else if (
data.__schema__ == 'event.datastudio.jupyter.com/kernel-message'
) {
listener.addCallback(
'event.datastudio.jupyter.com/kernel-message',
function (data: any) {
let path = '';
const kernel_id = data.kernel_id;
path = getPath(kernel_id);
if (!path) {
Expand All @@ -180,11 +182,17 @@ export const kernelStatusPlugin: JupyterFrontEndPlugin<void> = {
);
return;
}
// Log messages to the logger registry
const logger = loggerRegistry.getLogger(path);
logger.log({ type: 'text', data: data.message, level: logger.level });
logger.log({ type: 'text', data: data, level: 'debug' });
}
);

// Use kernel status messages to update the appropriate toolbar banner.
} else if (
data.__schema__ == 'event.datastudio.jupyter.com/kernel-status'
) {
listener.addCallback(
'event.datastudio.jupyter.com/kernel-status',
function (data: any) {
let path = '';
const kernel_id = data.kernel_id;
path = getPath(kernel_id);
if (!path) {
Expand Down Expand Up @@ -257,13 +265,13 @@ export const kernelStatusPlugin: JupyterFrontEndPlugin<void> = {
dialogProperty.set(panel, dialog);
dialog.launch();
}
// Log messages to the logger registry
const logger = loggerRegistry.getLogger(path);
logger.log({ type: 'text', data: data.message, level: logger.level });
logger.log({ type: 'text', data: data, level: 'debug' });
}

// Log messages to the logger registry
const logger = loggerRegistry.getLogger(path);
logger.log({ type: 'text', data: data.message, level: logger.level });
logger.log({ type: 'text', data: event.data, level: 'debug' });
});
);
listener.connect();
}
};

Expand Down
79 changes: 79 additions & 0 deletions src/telemetrylistener.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,79 @@
import { PageConfig, URLExt } from '@jupyterlab/coreutils';

/*
* TelemetryListener
*
* Singleton class that opens a websocket that listens to the
* `/subscribe` endpoint, attaches callbacks to specific events
* that will be broadcast across this channel, and fires these callbacks
* functions when the event occurs.
*/
export class TelemetryListener {
private static instance: TelemetryListener;
callbacks: Record<string, Function> = {};
url: string;

private constructor() {
this.url =
PageConfig.getOption('studioSubscribeURL') ||
URLExt.join(PageConfig.getWsUrl(), 'subscribe');
}
/**
* Get an instance of the singleton. If an instance doesn't
* already exists, create new one.
* @returns TelemetryListener
*/
static getInstance(): TelemetryListener {
if (!TelemetryListener.instance) {
TelemetryListener.instance = new TelemetryListener();
}
return TelemetryListener.instance;
}
/**
* Connect a (long-lived) websocket to the
* `/subscribe` endpoint.
*/
public connect() {
if (document.hidden) {
setTimeout(() => {
this.connect();
}, 1000);
return;
}
const ws = new WebSocket(this.url);
let listener = this;
ws.onclose = function (event: CloseEvent) {
// If the websocket is closed on purpose,
// don't create a new one.
if (event.code === 1000) {
return;
}
/* istanbul ignore next */
setTimeout(() => {
listener.connect();
}, 1000);
};
for (let name in this.callbacks) {
let callback = this.callbacks[name];
ws.addEventListener('message', event => {
callback(event);
});
}
}
/**
* Attach a callback to an event that will be triggered when
* that event comes across the `/subscribe` endpoint.
*
* @param {string} eventId - Telemtry event ID that should trigger this callback
* @param {Function} callback - The function to execute when a trigger happens.
*/
addCallback(eventId: string, callback: Function) {
const listener = function (event: MessageEvent) {
const data = JSON.parse(event.data);
if (data.__schema__ == eventId) {
callback(data);
}
};
this.callbacks[eventId] = listener;
}
}
63 changes: 63 additions & 0 deletions test/telemetrylistener.spec.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,63 @@
import WS from 'jest-websocket-mock';

import { TelemetryListener } from '../src/telemetrylistener';
import { PageConfig } from '@jupyterlab/coreutils';
import { sleep } from '@jupyterlab/testutils';

describe('telemetry listener', () => {
it('should test that the telemetry listener websocket listens to the /subscribe endpoint', () => {
const url = 'ws://localhost:5555';
PageConfig.setOption('studioSubscribeURL', url);
const server = new WS(url, { jsonProtocol: true });

let triggered = false;

/* Define a dummy callback function to verify it gets called */
const callback = function () {
triggered = true;
};

let listener = TelemetryListener.getInstance();
listener.addCallback('test-event', callback);
listener.connect();

server.send({ __schema__: 'test-event' });
expect(triggered).toEqual(true);

server.close();
WS.clean();
});

it('should test that the telemetry listener reconnects when the window is not hidden', async () => {
Object.defineProperty(document, 'hidden', {
value: true,
configurable: true
});

const url = 'ws://localhost:5555';
PageConfig.setOption('studioSubscribeURL', url);

const server = new WS(url, { jsonProtocol: true });
let listener = TelemetryListener.getInstance();

let triggered = false;
listener.addCallback('test-event', () => {
triggered = true;
});
listener.connect();
server.send({ __schema__: 'test-event' });
expect(triggered).toEqual(false);

Object.defineProperty(document, 'hidden', {
value: false,
configurable: true
});
await sleep(2000);

server.send({ __schema__: 'test-event' });
expect(triggered).toEqual(true);

server.close();
WS.clean();
});
});

0 comments on commit e7a73fa

Please sign in to comment.