Skip to content

Commit

Permalink
resolve restart when runtime is ready; support concurrently starting …
Browse files Browse the repository at this point in the history 
…while already restarting from ready state
  • Loading branch information
@seeM
seeM committed Dec 11, 2024
1 parent 72b8501 commit 7c081ad
Show file tree
Hide file tree
Showing 2 changed files with 152 additions and 69 deletions.
149 changes: 91 additions & 58 deletions src/vs/workbench/services/runtimeSession/common/runtimeSession.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -396,8 +396,8 @@ export class RuntimeSessionService extends Disposable implements IRuntimeSession
// are, return the promise that resolves when the session is ready to
// use. This makes it possible for multiple requests to start the same
// session to be coalesced.
const startingRuntimePromise = this._startingSessionsBySessionMapKey.get(
getSessionMapKey(sessionMode, runtimeId, notebookUri));
const sessionMapKey = getSessionMapKey(sessionMode, runtimeId, notebookUri);
const startingRuntimePromise = this._startingSessionsBySessionMapKey.get(sessionMapKey);
if (startingRuntimePromise && !startingRuntimePromise.isSettled) {
return startingRuntimePromise.p;
}
Expand Down Expand Up @@ -445,8 +445,9 @@ export class RuntimeSessionService extends Disposable implements IRuntimeSession
// are, return the promise that resolves when the session is ready to
// use. This makes it possible for multiple requests to start the same
// session to be coalesced.
const startingRuntimePromise = this._startingSessionsBySessionMapKey.get(
getSessionMapKey(sessionMetadata.sessionMode, runtimeMetadata.runtimeId, sessionMetadata.notebookUri));
const sessionMapKey = getSessionMapKey(
sessionMetadata.sessionMode, runtimeMetadata.runtimeId, sessionMetadata.notebookUri);
const startingRuntimePromise = this._startingSessionsBySessionMapKey.get(sessionMapKey);
if (startingRuntimePromise && !startingRuntimePromise.isSettled) {
return startingRuntimePromise.p.then(() => { });
}
Expand All @@ -468,8 +469,6 @@ export class RuntimeSessionService extends Disposable implements IRuntimeSession

// Create a promise that resolves when the runtime is ready to use.
const startPromise = new DeferredPromise<string>();
const sessionMapKey = getSessionMapKey(
sessionMetadata.sessionMode, runtimeMetadata.runtimeId, sessionMetadata.notebookUri);
this._startingSessionsBySessionMapKey.set(sessionMapKey, startPromise);

// It's possible that startPromise is never awaited, so we log any errors here
Expand Down Expand Up @@ -575,7 +574,92 @@ export class RuntimeSessionService extends Disposable implements IRuntimeSession
this._logService.info(
`Restarting session '` +
`${formatLanguageRuntimeSession(session)}' (Source: ${source})`);
await this.doRestartRuntime(session);

const state = session.getRuntimeState();
if (state === RuntimeState.Busy ||
state === RuntimeState.Idle ||
state === RuntimeState.Ready) {
// The runtime looks like it could handle a restart request, so send
// one over.
return this.doRestartRuntime(session);
} else if (state === RuntimeState.Uninitialized ||
state === RuntimeState.Exited) {
// The runtime has never been started, or is no longer running. Just
// tell it to start.
await this.startNewRuntimeSession(session.runtimeMetadata.runtimeId,
session.metadata.sessionName,
session.metadata.sessionMode,
session.metadata.notebookUri,
`'Restart Interpreter' command invoked`);
return;
} else if (state === RuntimeState.Starting ||
state === RuntimeState.Restarting) {
// The runtime is already starting or restarting. We could show an
// error, but this is probably just the result of a user mashing the
// restart when we already have one in flight.
return;
} else {
// The runtime is not in a state where it can be restarted.
throw new Error(`The ${session.runtimeMetadata.languageName} session is '${state}' ` +
`and cannot be restarted.`);
}
}

