Merge remote kernel finders.

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

@@ -23,7 +23,6 @@ import {
 import { IInterpreterService } from '../../../platform/interpreter/contracts';
 import { JupyterSessionManager } from '../session/jupyterSessionManager';
 import { JupyterSessionManagerFactory } from '../session/jupyterSessionManagerFactory';
-import { RemoteKernelFinder } from './remoteKernelFinder';
 import { ILocalKernelFinder } from '../../raw/types';
 import { ActiveKernelIdList, PreferredRemoteKernelIdProvider } from '../preferredRemoteKernelIdProvider';
 import {
@@ -46,11 +45,12 @@ import { KernelRankingHelper } from '../../../notebooks/controllers/kernelRankin
 import { IExtensions } from '../../../platform/common/types';
 import { takeTopRankKernel } from '../../raw/finder/localKernelFinder.unit.test';
 import { createEventHandler, TestEventHandler } from '../../../test/common';
+import { UniversalRemoteKernelFinder } from './universalRemoteKernelFinder';
 
 suite(`Remote Kernel Finder`, () => {
     let disposables: Disposable[] = [];
     let preferredRemoteKernelIdProvider: PreferredRemoteKernelIdProvider;
-    let remoteKernelFinder: RemoteKernelFinder;
+    let remoteKernelFinder: UniversalRemoteKernelFinder;
     let localKernelFinder: ILocalKernelFinder;
     let kernelFinder: KernelFinder;
     let kernelRankHelper: IKernelRankingHelper;
@@ -175,21 +175,22 @@ suite(`Remote Kernel Finder`, () => {
         disposables.push(kernelsChanged);
         kernelRankHelper = new KernelRankingHelper(preferredRemoteKernelIdProvider);
 
-        remoteKernelFinder = new RemoteKernelFinder(
+        remoteKernelFinder = new UniversalRemoteKernelFinder(
+            'currentremote',
+            'Local Kernels',
+            RemoteKernelSpecsCacheKey,
             instance(jupyterSessionManagerFactory),
             instance(interpreterService),
             instance(extensionChecker),
             instance(notebookProvider),
-            instance(serverUriStorage),
-            instance(connectionType),
             instance(memento),
             instance(env),
             instance(cachedRemoteKernelValidator),
             kernelFinder,
-            [],
             instance(kernelProvider),
             instance(extensions),
-            false
+            false,
+            serverEntry
         );
         remoteKernelFinder.activate().then(noop, noop);
     });

src/kernels/jupyter/finder/universalRemoteKernelFinder.ts

@@ -32,13 +32,12 @@ import { computeServerId } from '../jupyterUtils';
 import { PYTHON_LANGUAGE } from '../../../platform/common/constants';
 import { createPromiseFromCancellation } from '../../../platform/common/cancellation';
 import { DisplayOptions } from '../../displayOptions';
-import * as localize from '../../../platform/common/utils/localize';
 import { isArray } from '../../../platform/common/utils/sysTypes';
 import { deserializeKernelConnection } from '../../helpers';
 import { noop } from '../../../platform/common/utils/misc';
 import { IApplicationEnvironment } from '../../../platform/common/application/types';
 import { KernelFinder } from '../../kernelFinder';
-import { RemoteKernelSpecsCacheKey, removeOldCachedItems } from '../../common/commonFinder';
+import { removeOldCachedItems } from '../../common/commonFinder';
 import { ContributedKernelFinderKind, IContributedKernelFinderInfo } from '../../internalTypes';
 import { disposeAllDisposables } from '../../../platform/common/helpers';
 
@@ -53,8 +52,6 @@ export class UniversalRemoteKernelFinder implements IRemoteKernelFinder, IContri
      */
     private readonly kernelIdsToHide = new Set<string>();
     kind = ContributedKernelFinderKind.Remote;
-    id: string;
-    displayName: string;
     private _cacheUpdateCancelTokenSource: CancellationTokenSource | undefined;
     private cache: RemoteKernelConnectionMetadata[] = [];
 
@@ -69,6 +66,9 @@ export class UniversalRemoteKernelFinder implements IRemoteKernelFinder, IContri
     private wasPythonInstalledWhenFetchingKernels = false;
 
     constructor(
+        readonly id: string,
+        readonly displayName: string,
+        readonly cacheKey: string,
         private jupyterSessionManagerFactory: IJupyterSessionManagerFactory,
         private interpreterService: IInterpreterService,
         private extensionChecker: IPythonExtensionChecker,
@@ -82,14 +82,6 @@ export class UniversalRemoteKernelFinder implements IRemoteKernelFinder, IContri
         private isWebExtension: boolean,
         private readonly serverUri: IJupyterServerUriEntry
     ) {
-        // Register with remote-serverId as our ID
-        this.id = `${this.kind}-${serverUri.serverId}`;
-
-        // Create a reasonable display name for this kernel finder
-        this.displayName = localize.DataScience.universalRemoteKernelFinderDisplayName().format(
-            serverUri.displayName || serverUri.uri
-        );
-
         // When we register, add a disposable to clean ourselves up from the main kernel finder list
         // Unlike the Local kernel finder universal remote kernel finders will be added on the fly
         this.disposables.push(kernelFinder.registerKernelFinder(this));
@@ -160,7 +152,7 @@ export class UniversalRemoteKernelFinder implements IRemoteKernelFinder, IContri
         this.wasPythonInstalledWhenFetchingKernels = this.extensionChecker.isPythonExtensionInstalled;
     }
 
-    private async loadCache() {
+    public async loadCache() {
         traceInfoIfCI(`Remote Kernel Finder load cache Server: ${this.id}`);
 
         const kernelsFromCache = await this.getFromCache();
@@ -244,7 +236,7 @@ export class UniversalRemoteKernelFinder implements IRemoteKernelFinder, IContri
             traceVerbose('UniversalRemoteKernelFinder: get from cache');
 
             let results: RemoteKernelConnectionMetadata[] = this.cache;
-            const key = this.getCacheKey();
+            const key = this.cacheKey;
 
             // If not in memory, check memento
             if (!results || results.length === 0) {
@@ -389,18 +381,13 @@ export class UniversalRemoteKernelFinder implements IRemoteKernelFinder, IContri
         }
     }
 
-    private getCacheKey() {
-        // For Universal finders key each one per serverId
-        return `${RemoteKernelSpecsCacheKey}-${this.serverUri.serverId}`;
-    }
-
     private async writeToCache(values: RemoteKernelConnectionMetadata[]) {
         try {
             traceVerbose(
                 `UniversalRemoteKernelFinder: Writing ${values.length} remote kernel connection metadata to cache`
             );
 
-            const key = this.getCacheKey();
+            const key = this.cacheKey;
             this.cache = values;
             const serialized = values.map(serializeKernelConnection);
             await Promise.all([

src/kernels/jupyter/finder/universalRemoteKernelFinderController.ts

@@ -8,6 +8,7 @@ import { Memento } from 'vscode';
 import { IKernelFinder, IKernelProvider, INotebookProvider } from '../../types';
 import {
     GLOBAL_MEMENTO,
+    IConfigurationService,
     IDisposableRegistry,
     IExtensions,
     IMemento,
@@ -26,30 +27,36 @@ import { IApplicationEnvironment } from '../../../platform/common/application/ty
 import { KernelFinder } from '../../kernelFinder';
 import { IExtensionSingleActivationService } from '../../../platform/activation/types';
 import { UniversalRemoteKernelFinder } from './universalRemoteKernelFinder';
+import { ContributedKernelFinderKind } from '../../internalTypes';
+import * as localize from '../../../platform/common/utils/localize';
+import { RemoteKernelSpecsCacheKey } from '../../common/commonFinder';
 
-// This class creates RemoteKernelFinders for all registered Jupyter Server URIs
-@injectable()
-export class UniversalRemoteKernelFinderController implements IExtensionSingleActivationService {
+/** Strategy design */
+interface IRemoteKernelFinderRegistrationStrategy {
+    activate(): Promise<void>;
+    dispose(): void;
+}
+
+class MultiServerStrategy implements IRemoteKernelFinderRegistrationStrategy {
     private serverFinderMapping: Map<string, UniversalRemoteKernelFinder> = new Map<
         string,
         UniversalRemoteKernelFinder
     >();
 
     constructor(
-        @inject(IJupyterSessionManagerFactory) private jupyterSessionManagerFactory: IJupyterSessionManagerFactory,
-        @inject(IInterpreterService) private interpreterService: IInterpreterService,
-        @inject(IPythonExtensionChecker) private extensionChecker: IPythonExtensionChecker,
-        @inject(INotebookProvider) private readonly notebookProvider: INotebookProvider,
-        @inject(IJupyterServerUriStorage) private readonly serverUriStorage: IJupyterServerUriStorage,
-        @inject(IMemento) @named(GLOBAL_MEMENTO) private readonly globalState: Memento,
-        @inject(IApplicationEnvironment) private readonly env: IApplicationEnvironment,
-        @inject(IJupyterRemoteCachedKernelValidator)
+        private jupyterSessionManagerFactory: IJupyterSessionManagerFactory,
+        private interpreterService: IInterpreterService,
+        private extensionChecker: IPythonExtensionChecker,
+        private readonly notebookProvider: INotebookProvider,
+        private readonly serverUriStorage: IJupyterServerUriStorage,
+        private readonly globalState: Memento,
+        private readonly env: IApplicationEnvironment,
         private readonly cachedRemoteKernelValidator: IJupyterRemoteCachedKernelValidator,
-        @inject(IKernelFinder) private readonly kernelFinder: KernelFinder,
-        @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry,
-        @inject(IKernelProvider) private readonly kernelProvider: IKernelProvider,
-        @inject(IExtensions) private readonly extensions: IExtensions,
-        @inject(IsWebExtension) private isWebExtension: boolean
+        private readonly kernelFinder: KernelFinder,
+        private readonly disposables: IDisposableRegistry,
+        private readonly kernelProvider: IKernelProvider,
+        private readonly extensions: IExtensions,
+        private isWebExtension: boolean
     ) {}
 
     async activate(): Promise<void> {
@@ -66,11 +73,17 @@ export class UniversalRemoteKernelFinderController implements IExtensionSingleAc
 
     createRemoteKernelFinder(serverUri: IJupyterServerUriEntry) {
         if (!serverUri.isValidated) {
+            // TODO@rebornix, what if it's now validated?
             return;
         }
 
         if (!this.serverFinderMapping.has(serverUri.serverId)) {
             const finder = new UniversalRemoteKernelFinder(
+                `${ContributedKernelFinderKind.Remote}-${serverUri.serverId}`,
+                localize.DataScience.universalRemoteKernelFinderDisplayName().format(
+                    serverUri.displayName || serverUri.uri
+                ),
+                `${RemoteKernelSpecsCacheKey}-${serverUri.serverId}`,
                 this.jupyterSessionManagerFactory,
                 this.interpreterService,
                 this.extensionChecker,
@@ -100,4 +113,158 @@ export class UniversalRemoteKernelFinderController implements IExtensionSingleAc
             this.serverFinderMapping.delete(uri.serverId);
         });
     }
+
+    dispose() {
+        this.serverFinderMapping.forEach((finder) => finder.dispose());
+    }
+}
+
+class SingleServerStrategy implements IRemoteKernelFinderRegistrationStrategy {
+    private _activeServerFinder: { entry: IJupyterServerUriEntry; finder: UniversalRemoteKernelFinder } | undefined;
+    constructor(
+        private jupyterSessionManagerFactory: IJupyterSessionManagerFactory,
+        private interpreterService: IInterpreterService,
+        private extensionChecker: IPythonExtensionChecker,
+        private readonly notebookProvider: INotebookProvider,
+        private readonly serverUriStorage: IJupyterServerUriStorage,
+        private readonly globalState: Memento,
+        private readonly env: IApplicationEnvironment,
+        private readonly cachedRemoteKernelValidator: IJupyterRemoteCachedKernelValidator,
+        private readonly kernelFinder: KernelFinder,
+        private readonly disposables: IDisposableRegistry,
+        private readonly kernelProvider: IKernelProvider,
+        private readonly extensions: IExtensions,
+        private isWebExtension: boolean
+    ) {}
+
+    async activate(): Promise<void> {
+        this.disposables.push(
+            this.serverUriStorage.onDidChangeUri(() => {
+                this.updateRemoteKernelFinder().then(noop, noop);
+            })
+        );
+
+        this.updateRemoteKernelFinder().then(noop, noop);
+    }
+
+    async updateRemoteKernelFinder() {
+        if (this.serverUriStorage.isLocalLaunch) {
+            // no remote kernel finder needed
+            this._activeServerFinder?.finder.dispose();
+            this._activeServerFinder = undefined;
+            return;
+        }
+
+        const uri = await this.serverUriStorage.getRemoteUri();
+        // uri should not be local
+
+        if (!uri || !uri.isValidated) {
+            this._activeServerFinder?.finder.dispose();
+            this._activeServerFinder = undefined;
+            return;
+        }
+
+        if (this._activeServerFinder?.entry.serverId === uri.serverId) {
+            // no op
+            return;
+        }
+
+        this._activeServerFinder?.finder.dispose();
+        const finder = new UniversalRemoteKernelFinder(
+            'currentremote',
+            localize.DataScience.remoteKernelFinderDisplayName(),
+            RemoteKernelSpecsCacheKey,
+            this.jupyterSessionManagerFactory,
+            this.interpreterService,
+            this.extensionChecker,
+            this.notebookProvider,
+            this.globalState,
+            this.env,
+            this.cachedRemoteKernelValidator,
+            this.kernelFinder,
+            this.kernelProvider,
+            this.extensions,
+            this.isWebExtension,
+            uri
+        );
+
+        this._activeServerFinder = {
+            entry: uri,
+            finder
+        };
+
+        finder.activate().then(noop, noop);
+    }
+
+    dispose() {
+        this._activeServerFinder?.finder.dispose();
+    }
+}
+
+// This class creates RemoteKernelFinders for all registered Jupyter Server URIs
+@injectable()
+export class UniversalRemoteKernelFinderController implements IExtensionSingleActivationService {
+    private _strategy: IRemoteKernelFinderRegistrationStrategy;
+
+    constructor(
+        @inject(IConfigurationService) readonly configurationService: IConfigurationService,
+        @inject(IJupyterSessionManagerFactory)
+        private readonly jupyterSessionManagerFactory: IJupyterSessionManagerFactory,
+        @inject(IInterpreterService) private readonly interpreterService: IInterpreterService,
+        @inject(IPythonExtensionChecker) private readonly extensionChecker: IPythonExtensionChecker,
+        @inject(INotebookProvider) private readonly notebookProvider: INotebookProvider,
+        @inject(IJupyterServerUriStorage) private readonly serverUriStorage: IJupyterServerUriStorage,
+        @inject(IMemento) @named(GLOBAL_MEMENTO) private readonly globalState: Memento,
+        @inject(IApplicationEnvironment) private readonly env: IApplicationEnvironment,
+        @inject(IJupyterRemoteCachedKernelValidator)
+        private readonly cachedRemoteKernelValidator: IJupyterRemoteCachedKernelValidator,
+        @inject(IKernelFinder) private readonly kernelFinder: KernelFinder,
+        @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry,
+        @inject(IKernelProvider) private readonly kernelProvider: IKernelProvider,
+        @inject(IExtensions) private readonly extensions: IExtensions,
+        @inject(IsWebExtension) private readonly isWebExtension: boolean
+    ) {
+        this._strategy = this.getStrategy();
+        this.disposables.push(this._strategy);
+    }
+
+    private getStrategy(): IRemoteKernelFinderRegistrationStrategy {
+        if (this.configurationService.getSettings().kernelPickerType === 'Insiders') {
+            return new MultiServerStrategy(
+                this.jupyterSessionManagerFactory,
+                this.interpreterService,
+                this.extensionChecker,
+                this.notebookProvider,
+                this.serverUriStorage,
+                this.globalState,
+                this.env,
+                this.cachedRemoteKernelValidator,
+                this.kernelFinder,
+                this.disposables,
+                this.kernelProvider,
+                this.extensions,
+                this.isWebExtension
+            );
+        } else {
+            return new SingleServerStrategy(
+                this.jupyterSessionManagerFactory,
+                this.interpreterService,
+                this.extensionChecker,
+                this.notebookProvider,
+                this.serverUriStorage,
+                this.globalState,
+                this.env,
+                this.cachedRemoteKernelValidator,
+                this.kernelFinder,
+                this.disposables,
+                this.kernelProvider,
+                this.extensions,
+                this.isWebExtension
+            );
+        }
+    }
+
+    async activate(): Promise<void> {
+        this._strategy.activate().then(noop, noop);
+    }
 }

src/kernels/jupyter/launcher/serverUriStorage.ts

@@ -56,6 +56,9 @@ export class JupyterServerUriStorage implements IJupyterServerUriStorage, IServe
     public get onDidChange(): Event<void> {
         return this._onDidChangeUri.event;
     }
+    public get onDidChangeConnectionType(): Event<void> {
+        return this._onDidChangeUri.event;
+    }
     public get isLocalLaunch(): boolean {
         return this._localOnly;
     }