Skip to content

Commit

Permalink
Remove IAsyncDisposable from KernelProcess (#15661)
Browse files Browse the repository at this point in the history
* Remove IAsyncDisposable from KernelProcess

* Remove more instances of IAsyncDisposable
  • Loading branch information
DonJayamanne authored May 3, 2024
1 parent 77818ee commit 3656224
Show file tree
Hide file tree
Showing 14 changed files with 106 additions and 75 deletions.
3 changes: 1 addition & 2 deletions src/kernels/jupyter/connection/jupyterConnection.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@
// Licensed under the MIT License.

import { inject, injectable, optional } from 'inversify';
import { noop } from '../../../platform/common/utils/misc';
import { RemoteJupyterServerUriProviderError } from '../../errors/remoteJupyterServerUriProviderError';
import { BaseError } from '../../../platform/errors/types';
import { createJupyterConnectionInfo, handleExpiredCertsError, handleSelfCertsError } from '../jupyterUtils';
Expand Down Expand Up @@ -127,7 +126,7 @@ export class JupyterConnection {
} finally {
connection.dispose();
if (sessionManager) {
sessionManager.dispose().catch(noop);
sessionManager.dispose();
}
}
}
Expand Down
3 changes: 2 additions & 1 deletion src/kernels/jupyter/finder/remoteKernelFinder.ts
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ import { IFileSystem } from '../../../platform/common/platform/types';
import { computeServerId, generateIdFromRemoteProvider } from '../jupyterUtils';
import { RemoteKernelSpecCacheFileName } from '../constants';
import { JupyterLabHelper } from '../session/jupyterLabHelper';
import { disposeAsync } from '../../../platform/common/utils';

