Include windows registry perf improvement in pythonDiscoveryUsingWorkers experiment (#23075)

For #22146

Author: Kartik Raj

src/client/pythonEnvironments/base/locators/composite/resolverUtils.ts

@@ -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) {

src/client/pythonEnvironments/base/locators/lowLevel/condaLocator.ts

@@ -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;

src/client/pythonEnvironments/base/locators/lowLevel/windowsRegistryLocator.ts

@@ -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();
     }
 }
 
@@ -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.

src/client/pythonEnvironments/common/environmentManagers/conda.ts

@@ -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);
     }
@@ -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';
@@ -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));

src/client/pythonEnvironments/common/windowsUtils.ts

@@ -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']) {

src/test/pythonEnvironments/base/locators/lowLevel/condaLocator.testvirtualenvs.ts

@@ -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';
@@ -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);
 });