/**
* Internal method to restart a runtime session.
*
* @param session The runtime to restart.
*/
private async doRestartRuntime(session: ILanguageRuntimeSession): Promise<void> {
// If there is already a runtime starting for the session, return its promise.
const sessionMapKey = getSessionMapKey(
session.metadata.sessionMode, session.runtimeMetadata.runtimeId, session.metadata.notebookUri);
const startingRuntimePromise = this._startingSessionsBySessionMapKey.get(sessionMapKey);
if (startingRuntimePromise && !startingRuntimePromise.isSettled) {
return startingRuntimePromise.p.then(() => { });
}

const activeSession = this._activeSessionsBySessionId.get(session.sessionId);
if (!activeSession) {
throw new Error(`No active session '${session.sessionId}'`);
}

// Create a promise that resolves when the runtime is ready to use.
const startPromise = new DeferredPromise<string>();
this._startingSessionsBySessionMapKey.set(sessionMapKey, startPromise);

// Mark the session as starting.
this.setStartingSessionMaps(
session.metadata.sessionMode, session.runtimeMetadata, session.metadata.notebookUri);

// Mark the session as ready when it reaches the ready state,
// or after a timeout.
const disposables = new DisposableStore();
activeSession.register(disposables);
disposables.add(session.onDidChangeRuntimeState((state) => {
if (state === RuntimeState.Ready) {
disposables.dispose();
this.clearStartingSessionMaps(
session.metadata.sessionMode, session.runtimeMetadata, session.metadata.notebookUri);
startPromise.complete(session.sessionId);
}
}));
disposables.add(disposableTimeout(() => {
disposables.dispose();
startPromise.error(new Error(`Timed out waiting for runtime ` +
`${formatLanguageRuntimeSession(session)} to finish restarting.`));
}, 10_000));

// Ask the runtime to restart.
try {
await session.restart();
} catch (err) {
startPromise.error(err);
this.clearStartingSessionMaps(
session.metadata.sessionMode, session.runtimeMetadata, session.metadata.notebookUri);
}

return startPromise.p.then(() => { });
}

/**
Expand Down Expand Up @@ -1234,57 +1318,6 @@ export class RuntimeSessionService extends Disposable implements IRuntimeSession
}
}

/**
* Restarts a runtime session.
*
* @param session The runtime to restart.
*/
private async doRestartRuntime(session: ILanguageRuntimeSession): Promise<void> {
const state = session.getRuntimeState();
if (state === RuntimeState.Busy ||
state === RuntimeState.Idle ||
state === RuntimeState.Ready) {
// The runtime looks like it could handle a restart request, so send
// one over.

// Mark the session as starting until the restart sequence completes.
this.setStartingSessionMaps(
session.metadata.sessionMode, session.runtimeMetadata, session.metadata.notebookUri);
const disposable = this._register(session.onDidChangeRuntimeState((state) => {
if (state === RuntimeState.Ready) {
disposable.dispose();
this.clearStartingSessionMaps(
session.metadata.sessionMode, session.runtimeMetadata, session.metadata.notebookUri);
}
}));

// Restart the session.
return session.restart();
} else if (state === RuntimeState.Uninitialized ||
state === RuntimeState.Exited) {
// The runtime has never been started, or is no longer running. Just
// tell it to start.
await this.startNewRuntimeSession(session.runtimeMetadata.runtimeId,
session.metadata.sessionName,
session.metadata.sessionMode,
session.metadata.notebookUri,
`'Restart Interpreter' command invoked`);
return;
} else if (state === RuntimeState.Starting ||
state === RuntimeState.Restarting) {
// The runtime is already starting or restarting. We could show an
// error, but this is probably just the result of a user mashing the
// restart when we already have one in flight.
return;
} else {
// The runtime is not in a state where it can be restarted.
return Promise.reject(
new Error(`The ${session.runtimeMetadata.languageName} session is '${state}' ` +
`and cannot be restarted.`)
);
}
}

