From f943fb96d5ca1a38ac00ea77ebe60585188ea7b5 Mon Sep 17 00:00:00 2001 From: Wasim Lorgat Date: Wed, 3 Apr 2024 16:45:56 +0200 Subject: [PATCH] Use restored runtime sessions in notebooks and improve notebook runtime startup/shutdown stability (#2604) * use log output channel * extract notebook session service; fix concurrent start/shutdown calls; restore sessions * rename * fix copyright * remove debug log * docs * address review comments; renames * simplify start on execute * register `onDidCloseNotebookDocument` listener at extension-level * Apply suggestions from code review Co-authored-by: Nick Strayer Signed-off-by: Wasim Lorgat --------- Signed-off-by: Wasim Lorgat Co-authored-by: Nick Strayer --- .../src/extension.ts | 15 +- .../src/logging.ts | 14 -- .../positron-notebook-controllers/src/map.ts | 124 ++++++++++ .../src/notebookController.ts | 159 ++---------- .../src/notebookControllerManager.ts | 15 +- .../src/notebookSessionService.ts | 231 ++++++++++++++++++ src/positron-dts/positron.d.ts | 7 + .../positron/mainThreadLanguageRuntime.ts | 10 + .../positron/extHost.positron.api.impl.ts | 3 + .../positron/extHost.positron.protocol.ts | 1 + .../common/positron/extHostLanguageRuntime.ts | 12 + .../runtimeSession/common/runtimeSession.ts | 37 ++- .../common/runtimeSessionService.ts | 6 + .../runtimeStartup/common/runtimeStartup.ts | 13 +- 14 files changed, 490 insertions(+), 157 deletions(-) delete mode 100644 extensions/positron-notebook-controllers/src/logging.ts create mode 100644 extensions/positron-notebook-controllers/src/map.ts create mode 100644 extensions/positron-notebook-controllers/src/notebookSessionService.ts diff --git a/extensions/positron-notebook-controllers/src/extension.ts b/extensions/positron-notebook-controllers/src/extension.ts index 1c08cd948ca..9fbfa8e5d87 100644 --- a/extensions/positron-notebook-controllers/src/extension.ts +++ b/extensions/positron-notebook-controllers/src/extension.ts @@ -4,13 +4,22 @@ import * as positron from 'positron'; import * as vscode from 'vscode'; -import { initializeLogging } from './logging'; import { NotebookControllerManager } from './notebookControllerManager'; +import { NotebookSessionService } from './notebookSessionService'; + +export const log = vscode.window.createOutputChannel('Positron Notebook Controllers', { log: true }); export async function activate(context: vscode.ExtensionContext): Promise { - initializeLogging(); + const notebookSessionService = new NotebookSessionService(); + + // Shutdown any running sessions when a notebook is closed. + context.subscriptions.push(vscode.workspace.onDidCloseNotebookDocument(async (notebook) => { + if (notebookSessionService.hasStartingOrRunningNotebookSession(notebook.uri)) { + await notebookSessionService.shutdownRuntimeSession(notebook.uri); + } + })); - const manager = new NotebookControllerManager(); + const manager = new NotebookControllerManager(notebookSessionService); context.subscriptions.push(manager); // Register notebook controllers for newly registered runtimes. diff --git a/extensions/positron-notebook-controllers/src/logging.ts b/extensions/positron-notebook-controllers/src/logging.ts deleted file mode 100644 index 85d1247f914..00000000000 --- a/extensions/positron-notebook-controllers/src/logging.ts +++ /dev/null @@ -1,14 +0,0 @@ -/*--------------------------------------------------------------------------------------------- - * Copyright (C) 2023 Posit Software, PBC. All rights reserved. - *--------------------------------------------------------------------------------------------*/ - -import * as vscode from 'vscode'; - -let channel: vscode.OutputChannel | undefined; -export function initializeLogging() { - channel = vscode.window.createOutputChannel('Positron Notebook Controllers'); -} - -export function trace(message: string) { - channel?.appendLine(message); -} diff --git a/extensions/positron-notebook-controllers/src/map.ts b/extensions/positron-notebook-controllers/src/map.ts new file mode 100644 index 00000000000..36a93ec1154 --- /dev/null +++ b/extensions/positron-notebook-controllers/src/map.ts @@ -0,0 +1,124 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (C) 2024 Posit Software, PBC. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +import { Uri as URI } from 'vscode'; + +// The ResourceMap class and its dependencies are copied as is from the core project. + +interface ResourceMapKeyFn { + (resource: URI): string; +} + +class ResourceMapEntry { + constructor(readonly uri: URI, readonly value: T) { } +} + +function isEntries(arg: ResourceMap | ResourceMapKeyFn | readonly (readonly [URI, T])[] | undefined): arg is readonly (readonly [URI, T])[] { + return Array.isArray(arg); +} + +export class ResourceMap implements Map { + + private static readonly defaultToKey = (resource: URI) => resource.toString(); + + readonly [Symbol.toStringTag] = 'ResourceMap'; + + private readonly map: Map>; + private readonly toKey: ResourceMapKeyFn; + + /** + * + * @param toKey Custom uri identity function, e.g use an existing `IExtUri#getComparison`-util + */ + constructor(toKey?: ResourceMapKeyFn); + + /** + * + * @param other Another resource which this maps is created from + * @param toKey Custom uri identity function, e.g use an existing `IExtUri#getComparison`-util + */ + constructor(other?: ResourceMap, toKey?: ResourceMapKeyFn); + + /** + * + * @param other Another resource which this maps is created from + * @param toKey Custom uri identity function, e.g use an existing `IExtUri#getComparison`-util + */ + constructor(entries?: readonly (readonly [URI, T])[], toKey?: ResourceMapKeyFn); + + constructor(arg?: ResourceMap | ResourceMapKeyFn | readonly (readonly [URI, T])[], toKey?: ResourceMapKeyFn) { + if (arg instanceof ResourceMap) { + this.map = new Map(arg.map); + this.toKey = toKey ?? ResourceMap.defaultToKey; + } else if (isEntries(arg)) { + this.map = new Map(); + this.toKey = toKey ?? ResourceMap.defaultToKey; + + for (const [resource, value] of arg) { + this.set(resource, value); + } + } else { + this.map = new Map(); + this.toKey = arg ?? ResourceMap.defaultToKey; + } + } + + set(resource: URI, value: T): this { + this.map.set(this.toKey(resource), new ResourceMapEntry(resource, value)); + return this; + } + + get(resource: URI): T | undefined { + return this.map.get(this.toKey(resource))?.value; + } + + has(resource: URI): boolean { + return this.map.has(this.toKey(resource)); + } + + get size(): number { + return this.map.size; + } + + clear(): void { + this.map.clear(); + } + + delete(resource: URI): boolean { + return this.map.delete(this.toKey(resource)); + } + + forEach(clb: (value: T, key: URI, map: Map) => void, thisArg?: any): void { + if (typeof thisArg !== 'undefined') { + clb = clb.bind(thisArg); + } + for (const [_, entry] of this.map) { + clb(entry.value, entry.uri, this); + } + } + + *values(): IterableIterator { + for (const entry of this.map.values()) { + yield entry.value; + } + } + + *keys(): IterableIterator { + for (const entry of this.map.values()) { + yield entry.uri; + } + } + + *entries(): IterableIterator<[URI, T]> { + for (const entry of this.map.values()) { + yield [entry.uri, entry.value]; + } + } + + *[Symbol.iterator](): IterableIterator<[URI, T]> { + for (const [, entry] of this.map) { + yield [entry.uri, entry.value]; + } + } +} diff --git a/extensions/positron-notebook-controllers/src/notebookController.ts b/extensions/positron-notebook-controllers/src/notebookController.ts index dd4c6c8c91c..3794c359323 100644 --- a/extensions/positron-notebook-controllers/src/notebookController.ts +++ b/extensions/positron-notebook-controllers/src/notebookController.ts @@ -3,8 +3,8 @@ *--------------------------------------------------------------------------------------------*/ import * as vscode from 'vscode'; import * as positron from 'positron'; -import { trace } from './logging'; -import { DeferredPromise, delay } from './util'; +import { NotebookSessionService } from './notebookSessionService'; +import { log } from './extension'; /** * Wraps a vscode.NotebookController for a specific language, and manages a notebook runtime session @@ -17,15 +17,6 @@ export class NotebookController implements vscode.Disposable { /** The wrapped VSCode notebook controller. */ public readonly controller: vscode.NotebookController; - /** - * A map of sessions currently starting, keyed by notebook. Values are promises that resolve - * when the session has started and is ready to execute code. - */ - private readonly _startingSessionsByNotebook: Map> = new Map(); - - /** A map of the currently active sessions, keyed by notebook. */ - private readonly _activeSessionsByNotebook: Map = new Map(); - /** A map of the current execution order, keyed by session ID. */ private readonly _executionOrderBySessionId: Map = new Map(); @@ -33,12 +24,16 @@ export class NotebookController implements vscode.Disposable { private static _CELL_COUNTER = 0; /** - * @param languageId The language ID for which this controller is responsible. + * @param _languageId The language ID for which this controller is responsible. + * @param _notebookSessionService The notebook session service. */ - constructor(private readonly languageId: string) { + constructor( + private readonly _languageId: string, + private readonly _notebookSessionService: NotebookSessionService, + ) { // Create a VSCode notebook controller for this language. this.controller = vscode.notebooks.createNotebookController( - `positron-${languageId}`, + `positron-${_languageId}`, // The 'jupyter-notebook' notebook type is contributed via the built-in extension // extensions/ipynb. Registering our notebook controllers with the same type ensures // that they show up in the notebook UI's kernel picker for .ipynb files. @@ -56,15 +51,8 @@ export class NotebookController implements vscode.Disposable { this._disposables.push(this.controller); - this._disposables.push(vscode.workspace.onDidCloseNotebookDocument(async (notebook) => { - // Wait a few seconds before shutting down the runtime. If this was reached via a window - // reload, we want to give the runtime a chance to reconnect. - await delay(3000); - this.shutdownRuntimeSession(notebook); - })); - this._disposables.push(this.controller.onDidChangeSelectedNotebooks(async (e) => { - await this.shutdownRuntimeSession(e.notebook); + log.debug(`Notebook ${e.notebook.uri}, controller ${this._languageId}, selected ${e.selected}`); // Has this controller been selected for a notebook? if (e.selected) { @@ -72,61 +60,34 @@ export class NotebookController implements vscode.Disposable { // already selected. await Promise.all([ - updateNotebookLanguage(e.notebook, this.languageId), - this.startNewRuntimeSession(e.notebook), + updateNotebookLanguage(e.notebook, this._languageId), + this.startRuntimeSession(e.notebook), ]); + } else { + await this._notebookSessionService.shutdownRuntimeSession(e.notebook.uri); } })); } /** The human-readable label of the controller. */ public get label(): string { - return `${this.languageId[0].toUpperCase()}${this.languageId.slice(1)}`; + return `${this._languageId[0].toUpperCase()}${this._languageId.slice(1)}`; } /** - * Shutdown the runtime session for a notebook. - * - * @param notebook Notebook whose runtime to shutdown. - * @returns Promise that resolves when the runtime shutdown sequence has been started. - */ - private async shutdownRuntimeSession(notebook: vscode.NotebookDocument): Promise { - // Get the notebook's session. - let session: positron.LanguageRuntimeSession | undefined; - const startingSessionPromise = this._startingSessionsByNotebook.get(notebook); - if (startingSessionPromise && !startingSessionPromise.isSettled) { - // If the runtime is still starting, wait for it to be ready. - session = await startingSessionPromise.p; - } else { - session = this._activeSessionsByNotebook.get(notebook); - } - - if (!session) { - trace(`Tried to shutdown runtime for notebook without a running runtime: ${notebook.uri.path}`); - return; - } - - await session.shutdown(positron.RuntimeExitReason.Shutdown); - session.dispose(); - this._activeSessionsByNotebook.delete(notebook); - this._executionOrderBySessionId.delete(session.metadata.sessionId); - trace(`Shutdown runtime ${session.runtimeMetadata.runtimeName} for notebook ${notebook.uri.path}`); - } - - /** - * Start a new runtime session for a notebook. + * Start a runtime session for a notebook. * * @param notebook The notebook to start a runtime for. * @returns Promise that resolves when the runtime has started. */ - public async startNewRuntimeSession(notebook: vscode.NotebookDocument): Promise { + public async startRuntimeSession(notebook: vscode.NotebookDocument): Promise { try { - return await this.doStartNewRuntimeSession(notebook); + return await this._notebookSessionService.startRuntimeSession(notebook.uri, this._languageId); } catch (err) { - const retry = vscode.l10n.t(`Retry`); + const retry = vscode.l10n.t('Retry'); const selection = await vscode.window.showErrorMessage( vscode.l10n.t( - "Starting {0} interpreter for '{1}' failed. Reason: {2}", + 'Starting {0} interpreter for "{1}" failed. Reason: {2}', this.label, notebook.uri.path, err @@ -134,72 +95,13 @@ export class NotebookController implements vscode.Disposable { retry, ); if (selection === retry) { - return vscode.window.withProgress(this.startProgressOptions(notebook), () => this.startNewRuntimeSession(notebook)); + return vscode.window.withProgress(this.startProgressOptions(notebook), () => this.startRuntimeSession(notebook)); } throw err; } } - private async doStartNewRuntimeSession(notebook: vscode.NotebookDocument): Promise { - if (this._activeSessionsByNotebook.has(notebook)) { - throw new Error(`Tried to start a runtime for a notebook that already has one: ${notebook.uri.path}`); - } - - const startingSessionPromise = this._startingSessionsByNotebook.get(notebook); - if (startingSessionPromise && !startingSessionPromise.isSettled) { - return startingSessionPromise.p; - } - - // Update the starting sessions map. This needs to be set before any awaits. - // When a user executes code without a controller selected, they will be presented - // with a quickpick. Once they make a selection, this is event is fired, and - // the execute handler is called immediately after. We need to ensure that the map - // is updated before that happens. - const startPromise = new DeferredPromise(); - this._startingSessionsByNotebook.set(notebook, startPromise); - - // Get the preferred runtime for this language. - let preferredRuntime: positron.LanguageRuntimeMetadata; - try { - preferredRuntime = await positron.runtime.getPreferredRuntime(this.languageId); - } catch (err) { - trace(`Getting preferred runtime for language '${this.languageId}' failed. Reason: ${err}`); - startPromise.error(err); - this._startingSessionsByNotebook.delete(notebook); - - throw err; - } - - // Start a session for the preferred runtime. - let session: positron.LanguageRuntimeSession; - try { - session = await positron.runtime.startLanguageRuntime( - preferredRuntime.runtimeId, - notebook.uri.path, // Use the notebook's path as the session name. - notebook.uri); - trace( - `Starting session for language runtime ${session.metadata.sessionId} ` - + `(language: ${this.label}, name: ${session.runtimeMetadata.runtimeName}, ` - + `version: ${session.runtimeMetadata.runtimeVersion}, notebook: ${notebook.uri.path})` - ); - } catch (err) { - trace(`Starting session for language runtime ${preferredRuntime.runtimeName} failed. Reason: ${err}`); - startPromise.error(err); - this._startingSessionsByNotebook.delete(notebook); - - throw err; - } - - this._activeSessionsByNotebook.set(notebook, session); - this._executionOrderBySessionId.set(session.metadata.sessionId, 0); - this._startingSessionsByNotebook.delete(notebook); - startPromise.complete(session); - trace(`Session ${session.metadata.sessionId} is ready`); - - return session; - } - /** * Notebook controller execute handler. * @@ -210,18 +112,11 @@ export class NotebookController implements vscode.Disposable { */ private async executeCells(cells: vscode.NotebookCell[], notebook: vscode.NotebookDocument, _controller: vscode.NotebookController) { // Get the notebook's session. - let session: positron.LanguageRuntimeSession | undefined; - const startingSessionPromise = this._startingSessionsByNotebook.get(notebook); - if (startingSessionPromise && !startingSessionPromise.isSettled) { - // If the runtime is still starting, wait for it to be ready. - session = await vscode.window.withProgress(this.startProgressOptions(notebook), () => startingSessionPromise.p); - } else { - session = this._activeSessionsByNotebook.get(notebook); - } + let session = this._notebookSessionService.getNotebookSession(notebook.uri); // No session has been started for this notebook, start one. if (!session) { - session = await vscode.window.withProgress(this.startProgressOptions(notebook), () => this.startNewRuntimeSession(notebook)); + session = await vscode.window.withProgress(this.startProgressOptions(notebook), () => this.startRuntimeSession(notebook)); } for (const cell of cells) { @@ -236,12 +131,8 @@ export class NotebookController implements vscode.Disposable { * @returns Promise that resolves when the runtime has finished executing the cell. */ private async executeCell(cell: vscode.NotebookCell, session: positron.LanguageRuntimeSession): Promise { - // Get the execution order for the session. - let executionOrder = this._executionOrderBySessionId.get(session.metadata.sessionId); - if (executionOrder === undefined) { - trace(`No execution order for session ${session.metadata.sessionId}, resetting to 0`); - executionOrder = 0; - } + // Get the execution order for the session, default to 0 for the first execution. + let executionOrder = this._executionOrderBySessionId.get(session.metadata.sessionId) ?? 0; // Create a cell execution. const currentExecution = this.controller.createNotebookCellExecution(cell); diff --git a/extensions/positron-notebook-controllers/src/notebookControllerManager.ts b/extensions/positron-notebook-controllers/src/notebookControllerManager.ts index f46fdb45351..53548128373 100644 --- a/extensions/positron-notebook-controllers/src/notebookControllerManager.ts +++ b/extensions/positron-notebook-controllers/src/notebookControllerManager.ts @@ -3,8 +3,9 @@ *--------------------------------------------------------------------------------------------*/ import * as vscode from 'vscode'; -import { trace } from './logging'; +import { log } from './extension'; import { NotebookController } from './notebookController'; +import { NotebookSessionService } from './notebookSessionService'; /** * Manages notebook controllers. @@ -13,6 +14,12 @@ export class NotebookControllerManager implements vscode.Disposable { /** Notebook controllers keyed by languageId. */ public readonly controllers = new Map(); + /** + * + * @param _notebookSessionService The notebook session service. + */ + constructor(private readonly _notebookSessionService: NotebookSessionService) { } + /** * Create a notebook controller for a language. * @@ -22,9 +29,9 @@ export class NotebookControllerManager implements vscode.Disposable { if (this.controllers.has(languageId)) { throw new Error(`Notebook controller already exists for language: ${languageId}`); } - const controller = new NotebookController(languageId); + const controller = new NotebookController(languageId, this._notebookSessionService); this.controllers.set(languageId, controller); - trace(`Registered notebook controller for language: ${languageId}`); + log.info(`Registered notebook controller for language: ${languageId}`); } /** @@ -79,7 +86,7 @@ export class NotebookControllerManager implements vscode.Disposable { ? vscode.NotebookControllerAffinity.Preferred : vscode.NotebookControllerAffinity.Default; controller.controller.updateNotebookAffinity(notebook, affinity); - trace(`Updated notebook affinity for language: ${languageId}, notebook: ${notebook.uri.path}, affinity: ${affinity}`); + log.debug(`Updated notebook affinity for controller: ${controller.label}, notebook: ${notebook.uri.path}, affinity: ${affinity}`); } } diff --git a/extensions/positron-notebook-controllers/src/notebookSessionService.ts b/extensions/positron-notebook-controllers/src/notebookSessionService.ts new file mode 100644 index 00000000000..58a7bc352ee --- /dev/null +++ b/extensions/positron-notebook-controllers/src/notebookSessionService.ts @@ -0,0 +1,231 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (C) 2024 Posit Software, PBC. All rights reserved. + *--------------------------------------------------------------------------------------------*/ + +import * as positron from 'positron'; +import { Uri } from 'vscode'; +import { log } from './extension'; +import { ResourceMap } from './map'; +import { DeferredPromise } from './util'; + +/** + * The notebook session service is the main interface for interacting with + * runtime sessions; it manages the set of active sessions and provides + * facilities for starting, stopping, and interacting with them. + * + * TODO(seem): Most of this code is copied from the runtime session service. We should bring what's + * required into the runtime session service and expose what's needed via the Positron Extensions + * API. + */ +export class NotebookSessionService { + + /** + * A map of sessions currently starting, keyed by notebook URI. Values are promises that resolve + * when the session has started and is ready to execute code. + */ + private readonly _startingSessionsByNotebookUri = new ResourceMap>(); + + /** + * A map of sessions currently shutting down, keyed by notebook URI. Values are promises that resolve + * when the session has completed the shutdown sequence. + */ + private readonly _shuttingDownSessionsByNotebookUri = new ResourceMap>(); + + /** A map of the currently active notebook sessions, keyed by notebook URI. */ + private readonly _notebookSessionsByNotebookUri = new ResourceMap(); + + /** + * Checks for a starting or running notebook for the given notebook URI. + * + * @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: Uri): boolean { + return this._startingSessionsByNotebookUri.has(notebookUri) || this._notebookSessionsByNotebookUri.has(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: Uri): positron.LanguageRuntimeSession | undefined { + return this._notebookSessionsByNotebookUri.get(notebookUri); + } + + /** + * Start a new runtime session for a notebook. + * + * @param notebookUri The notebook URI to start a runtime for. + * @returns Promise that resolves when the runtime startup sequence has been started. + */ + async startRuntimeSession(notebookUri: Uri, languageId: string): Promise { + // Return the existing promise, if there is one. + const startingSessionPromise = this._startingSessionsByNotebookUri.get(notebookUri); + if (startingSessionPromise && !startingSessionPromise.isSettled) { + return startingSessionPromise.p; + } + + // Update the starting sessions map. This needs to be set before any awaits in case another + // caller tries to start a runtime or access the start promise concurrently. + const startPromise = new DeferredPromise(); + this._startingSessionsByNotebookUri.set(notebookUri, startPromise); + + // If the notebook has a session that is still shutting down, wait for it to finish. + const shuttingDownSessionPromise = this._shuttingDownSessionsByNotebookUri.get(notebookUri); + if (shuttingDownSessionPromise && !shuttingDownSessionPromise.isSettled) { + await shuttingDownSessionPromise.p; + } + + // Ensure that we don't start a runtime for a notebook that already has one. + if (this._notebookSessionsByNotebookUri.has(notebookUri)) { + const err = new Error(`Tried to start a runtime for a notebook that already has one: ${notebookUri.path}`); + startPromise.error(err); + this._startingSessionsByNotebookUri.delete(notebookUri); + throw err; + } + + // If there's already a session for this runtime e.g. one restored after a window reload, use it. + let session: positron.LanguageRuntimeSession | undefined; + try { + session = await positron.runtime.getNotebookSession(notebookUri); + if (session) { + 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})` + ); + } + } catch (err) { + log.error( + `Getting existing session for notebook ${notebookUri.path}' failed. Reason: ${err}` + ); + startPromise.error(err); + this._startingSessionsByNotebookUri.delete(notebookUri); + throw err; + } + + // 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. + + // If we couldn't restore a session, start a new one. + if (!session) { + // Get the preferred runtime for this language. + let preferredRuntime: positron.LanguageRuntimeMetadata; + try { + preferredRuntime = await positron.runtime.getPreferredRuntime(languageId); + } catch (err) { + log.error(`Getting preferred runtime for language '${languageId}' failed. Reason: ${err}`); + startPromise.error(err); + this._startingSessionsByNotebookUri.delete(notebookUri); + throw err; + } + + try { + session = await positron.runtime.startLanguageRuntime( + preferredRuntime.runtimeId, + notebookUri.path, // Use the notebook's path as the session name. + notebookUri); + log.info( + `Starting session for language runtime ${session.metadata.sessionId} ` + + `(language: ${session.runtimeMetadata.languageName}, name: ${session.runtimeMetadata.runtimeName}, ` + + `version: ${session.runtimeMetadata.runtimeVersion}, notebook: ${notebookUri.path})` + ); + } catch (err) { + log.error(`Starting session for language runtime ${preferredRuntime.runtimeName} failed. Reason: ${err}`); + startPromise.error(err); + this._startingSessionsByNotebookUri.delete(notebookUri); + throw err; + } + } + + // Complete the promise and update the session maps. + this._notebookSessionsByNotebookUri.set(notebookUri, session); + this._startingSessionsByNotebookUri.delete(notebookUri); + startPromise.complete(session); + log.info(`Session ${session.metadata.sessionId} is ready`); + + return session; + } + + /** + * Shutdown the runtime session for a notebook. + * + * @param notebookUri The notebook URI whose runtime to shutdown. + * @returns Promise that resolves when the runtime shutdown sequence has been started. + */ + async shutdownRuntimeSession(notebookUri: Uri): Promise { + // Return the existing promise, if there is one. + const shuttingDownSessionPromise = this._shuttingDownSessionsByNotebookUri.get(notebookUri); + if (shuttingDownSessionPromise && !shuttingDownSessionPromise.isSettled) { + return shuttingDownSessionPromise.p; + } + + // Update the shutting down sessions map. This needs to be set before any awaits in case + // another caller tries to shutdown a runtime or access the shutdown promise concurrently. + const shutDownPromise = new DeferredPromise(); + this._shuttingDownSessionsByNotebookUri.set(notebookUri, shutDownPromise); + + // Get the notebook's session. + let session = this._notebookSessionsByNotebookUri.get(notebookUri); + + // If the runtime is still starting, wait for it to be ready. + if (!session) { + const startingSessionPromise = this._startingSessionsByNotebookUri.get(notebookUri); + if (startingSessionPromise && !startingSessionPromise.isSettled) { + session = await startingSessionPromise.p; + } + } + + // Ensure that we have a session. + if (!session) { + const err = new Error(`Tried to shutdown runtime for notebook without a running runtime: ${notebookUri.path}`); + this._shuttingDownSessionsByNotebookUri.delete(notebookUri); + shutDownPromise.error(err); + throw err; + } + + // Start the shutdown sequence. + try { + await session.shutdown(positron.RuntimeExitReason.Shutdown); + log.info(`Shutting down runtime ${session.runtimeMetadata.runtimeName} for notebook ${notebookUri.path}`); + } catch (err) { + log.error(`Shutting down runtime ${session.runtimeMetadata.runtimeName} for notebook ${notebookUri.path} failed. Reason: ${err}`); + this._shuttingDownSessionsByNotebookUri.delete(notebookUri); + this._notebookSessionsByNotebookUri.delete(notebookUri); + shutDownPromise.error(err); + throw err; + } + + // Wait for the session to end. This is necessary so that we know when to start the next + // session for the notebook, since at most one session can exist per notebook. + try { + const timeout = new Promise((_, reject) => { + setTimeout(() => { + reject(new Error(`Shutting down runtime ${session.runtimeMetadata.runtimeName} for notebook ${notebookUri.path} timed out`)); + }, 5000); + }); + const promise = new Promise(resolve => { + const disposable = session.onDidEndSession(() => { + disposable.dispose(); + resolve(); + }); + }); + await Promise.race([promise, timeout]); + } catch (err) { + log.error(err); + this._shuttingDownSessionsByNotebookUri.delete(notebookUri); + this._notebookSessionsByNotebookUri.delete(notebookUri); + shutDownPromise.error(err); + throw err; + } + + // Complete the promise and update the session maps. + this._shuttingDownSessionsByNotebookUri.delete(notebookUri); + this._notebookSessionsByNotebookUri.delete(notebookUri); + shutDownPromise.complete(); + log.info(`Session ${session.metadata.sessionId} shutdown completed`); + } + +} diff --git a/src/positron-dts/positron.d.ts b/src/positron-dts/positron.d.ts index 010b9a6d774..960e3f89564 100644 --- a/src/positron-dts/positron.d.ts +++ b/src/positron-dts/positron.d.ts @@ -1155,6 +1155,13 @@ declare module 'positron' { */ export function getForegroundSession(): Thenable; + /** + * Get the session corresponding to a notebook, if any. + * + * @param notebookUri The URI of the notebook. + */ + export function getNotebookSession(notebookUri: vscode.Uri): Thenable; + /** * Select and start a runtime previously registered with Positron. Any * previously active runtimes for the language will be shut down. diff --git a/src/vs/workbench/api/browser/positron/mainThreadLanguageRuntime.ts b/src/vs/workbench/api/browser/positron/mainThreadLanguageRuntime.ts index 56429f264b1..f97fa2cc2cd 100644 --- a/src/vs/workbench/api/browser/positron/mainThreadLanguageRuntime.ts +++ b/src/vs/workbench/api/browser/positron/mainThreadLanguageRuntime.ts @@ -1074,6 +1074,16 @@ export class MainThreadLanguageRuntime return Promise.resolve(this._runtimeSessionService.foregroundSession?.sessionId); } + $getNotebookSession(notebookUri: URI): Promise { + // Revive the URI from the serialized form + const uri = URI.revive(notebookUri); + + // Get the session for the notebook URI + const session = this._runtimeSessionService.getNotebookSessionForNotebookUri(uri); + + return Promise.resolve(session?.sessionId); + } + // Called by the extension host to select a previously registered language runtime $selectLanguageRuntime(runtimeId: string): Promise { return this._runtimeSessionService.selectRuntime( diff --git a/src/vs/workbench/api/common/positron/extHost.positron.api.impl.ts b/src/vs/workbench/api/common/positron/extHost.positron.api.impl.ts index c7be0e8b059..31425064f37 100644 --- a/src/vs/workbench/api/common/positron/extHost.positron.api.impl.ts +++ b/src/vs/workbench/api/common/positron/extHost.positron.api.impl.ts @@ -84,6 +84,9 @@ export function createPositronApiFactoryAndRegisterActors(accessor: ServicesAcce getForegroundSession(): Thenable { return extHostLanguageRuntime.getForegroundSession(); }, + getNotebookSession(notebookUri: vscode.Uri): Thenable { + return extHostLanguageRuntime.getNotebookSession(notebookUri); + }, selectLanguageRuntime(runtimeId: string): Thenable { return extHostLanguageRuntime.selectLanguageRuntime(runtimeId); }, diff --git a/src/vs/workbench/api/common/positron/extHost.positron.protocol.ts b/src/vs/workbench/api/common/positron/extHost.positron.protocol.ts index 34e136b4363..099747dcd72 100644 --- a/src/vs/workbench/api/common/positron/extHost.positron.protocol.ts +++ b/src/vs/workbench/api/common/positron/extHost.positron.protocol.ts @@ -38,6 +38,7 @@ export interface MainThreadLanguageRuntimeShape extends IDisposable { $executeCode(languageId: string, code: string, focus: boolean, allowIncomplete?: boolean): Promise; $getPreferredRuntime(languageId: string): Promise; $getForegroundSession(): Promise; + $getNotebookSession(notebookUri: URI): Promise; $restartSession(handle: number): Promise; $emitLanguageRuntimeMessage(handle: number, message: ILanguageRuntimeMessage): void; $emitLanguageRuntimeState(handle: number, clock: number, state: RuntimeState): void; diff --git a/src/vs/workbench/api/common/positron/extHostLanguageRuntime.ts b/src/vs/workbench/api/common/positron/extHostLanguageRuntime.ts index 66b7aefbb4b..8b9728e4f16 100644 --- a/src/vs/workbench/api/common/positron/extHostLanguageRuntime.ts +++ b/src/vs/workbench/api/common/positron/extHostLanguageRuntime.ts @@ -545,6 +545,18 @@ export class ExtHostLanguageRuntime implements extHostProtocol.ExtHostLanguageRu return session; } + public async getNotebookSession(notebookUri: URI): Promise { + const sessionId = await this._proxy.$getNotebookSession(notebookUri); + if (!sessionId) { + return; + } + const session = this._runtimeSessions.find(session => session.metadata.sessionId === sessionId); + if (!session) { + throw new Error(`Session ID '${sessionId}' exists for notebook '${notebookUri.toString()}', but is not known to the extension host.`); + } + return session; + } + /** * Registers a new language runtime manager with the extension host. * diff --git a/src/vs/workbench/services/runtimeSession/common/runtimeSession.ts b/src/vs/workbench/services/runtimeSession/common/runtimeSession.ts index acc4775b12c..91de4d94fa3 100644 --- a/src/vs/workbench/services/runtimeSession/common/runtimeSession.ts +++ b/src/vs/workbench/services/runtimeSession/common/runtimeSession.ts @@ -18,6 +18,7 @@ import { IModalDialogPromptInstance, IPositronModalDialogsService } from 'vs/wor import { IUiClientMessageInput, IUiClientMessageOutput, UiClientInstance } from 'vs/workbench/services/languageRuntime/common/languageRuntimeUiClient'; import { UiFrontendEvent } from 'vs/workbench/services/languageRuntime/common/positronUiComm'; import { ILanguageService } from 'vs/editor/common/languages/language'; +import { ResourceMap } from 'vs/base/common/map'; /** * Utility class for tracking state changes in a language runtime session. @@ -60,6 +61,10 @@ export class RuntimeSessionService extends Disposable implements IRuntimeSession // (metadata.languageId) of the runtime owning the session. private readonly _startingConsolesByLanguageId = new Map(); + // A map of the starting notebooks. This is keyed by the notebook URI + // owning the session. + private readonly _startingNotebooksByNotebookUri = new ResourceMap(); + // A map of runtimes currently starting to promises that resolve when the runtime // is ready to use. This is keyed by the runtimeId (metadata.runtimeId) of the runtime. private readonly _startingRuntimesByRuntimeId = new Map>(); @@ -69,6 +74,10 @@ export class RuntimeSessionService extends Disposable implements IRuntimeSession // languageId (metadata.languageId) of the session. private readonly _consoleSessionsByLanguageId = new Map(); + // A map of the currently active notebook sessions. This is keyed by the notebook URI + // owning the session. + private readonly _notebookSessionsByNotebookUri = new ResourceMap(); + // The event emitter for the onWillStartRuntime event. private readonly _onWillStartRuntimeEmitter = this._register(new Emitter); @@ -217,6 +226,17 @@ export class RuntimeSessionService extends Disposable implements IRuntimeSession return this._consoleSessionsByLanguageId.get(runtimeId); } + /** + * Gets the notebook session for a notebook URI, if one exists. + * + * @param notebookUri The notebook URI of the session to retrieve. + * @returns The notebook session with the given notebook URI, or undefined if + * no notebook session with the given notebook URI exists. + */ + getNotebookSessionForNotebookUri(notebookUri: URI): ILanguageRuntimeSession | undefined { + return this._notebookSessionsByNotebookUri.get(notebookUri); + } + /** * Selects and starts a new runtime session, after shutting down any currently active * sessions for the language. @@ -448,7 +468,7 @@ export class RuntimeSessionService extends Disposable implements IRuntimeSession /** * Restarts a runtime session. * - * @param sessionId The session ID of the runtime to interrupt. + * @param sessionId The session ID of the runtime to restart. * @param source The source of the request to restart the runtime. */ async restartSession(sessionId: string, source: string): Promise { @@ -740,6 +760,14 @@ export class RuntimeSessionService extends Disposable implements IRuntimeSession if (session.metadata.sessionMode === LanguageRuntimeSessionMode.Console) { this._startingConsolesByLanguageId.delete(session.runtimeMetadata.languageId); this._consoleSessionsByLanguageId.set(session.runtimeMetadata.languageId, session); + } else if (session.metadata.sessionMode === LanguageRuntimeSessionMode.Notebook) { + if (session.metadata.notebookUri) { + this._startingNotebooksByNotebookUri.delete(session.metadata.notebookUri); + this._notebookSessionsByNotebookUri.set(session.metadata.notebookUri, session); + } else { + this._logService.error(`Notebook session ${formatLanguageRuntimeSession(session)} ` + + `does not have a notebook URI.`); + } } // Fire the onDidStartRuntime event. @@ -811,6 +839,13 @@ export class RuntimeSessionService extends Disposable implements IRuntimeSession this._startingConsolesByLanguageId.delete(session.runtimeMetadata.languageId); if (session.metadata.sessionMode === LanguageRuntimeSessionMode.Console) { this._consoleSessionsByLanguageId.delete(session.runtimeMetadata.languageId); + } else if (session.metadata.sessionMode === LanguageRuntimeSessionMode.Notebook) { + if (session.metadata.notebookUri) { + this._notebookSessionsByNotebookUri.delete(session.metadata.notebookUri); + } else { + this._logService.error(`Notebook session ${formatLanguageRuntimeSession(session)} ` + + `does not have a notebook URI.`); + } } break; } diff --git a/src/vs/workbench/services/runtimeSession/common/runtimeSessionService.ts b/src/vs/workbench/services/runtimeSession/common/runtimeSessionService.ts index b8e7085a331..606d5b8d989 100644 --- a/src/vs/workbench/services/runtimeSession/common/runtimeSessionService.ts +++ b/src/vs/workbench/services/runtimeSession/common/runtimeSessionService.ts @@ -277,6 +277,12 @@ export interface IRuntimeSessionService { */ getConsoleSessionForLanguage(languageId: string): ILanguageRuntimeSession | undefined; + /** + * Gets a specific notebook session by notebook URI. Currently, only one + * notebook session can exist per notebook URI. + */ + getNotebookSessionForNotebookUri(notebookUri: URI): ILanguageRuntimeSession | undefined; + /** * Checks for a starting or running console for the given language ID. * diff --git a/src/vs/workbench/services/runtimeStartup/common/runtimeStartup.ts b/src/vs/workbench/services/runtimeStartup/common/runtimeStartup.ts index 4f72aceba15..d4ad895033f 100644 --- a/src/vs/workbench/services/runtimeStartup/common/runtimeStartup.ts +++ b/src/vs/workbench/services/runtimeStartup/common/runtimeStartup.ts @@ -20,6 +20,7 @@ import { ExtensionIdentifier } from 'vs/platform/extensions/common/extensions'; import { ILifecycleService, ShutdownReason, StartupKind } from 'vs/workbench/services/lifecycle/common/lifecycle'; import { INotificationService } from 'vs/platform/notification/common/notification'; import { IWorkspaceTrustManagementService } from 'vs/platform/workspace/common/workspaceTrust'; +import { URI } from 'vs/base/common/uri'; interface ILanguageRuntimeProviderMetadata { languageId: string; @@ -634,7 +635,17 @@ export class RuntimeStartupService extends Disposable implements IRuntimeStartup StorageScope.WORKSPACE); if (storedSessions) { try { - const sessions = JSON.parse(storedSessions) as SerializedSessionMetadata[]; + const serializedSessions = JSON.parse(storedSessions) as SerializedSessionMetadata[]; + + // Revive the URIs in the session metadata. + const sessions: SerializedSessionMetadata[] = serializedSessions.map(session => ({ + ...session, + metadata: { + ...session.metadata, + notebookUri: URI.revive(session.metadata.notebookUri), + }, + })); + if (sessions.length > 0) { // If this workspace has sessions, attempt to reconnect to // them.