Use restored runtime sessions in notebooks and improve notebook runtime startup/shutdown stability

extensions/positron-notebook-controllers/src/extension.ts

@@ -4,13 +4,14 @@
 
 import * as positron from 'positron';
 import * as vscode from 'vscode';
-import { initializeLogging } from './logging';
 import { NotebookControllerManager } from './notebookControllerManager';
+import { NotebookSessionService } from './notebookSessionService';
 
-export async function activate(context: vscode.ExtensionContext): Promise<void> {
-	initializeLogging();
+export const log = vscode.window.createOutputChannel('Positron Notebook Controllers', { log: true });
 
-	const manager = new NotebookControllerManager();
+export async function activate(context: vscode.ExtensionContext): Promise<void> {
+	const notebookSessionService = new NotebookSessionService();
+	const manager = new NotebookControllerManager(notebookSessionService);
 	context.subscriptions.push(manager);
 
 	// Register notebook controllers for newly registered runtimes.

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<T> {
+	constructor(readonly uri: URI, readonly value: T) { }
+}
+
+function isEntries<T>(arg: ResourceMap<T> | ResourceMapKeyFn | readonly (readonly [URI, T])[] | undefined): arg is readonly (readonly [URI, T])[] {
+	return Array.isArray(arg);
+}
+
+export class ResourceMap<T> implements Map<URI, T> {
+
+	private static readonly defaultToKey = (resource: URI) => resource.toString();
+
+	readonly [Symbol.toStringTag] = 'ResourceMap';
+
+	private readonly map: Map<string, ResourceMapEntry<T>>;
+	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<T>, 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<T> | 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<URI, T>) => void, thisArg?: any): void {
+		if (typeof thisArg !== 'undefined') {
+			clb = clb.bind(thisArg);
+		}
+		for (const [_, entry] of this.map) {
+			clb(entry.value, entry.uri, <any>this);
+		}
+	}
+
+	*values(): IterableIterator<T> {
+		for (const entry of this.map.values()) {
+			yield entry.value;
+		}
+	}
+
+	*keys(): IterableIterator<URI> {
+		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];
+		}
+	}
+}

extensions/positron-notebook-controllers/src/notebookController.ts

@@ -3,8 +3,7 @@
  *--------------------------------------------------------------------------------------------*/
 import * as vscode from 'vscode';
 import * as positron from 'positron';
-import { trace } from './logging';
-import { DeferredPromise, delay } from './util';
+import { INotebookSessionService } from './notebookSessionService';
 
 /**
  * Wraps a vscode.NotebookController for a specific language, and manages a notebook runtime session
@@ -17,15 +16,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<vscode.NotebookDocument, DeferredPromise<positron.LanguageRuntimeSession>> = new Map();
-
-	/** A map of the currently active sessions, keyed by notebook. */
-	private readonly _activeSessionsByNotebook: Map<vscode.NotebookDocument, positron.LanguageRuntimeSession> = new Map();
-
 	/** A map of the current execution order, keyed by session ID. */
 	private readonly _executionOrderBySessionId: Map<string, number> = new Map();
 
@@ -34,8 +24,12 @@ export class NotebookController implements vscode.Disposable {
 
 	/**
 	 * @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: INotebookSessionService,
+	) {
 		// Create a VSCode notebook controller for this language.
 		this.controller = vscode.notebooks.createNotebookController(
 			`positron-${languageId}`,
@@ -57,24 +51,23 @@ 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);
+			if (notebookSessionService.hasStartingOrRunningNotebookSession(notebook.uri)) {
+				await notebookSessionService.shutdownRuntimeSession(notebook.uri);
+			}
 		}));
 
 		this._disposables.push(this.controller.onDidChangeSelectedNotebooks(async (e) => {
-			await this.shutdownRuntimeSession(e.notebook);
-
 			// Has this controller been selected for a notebook?
 			if (e.selected) {
 				// Note that this is also reached when a notebook is opened, if this controller was
 				// already selected.
 
 				await Promise.all([
 					updateNotebookLanguage(e.notebook, this.languageId),
-					this.startNewRuntimeSession(e.notebook),
+					this.startRuntimeSession(e.notebook),
 				]);
+			} else {
+				await notebookSessionService.shutdownRuntimeSession(e.notebook.uri);
 			}
 		}));
 	}
@@ -85,45 +78,16 @@ export class NotebookController implements vscode.Disposable {
 	}
 
 	/**
-	 * 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<void> {
-		// 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<positron.LanguageRuntimeSession> {
+	public async startRuntimeSession(notebook: vscode.NotebookDocument): Promise<positron.LanguageRuntimeSession> {
 		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}",
@@ -134,72 +98,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<positron.LanguageRuntimeSession> {
-		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<positron.LanguageRuntimeSession>();
-		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.
 	 *
@@ -211,17 +116,17 @@ 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);
+		const startingSessionPromise = this.notebookSessionService.getStartingNotebookSessionPromise(notebook.uri);
 		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);
+			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 +141,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<void> {
-		// 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);