Skip to content

Commit

Permalink
Include windows registry perf improvement in `pythonDiscoveryUsingWor…
Browse files Browse the repository at this point in the history
  • Loading branch information
Kartik Raj authored and wesm committed Apr 8, 2024
1 parent 6d74b7a commit 27a1f4d
Show file tree
Hide file tree
Showing 6 changed files with 21 additions and 32 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -109,7 +109,7 @@ async function updateEnvUsingRegistry(env: PythonEnvInfo): Promise<void> {
let interpreters = getRegistryInterpretersSync();
if (!interpreters) {
traceError('Expected registry interpreter cache to be initialized already');
interpreters = await getRegistryInterpreters(true);
interpreters = await getRegistryInterpreters();
}
const data = interpreters.find((i) => arePathsSame(i.interpreterPath, env.executable.filename));
if (data) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,8 +19,8 @@ export class CondaEnvironmentLocator extends FSWatchingLocator {
}

// eslint-disable-next-line class-methods-use-this
public async *doIterEnvs(_: unknown, useWorkerThreads = false): IPythonEnvsIterator<BasicEnvInfo> {
const conda = await Conda.getConda(undefined, useWorkerThreads);
public async *doIterEnvs(_: unknown): IPythonEnvsIterator<BasicEnvInfo> {
const conda = await Conda.getConda();
if (conda === undefined) {
traceVerbose(`Couldn't locate the conda binary.`);
return;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -9,26 +9,31 @@ import { getRegistryInterpreters } from '../../../common/windowsUtils';
import { traceError, traceVerbose } from '../../../../logging';
import { isMicrosoftStoreDir } from '../../../common/environmentManagers/microsoftStoreEnv';
import { PythonEnvsChangedEvent } from '../../watcher';
import { DiscoveryUsingWorkers } from '../../../../common/experiments/groups';
import { inExperiment } from '../../../common/externalDependencies';

export const WINDOWS_REG_PROVIDER_ID = 'windows-registry';

