Skip to content

Commit

Permalink
use positron.runtime.restartSession
Browse files Browse the repository at this point in the history
  • Loading branch information
@seeM
seeM committed Dec 11, 2024
1 parent b6691cf commit 72b8501
Show file tree
Hide file tree
Showing 4 changed files with 7 additions and 111 deletions.
8 changes: 4 additions & 4 deletions extensions/positron-notebook-controllers/src/commands.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,13 +3,13 @@
* Licensed under the Elastic License 2.0. See LICENSE.txt for license information.
*--------------------------------------------------------------------------------------------*/

import * as positron from 'positron';
import * as vscode from 'vscode';
import { NotebookSessionService } from './notebookSessionService';
import { getNotebookSession, isActiveNotebookEditorUri } from './utils';
import { setHasRunningNotebookSessionContext } from './extension';

export function registerCommands(context: vscode.ExtensionContext, notebookSessionService: NotebookSessionService): void {
context.subscriptions.push(vscode.commands.registerCommand('positron.restartKernel', async () => {
export function registerCommands(disposables: vscode.Disposable[]): void {
disposables.push(vscode.commands.registerCommand('positron.restartKernel', async () => {
// Get the active notebook.
const notebook = vscode.window.activeNotebookEditor?.notebook;
if (!notebook) {
Expand All @@ -32,7 +32,7 @@ export function registerCommands(context: vscode.ExtensionContext, notebookSessi
await vscode.window.withProgress({
location: vscode.ProgressLocation.Notification,
title: vscode.l10n.t("Restarting {0} interpreter for '{1}'", session.runtimeMetadata.runtimeName, notebook.uri.path),
}, () => notebookSessionService.restartRuntimeSession(notebook.uri));
}, () => positron.runtime.restartSession(session.metadata.sessionId));

// Enable the hasRunningNotebookSession context.
if (isActiveNotebookEditorUri(notebook.uri)) {
Expand Down
2 changes: 1 addition & 1 deletion extensions/positron-notebook-controllers/src/extension.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -99,7 +99,7 @@ export async function activate(context: vscode.ExtensionContext): Promise<void>
}
}));

registerCommands(context, notebookSessionService);
registerCommands(context.subscriptions);

registerExecutionInfoStatusBar(context.subscriptions, manager);
}
Expand Down
3 changes: 0 additions & 3 deletions extensions/positron-notebook-controllers/src/notebookController.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -276,9 +276,6 @@ export class NotebookController implements vscode.Disposable {
// 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 = await getNotebookSession(notebook.uri, this._runtimeMetadata.runtimeId);

Expand Down
105 changes: 2 additions & 103 deletions extensions/positron-notebook-controllers/src/notebookSessionService.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,12 +32,6 @@ export class NotebookSessionService {
*/
private readonly _shuttingDownSessionsByNotebookUri = new ResourceMap<Promise<void>>();

/**
* A map of sessions currently restarting, keyed by notebook URI. Values are promises that resolve
* when the session has completed the restart sequence.
*/
private readonly _restartingSessionsByNotebookUri = new ResourceMap<Promise<positron.LanguageRuntimeSession>>();

/**
* Wait for a notebook session to complete a shutdown sequence.
*
Expand All @@ -48,16 +42,6 @@ export class NotebookSessionService {
await this._shuttingDownSessionsByNotebookUri.get(notebookUri);
}

/**
* Wait for a notebook session to complete a restart sequence.
*
* @param notebookUri The notebook URI to wait for.
* @returns A promise that resolves when the session has completed the restart sequence.
*/
async waitForNotebookSessionToRestart(notebookUri: vscode.Uri): Promise<void> {
await this._restartingSessionsByNotebookUri.get(notebookUri);
}

/**
* Start a new runtime session for a notebook.
*
Expand All @@ -67,8 +51,7 @@ export class NotebookSessionService {
*/
async startRuntimeSession(notebookUri: vscode.Uri, runtimeId: string): Promise<positron.LanguageRuntimeSession> {
// Return the existing promise, if there is one.
const startingSessionPromise = this._startingSessionsByNotebookUri.get(notebookUri) ||
this._restartingSessionsByNotebookUri.get(notebookUri);
const startingSessionPromise = this._startingSessionsByNotebookUri.get(notebookUri);
if (startingSessionPromise) {
return startingSessionPromise;
}
Expand Down Expand Up @@ -187,8 +170,7 @@ export class NotebookSessionService {
}

// Check for a pending session.
const pendingSessionPromise = this._startingSessionsByNotebookUri.get(notebookUri) ||
this._restartingSessionsByNotebookUri.get(notebookUri);
const pendingSessionPromise = this._startingSessionsByNotebookUri.get(notebookUri);
if (pendingSessionPromise) {
try {
return await pendingSessionPromise;
Expand All @@ -200,87 +182,4 @@ export class NotebookSessionService {
// There is no existing or pending session for the notebook.
return undefined;
}

/**
* Restart a runtime session for a notebook.
*
* @param notebookUri The notebook URI to restart a runtime for.
* @returns Promise that resolves when the runtime restart sequence has completed and the
* session is enters the ready state.
*/
async restartRuntimeSession(notebookUri: vscode.Uri): Promise<positron.LanguageRuntimeSession> {
// Return the existing promise, if there is one.
const startingSessionPromise = this._startingSessionsByNotebookUri.get(notebookUri) ||
this._restartingSessionsByNotebookUri.get(notebookUri);
if (startingSessionPromise) {
return startingSessionPromise;
}

// Construct a wrapping promise that resolves/rejects after the session maps have been updated.
const restartPromise = (async () => {
try {
const session = await this.doRestartRuntimeSession(notebookUri);
this._restartingSessionsByNotebookUri.delete(notebookUri);
log.info(`Session ${session.metadata.sessionId} is restarted`);
return session;
} catch (err) {
this._restartingSessionsByNotebookUri.delete(notebookUri);
throw err;
}
})();

this._restartingSessionsByNotebookUri.set(notebookUri, restartPromise);

return restartPromise;
}

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

// If the notebook's session is still shutting down, wait for it to finish.
const shuttingDownSessionPromise = this._shuttingDownSessionsByNotebookUri.get(notebookUri);
if (shuttingDownSessionPromise) {
try {
await shuttingDownSessionPromise;
} catch (err) {
log.error(`Waiting for notebook runtime to shutdown before starting failed. Reason ${err}`);
throw err;
}
}

// Create a promise that resolves when the session is ready.
const promise = new Promise<void>((resolve) => {
const disposable = session.onDidChangeRuntimeState((state) => {
if (state === positron.RuntimeState.Ready) {
disposable.dispose();
resolve();
}
});
});

// Start the restart sequence.
try {
log.info(`Restarting session ${session.metadata.sessionId} for notebook ${notebookUri.path}`);
await positron.runtime.restartSession(session.metadata.sessionId);
} catch (err) {
log.error(`Restarting session ${session.metadata.sessionId} for notebook ${notebookUri.path} failed. Reason: ${err}`);
throw err;
}

// Wait for the session to be ready, or for a timeout.
const timeout = new Promise<void>((_, reject) =>
setTimeout(() => reject(new Error('Timeout waiting for runtime to restart')), 5000));
try {
await Promise.race([promise, timeout]);
} catch (err) {
log.error(err);
throw err;
}

return session;
}
}

0 comments on commit 72b8501

Please sign in to comment.