Skip to content

Commit

Permalink
Retrieve notebook sessions via the positron.runtime API (#5632)
Browse files Browse the repository at this point in the history 
The goal of this PR is to stop relying on the
positron-notebook-controller extension's notebook session service
tracking the running notebook sessions and instead use the existing
`positron.runtime.getNotebookSession` method. It's another step toward
removing the notebook session service altogether.

Related to #2671.

### QA Notes

Tests should continue to pass and opening/closing/restarting/executing
notebooks should feel stable in general.
  • Loading branch information
@seeM
seeM authored Dec 10, 2024
1 parent 54d791d commit 0a51167
Show file tree
Hide file tree
Showing 8 changed files with 85 additions and 94 deletions.
3 changes: 2 additions & 1 deletion extensions/positron-notebook-controllers/src/commands.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@

import * as vscode from 'vscode';
import { NotebookSessionService } from './notebookSessionService';
import { getNotebookSession } from './utils';

export function registerCommands(context: vscode.ExtensionContext, notebookSessionService: NotebookSessionService): void {
context.subscriptions.push(vscode.commands.registerCommand('positron.restartKernel', async () => {
Expand All @@ -15,7 +16,7 @@ export function registerCommands(context: vscode.ExtensionContext, notebookSessi
}

// Get the session for the active notebook.
const session = notebookSessionService.getNotebookSession(notebook.uri);
const session = await getNotebookSession(notebook.uri);
if (!session) {
throw new Error('No session found for active notebook. This command should only be available when a session is running.');
}
Expand Down
15 changes: 7 additions & 8 deletions extensions/positron-notebook-controllers/src/extension.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import { NotebookSessionService } from './notebookSessionService';
import { registerCommands } from './commands';
import { JUPYTER_NOTEBOOK_TYPE } from './constants';
import { registerExecutionInfoStatusBar } from './statusBar';
import { getNotebookSession } from './utils';

export const log = vscode.window.createOutputChannel('Positron Notebook Controllers', { log: true });

Expand All @@ -20,9 +21,7 @@ export async function activate(context: vscode.ExtensionContext): Promise<void>
// Shutdown any running sessions when a notebook is closed.
context.subscriptions.push(vscode.workspace.onDidCloseNotebookDocument(async (notebook) => {
log.debug(`Notebook closed: ${notebook.uri.path}`);
if (notebookSessionService.hasStartingOrRunningNotebookSession(notebook.uri)) {
await notebookSessionService.shutdownRuntimeSession(notebook.uri);
}
await notebookSessionService.shutdownRuntimeSession(notebook.uri);
}));

const manager = new NotebookControllerManager(notebookSessionService);
Expand Down Expand Up @@ -53,15 +52,15 @@ export async function activate(context: vscode.ExtensionContext): Promise<void>
}

// Set the hasRunningNotebookSession context when the active notebook editor changes.
context.subscriptions.push(vscode.window.onDidChangeActiveNotebookEditor((editor) => {
const value = notebookSessionService.hasRunningNotebookSession(editor?.notebook.uri);
setHasRunningNotebookSessionContext(value);
context.subscriptions.push(vscode.window.onDidChangeActiveNotebookEditor(async (editor) => {
const session = editor && await getNotebookSession(editor.notebook.uri);
setHasRunningNotebookSessionContext(Boolean(session));
}));

// Set the hasRunningNotebookSession context when a session is started/shutdown for the active notebook.
context.subscriptions.push(notebookSessionService.onDidChangeNotebookSession((e) => {
if (e.notebookUri === vscode.window.activeNotebookEditor?.notebook.uri) {
setHasRunningNotebookSessionContext(!!e.session);
if (e.notebookUri.toString() === vscode.window.activeNotebookEditor?.notebook.uri.toString()) {
setHasRunningNotebookSessionContext(Boolean(e.session));
}
}));

Expand Down
14 changes: 8 additions & 6 deletions extensions/positron-notebook-controllers/src/notebookController.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import { NotebookSessionService } from './notebookSessionService';
import { JUPYTER_NOTEBOOK_TYPE } from './constants';
import { log } from './extension';
import { ResourceMap } from './map';
import { getNotebookSession } from './utils';

/** The type of a Jupyter notebook cell output. */
enum NotebookCellOutputType {
Expand Down Expand Up @@ -121,9 +122,7 @@ export class NotebookController implements vscode.Disposable {

private async selectRuntimeSession(notebook: vscode.NotebookDocument): Promise<void> {
// If there's an existing session from another runtime, shut it down.
if (this._notebookSessionService.hasStartingOrRunningNotebookSession(notebook.uri)) {
await this._notebookSessionService.shutdownRuntimeSession(notebook.uri);
}
await this._notebookSessionService.shutdownRuntimeSession(notebook.uri);

// Start the new session.
await this.startRuntimeSession(notebook);
Expand Down Expand Up @@ -166,7 +165,7 @@ export class NotebookController implements vscode.Disposable {
*/
private async interruptRuntimeSession(notebook: vscode.NotebookDocument): Promise<void> {
// If the notebook has a session, interrupt it.
const session = this._notebookSessionService.getNotebookSession(notebook.uri);
const session = await getNotebookSession(notebook.uri);
if (session) {
await session.interrupt();
return;
Expand Down Expand Up @@ -207,7 +206,7 @@ export class NotebookController implements vscode.Disposable {
try {
await Promise.all(cells.map(cell => this.queueCellExecution(cell, notebook, tokenSource.token)));
} catch (err) {
log.debug(`Error executing cells: ${err}`);
log.debug(`Error executing cells: ${err.stack ?? err}`);
} finally {
// Clean up the cancellation token source for this execution.
if (this._executionTokenSourceByNotebookUri.get(notebook.uri) === tokenSource) {
Expand Down Expand Up @@ -260,11 +259,14 @@ export class NotebookController implements vscode.Disposable {
return;
}

// If a session is shutting down for this notebook, wait for it to finish.
await this._notebookSessionService.waitForNotebookSessionToShutdown(notebook.uri);

// If a session is restarting for this notebook, wait for it to finish.
await this._notebookSessionService.waitForNotebookSessionToRestart(notebook.uri);

// Get the notebook's session.
let session = this._notebookSessionService.getNotebookSession(notebook.uri);
let session = await getNotebookSession(notebook.uri, this._runtimeMetadata.runtimeId);

// No session has been started for this notebook, start one.
if (!session) {
Expand Down
88 changes: 13 additions & 75 deletions extensions/positron-notebook-controllers/src/notebookSessionService.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import * as vscode from 'vscode';
import * as path from 'path';
import { log } from './extension';
import { ResourceMap } from './map';
import { getNotebookSession } from './utils';

export interface INotebookSessionDidChangeEvent {
/** The URI of the notebook corresponding to the session. */
Expand Down Expand Up @@ -47,9 +48,6 @@ export class NotebookSessionService implements vscode.Disposable {
*/
private readonly _restartingSessionsByNotebookUri = new ResourceMap<Promise<positron.LanguageRuntimeSession>>();

/** A map of the currently active notebook sessions, keyed by notebook URI. */
private readonly _notebookSessionsByNotebookUri = new ResourceMap<positron.LanguageRuntimeSession>();

/** The event emitter for the onDidChangeNotebookSession event. */
private readonly _onDidChangeNotebookSession = this._register(new vscode.EventEmitter<INotebookSessionDidChangeEvent>);

Expand All @@ -62,28 +60,13 @@ export class NotebookSessionService implements vscode.Disposable {
}

/**
* Checks for a starting or running notebook for the given notebook URI.
* Wait for a notebook session to complete a shutdown sequence.
*
* @param notebookUri The notebook URI to check for.
* @returns True if a starting or running notebook session exists for the given notebook URI.
*/
hasStartingOrRunningNotebookSession(notebookUri: vscode.Uri): boolean {
return this._startingSessionsByNotebookUri.has(notebookUri) ||
this._restartingSessionsByNotebookUri.has(notebookUri) ||
this._notebookSessionsByNotebookUri.has(notebookUri);
}

/**
* Checks for a running notebook for the given notebook URI.
*
* @param notebookUri The notebook URI to check for.
* @returns True if a running notebook session exists for the given notebook URI.
* @param notebookUri The notebook URI to wait for.
* @returns A promise that resolves when the session has completed the shutdown sequence.
*/
hasRunningNotebookSession(notebookUri: vscode.Uri | undefined): boolean {
if (!notebookUri) {
return false;
}
return this._notebookSessionsByNotebookUri.has(notebookUri);
async waitForNotebookSessionToShutdown(notebookUri: vscode.Uri): Promise<void> {
await this._shuttingDownSessionsByNotebookUri.get(notebookUri);
}

/**
Expand All @@ -96,31 +79,6 @@ export class NotebookSessionService implements vscode.Disposable {
await this._restartingSessionsByNotebookUri.get(notebookUri);
}

/**
* Get the running notebook session for the given notebook URI, if one exists.
*
* @param notebookUri The notebook URI of the session to retrieve.
* @returns The running notebook session for the given notebook URI, if one exists.
*/
getNotebookSession(notebookUri: vscode.Uri): positron.LanguageRuntimeSession | undefined {
return this._notebookSessionsByNotebookUri.get(notebookUri);
}

/**
* Set a notebook session for a notebook URI.
*
* @param notebookUri The notebook URI of the session to set.
* @param session The session to set for the notebook URI, or undefined to delete the session.
*/
setNotebookSession(notebookUri: vscode.Uri, session: positron.LanguageRuntimeSession | undefined): void {
if (session) {
this._notebookSessionsByNotebookUri.set(notebookUri, session);
} else {
this._notebookSessionsByNotebookUri.delete(notebookUri);
}
this._onDidChangeNotebookSession.fire({ notebookUri, session });
}

/**
* Start a new runtime session for a notebook.
*
Expand All @@ -141,7 +99,7 @@ export class NotebookSessionService implements vscode.Disposable {
try {
const session = await this.doStartRuntimeSession(notebookUri, runtimeId);
this._startingSessionsByNotebookUri.delete(notebookUri);
this.setNotebookSession(notebookUri, session);
this._onDidChangeNotebookSession.fire({ notebookUri, session });
log.info(`Session ${session.metadata.sessionId} is started`);
return session;
} catch (err) {
Expand All @@ -167,27 +125,7 @@ export class NotebookSessionService implements vscode.Disposable {
}
}

// If there's already a session for this runtime e.g. one restored after a window reload, use it.
try {
const session = await positron.runtime.getNotebookSession(notebookUri);
if (session) {
// TODO: If it isn't running, log an error and start a new one.
// TODO: If it doesn't match the runtime ID, log an error, shut it down, and start a new one.
log.info(
`Restored session for language runtime ${session.metadata.sessionId} `
+ `(language: ${session.runtimeMetadata.languageName}, name: ${session.runtimeMetadata.runtimeName}, `
+ `version: ${session.runtimeMetadata.runtimeVersion}, notebook: ${notebookUri.path})`
);
return session;
}
} catch (err) {
log.error(
`Getting existing session for notebook ${notebookUri.path}' failed. Reason: ${err}`
);
throw err;
}

// If we couldn't restore a session, start a new one.
// Start the session.
try {
const session = await positron.runtime.startLanguageRuntime(
runtimeId,
Expand Down Expand Up @@ -223,7 +161,7 @@ export class NotebookSessionService implements vscode.Disposable {
try {
await this.doShutdownRuntimeSession(notebookUri);
this._shuttingDownSessionsByNotebookUri.delete(notebookUri);
this.setNotebookSession(notebookUri, undefined);
this._onDidChangeNotebookSession.fire({ notebookUri, session: undefined });
} catch (err) {
this._startingSessionsByNotebookUri.delete(notebookUri);
throw err;
Expand Down Expand Up @@ -277,7 +215,7 @@ export class NotebookSessionService implements vscode.Disposable {

private async getExistingOrPendingSession(notebookUri: vscode.Uri): Promise<positron.LanguageRuntimeSession | undefined> {
// Check for an active session first.
const activeSession = this._notebookSessionsByNotebookUri.get(notebookUri);
const activeSession = await getNotebookSession(notebookUri);
if (activeSession) {
return activeSession;
}
Expand Down Expand Up @@ -317,7 +255,7 @@ export class NotebookSessionService implements vscode.Disposable {
try {
const session = await this.doRestartRuntimeSession(notebookUri);
this._restartingSessionsByNotebookUri.delete(notebookUri);
this.setNotebookSession(notebookUri, session);
this._onDidChangeNotebookSession.fire({ notebookUri, session });
log.info(`Session ${session.metadata.sessionId} is restarted`);
return session;
} catch (err) {
Expand All @@ -333,14 +271,14 @@ export class NotebookSessionService implements vscode.Disposable {

async doRestartRuntimeSession(notebookUri: vscode.Uri): Promise<positron.LanguageRuntimeSession> {
// Get the notebook's session.
const session = this._notebookSessionsByNotebookUri.get(notebookUri);
const session = await getNotebookSession(notebookUri);
if (!session) {
throw new Error(`Tried to restart runtime for notebook without a running runtime: ${notebookUri.path}`);
}

// Remove the session from the map of active notebooks in case it's accessed while we're
// restarting.
this.setNotebookSession(notebookUri, undefined);
this._onDidChangeNotebookSession.fire({ notebookUri, session: undefined });

// If the notebook's session is still shutting down, wait for it to finish.
const shuttingDownSessionPromise = this._shuttingDownSessionsByNotebookUri.get(notebookUri);
Expand Down
16 changes: 12 additions & 4 deletions extensions/positron-notebook-controllers/src/test/notebookController.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ suite('NotebookController', () => {
let notebook: vscode.NotebookDocument;
let cells: vscode.NotebookCell[];
let session: TestLanguageRuntimeSession;
let getNotebookSessionStub: sinon.SinonStub;
let executions: TestNotebookCellExecution[];
let onDidCreateNotebookCellExecution: vscode.EventEmitter<TestNotebookCellExecution>;

Expand Down Expand Up @@ -79,9 +80,10 @@ suite('NotebookController', () => {
} as vscode.NotebookCell];

// Create a test session.
session = new TestLanguageRuntimeSession();
session = new TestLanguageRuntimeSession(runtime);
disposables.push(session);
notebookSessionService.getNotebookSession.withArgs(notebook.uri).returns(session as any);
getNotebookSessionStub = sinon.stub(positron.runtime, 'getNotebookSession')
.withArgs(notebook.uri).resolves(session as any);

// Stub the notebook controller to return a test cell execution.
executions = [];
Expand Down Expand Up @@ -244,13 +246,19 @@ suite('NotebookController', () => {
const executionEndedPromise = executeNotebook([0]);
await executionStartedPromise;

const sessionInterruptSpy = sinon.spy(session, 'interrupt');

// Simulate the session exiting.
notebookSessionService.getNotebookSession.withArgs(notebook.uri).returns(undefined);
getNotebookSessionStub.withArgs(notebook.uri).resolves(undefined);

// Interrupt and wait for the execution to end.
// Interrupt and wait for the execution to end (it should actually end!).
await interruptNotebook();
await executionEndedPromise;

// session.interrupt() should not be called.
sinon.assert.notCalled(sessionInterruptSpy);

// The execution should still end unsuccessfully.
assert.equal(executions.length, 1);
executions[0].assertDidEndUnsuccessfully();
});
Expand Down
4 changes: 4 additions & 0 deletions extensions/positron-notebook-controllers/src/test/testLanguageRuntimeSession.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,10 @@ export class TestLanguageRuntimeSession implements Partial<positron.LanguageRunt
sessionId: 'test-session',
} as positron.RuntimeSessionMetadata;

constructor(
public readonly runtimeMetadata: positron.LanguageRuntimeMetadata,
) { }

execute(_code: string, id: string, _mode: positron.RuntimeCodeExecutionMode, _errorBehavior: positron.RuntimeErrorBehavior) {
this._lastExecutionId = id;
this._onDidExecute.fire(id);
Expand Down
30 changes: 30 additions & 0 deletions extensions/positron-notebook-controllers/src/utils.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,10 @@
* Licensed under the Elastic License 2.0. See LICENSE.txt for license information.
*--------------------------------------------------------------------------------------------*/

import * as positron from 'positron';
import * as vscode from 'vscode';
import { log } from './extension';

export function delay(ms: number) {
return new Promise(resolve => setTimeout(resolve, ms));
}
Expand All @@ -13,3 +17,29 @@ export function formatCount(count: number, unit: string): string {
}
return `${count} ${unit}s`;
}

/**
* Get the language runtime session for a notebook.
*
* @param notebookUri The URI of the notebook.
* @param runtimeId Optional runtime ID to filter the session by.
* @returns Promise that resolves with the language runtime session, or `undefined` if no session is found.
*/
export async function getNotebookSession(
notebookUri: vscode.Uri, runtimeId?: string,
): Promise<positron.LanguageRuntimeSession | undefined> {
// Get the session for the notebook.
const session = await positron.runtime.getNotebookSession(notebookUri);
if (!session) {
return undefined;
}

// Ensure that the session is for the requested runtime.
if (runtimeId && session.runtimeMetadata.runtimeId !== runtimeId) {
log.warn(`Expected session for notebook ${notebookUri} to be for runtime ${runtimeId}, ` +
`but it is for runtime ${session.runtimeMetadata.runtimeId}`);
return undefined;
}

return session;
}
9 changes: 9 additions & 0 deletions src/positron-dts/positron.d.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -93,11 +93,20 @@ declare module 'positron' {
/** The runtime is busy executing code. */
Busy = 'busy',

/** The runtime is in the process of restarting. */
Restarting = 'restarting',

/** The runtime is in the process of shutting down. */
Exiting = 'exiting',

/** The runtime's host process has ended. */
Exited = 'exited',

/** The runtime is not responding to heartbeats and is presumed offline. */
Offline = 'offline',

/** The user has interrupted a busy runtime, but the runtime is not idle yet. */
Interrupting = 'interrupting',
}

/**
Expand Down

0 comments on commit 0a51167

Please sign in to comment.