export class WindowsRegistryLocator extends Locator<BasicEnvInfo> {
public readonly providerId: string = WINDOWS_REG_PROVIDER_ID;

// eslint-disable-next-line class-methods-use-this
public iterEnvs(query?: PythonLocatorQuery, useWorkerThreads = false): IPythonEnvsIterator<BasicEnvInfo> {
public iterEnvs(
query?: PythonLocatorQuery,
useWorkerThreads = inExperiment(DiscoveryUsingWorkers.experiment),
): IPythonEnvsIterator<BasicEnvInfo> {
if (useWorkerThreads) {
/**
* Windows registry is slow and often not necessary, so notify completion immediately, but use watcher
* change events to signal for any new envs which are found.
*/
if (query?.providerId === this.providerId) {
// Query via change event, so iterate all envs.
return iterateEnvs(true);
return iterateEnvs();
}
return iterateEnvsLazily(this.emitter);
}
return iterateEnvs(false);
return iterateEnvs();
}
}

Expand All @@ -38,13 +43,13 @@ async function* iterateEnvsLazily(changed: IEmitter<PythonEnvsChangedEvent>): IP

async function loadAllEnvs(changed: IEmitter<PythonEnvsChangedEvent>) {
traceVerbose('Searching for windows registry interpreters');
await getRegistryInterpreters(true);
await getRegistryInterpreters();
changed.fire({ providerId: WINDOWS_REG_PROVIDER_ID });
traceVerbose('Finished searching for windows registry interpreters');
}

async function* iterateEnvs(useWorkerThreads: boolean): IPythonEnvsIterator<BasicEnvInfo> {
const interpreters = await getRegistryInterpreters(useWorkerThreads);
async function* iterateEnvs(): IPythonEnvsIterator<BasicEnvInfo> {
const interpreters = await getRegistryInterpreters(); // Value should already be loaded at this point, so this returns immediately.
for (const interpreter of interpreters) {
try {
// Filter out Microsoft Store app directories. We have a store app locator that handles this.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -280,9 +280,9 @@ export class Conda {
});
}

public static async getConda(shellPath?: string, useWorkerThreads?: boolean): Promise<Conda | undefined> {
public static async getConda(shellPath?: string): Promise<Conda | undefined> {
if (Conda.condaPromise.get(shellPath) === undefined || isTestExecution()) {
Conda.condaPromise.set(shellPath, Conda.locate(shellPath, useWorkerThreads));
Conda.condaPromise.set(shellPath, Conda.locate(shellPath));
}
return Conda.condaPromise.get(shellPath);
}
Expand All @@ -293,11 +293,7 @@ export class Conda {
*
* @return A Conda instance corresponding to the binary, if successful; otherwise, undefined.
*/
private static async locate(shellPath?: string, useWorkerThread?: boolean): Promise<Conda | undefined> {
let useWorkerThreads: boolean;
if (useWorkerThread === undefined) {
useWorkerThreads = false;
}
private static async locate(shellPath?: string): Promise<Conda | undefined> {
traceVerbose(`Searching for conda.`);
const home = getUserHomeDir();
let customCondaPath: string | undefined = 'conda';
Expand All @@ -324,7 +320,7 @@ export class Conda {
}

async function* getCandidatesFromRegistry() {
const interps = await getRegistryInterpreters(useWorkerThreads);
const interps = await getRegistryInterpreters();
const candidates = interps
.filter((interp) => interp.interpreterPath && interp.distroOrgName === 'ContinuumAnalytics')
.map((interp) => path.join(path.win32.dirname(interp.interpreterPath), suffix));
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -126,15 +126,15 @@ export function getRegistryInterpretersSync(): IRegistryInterpreterData[] | unde

let registryInterpretersPromise: Promise<IRegistryInterpreterData[]> | undefined;

export async function getRegistryInterpreters(useWorkerThreads: boolean): Promise<IRegistryInterpreterData[]> {
export async function getRegistryInterpreters(): Promise<IRegistryInterpreterData[]> {
if (!isTestExecution() && registryInterpretersPromise !== undefined) {
return registryInterpretersPromise;
}
registryInterpretersPromise = getRegistryInterpretersImpl(useWorkerThreads);
registryInterpretersPromise = getRegistryInterpretersImpl();
return registryInterpretersPromise;
}

async function getRegistryInterpretersImpl(useWorkerThreads: boolean): Promise<IRegistryInterpreterData[]> {
async function getRegistryInterpretersImpl(useWorkerThreads = false): Promise<IRegistryInterpreterData[]> {
let registryData: IRegistryInterpreterData[] = [];

for (const arch of ['x64', 'x86']) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,6 @@ import { PythonEnvsChangedEvent } from '../../../../../client/pythonEnvironments
import { EXTENSION_ROOT_DIR_FOR_TESTS, TEST_TIMEOUT } from '../../../../constants';
import { traceWarn } from '../../../../../client/logging';
import { TEST_LAYOUT_ROOT } from '../../../common/commonTestConstants';
import { getEnvs } from '../../common';
import { assertBasicEnvsEqual } from '../envTestUtils';
import { PYTHON_VIRTUAL_ENVS_LOCATION } from '../../../../ciConstants';
import { isCI } from '../../../../../client/common/constants';
import * as externalDependencies from '../../../../../client/pythonEnvironments/common/externalDependencies';
Expand Down Expand Up @@ -131,14 +129,4 @@ suite('Conda Env Locator', async () => {

assert.deepEqual(actualEvent!, expectedEvent, 'Unexpected event emitted');
});

test('Worker thread to fetch conda environments is working', async () => {
locator = new CondaEnvironmentLocator();
const items = await getEnvs(locator.doIterEnvs(undefined, false));
const workerItems = await getEnvs(locator.doIterEnvs(undefined, true));
console.log('Number of items Conda locator returned:', items.length);
// Make sure items returned when using worker threads v/s not are the same.
assertBasicEnvsEqual(items, workerItems);
assert(workerItems.length > 0, 'No environments found');
}).timeout(TEST_TIMEOUT * 2);
});

0 comments on commit 27a1f4d

Please sign in to comment.