Faster access of Jupyter Uri from storage (#14088)

src/kernels/jupyter/connection/jupyterUriProviderRegistration.ts

@@ -16,7 +16,7 @@ import {
  JupyterServerProviderHandle
 } from '../types';
 import { sendTelemetryEvent } from '../../../telemetry';
-import { traceError } from '../../../platform/logging';
+import { traceError, traceVerbose } from '../../../platform/logging';
 import { IJupyterServerUri, IJupyterUriProvider } from '../../../api';
 import { Disposables } from '../../../platform/common/utils';
 import { IServiceContainer } from '../../../platform/ioc/types';
@@ -103,11 +103,26 @@ export class JupyterUriProviderRegistration
  const provider = this._providers.get(id);
  if (!provider) {
  throw new Error(
- `${localize.DataScience.unknownServerUri}. Provider Id=${id} and handle=${providerHandle.handle}`
+ `Provider Id=${id} and handle=${providerHandle.handle} not registered. Did you uninstall the extension ${providerHandle.extensionId}?`
  );
  }
  return provider.getServerUri(providerHandle.handle, doNotPromptForAuthInfo);
  }
+ public async getDisplayNameIfProviderIsLoaded(
+ providerHandle: JupyterServerProviderHandle
+ ): Promise<string | undefined> {
+ await this.loadExtension(providerHandle.extensionId, providerHandle.id);
+ const id = getProviderId(providerHandle.extensionId, providerHandle.id);
+ const provider = this._providers.get(id);
+ if (!provider) {
+ traceVerbose(
+ `Provider Id=${id} and handle=${providerHandle.handle} not registered. Extension ${providerHandle.extensionId} may not have yet loaded or provider not yet registered?`
+ );
+ return;
+ }
+ const server = await provider.getServerUri(providerHandle.handle, true);
+ return server.displayName;
+ }
  private async loadExtension(extensionId: string, providerId: string) {
  if (extensionId === JVSC_EXTENSION_ID) {
  return;

src/kernels/jupyter/finder/remoteKernelFinderController.ts

@@ -81,8 +81,13 @@ export class RemoteKernelFinderController implements IExtensionSyncActivationSer
  private async validateAndCreateFinder(serverUri: IJupyterServerUriEntry) {
  const serverId = generateIdFromRemoteProvider(serverUri.provider);
  if (!this.serverFinderMapping.has(serverId)) {
- const info = await this.jupyterPickerRegistration.getJupyterServerUri(serverUri.provider, true);
- this.createRemoteKernelFinder(serverUri.provider, info.displayName);
+ const displayName = await this.jupyterPickerRegistration.getDisplayNameIfProviderIsLoaded(
+ serverUri.provider
+ );
+ // If display name is empty/undefined, then the extension has not yet loaded or provider not yet registered.
+ if (displayName) {
+ this.createRemoteKernelFinder(serverUri.provider, displayName);
+ }
  }
  }
 

src/kernels/jupyter/types.ts

@@ -166,6 +166,11 @@ export interface IJupyterUriProviderRegistration {
  readonly providers: ReadonlyArray<IInternalJupyterUriProvider>;
  getProvider(extensionId: string, id: string): Promise<IInternalJupyterUriProvider | undefined>;
  registerProvider(provider: IJupyterUriProvider, extensionId: string): IDisposable;
+ /**
+ * Temporary, until the new API is finalized.
+ * We need a way to get the displayName of the Server.
+ */
+ getDisplayNameIfProviderIsLoaded(providerHandle: JupyterServerProviderHandle): Promise<string | undefined>;
  getJupyterServerUri(
  serverHandle: JupyterServerProviderHandle,
  doNotPromptForAuthInfo?: boolean

src/kernels/portAttributeProvider.node.ts

@@ -24,13 +24,11 @@ export class PortAttributesProviders implements PortAttributesProvider, IExtensi
  }
  }
  providePortAttributes(
- port: number,
- _pid: number | undefined,
- _commandLine: string | undefined,
+ attributes: { port: number; pid?: number; commandLine?: string },
  _token: CancellationToken
  ): PortAttributes | undefined {
  try {
- if (UsedPorts.has(port)) {
+ if (UsedPorts.has(attributes.port)) {
  return new PortAttributes(PortAutoForwardAction.Ignore);
  }
  } catch (ex) {

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

@@ -111,8 +111,7 @@ suite('kernel Launcher', () => {
 
  for (const port of UsedPorts) {
  assert.equal(
- portAttributeProvider.providePortAttributes(port, undefined, undefined, cancellation.token)
- ?.autoForwardAction,
+ portAttributeProvider.providePortAttributes({ port }, cancellation.token)?.autoForwardAction,
  PortAutoForwardAction.Ignore,
  'Kernel Port should not be forwarded'
  );

src/platform/common/utils/localize.ts

@@ -80,9 +80,6 @@ export namespace DataScience {
  export const pythonExtensionInstalled = l10n.t(
  'Python Extension is now installed. Some features might not be available until a notebook or interactive window session is restarted.'
  );
- export const unknownServerUri = l10n.t(
- 'Server URL cannot be used. Did you uninstall an extension that provided a Jupyter server connection?'
- );
  export const uriProviderDescriptionFormat = (description: string, extensionId: string) =>
  l10n.t('{0} (From {1} extension)', description, extensionId);
  export const unknownPackage = l10n.t('unknown');

src/standalone/userJupyterServer/userServerUriProvider.unit.test.ts

@@ -249,12 +249,18 @@ suite('User Uri Provider', () => {
  verify(serverUriStorage.add(anything(), anything())).atLeast(2);
 
  // Verify the items were added into both of the stores.
- const serversInNewStorage = await provider.newStorage.getServers();
-
+ const [serversInNewStorage, serversInNewStorage2] = await Promise.all([
+ provider.newStorage.getServers(false),
+ provider.newStorage.getServers(true)
+ ]);
  assert.deepEqual(
  serversInNewStorage.map((s) => s.serverInfo.displayName),
  ['Hello World', 'Foo Bar']
  );
+ assert.deepEqual(
+ serversInNewStorage2.map((s) => s.serverInfo.displayName),
+ ['Hello World', 'Foo Bar']
+ );
  }
  test('Migrate Old Urls', async () => testMigration());
  test('Migrate display names from Uri Storage', async () => {
@@ -316,11 +322,18 @@ suite('User Uri Provider', () => {
  );
 
  // Verify the of the servers have the actual names in the stores.
- const serversInNewStorage = await provider.newStorage.getServers();
+ const [serversInNewStorage, serversInNewStorage2] = await Promise.all([
+ provider.newStorage.getServers(false),
+ provider.newStorage.getServers(true)
+ ]);
  assert.deepEqual(
  serversInNewStorage.map((s) => s.serverInfo.displayName),
  ['Azure ML', 'My Remote Server Name']
  );
+ assert.deepEqual(
+ serversInNewStorage2.map((s) => s.serverInfo.displayName),
+ ['Azure ML', 'My Remote Server Name']
+ );
  });
  test('Add a new Url and verify it is in the storage', async () => {
  await testMigration();
@@ -341,9 +354,12 @@ suite('User Uri Provider', () => {
  assert.isAtLeast(handles.length, 3, '2 migrated urls and one entered');
  assert.include(handles, handle);
 
- const serversInNewStorage = await provider.newStorage.getServers();
-
+ const [serversInNewStorage, serversInNewStorage2] = await Promise.all([
+ provider.newStorage.getServers(false),
+ provider.newStorage.getServers(true)
+ ]);
  assert.strictEqual(serversInNewStorage.length, 3);
+ assert.strictEqual(serversInNewStorage2.length, 3);
  });
  test('When adding a HTTP url (without pwd, and without a token) prompt user to use insecure sites (in new pwd manager)', async function () {
  await testMigration();
@@ -371,9 +387,12 @@ suite('User Uri Provider', () => {
  assert.isAtLeast(handles.length, 3, '2 migrated urls and one entered');
  assert.include(handles, handle);
 
- const serversInNewStorage = await provider.newStorage.getServers();
-
+ const [serversInNewStorage, serversInNewStorage2] = await Promise.all([
+ provider.newStorage.getServers(false),
+ provider.newStorage.getServers(true)
+ ]);
  assert.strictEqual(serversInNewStorage.length, 3);
+ assert.strictEqual(serversInNewStorage2.length, 3);
  });
  test('When prompted to use insecure sites and ignored/cancelled, then do not add the url', async function () {
  await testMigration();
@@ -400,8 +419,12 @@ suite('User Uri Provider', () => {
  const handles = await provider.getHandles();
  assert.isAtLeast(handles.length, 2, '2 migrated urls');
 
- const serversInNewStorage = await provider.newStorage.getServers();
+ const [serversInNewStorage, serversInNewStorage2] = await Promise.all([
+ provider.newStorage.getServers(false),
+ provider.newStorage.getServers(true)
+ ]);
  assert.strictEqual(serversInNewStorage.length, 2);
+ assert.strictEqual(serversInNewStorage2.length, 2);
  });
  test('When adding a HTTP url (with a pwd, and without a token) do not prompt user to use insecure sites (in new pwd manager)', async function () {
  await testMigration();
@@ -432,8 +455,11 @@ suite('User Uri Provider', () => {
  assert.isAtLeast(handles.length, 3, '2 migrated urls and one entered');
  assert.include(handles, handle);
 
- const serversInNewStorage = await provider.newStorage.getServers();
-
+ const [serversInNewStorage, serversInNewStorage2] = await Promise.all([
+ provider.newStorage.getServers(false),
+ provider.newStorage.getServers(true)
+ ]);
  assert.strictEqual(serversInNewStorage.length, 3);
+ assert.strictEqual(serversInNewStorage2.length, 3);
  });
 });

src/standalone/userJupyterServer/userServerUrlProvider.ts

@@ -151,7 +151,7 @@ export class UserJupyterServerUrlProvider
  onDidChangeServers = this._onDidChangeServers.event;
  async getJupyterServers(_token: CancellationToken): Promise<JupyterServer[]> {
  await this.initializeServers();
- const servers = await this.newStorage.getServers();
+ const servers = await this.newStorage.getServers(false);
  return servers.map((s) => {
  return {
  id: s.handle,
@@ -549,7 +549,7 @@ export class UserJupyterServerUrlProvider
  }
  async getServerUriWithoutAuthInfo(handle: string): Promise<IJupyterServerUri> {
  await this.initializeServers();
- const servers = await this.newStorage.getServers();
+ const servers = await this.newStorage.getServers(false);
  const server = servers.find((s) => s.handle === handle);
  if (!server) {
  throw new Error('Server not found');
@@ -565,7 +565,7 @@ export class UserJupyterServerUrlProvider
  }
  async getServerUri(handle: string): Promise<IJupyterServerUri> {
  await this.initializeServers();
- const servers = await this.newStorage.getServers();
+ const servers = await this.newStorage.getServers(false);
  const server = servers.find((s) => s.handle === handle);
  if (!server) {
  throw new Error('Server not found');
@@ -590,7 +590,7 @@ export class UserJupyterServerUrlProvider
  }
  async getHandles(): Promise<string[]> {
  await this.initializeServers();
- const servers = await this.newStorage.getServers();
+ const servers = await this.newStorage.getServers(false);
  return servers.map((s) => s.handle);
  }
 
@@ -732,6 +732,7 @@ function storageFormatToServers(items: StorageItem[]) {
 export class NewStorage {
  private readonly _migrationDone: Deferred<void>;
  private updatePromise = Promise.resolve();
+ private servers?: { handle: string; uri: string; serverInfo: IJupyterServerUri }[];
  public get migrationDone(): Promise<void> {
  return this._migrationDone.promise;
  }
@@ -815,7 +816,12 @@ export class NewStorage {
  JSON.stringify(serverToStorageFormat(userServers))
  );
  }
- public async getServers(): Promise<{ handle: string; uri: string; serverInfo: IJupyterServerUri }[]> {
+ public async getServers(
+ ignoreCache: boolean
+ ): Promise<{ handle: string; uri: string; serverInfo: IJupyterServerUri }[]> {
+ if (this.servers && !ignoreCache) {
+ return this.servers;
+ }
  const data = await this.encryptedStorage.retrieve(
  Settings.JupyterServerRemoteLaunchService,
  UserJupyterServerUriListKeyV2
@@ -824,16 +830,20 @@ export class NewStorage {
  return [];
  }
  try {
- return storageFormatToServers(JSON.parse(data));
+ return (this.servers = storageFormatToServers(JSON.parse(data)));
  } catch {
  return [];
  }
  }
 
  public async add(server: { handle: string; uri: string; serverInfo: IJupyterServerUri }) {
+ if (this.servers) {
+ this.servers = this.servers.filter((s) => s.handle !== server.handle).concat(server);
+ }
  await (this.updatePromise = this.updatePromise
  .then(async () => {
- const servers = (await this.getServers()).concat(server);
+ const servers = (await this.getServers(true)).concat(server);
+ this.servers = servers;
  await this.encryptedStorage.store(
  Settings.JupyterServerRemoteLaunchService,
  UserJupyterServerUriListKeyV2,
@@ -843,9 +853,13 @@ export class NewStorage {
  .catch(noop));
  }
  public async remove(handle: string) {
+ if (this.servers) {
+ this.servers = this.servers.filter((s) => s.handle !== handle);
+ }
  await (this.updatePromise = this.updatePromise
  .then(async () => {
- const servers = (await this.getServers()).filter((s) => s.handle !== handle);
+ const servers = (await this.getServers(true)).filter((s) => s.handle !== handle);
+ this.servers = servers;
  return this.encryptedStorage.store(
  Settings.JupyterServerRemoteLaunchService,
  UserJupyterServerUriListKeyV2,
@@ -855,8 +869,10 @@ export class NewStorage {
  .catch(noop));
  }
  public async clear() {
+ this.servers = [];
  await (this.updatePromise = this.updatePromise
  .then(async () => {
+ this.servers = [];
  await this.encryptedStorage.store(
  Settings.JupyterServerRemoteLaunchService,
  UserJupyterServerUriListKeyV2,