Create finders only for Servers that were used (#14123)

* Create finders only for Servers that were used

* Update title

* Updates

src/kernels/jupyter/connection/jupyterServerProviderRegistry.ts

@@ -67,6 +67,7 @@ class JupyterUriProviderAdaptor extends Disposables implements IJupyterUriProvid
     private providerChanges: IDisposable[] = [];
     removeHandle?(handle: string): Promise<void>;
     getServerUriWithoutAuthInfo?(handle: string): Promise<IJupyterServerUri>;
+    private commands = new Map<string, JupyterServerCommand>();
     constructor(
         private readonly provider: JupyterServerCollection,
         public readonly extensionId: string
@@ -98,7 +99,6 @@ class JupyterUriProviderAdaptor extends Disposables implements IJupyterUriProvid
             );
         }
     }
-    private commands = new Map<string, JupyterServerCommand>();
     async getQuickPickEntryItems(value?: string): Promise<(QuickPickItem & { default?: boolean | undefined })[]> {
         if (!this.provider.commandProvider) {
             throw new Error(`No Jupyter Server Command Provider for ${this.provider.extensionId}#${this.provider.id}`);
@@ -280,14 +280,14 @@ export class JupyterServerProviderRegistry extends Disposables implements IJupyt
         const serverProvider = new JupyterServerCollectionImpl(extensionId, id, label);
         this._serverProviders.set(extId, serverProvider);
         let uriRegistration: IDisposable | undefined;
+        let adapter: JupyterUriProviderAdaptor | undefined;
         serverProvider.onDidChangeProvider(
             () => {
                 if (serverProvider.serverProvider) {
+                    adapter?.dispose();
                     uriRegistration?.dispose();
-                    uriRegistration = this.jupyterUriProviderRegistration.registerProvider(
-                        new JupyterUriProviderAdaptor(serverProvider, extensionId),
-                        extensionId
-                    );
+                    adapter = new JupyterUriProviderAdaptor(serverProvider, extensionId);
+                    uriRegistration = this.jupyterUriProviderRegistration.registerProvider(adapter, extensionId);
                     this.disposables.push(uriRegistration);
                     this._onDidChangeProviders.fire();
                 }
@@ -298,6 +298,7 @@ export class JupyterServerProviderRegistry extends Disposables implements IJupyt
 
         serverProvider.onDidDispose(
             () => {
+                adapter?.dispose();
                 uriRegistration?.dispose();
                 this._serverProviders.delete(extId);
                 this._onDidChangeProviders.fire();

src/kernels/jupyter/connection/serverUriStorage.ts

@@ -42,6 +42,10 @@ export type StorageMRUItem = {
  */
 @injectable()
 export class JupyterServerUriStorage extends Disposables implements IJupyterServerUriStorage {
+    private _onDidLoad = new EventEmitter<void>();
+    public get onDidLoad() {
+        return this._onDidLoad.event;
+    }
     private _onDidChangeUri = new EventEmitter<void>();
     public get onDidChange() {
         return this._onDidChangeUri.event;
@@ -57,6 +61,10 @@ export class JupyterServerUriStorage extends Disposables implements IJupyterServ
     private readonly oldStorage: OldStorage;
     private readonly newStorage: NewStorage;
     private storageEventsHooked?: boolean;
+    private _all: IJupyterServerUriEntry[] = [];
+    public get all() {
+        return this._all;
+    }
     constructor(
         @inject(IEncryptedStorage) encryptedStorage: IEncryptedStorage,
         @inject(IMemento) @named(GLOBAL_MEMENTO) globalMemento: Memento,
@@ -74,6 +82,7 @@ export class JupyterServerUriStorage extends Disposables implements IJupyterServ
         this.oldStorage = new OldStorage(encryptedStorage, globalMemento);
         // eslint-disable-next-line @typescript-eslint/no-use-before-define
         this.newStorage = new NewStorage(fs, storageFile, this.oldStorage, globalMemento);
+        this.disposables.push(this._onDidLoad);
         this.disposables.push(this._onDidAddUri);
         this.disposables.push(this._onDidChangeUri);
         this.disposables.push(this._onDidRemoveUris);
@@ -91,12 +100,19 @@ export class JupyterServerUriStorage extends Disposables implements IJupyterServ
     public async getAll(): Promise<IJupyterServerUriEntry[]> {
         this.hookupStorageEvents();
         await this.newStorage.migrateMRU();
-        return this.newStorage.getAll();
+        const previous = this._all;
+        this._all = await this.newStorage.getAll();
+        if (previous.length !== this._all.length || JSON.stringify(this._all) !== JSON.stringify(previous)) {
+            this._onDidLoad.fire();
+        }
+        return this._all;
     }
     public async clear(): Promise<void> {
         this.hookupStorageEvents();
         await this.newStorage.migrateMRU();
         await Promise.all([this.oldStorage.clear(), this.newStorage.clear()]);
+        this._all = [];
+        this._onDidLoad.fire();
     }
     public async add(jupyterHandle: JupyterServerProviderHandle, options?: { time: number }): Promise<void> {
         this.hookupStorageEvents();
@@ -109,16 +125,19 @@ export class JupyterServerUriStorage extends Disposables implements IJupyterServ
         };
 
         await this.newStorage.add(entry);
+        this.getAll().catch(noop);
     }
     public async update(server: JupyterServerProviderHandle) {
         this.hookupStorageEvents();
         await this.newStorage.migrateMRU();
         await this.newStorage.update(server);
+        this.getAll().catch(noop);
     }
     public async remove(server: JupyterServerProviderHandle) {
         this.hookupStorageEvents();
         await this.newStorage.migrateMRU();
         await this.newStorage.remove(server);
+        this.getAll().catch(noop);
     }
 }
 

src/kernels/jupyter/finder/remoteKernelFinderController.ts

@@ -11,7 +11,8 @@ import {
     IJupyterServerUriEntry,
     IJupyterUriProviderRegistration,
     JupyterServerProviderHandle,
-    IJupyterServerProviderRegistry
+    IJupyterServerProviderRegistry,
+    IRemoteKernelFinder
 } from '../types';
 import { noop } from '../../../platform/common/utils/misc';
 import { IApplicationEnvironment } from '../../../platform/common/application/types';
@@ -26,9 +27,10 @@ import { swallowExceptions } from '../../../platform/common/utils/decorators';
 import { IJupyterUriProvider, JupyterServerCollection, JupyterServerProvider } from '../../../api';
 import { CancellationTokenSource } from 'vscode';
 import { traceError } from '../../../platform/logging';
+import { IRemoteKernelFinderController } from './types';
 
 @injectable()
-export class RemoteKernelFinderController implements IExtensionSyncActivationService {
+export class RemoteKernelFinderController implements IRemoteKernelFinderController, IExtensionSyncActivationService {
     private serverFinderMapping: Map<string, RemoteKernelFinder> = new Map<string, RemoteKernelFinder>();
 
     constructor(
@@ -55,6 +57,7 @@ export class RemoteKernelFinderController implements IExtensionSyncActivationSer
     activate() {
         this.serverUriStorage.onDidAdd((server) => this.validateAndCreateFinder(server), this, this.disposables);
         this.serverUriStorage.onDidChange(this.buildListOfFinders, this, this.disposables);
+        this.serverUriStorage.onDidLoad(this.handleProviderChanges, this, this.disposables);
         // Possible some extensions register their providers later.
         // And also possible they load their old server later, hence we need to go through the
         // MRU list again and try to build the finders, as the servers might now exist.
@@ -92,10 +95,14 @@ export class RemoteKernelFinderController implements IExtensionSyncActivationSer
         if (!this.features.features.enableProposedJupyterServerProviderApi) {
             return;
         }
+        if (!this.serverUriStorage.all.length) {
+            // We do not have any of the previously used servers, or the data has not yet loaded.
+            return;
+        }
         const token = new CancellationTokenSource();
         this.disposables.push(token);
         await Promise.all(
-            this.jupyterServerProviderRegistry.providers.map(async (collection) => {
+            this.jupyterServerProviderRegistry.providers.map((collection) => {
                 const serverProvider = collection.serverProvider;
                 if (!serverProvider || this.mappedProviders.has(serverProvider)) {
                     return;
@@ -108,13 +115,23 @@ export class RemoteKernelFinderController implements IExtensionSyncActivationSer
                         this.disposables
                     );
                 }
-                await this.lookForServersInCollection(collection).catch(noop);
+                this.serverUriStorage.onDidLoad(
+                    () => this.lookForServersInCollection(collection),
+                    this,
+                    this.disposables
+                );
+                return this.lookForServersInCollection(collection).catch(noop);
             })
         );
         token.dispose();
     }
     @swallowExceptions('Check Servers in Jupyter Server Provider')
     private async lookForServersInCollection(collection: JupyterServerCollection) {
+        if (!this.serverUriStorage.all.length) {
+            // We do not have any of the previously used servers, or the data has not yet loaded.
+            return;
+        }
+        const usedServers = new Set(this.serverUriStorage.all.map((s) => generateIdFromRemoteProvider(s.provider)));
         const serverProvider = collection.serverProvider;
         if (!serverProvider) {
             return;
@@ -123,17 +140,18 @@ export class RemoteKernelFinderController implements IExtensionSyncActivationSer
         try {
             const servers = await serverProvider.getJupyterServers(token.token);
             servers.forEach((server) => {
-                const serverId = `${collection.extensionId}#${collection.id}#${server.id}`;
-                if (this.mappedServers.has(serverId)) {
-                    return;
-                }
-                this.mappedServers.add(serverId);
-                const info = {
+                const serverProviderHandle = {
                     extensionId: collection.extensionId,
                     handle: server.id,
                     id: collection.id
                 };
-                this.createRemoteKernelFinder(info, server.label);
+                const serverId = generateIdFromRemoteProvider(serverProviderHandle);
+                // If this sever was never used in the past, then no need to create a finder for this.
+                if (this.mappedServers.has(serverId) || !usedServers.has(serverId)) {
+                    return;
+                }
+                this.mappedServers.add(serverId);
+                this.createRemoteKernelFinder(serverProviderHandle, server.label);
             });
         } catch (ex) {
             traceError(`Failed to get servers for Collection ${collection.id} in ${collection.extensionId}`, ex);
@@ -155,7 +173,10 @@ export class RemoteKernelFinderController implements IExtensionSyncActivationSer
         }
     }
 
-    createRemoteKernelFinder(serverProviderHandle: JupyterServerProviderHandle, displayName: string) {
+    public getOrCreateRemoteKernelFinder(
+        serverProviderHandle: JupyterServerProviderHandle,
+        displayName: string
+    ): IRemoteKernelFinder {
         const serverId = generateIdFromRemoteProvider(serverProviderHandle);
         if (!this.serverFinderMapping.has(serverId)) {
             const finder = new RemoteKernelFinder(
@@ -177,6 +198,10 @@ export class RemoteKernelFinderController implements IExtensionSyncActivationSer
 
             finder.activate().then(noop, noop);
         }
+        return this.serverFinderMapping.get(serverId)!;
+    }
+    createRemoteKernelFinder(serverProviderHandle: JupyterServerProviderHandle, displayName: string) {
+        this.getOrCreateRemoteKernelFinder(serverProviderHandle, displayName);
     }
 
     // When a URI is removed, dispose the kernel finder for it

src/kernels/jupyter/finder/types.ts

@@ -0,0 +1,12 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+import { JupyterServerProviderHandle, IRemoteKernelFinder } from '../types';
+
+export const IRemoteKernelFinderController = Symbol('RemoteKernelFinderController');
+export interface IRemoteKernelFinderController {
+    getOrCreateRemoteKernelFinder(
+        serverProviderHandle: JupyterServerProviderHandle,
+        displayName: string
+    ): IRemoteKernelFinder;
+}

src/kernels/jupyter/serviceRegistry.node.ts

@@ -57,6 +57,7 @@ import { KernelSessionFactory } from '../common/kernelSessionFactory';
 import { OldJupyterKernelSessionFactory } from './session/oldJupyterKernelSessionFactory';
 import { JupyterKernelSessionFactory } from './session/jupyterKernelSessionFactory';
 import { JupyterServerProviderRegistry } from './connection/jupyterServerProviderRegistry';
+import { IRemoteKernelFinderController } from './finder/types';
 
 export function registerTypes(serviceManager: IServiceManager, _isDevMode: boolean) {
     serviceManager.add<IJupyterCommandFactory>(IJupyterCommandFactory, JupyterCommandFactory);
@@ -132,10 +133,11 @@ export function registerTypes(serviceManager: IServiceManager, _isDevMode: boole
     );
     serviceManager.addSingleton<JupyterDetectionTelemetry>(IExtensionSyncActivationService, JupyterDetectionTelemetry);
     serviceManager.addSingleton<IDataScienceErrorHandler>(IDataScienceErrorHandler, DataScienceErrorHandlerNode);
-    serviceManager.addSingleton<IExtensionSyncActivationService>(
-        IExtensionSyncActivationService,
+    serviceManager.addSingleton<IRemoteKernelFinderController>(
+        IRemoteKernelFinderController,
         RemoteKernelFinderController
     );
+    serviceManager.addBinding(IRemoteKernelFinderController, IExtensionSyncActivationService);
     serviceManager.addSingleton<IJupyterServerProviderRegistry>(
         IJupyterServerProviderRegistry,
         JupyterServerProviderRegistry

src/kernels/jupyter/serviceRegistry.web.ts

@@ -34,6 +34,7 @@ import { KernelSessionFactory } from '../common/kernelSessionFactory';
 import { OldJupyterKernelSessionFactory } from './session/oldJupyterKernelSessionFactory';
 import { JupyterKernelSessionFactory } from './session/jupyterKernelSessionFactory';
 import { JupyterServerProviderRegistry } from './connection/jupyterServerProviderRegistry';
+import { IRemoteKernelFinderController } from './finder/types';
 
 export function registerTypes(serviceManager: IServiceManager, _isDevMode: boolean) {
     serviceManager.addSingleton<IOldJupyterSessionManagerFactory>(
@@ -68,10 +69,11 @@ export function registerTypes(serviceManager: IServiceManager, _isDevMode: boole
         JupyterRemoteCachedKernelValidator
     );
     serviceManager.addSingleton<IDataScienceErrorHandler>(IDataScienceErrorHandler, DataScienceErrorHandlerWeb);
-    serviceManager.addSingleton<IExtensionSyncActivationService>(
-        IExtensionSyncActivationService,
+    serviceManager.addSingleton<IRemoteKernelFinderController>(
+        IRemoteKernelFinderController,
         RemoteKernelFinderController
     );
+    serviceManager.addBinding(IRemoteKernelFinderController, IExtensionSyncActivationService);
     serviceManager.addSingleton<IJupyterServerProviderRegistry>(
         IJupyterServerProviderRegistry,
         JupyterServerProviderRegistry

src/kernels/jupyter/types.ts

@@ -196,13 +196,23 @@ export interface IJupyterServerUriEntry {
 
 export const IJupyterServerUriStorage = Symbol('IJupyterServerUriStorage');
 export interface IJupyterServerUriStorage {
+    /**
+     * Temporary event, the old API was async, the new API is sync.
+     * However until we migrate everything we will never know whether
+     * the data has been successfully loaded into memory.
+     */
+    readonly onDidLoad: Event<void>;
     readonly onDidChange: Event<void>;
     readonly onDidRemove: Event<IJupyterServerUriEntry[]>;
     readonly onDidAdd: Event<IJupyterServerUriEntry>;
+    readonly all: readonly IJupyterServerUriEntry[];
     /**
      * Updates MRU list marking this server as the most recently used.
      */
     update(serverProviderHandle: JupyterServerProviderHandle): Promise<void>;
+    /**
+     * @deprecated Use `all` and `onDidLoad` instead.
+     */
     getAll(): Promise<IJupyterServerUriEntry[]>;
     remove(serverProviderHandle: JupyterServerProviderHandle): Promise<void>;
     clear(): Promise<void>;
@@ -301,6 +311,6 @@ export interface IRemoteKernelFinder
 export const IJupyterServerProviderRegistry = Symbol('IJupyterServerProviderRegistry');
 export interface IJupyterServerProviderRegistry {
     onDidChangeProviders: Event<void>;
-    providers: readonly JupyterServerCollection[];
+    readonly providers: readonly JupyterServerCollection[];
     createJupyterServerCollection(extensionId: string, id: string, label: string): JupyterServerCollection;
 }