Remove IAsyncDisposable from KernelProcess (#15661)

* Remove IAsyncDisposable from KernelProcess

* Remove more instances of IAsyncDisposable

src/kernels/jupyter/connection/jupyterConnection.ts

@@ -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';
@@ -127,7 +126,7 @@ export class JupyterConnection {
         } finally {
             connection.dispose();
             if (sessionManager) {
-                sessionManager.dispose().catch(noop);
+                sessionManager.dispose();
             }
         }
     }

src/kernels/jupyter/finder/remoteKernelFinder.ts

@@ -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.
@@ -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([

src/kernels/jupyter/finder/remoteKernelFinder.unit.test.ts

@@ -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 {
@@ -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();

src/kernels/jupyter/launcher/jupyterServerHelper.node.ts

@@ -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,
@@ -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);
 
@@ -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) {
@@ -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)
@@ -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);
@@ -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;

src/kernels/jupyter/session/jupyterKernelSessionFactory.ts

@@ -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';
@@ -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 {
@@ -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));
 
@@ -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({

src/kernels/jupyter/session/jupyterLabHelper.ts

@@ -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
@@ -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({
@@ -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[]> {

src/kernels/jupyter/types.ts

@@ -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 {
@@ -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>;

src/kernels/raw/finder/pythonKernelInterruptDaemon.node.ts

@@ -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';
@@ -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>;
 };
@@ -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

src/kernels/raw/launcher/kernelProcess.node.ts

@@ -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;
@@ -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: '' });
@@ -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) {

src/kernels/raw/launcher/kernelProcess.node.unit.test.ts

@@ -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);
@@ -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
                     });