// Even after shutting down a kernel, the server API still returns the old information.
// Re-query after 2 seconds to ensure we don't get stale information.
Expand Down Expand Up @@ -327,7 +328,7 @@ export class RemoteKernelFinder extends ObservableDisposable implements IRemoteK
const disposables: IAsyncDisposable[] = [];
try {
const sessionManager = JupyterLabHelper.create(connInfo.settings);
disposables.push(sessionManager);
disposables.push({ dispose: () => disposeAsync(sessionManager) });

// Get running and specs at the same time
const [running, specs, sessions, serverId] = await Promise.all([
Expand Down
9 changes: 7 additions & 2 deletions src/kernels/jupyter/finder/remoteKernelFinder.unit.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ import type { Session } from '@jupyterlab/services';
import { assert } from 'chai';
import { anything, instance, mock, verify, when } from 'ts-mockito';
import { getDisplayNameOrNameOfKernelConnection } from '../../helpers';
import { Disposable, Uri } from 'vscode';
import { Disposable, EventEmitter, Uri } from 'vscode';
import { CryptoUtils } from '../../../platform/common/crypto';
import { noop, sleep } from '../../../test/core';
import {
Expand Down Expand Up @@ -124,7 +124,12 @@ suite(`Remote Kernel Finder`, () => {
return Promise.resolve(d.toLowerCase());
});
jupyterSessionManager = mock<JupyterLabHelper>();
when(jupyterSessionManager.dispose()).thenResolve();
const onDidDispose = new EventEmitter<void>();
disposables.push(onDidDispose);
when(jupyterSessionManager.onDidDispose).thenReturn(onDidDispose.event);
when(jupyterSessionManager.dispose()).thenCall(() => {
onDidDispose.fire();
});
sinon.stub(JupyterLabHelper, 'create').callsFake(() => resolvableInstance(jupyterSessionManager));
when(fs.delete(anything())).thenResolve();
when(fs.createDirectory(uriEquals(globalStorageUri))).thenResolve();
Expand Down
38 changes: 23 additions & 15 deletions src/kernels/jupyter/launcher/jupyterServerHelper.node.ts
Original file line number Diff line number Diff line change
Expand Up @@ -24,16 +24,17 @@ import { expandWorkingDir } from '../jupyterUtils';
import { noop } from '../../../platform/common/utils/misc';
import { getRootFolder } from '../../../platform/common/application/workspace.base';
import { computeWorkingDirectory } from '../../../platform/common/application/workspace.node';
import { disposeAsync } from '../../../platform/common/utils';
import { ObservableDisposable } from '../../../platform/common/utils/lifecycle';

/**
* Jupyter server implementation that uses the JupyterExecutionBase class to launch Jupyter.
*/
@injectable()
export class JupyterServerHelper implements IJupyterServerHelper {
export class JupyterServerHelper extends ObservableDisposable implements IJupyterServerHelper {
private usablePythonInterpreter: PythonEnvironment | undefined;
private cache?: Promise<IJupyterConnection>;
private disposed: boolean = false;
private _disposed = false;
private _isDisposing = false;
constructor(
@inject(IInterpreterService) private readonly interpreterService: IInterpreterService,
@inject(IDisposableRegistry) private readonly disposableRegistry: IDisposableRegistry,
Expand All @@ -44,6 +45,7 @@ export class JupyterServerHelper implements IJupyterServerHelper {
@optional()
private readonly jupyterInterpreterService: IJupyterSubCommandExecutionService | undefined
) {
super();
this.disposableRegistry.push(this.interpreterService.onDidChangeInterpreter(() => this.onSettingsChanged()));
this.disposableRegistry.push(this);

Expand All @@ -57,21 +59,27 @@ export class JupyterServerHelper implements IJupyterServerHelper {
this,
this.disposableRegistry
);
asyncRegistry.push(this);
asyncRegistry.push({ dispose: () => disposeAsync(this) });
}

public async dispose(): Promise<void> {
if (!this._disposed) {
this._disposed = true;
this.disposed = true;

// Cleanup on dispose. We are going away permanently
await this.cache?.then((s) => s.dispose()).catch(noop);
public override dispose() {
if (!this._isDisposing) {
this._isDisposing = true;

if (this.cache) {
// Cleanup on dispose. We are going away permanently
this.cache
.then((s) => s.dispose())
.catch(noop)
.finally(() => super.dispose());
} else {
super.dispose();
}
}
}

public async startServer(resource: Resource, cancelToken: CancellationToken): Promise<IJupyterConnection> {
if (this._disposed) {
if (this.isDisposed || this._isDisposing) {
throw new Error('Notebook server is disposed');
}
if (!this.cache) {
Expand Down Expand Up @@ -103,7 +111,7 @@ export class JupyterServerHelper implements IJupyterServerHelper {

public async getUsableJupyterPython(cancelToken?: CancellationToken): Promise<PythonEnvironment | undefined> {
// Only try to compute this once.
if (!this.usablePythonInterpreter && !this.disposed && this.jupyterInterpreterService) {
if (!this.usablePythonInterpreter && !this.isDisposed && !this._isDisposing && this.jupyterInterpreterService) {
this.usablePythonInterpreter = await raceCancellationError(
cancelToken,
this.jupyterInterpreterService!.getSelectedInterpreter(cancelToken)
Expand All @@ -121,7 +129,7 @@ export class JupyterServerHelper implements IJupyterServerHelper {
let tryCount = 1;
const maxTries = Math.max(1, this.configuration.getSettings(undefined).jupyterLaunchRetries);
let lastTryError: Error;
while (tryCount <= maxTries && !this.disposed) {
while (tryCount <= maxTries && !this.isDisposed && !this._isDisposing) {
try {
// Start or connect to the process
connection = await this.startImpl(resource, cancelToken);
Expand All @@ -140,7 +148,7 @@ export class JupyterServerHelper implements IJupyterServerHelper {
tryCount += 1;
} else if (connection) {
// If this is occurring during shutdown, don't worry about it.
if (this.disposed) {
if (this.isDisposed || this._isDisposing) {
throw err;
}
throw err;
Expand Down
9 changes: 5 additions & 4 deletions src/kernels/jupyter/session/jupyterKernelSessionFactory.ts
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

import { CancellationError, CancellationToken, Disposable } from 'vscode';
import { CancellationError, CancellationToken } from 'vscode';
import { Cancellation, raceCancellationError } from '../../../platform/common/cancellation';
import uuid from 'uuid/v4';
import * as urlPath from '../../../platform/vscode-path/resources';
Expand Down Expand Up @@ -45,6 +45,7 @@ import { getNameOfKernelConnection, jvscIdentifier } from '../../helpers';
import { waitForCondition } from '../../../platform/common/utils/async';
import { JupyterLabHelper } from './jupyterLabHelper';
import { JupyterSessionWrapper, getRemoteSessionOptions } from './jupyterSession';
import { disposeAsync } from '../../../platform/common/utils';

@injectable()
export class JupyterKernelSessionFactory implements IKernelSessionFactory {
Expand Down Expand Up @@ -97,8 +98,8 @@ export class JupyterKernelSessionFactory implements IKernelSessionFactory {
await raceCancellationError(options.token, this.validateLocalKernelDependencies(options));

const sessionManager = JupyterLabHelper.create(connection.settings);
this.asyncDisposables.push(sessionManager);
disposablesIfAnyErrors.push(new Disposable(() => sessionManager.dispose().catch(noop)));
this.asyncDisposables.push({ dispose: () => disposeAsync(sessionManager) });
disposablesIfAnyErrors.push(sessionManager);

await raceCancellationError(options.token, this.validateRemoteServer(options, sessionManager));

Expand Down Expand Up @@ -130,7 +131,7 @@ export class JupyterKernelSessionFactory implements IKernelSessionFactory {
);
const disposed = session.disposed;
const onDidDisposeSession = () => {
sessionManager.dispose().catch(noop);
sessionManager.dispose();
disposed.disconnect(onDidDisposeSession);
};
this.asyncDisposables.push({
Expand Down
66 changes: 34 additions & 32 deletions src/kernels/jupyter/session/jupyterLabHelper.ts
Original file line number Diff line number Diff line change
Expand Up @@ -20,21 +20,17 @@ import { JupyterKernelSpec } from '../jupyterKernelSpec';
import { createDeferred, raceTimeout } from '../../../platform/common/utils/async';
import { IJupyterKernel } from '../types';
import { sendTelemetryEvent, Telemetry } from '../../../telemetry';
import { dispose } from '../../../platform/common/utils/lifecycle';
import { ObservableDisposable, dispose } from '../../../platform/common/utils/lifecycle';
import { StopWatch } from '../../../platform/common/utils/stopWatch';
import type { ISpecModel } from '@jupyterlab/services/lib/kernelspec/kernelspec';
import { noop } from '../../../platform/common/utils/misc';

export class JupyterLabHelper {
export class JupyterLabHelper extends ObservableDisposable {
public sessionManager: SessionManager;
public kernelSpecManager: KernelSpecManager;
public kernelManager: KernelManager;
public contentsManager: ContentsManager;
private _jupyterlab?: typeof import('@jupyterlab/services');
private disposed?: boolean;
public get isDisposed() {
return this.disposed === true;
}
private get jupyterlab(): typeof import('@jupyterlab/services') {
if (!this._jupyterlab) {
// eslint-disable-next-line @typescript-eslint/no-require-imports
Expand All @@ -43,6 +39,7 @@ export class JupyterLabHelper {
return this._jupyterlab!;
}
private constructor(private readonly serverSettings: ServerConnection.ISettings) {
super();
this.kernelSpecManager = new this.jupyterlab.KernelSpecManager({ serverSettings: this.serverSettings });
this.kernelManager = new this.jupyterlab.KernelManager({ serverSettings: this.serverSettings });
this.sessionManager = new this.jupyterlab.SessionManager({
Expand All @@ -55,36 +52,41 @@ export class JupyterLabHelper {
return new JupyterLabHelper(serverSettings);
}

public async dispose() {
if (this.disposed) {
private _isDisposing = false;
public override dispose() {
if (this.isDisposed || this._isDisposing) {
return;
}
this.disposed = true;
logger.trace(`Disposing Jupyter Lab Helper`);
try {
if (this.contentsManager) {
logger.trace('SessionManager - dispose contents manager');
this.contentsManager.dispose();
}
if (this.sessionManager && !this.sessionManager.isDisposed) {
logger.trace('ShutdownSessionAndConnection - dispose session manager');
// Make sure it finishes startup.
await raceTimeout(10_000, this.sessionManager.ready);
this._isDisposing = true;
(async () => {
logger.trace(`Disposing Jupyter Lab Helper`);
try {
if (this.contentsManager) {
logger.trace('SessionManager - dispose contents manager');
this.contentsManager.dispose();
}
if (this.sessionManager && !this.sessionManager.isDisposed) {
logger.trace('ShutdownSessionAndConnection - dispose session manager');
// Make sure it finishes startup.
await raceTimeout(10_000, this.sessionManager.ready);

// eslint-disable-next-line @typescript-eslint/no-explicit-any
this.sessionManager.dispose(); // Note, shutting down all will kill all kernels on the same connection. We don't want that.
}
if (!this.kernelManager?.isDisposed) {
this.kernelManager?.dispose();
}
if (!this.kernelSpecManager?.isDisposed) {
this.kernelSpecManager?.dispose();
// eslint-disable-next-line @typescript-eslint/no-explicit-any
this.sessionManager.dispose(); // Note, shutting down all will kill all kernels on the same connection. We don't want that.
}
if (!this.kernelManager?.isDisposed) {
this.kernelManager?.dispose();
}
if (!this.kernelSpecManager?.isDisposed) {
this.kernelSpecManager?.dispose();
}
} catch (e) {
logger.error(`Exception on Jupyter Lab Helper shutdown: `, e);
} finally {
logger.trace('Finished disposing Jupyter Lab Helper');
}
} catch (e) {
logger.error(`Exception on Jupyter Lab Helper shutdown: `, e);
} finally {
logger.trace('Finished disposing Jupyter Lab Helper');
}
})()
.catch(noop)
.finally(() => super.dispose());
}

public async getRunningSessions(): Promise<Session.IModel[]> {
Expand Down
4 changes: 2 additions & 2 deletions src/kernels/jupyter/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ import type WebSocketIsomorphic from 'isomorphic-ws';
import { CancellationToken, Disposable, Event, type NotebookCellData } from 'vscode';
import { SemVer } from 'semver';
import { Uri } from 'vscode';
import { IAsyncDisposable, IDisplayOptions, IDisposable, Resource } from '../../platform/common/types';
import { IDisplayOptions, IDisposable, Resource } from '../../platform/common/types';
import { JupyterInstallError } from '../../platform/errors/jupyterInstallError';
import { PythonEnvironment } from '../../platform/pythonEnvironments/info';
import {
Expand Down Expand Up @@ -42,7 +42,7 @@ export enum JupyterInterpreterDependencyResponse {
}

export const IJupyterServerHelper = Symbol('JupyterServerHelper');
export interface IJupyterServerHelper extends IAsyncDisposable {
export interface IJupyterServerHelper {
isJupyterServerSupported(cancelToken?: CancellationToken): Promise<boolean>;
startServer(resource: Resource, cancelToken?: CancellationToken): Promise<IJupyterConnection>;
getUsableJupyterPython(cancelToken?: CancellationToken): Promise<PythonEnvironment | undefined>;
Expand Down
8 changes: 4 additions & 4 deletions src/kernels/raw/finder/pythonKernelInterruptDaemon.node.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ import { logger } from '../../../platform/logging';
import { ObservableExecutionResult } from '../../../platform/common/process/types.node';
import { EnvironmentType, PythonEnvironment } from '../../../platform/pythonEnvironments/info';
import { inject, injectable } from 'inversify';
import { IAsyncDisposable, IDisposableRegistry, IExtensionContext, Resource } from '../../../platform/common/types';
import { IDisposableRegistry, IExtensionContext, Resource, type IDisposable } from '../../../platform/common/types';
import { createDeferred, Deferred } from '../../../platform/common/utils/async';
import { Disposable, Uri } from 'vscode';
import { EOL } from 'os';
Expand Down Expand Up @@ -57,7 +57,7 @@ type Command =
| { command: 'INITIALIZE_INTERRUPT' }
| { command: 'INTERRUPT'; handle: InterruptHandle }
| { command: 'DISPOSE_INTERRUPT_HANDLE'; handle: InterruptHandle };
export type Interrupter = IAsyncDisposable & {
export type Interrupter = IDisposable & {
handle: InterruptHandle;
interrupt: () => Promise<void>;
};
Expand Down Expand Up @@ -103,8 +103,8 @@ export class PythonKernelInterruptDaemon {
interrupt: async () => {
await this.sendCommand({ command: 'INTERRUPT', handle: interruptHandle }, pythonEnvironment, resource);
},
dispose: async () => {
await this.sendCommand(
dispose: () => {
void this.sendCommand(
{ command: 'DISPOSE_INTERRUPT_HANDLE', handle: interruptHandle },
pythonEnvironment,
resource
Expand Down
8 changes: 4 additions & 4 deletions src/kernels/raw/launcher/kernelProcess.node.ts
Original file line number Diff line number Diff line change
Expand Up @@ -332,9 +332,9 @@ export class KernelProcess extends ObservableDisposable implements IKernelProces
}
}

public override async dispose(): Promise<void> {
public override dispose() {
if (this._disposingPromise) {
return this._disposingPromise;
return;
}
if (this.isDisposed) {
return;
Expand All @@ -347,7 +347,7 @@ export class KernelProcess extends ObservableDisposable implements IKernelProces
this.killChildProcesses(this._process?.pid).catch(noop)
);
try {
this.interrupter?.dispose().catch(noop);
this.interrupter?.dispose();
this._process?.kill(); // NOSONAR
if (!this.exitEventFired) {
this.exitEvent.fire({ stderr: '' });
Expand All @@ -364,7 +364,7 @@ export class KernelProcess extends ObservableDisposable implements IKernelProces
}
logger.debug(`Disposed Kernel process ${pid}.`);
})();
super.dispose();
void this._disposingPromise.finally(() => super.dispose()).catch(noop);
}

private async killChildProcesses(pid?: number) {
Expand Down
4 changes: 2 additions & 2 deletions src/kernels/raw/launcher/kernelProcess.node.unit.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -137,7 +137,7 @@ suite('kernel Process', () => {
});
const interrupter = {
handle: 1,
dispose: () => Promise.resolve(),
dispose: noop,
interrupt: () => Promise.resolve()
};
when(daemon.createInterrupter(anything(), anything())).thenResolve(interrupter);
Expand Down Expand Up @@ -516,7 +516,7 @@ suite('Kernel Process', () => {
when(settings.enablePythonKernelLogging).thenReturn(false);
const interruptDaemon = mock<PythonKernelInterruptDaemon>();
when(interruptDaemon.createInterrupter(anything(), anything())).thenResolve({
dispose: () => Promise.resolve(),
dispose: noop,
interrupt: () => Promise.resolve(),
handle: 1
});
Expand Down
Loading

0 comments on commit 3656224

Please sign in to comment.