Skip to content

Commit

Permalink
extract notebook session service; fix concurrent start/shutdown calls…
Browse files Browse the repository at this point in the history
…; restore sessions
  • Loading branch information
seeM committed Apr 2, 2024
1 parent af65545 commit 53b55c9
Show file tree
Hide file tree
Showing 13 changed files with 501 additions and 121 deletions.
4 changes: 3 additions & 1 deletion extensions/positron-notebook-controllers/src/extension.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,11 +5,13 @@
import * as positron from 'positron';
import * as vscode from 'vscode';
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<void> {
const manager = new NotebookControllerManager();
const notebookSessionService = new NotebookSessionService();
const manager = new NotebookControllerManager(notebookSessionService);
context.subscriptions.push(manager);

// Register notebook controllers for newly registered runtimes.
Expand Down
123 changes: 123 additions & 0 deletions extensions/positron-notebook-controllers/src/map.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,123 @@
/*---------------------------------------------------------------------------------------------
* Copyright (c) Microsoft Corporation. All rights reserved.
* Licensed under the MIT License. See License.txt in the project root for license information.
*--------------------------------------------------------------------------------------------*/

import { Uri as URI } from 'vscode';

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];
}
}
}
133 changes: 17 additions & 116 deletions extensions/positron-notebook-controllers/src/notebookController.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,7 @@
*--------------------------------------------------------------------------------------------*/
import * as vscode from 'vscode';
import * as positron from 'positron';
import { log } from './extension';
import { DeferredPromise, delay } from './util';
import { INotebookSessionService } from './notebookSessionService';

/**
* Wraps a vscode.NotebookController for a specific language, and manages a notebook runtime session
Expand All @@ -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();

Expand All @@ -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}`,
Expand All @@ -57,15 +51,12 @@ 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
Expand All @@ -75,6 +66,8 @@ export class NotebookController implements vscode.Disposable {
updateNotebookLanguage(e.notebook, this.languageId),
this.startNewRuntimeSession(e.notebook),
]);
} else {
await notebookSessionService.shutdownRuntimeSession(e.notebook.uri);
}
}));
}
Expand All @@ -84,35 +77,6 @@ export class NotebookController implements vscode.Disposable {
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<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) {
log.warn(`[${this.languageId}] 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);
log.info(`Shutdown runtime ${session.runtimeMetadata.runtimeName} for notebook ${notebook.uri.path}`);
}

/**
* Start a new runtime session for a notebook.
*
Expand All @@ -121,9 +85,9 @@ export class NotebookController implements vscode.Disposable {
*/
public async startNewRuntimeSession(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}",
Expand All @@ -141,65 +105,6 @@ export class NotebookController implements vscode.Disposable {
}
}

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) {
log.error(`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);
log.info(`Session ${session.metadata.sessionId} is ready`);

return session;
}

/**
* Notebook controller execute handler.
*
Expand All @@ -211,12 +116,12 @@ 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.
Expand All @@ -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) {
log.error(`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);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
import * as vscode from 'vscode';
import { log } from './extension';
import { NotebookController } from './notebookController';
import { INotebookSessionService } from './notebookSessionService';

/**
* Manages notebook controllers.
Expand All @@ -13,6 +14,12 @@ export class NotebookControllerManager implements vscode.Disposable {
/** Notebook controllers keyed by languageId. */
public readonly controllers = new Map<string, NotebookController>();

/**
*
* @param notebookSessionService The notebook session service.
*/
constructor(private readonly notebookSessionService: INotebookSessionService) { }

/**
* Create a notebook controller for a language.
*
Expand All @@ -22,7 +29,7 @@ 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);
log.info(`Registered notebook controller for language: ${languageId}`);
}
Expand Down Expand Up @@ -79,7 +86,7 @@ export class NotebookControllerManager implements vscode.Disposable {
? vscode.NotebookControllerAffinity.Preferred
: vscode.NotebookControllerAffinity.Default;
controller.controller.updateNotebookAffinity(notebook, affinity);
log.info(`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}`);
}
}

Expand Down
Loading

0 comments on commit 53b55c9

Please sign in to comment.