/**
* Waits for the runtime to report that interrupt processing is complete (by
* returning to the idle state). If the runtime does not return to the idle
Expand Down
72 changes: 61 additions & 11 deletions src/vs/workbench/services/runtimeSession/test/common/runtimeSession.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -681,7 +681,7 @@ suite('Positron - RuntimeSessionService', () => {
const clock = sinon.useFakeTimers();
const promise = assert.rejects(selectRuntime(anotherRuntime), new Error(`Timed out waiting for runtime ` +
`${formatLanguageRuntimeSession(session)} to finish exiting.`));
await clock.tickAsync(10_000);
await clock.tickAsync(5_000);
await promise;
});

Expand Down Expand Up @@ -728,12 +728,57 @@ suite('Positron - RuntimeSessionService', () => {

await restartSession(session.sessionId);

assert.strictEqual(session.getRuntimeState(), RuntimeState.Restarting);
assertSingleSessionIsRestarting(session);
assert.strictEqual(session.getRuntimeState(), RuntimeState.Ready);
assertSingleSessionIsReady(session);
});

test(`restart ${mode} in '${state}' state encounters session.restart() error`, async () => {
// Start the session and wait for it to be ready.
const session = await start();
await waitForRuntimeState(session, RuntimeState.Ready);

// Set the state to the desired state.
if (session.getRuntimeState() !== state) {
session.setRuntimeState(state);
}

// Stub session.restart() to reject.
const restartStub = sinon.stub(session, 'restart').rejects(new Error('Session failed to restart'));

// Restart the session. It should error.
await assert.rejects(restartSession(session.sessionId));

// The session's state should not have changed.
assert.strictEqual(session.getRuntimeState(), state);

// If we restart now, without session.restart() rejecting, it should work.
restartStub.restore();
await restartSession(session.sessionId);

assert.strictEqual(session.getRuntimeState(), RuntimeState.Ready);
assertSingleSessionIsReady(session);
});

test(`restart ${mode} in '${state}' state and session never reaches ready state`, async () => {
// Start the session and wait for it to be ready.
const session = await start();
await waitForRuntimeState(session, RuntimeState.Ready);

// Set the state to the desired state.
if (session.getRuntimeState() !== state) {
session.setRuntimeState(state);
}

// Stub onDidChangeRuntimeState to never fire, causing the restart to time out.
sinon.stub(session, 'onDidChangeRuntimeState').returns({ dispose: () => { } });

// Use a fake timer to avoid actually having to wait for the timeout.
const clock = sinon.useFakeTimers();
const promise = assert.rejects(restartSession(session.sessionId), new Error(`Timed out waiting for runtime ` +
`${formatLanguageRuntimeSession(session)} to finish restarting.`));
await clock.tickAsync(10_000);
await promise;
});
}

for (const state of [RuntimeState.Uninitialized, RuntimeState.Exited]) {
Expand All @@ -746,7 +791,7 @@ suite('Positron - RuntimeSessionService', () => {

await restartSession(session.sessionId);

// The existing sessino should remain exited.
// The existing session should remain exited.
assert.strictEqual(session.getRuntimeState(), RuntimeState.Exited);

// A new session should be starting.
Expand Down Expand Up @@ -820,36 +865,41 @@ suite('Positron - RuntimeSessionService', () => {
restartSession(session.sessionId),
]);

assertSingleSessionIsRestarting(session);
assertSingleSessionIsReady(session);

sinon.assert.calledOnce(target);
});

test(`restart ${mode} successively`, async () => {
const session = await start();
await waitForRuntimeState(session, RuntimeState.Ready);

const target = sinon.spy(session, 'restart');

await waitForRuntimeState(session, RuntimeState.Ready);
await restartSession(session.sessionId);
await waitForRuntimeState(session, RuntimeState.Ready);
await restartSession(session.sessionId);
await waitForRuntimeState(session, RuntimeState.Ready);
await restartSession(session.sessionId);

assertSingleSessionIsRestarting(session);
assertSingleSessionIsReady(session);

sinon.assert.calledThrice(target);
});

test(`restart ${mode} while ready -> start`, async () => {
test(`restart ${mode} while ready, then start successively`, async () => {
const session = await start();
await waitForRuntimeState(session, RuntimeState.Ready);

await restartSession(session.sessionId);
const newSession = await start();

assertSingleSessionIsReady(newSession);
});

test(`restart ${mode} while ready, then start concurrently`, async () => {
const session = await start();
await waitForRuntimeState(session, RuntimeState.Ready);

const newSession = await start();
const [, newSession] = await Promise.all([restartSession(session.sessionId), start()]);

assertSingleSessionIsReady(newSession);
});
Expand Down

0 comments on commit 7c081ad

Please sign in to comment.