Use new port forwarding API (#15243)

* Use new port forwarding API

* Fix tests

src/kernels/common/usedPorts.ts

@@ -1,5 +1,45 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
+import {
+    PortAttributes,
+    type CancellationToken,
+    type PortAttributesProvider,
+    PortAutoForwardAction,
+    type PortAttributesSelector,
+    workspace
+} from 'vscode';
+import { DisposableStore } from '../../platform/common/utils/lifecycle';
+import { traceError } from '../../platform/logging';
+
 // Keeps track of all ports used by Kernels and other processes spawned by Kernels and related code
 export const UsedPorts = new Set<number>();
+
+export function ignorePortForwarding(...ports: number[]) {
+    const disposableStore = new DisposableStore();
+    try {
+        const provider = new (class implements PortAttributesProvider {
+            async providePortAttributes(
+                attributes: { port: number; pid?: number; commandLine?: string },
+                _token: CancellationToken
+            ) {
+                if (ports.includes(attributes.port)) {
+                    return new PortAttributes(PortAutoForwardAction.Ignore);
+                }
+                return undefined;
+            }
+        })();
+
+        for (const port of ports) {
+            const portSelector: PortAttributesSelector = {
+                portRange: port
+            };
+            disposableStore.add(workspace.registerPortAttributesProvider(portSelector, provider));
+        }
+    } catch (ex) {
+        // In case proposed API changes.
+        traceError('Failure in registering port attributes', ex);
+    }
+
+    return disposableStore;
+}

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

@@ -28,7 +28,7 @@ import { getFilePath } from '../../../platform/common/platform/fs-paths';
 import { JupyterNotebookNotInstalled } from '../../../platform/errors/jupyterNotebookNotInstalled';
 import { JupyterCannotBeLaunchedWithRootError } from '../../../platform/errors/jupyterCannotBeLaunchedWithRootError';
 import { noop } from '../../../platform/common/utils/misc';
-import { UsedPorts } from '../../common/usedPorts';
+import { UsedPorts, ignorePortForwarding } from '../../common/usedPorts';
 
 /**
  * Responsible for starting a notebook.
@@ -132,8 +132,12 @@ export class JupyterServerStarter implements IJupyterServerStarter {
             try {
                 const port = parseInt(new URL(connection.baseUrl).port || '0', 10);
                 if (port && !isNaN(port)) {
+                    const disposable = ignorePortForwarding(port);
                     if (launchResult.proc) {
-                        launchResult.proc.on('exit', () => UsedPorts.delete(port));
+                        launchResult.proc.on('exit', () => {
+                            UsedPorts.delete(port);
+                            disposable.dispose();
+                        });
                     }
                     UsedPorts.add(port);
                 }

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

@@ -38,8 +38,9 @@ import { StopWatch } from '../../../platform/common/utils/stopWatch';
 import { getResourceType } from '../../../platform/common/utils';
 import { format, splitLines } from '../../../platform/common/helpers';
 import { IPythonExecutionFactory } from '../../../platform/interpreter/types.node';
-import { UsedPorts } from '../../common/usedPorts';
+import { UsedPorts, ignorePortForwarding } from '../../common/usedPorts';
 import { isPythonKernelConnection } from '../../helpers';
+import { once } from '../../../platform/common/utils/events';
 
 const PortFormatString = `kernelLauncherPortStart_{0}.tmp`;
 // Launches and returns a kernel process given a resource or python interpreter.
@@ -206,7 +207,13 @@ export class KernelLauncher implements IKernelLauncher {
                 ? window.createOutputChannel(DataScience.kernelConsoleOutputChannel(baseName), 'log')
                 : undefined;
         outputChannel?.clear();
-
+        const portAttributeProvider = ignorePortForwarding(
+            connection.control_port,
+            connection.hb_port,
+            connection.iopub_port,
+            connection.shell_port,
+            connection.stdin_port
+        );
         // Create the process
         const kernelProcess = new KernelProcess(
             this.processExecutionFactory,
@@ -223,8 +230,8 @@ export class KernelLauncher implements IKernelLauncher {
             this.pythonKernelInterruptDaemon,
             this.platformService
         );
-
-        kernelProcess.exited(() => outputChannel?.dispose(), this, this.disposables);
+        once(kernelProcess.onDidDispose)(() => portAttributeProvider.dispose(), this, this.disposables);
+        once(kernelProcess.exited)(() => outputChannel?.dispose(), this, this.disposables);
         try {
             await raceCancellationError(cancelToken, kernelProcess.launch(workingDirectory, timeout, cancelToken));
         } catch (ex) {
@@ -233,7 +240,7 @@ export class KernelLauncher implements IKernelLauncher {
             throw ex;
         }
 
-        const disposable = kernelProcess.exited(
+        const disposable = once(kernelProcess.exited)(
             ({ exitCode, reason }) => {
                 sendKernelTelemetryEvent(
                     resource,

src/kernels/raw/launcher/kernelLauncher.unit.test.ts

@@ -17,11 +17,18 @@ import { dispose } from '../../../platform/common/utils/lifecycle';
 import { anything, instance, mock, when } from 'ts-mockito';
 import { IFileSystemNode } from '../../../platform/common/platform/types.node';
 import { PythonKernelConnectionMetadata } from '../../types';
-import { CancellationTokenSource, Disposable, EventEmitter, PortAutoForwardAction, Uri } from 'vscode';
+import {
+    CancellationTokenSource,
+    Disposable,
+    EventEmitter,
+    PortAutoForwardAction,
+    Uri,
+    type PortAttributesProvider
+} from 'vscode';
 import { KernelProcess } from './kernelProcess.node';
-import { PortAttributesProviders } from '../../portAttributeProvider.node';
 import { IPythonExecutionFactory, IPythonExecutionService } from '../../../platform/interpreter/types.node';
 import { UsedPorts } from '../../common/usedPorts';
+import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../../test/vscode-mock';
 
 suite('kernel Launcher', () => {
     let disposables: IDisposable[] = [];
@@ -67,7 +74,10 @@ suite('kernel Launcher', () => {
             instance(platform)
         );
     });
-    teardown(() => (disposables = dispose(disposables)));
+    teardown(() => {
+        disposables = dispose(disposables);
+        resetVSCodeMocks();
+    });
     async function launchKernel() {
         const kernelSpec = PythonKernelConnectionMetadata.create({
             id: '1',
@@ -103,17 +113,25 @@ suite('kernel Launcher', () => {
         assert.notDeepEqual(Array.from(UsedPorts), oldPorts, 'Ports not updated');
     });
     test('Verify Kernel ports are not forwarded', async () => {
-        await launchKernel();
+        const oldPorts = new Set(UsedPorts);
+        const providers: PortAttributesProvider[] = [];
+        when(mockedVSCodeNamespaces.workspace.registerPortAttributesProvider(anything(), anything())).thenCall(
+            (_, provider) => providers.push(provider)
+        );
 
-        const portAttributeProvider = new PortAttributesProviders(disposables);
+        await launchKernel();
         const cancellation = new CancellationTokenSource();
+        disposables.push(cancellation);
 
         for (const port of UsedPorts) {
-            assert.equal(
-                portAttributeProvider.providePortAttributes({ port }, cancellation.token)?.autoForwardAction,
-                PortAutoForwardAction.Ignore,
-                'Kernel Port should not be forwarded'
+            if (oldPorts.has(port)) {
+                continue;
+            }
+            const results = await Promise.all(
+                providers.map((p) => Promise.resolve(p.providePortAttributes({ port }, cancellation.token)))
             );
+            const portForwardingIgnored = results.some((r) => r?.autoForwardAction === PortAutoForwardAction.Ignore);
+            assert.isTrue(portForwardingIgnored, `Kernel Port ${port} should not be forwarded`);
         }
     });
 });

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

@@ -61,14 +61,15 @@ import { IPythonExecutionFactory } from '../../../platform/interpreter/types.nod
 import { getDisplayPath } from '../../../platform/common/platform/fs-paths';
 import { StopWatch } from '../../../platform/common/utils/stopWatch';
 import { ServiceContainer } from '../../../platform/ioc/container';
+import { ObservableDisposable } from '../../../platform/common/utils/lifecycle';
 
 const kernelOutputWithConnectionFile = 'To connect another client to this kernel, use:';
 const kernelOutputToNotLog =
     'NOTE: When using the `ipython kernel` entry point, Ctrl-C will not work.\n\nTo exit, you will have to explicitly quit this process, by either sending\n"quit" from a client, or using Ctrl-\\ in UNIX-like environments.\n\nTo read more about this, see https://github.com/ipython/ipython/issues/2049\n\n\n';
 
 // Launches and disposes a kernel process given a kernelspec and a resource or python interpreter.
 // Exposes connection information and the process itself.
-export class KernelProcess implements IKernelProcess {
+export class KernelProcess extends ObservableDisposable implements IKernelProcess {
     private _pid?: number;
     private _disposingPromise?: Promise<void>;
     public get pid() {
@@ -100,7 +101,6 @@ export class KernelProcess implements IKernelProcess {
     private _process?: ChildProcess;
     private exitEvent = new EventEmitter<{ exitCode?: number; reason?: string }>();
     private launchedOnce?: boolean;
-    private disposed?: boolean;
     private connectionFile?: Uri;
     private _launchKernelSpec?: IJupyterKernelSpec;
     private interrupter?: Interrupter;
@@ -122,6 +122,7 @@ export class KernelProcess implements IKernelProcess {
         private readonly pythonKernelInterruptDaemon: PythonKernelInterruptDaemon,
         private readonly platform: IPlatformService
     ) {
+        super();
         this._kernelConnectionMetadata = kernelConnectionMetadata;
     }
     public async interrupt(): Promise<void> {
@@ -176,7 +177,7 @@ export class KernelProcess implements IKernelProcess {
         if (proc) {
             proc.on('exit', (exitCode) => {
                 exitCode = exitCode || providedExitCode;
-                if (this.disposed) {
+                if (this.isDisposed) {
                     traceVerbose(`KernelProcess Exited, Exit Code - ${exitCode}`);
                     return;
                 }
@@ -248,7 +249,7 @@ export class KernelProcess implements IKernelProcess {
             this.sendToOutput(output.out);
         });
         exeObs.out.done.catch((error) => {
-            if (this.disposed) {
+            if (this.isDisposed) {
                 traceWarning('Kernel died', error, stderr);
                 return;
             }
@@ -337,14 +338,13 @@ export class KernelProcess implements IKernelProcess {
         }
     }
 
-    public async dispose(): Promise<void> {
+    public override async dispose(): Promise<void> {
         if (this._disposingPromise) {
             return this._disposingPromise;
         }
-        if (this.disposed) {
+        if (this.isDisposed) {
             return;
         }
-        this.disposed = true;
         const pid = this._process?.pid;
         traceInfo(`Dispose Kernel process ${pid}.`);
         this._disposingPromise = (async () => {
@@ -370,6 +370,7 @@ export class KernelProcess implements IKernelProcess {
             });
             traceVerbose(`Disposed Kernel process ${pid}.`);
         })();
+        super.dispose();
     }
 
     private async killChildProcesses(pid?: number) {

src/kernels/serviceRegistry.node.ts

@@ -22,7 +22,6 @@ import { KernelRefreshIndicator } from './kernelRefreshIndicator.node';
 import { KernelStartupCodeProviders } from './kernelStartupCodeProviders.node';
 import { KernelStartupTelemetry } from './kernelStartupTelemetry.node';
 import { KernelStatusProvider } from './kernelStatusProvider';
-import { PortAttributesProviders } from './portAttributeProvider.node';
 import { ContributedLocalKernelSpecFinder } from './raw/finder/contributedLocalKernelSpecFinder.node';
 import { JupyterPaths } from './raw/finder/jupyterPaths.node';
 import { LocalKnownPathKernelSpecFinder } from './raw/finder/localKnownPathKernelSpecFinder.node';
@@ -54,10 +53,6 @@ import { KernelChatStartupCodeProvider } from './chat/kernelStartupCodeProvider'
 export function registerTypes(serviceManager: IServiceManager, isDevMode: boolean) {
     serviceManager.addSingleton<IExtensionSyncActivationService>(IExtensionSyncActivationService, Activation);
     serviceManager.addSingleton<IExtensionSyncActivationService>(IExtensionSyncActivationService, ServerPreload);
-    serviceManager.addSingleton<IExtensionSyncActivationService>(
-        IExtensionSyncActivationService,
-        PortAttributesProviders
-    );
     serviceManager.addSingleton<IRawNotebookSupportedService>(
         IRawNotebookSupportedService,
         RawNotebookSupportedService