Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Faster access of Jupyter Uri from storage #14088

Merged
merged 3 commits into from
Aug 9, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 17 additions & 2 deletions src/kernels/jupyter/connection/jupyterUriProviderRegistration.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -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;
Expand Down
9 changes: 7 additions & 2 deletions src/kernels/jupyter/finder/remoteKernelFinderController.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
}
}

Expand Down
5 changes: 5 additions & 0 deletions src/kernels/jupyter/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
6 changes: 2 additions & 4 deletions src/kernels/portAttributeProvider.node.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
3 changes: 1 addition & 2 deletions src/kernels/raw/launcher/kernelLauncher.unit.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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'
);
Expand Down
3 changes: 0 additions & 3 deletions src/platform/common/utils/localize.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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');
Expand Down
46 changes: 36 additions & 10 deletions src/standalone/userJupyterServer/userServerUriProvider.unit.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 () => {
Expand Down Expand Up @@ -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();
Expand All @@ -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();
Expand Down Expand Up @@ -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();
Expand All @@ -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();
Expand Down Expand Up @@ -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);
});
});
32 changes: 24 additions & 8 deletions src/standalone/userJupyterServer/userServerUrlProvider.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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');
Expand All @@ -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');
Expand All @@ -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);
}

Expand Down Expand Up @@ -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;
}
Expand Down Expand Up @@ -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
Expand All @@ -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,
Expand All @@ -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,
Expand All @@ -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,
Expand Down
Loading