Use worker threads for fetching Windows Registry interpreters (#22479)

For #22146

Author: Kartik Raj

.github/workflows/pr-check.yml

@@ -293,7 +293,7 @@ jobs:
         with:
           run: npm run testSingleWorkspace
           working-directory: ${{ env.special-working-directory }}
-        if: matrix.test-suite == 'venv' && matrix.os == 'ubuntu-latest'
+        if: matrix.test-suite == 'venv'
 
       - name: Run single-workspace tests
         env:

package.json

@@ -538,6 +538,7 @@
                             "pythonSurveyNotification",
                             "pythonPromptNewToolsExt",
                             "pythonTerminalEnvVarActivation",
+                            "pythonDiscoveryUsingWorkers",
                             "pythonTestAdapter",
                             "pythonREPLSmartSend",
                             "pythonRecommendTensorboardExt"
@@ -547,6 +548,7 @@
                             "%python.experiments.pythonSurveyNotification.description%",
                             "%python.experiments.pythonPromptNewToolsExt.description%",
                             "%python.experiments.pythonTerminalEnvVarActivation.description%",
+                            "%python.experiments.pythonDiscoveryUsingWorkers.description%",
                             "%python.experiments.pythonTestAdapter.description%",
                             "%python.experiments.pythonREPLSmartSend.description%",
                             "%python.experiments.pythonRecommendTensorboardExt.description%"
@@ -565,6 +567,7 @@
                             "pythonSurveyNotification",
                             "pythonPromptNewToolsExt",
                             "pythonTerminalEnvVarActivation",
+                            "pythonDiscoveryUsingWorkers",
                             "pythonTestAdapter",
                             "pythonREPLSmartSend"
                         ],
@@ -573,6 +576,7 @@
                             "%python.experiments.pythonSurveyNotification.description%",
                             "%python.experiments.pythonPromptNewToolsExt.description%",
                             "%python.experiments.pythonTerminalEnvVarActivation.description%",
+                            "%python.experiments.pythonDiscoveryUsingWorkers.description%",
                             "%python.experiments.pythonTestAdapter.description%",
                             "%python.experiments.pythonREPLSmartSend.description%"
                         ]

package.nls.json

@@ -40,6 +40,7 @@
     "python.experiments.pythonSurveyNotification.description": "Denotes the Python Survey Notification experiment.",
     "python.experiments.pythonPromptNewToolsExt.description": "Denotes the Python Prompt New Tools Extension experiment.",
     "python.experiments.pythonTerminalEnvVarActivation.description": "Enables use of environment variables to activate terminals instead of sending activation commands.",
+    "python.experiments.pythonDiscoveryUsingWorkers.description": "Enables use of worker threads to do heavy computation when discovering interpreters.",
     "python.experiments.pythonTestAdapter.description": "Denotes the Python Test Adapter experiment.",
     "python.experiments.pythonREPLSmartSend.description": "Denotes the Python REPL Smart Send experiment.",
     "python.experiments.pythonRecommendTensorboardExt.description": "Denotes the Tensorboard Extension recommendation experiment.",

src/client/common/experiments/groups.ts

@@ -11,6 +11,10 @@ export enum TerminalEnvVarActivation {
     experiment = 'pythonTerminalEnvVarActivation',
 }
 
+export enum DiscoveryUsingWorkers {
+    experiment = 'pythonDiscoveryUsingWorkers',
+}
+
 // Experiment to enable the new testing rewrite.
 export enum EnableTestAdapterRewrite {
     experiment = 'pythonTestAdapter',

src/client/pythonEnvironments/base/locator.ts

@@ -178,7 +178,7 @@ export interface ILocator<I = PythonEnvInfo, E = PythonEnvsChangedEvent> extends
      * @param query - if provided, the locator will limit results to match
      * @returns - the fast async iterator of Python envs, which may have incomplete info
      */
-    iterEnvs(query?: QueryForEvent<E>): IPythonEnvsIterator<I>;
+    iterEnvs(query?: QueryForEvent<E>, useWorkerThreads?: boolean): IPythonEnvsIterator<I>;
 }
 
 export type ICompositeLocator<I = PythonEnvInfo, E = PythonEnvsChangedEvent> = Omit<ILocator<I, E>, 'providerId'>;

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();
+        interpreters = await getRegistryInterpreters(true);
     }
     const data = interpreters.find((i) => arePathsSame(i.interpreterPath, env.executable.filename));
     if (data) {

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

@@ -12,10 +12,10 @@ export class WindowsRegistryLocator extends Locator<BasicEnvInfo> {
     public readonly providerId: string = 'windows-registry';
 
     // eslint-disable-next-line class-methods-use-this
-    public iterEnvs(): IPythonEnvsIterator<BasicEnvInfo> {
+    public iterEnvs(_?: unknown, useWorkerThreads = true): IPythonEnvsIterator<BasicEnvInfo> {
         const iterator = async function* () {
             traceVerbose('Searching for windows registry interpreters');
-            const interpreters = await getRegistryInterpreters();
+            const interpreters = await getRegistryInterpreters(useWorkerThreads);
             for (const interpreter of interpreters) {
                 try {
                     // Filter out Microsoft Store app directories. We have a store app locator that handles this.

src/client/pythonEnvironments/base/locators/wrappers.ts

@@ -3,10 +3,13 @@
 
 // eslint-disable-next-line max-classes-per-file
 import { Uri } from 'vscode';
+import { DiscoveryUsingWorkers } from '../../../common/experiments/groups';
 import { IDisposable } from '../../../common/types';
 import { iterEmpty } from '../../../common/utils/async';
 import { getURIFilter } from '../../../common/utils/misc';
 import { Disposables } from '../../../common/utils/resourceLifecycle';
+import { traceLog } from '../../../logging';
+import { inExperiment } from '../../common/externalDependencies';
 import { PythonEnvInfo } from '../info';
 import { BasicEnvInfo, ILocator, IPythonEnvsIterator, PythonLocatorQuery } from '../locator';
 import { combineIterators, Locators } from '../locators';
@@ -17,6 +20,8 @@ import { LazyResourceBasedLocator } from './common/resourceBasedLocator';
  */
 
 export class ExtensionLocators<I = PythonEnvInfo> extends Locators<I> {
+    private readonly useWorkerThreads: boolean;
+
     constructor(
         // These are expected to be low-level locators (e.g. system).
         private readonly nonWorkspace: ILocator<I>[],
@@ -25,6 +30,10 @@ export class ExtensionLocators<I = PythonEnvInfo> extends Locators<I> {
         private readonly workspace: ILocator<I>,
     ) {
         super([...nonWorkspace, workspace]);
+        this.useWorkerThreads = inExperiment(DiscoveryUsingWorkers.experiment);
+        if (this.useWorkerThreads) {
+            traceLog('Using worker threads for discovery...');
+        }
     }
 
     public iterEnvs(query?: PythonLocatorQuery): IPythonEnvsIterator<I> {
@@ -33,7 +42,7 @@ export class ExtensionLocators<I = PythonEnvInfo> extends Locators<I> {
             const nonWorkspace = query?.providerId
                 ? this.nonWorkspace.filter((locator) => query.providerId === locator.providerId)
                 : this.nonWorkspace;
-            iterators.push(...nonWorkspace.map((loc) => loc.iterEnvs(query)));
+            iterators.push(...nonWorkspace.map((loc) => loc.iterEnvs(query, this.useWorkerThreads)));
         }
         return combineIterators(iterators);
     }

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

@@ -307,7 +307,7 @@ export class Conda {
         }
 
         async function* getCandidatesFromRegistry() {
-            const interps = await getRegistryInterpreters();
+            const interps = await getRegistryInterpreters(true);
             const candidates = interps
                 .filter((interp) => interp.interpreterPath && interp.distroOrgName === 'ContinuumAnalytics')
                 .map((interp) => path.join(path.win32.dirname(interp.interpreterPath), suffix));

src/client/pythonEnvironments/common/externalDependencies.ts

@@ -3,10 +3,11 @@
 
 import * as fsapi from 'fs-extra';
 import * as path from 'path';
+import { Worker } from 'worker_threads';
 import * as vscode from 'vscode';
 import { IWorkspaceService } from '../../common/application/types';
 import { ExecutionResult, IProcessServiceFactory, ShellOptions, SpawnOptions } from '../../common/process/types';
-import { IDisposable, IConfigurationService } from '../../common/types';
+import { IDisposable, IConfigurationService, IExperimentService } from '../../common/types';
 import { chain, iterable } from '../../common/utils/async';
 import { getOSType, OSType } from '../../common/utils/platform';
 import { IServiceContainer } from '../../ioc/types';
@@ -29,6 +30,11 @@ export async function exec(file: string, args: string[], options: SpawnOptions =
     return service.exec(file, args, options);
 }
 
+export function inExperiment(experimentName: string): boolean {
+    const service = internalServiceContainer.get<IExperimentService>(IExperimentService);
+    return service.inExperimentSync(experimentName);
+}
+
 // Workspace
 
 export function isVirtualWorkspace(): boolean {
@@ -195,3 +201,25 @@ export function onDidChangePythonSetting(name: string, callback: () => void, roo
         }
     });
 }
+
+// eslint-disable-next-line @typescript-eslint/no-explicit-any, @typescript-eslint/explicit-module-boundary-types
+export async function executeWorkerFile(workerFileName: string, workerData: any): Promise<any> {
+    return new Promise((resolve, reject) => {
+        const worker = new Worker(workerFileName, { workerData });
+        worker.on('message', (res: { err: Error; res: unknown }) => {
+            if (res.err) {
+                reject(res.err);
+            }
+            resolve(res.res);
+        });
+        worker.on('error', (ex: Error) => {
+            traceError(`Error in worker ${workerFileName}`, ex);
+            reject(ex);
+        });
+        worker.on('exit', (code) => {
+            if (code !== 0) {
+                reject(new Error(`Worker ${workerFileName} stopped with exit code ${code}`));
+            }
+        });
+    });
+}

src/client/pythonEnvironments/common/registryKeysWorker.ts

@@ -0,0 +1,24 @@
+import { Registry } from 'winreg';
+import { parentPort, workerData } from 'worker_threads';
+import { IRegistryKey } from './windowsRegistry';
+
+const WinReg = require('winreg');
+
+const regKey = new WinReg(workerData);
+
+function copyRegistryKeys(keys: IRegistryKey[]): IRegistryKey[] {
+    // Use the map function to create a new array with copies of the specified properties.
+    return keys.map((key) => ({
+        hive: key.hive,
+        arch: key.arch,
+        key: key.key,
+    }));
+}
+
+regKey.keys((err: Error, res: Registry[]) => {
+    if (!parentPort) {
+        throw new Error('Not in a worker thread');
+    }
+    const messageRes = copyRegistryKeys(res);
+    parentPort.postMessage({ err, res: messageRes });
+});

src/client/pythonEnvironments/common/registryValuesWorker.ts

@@ -0,0 +1,27 @@
+import { RegistryItem } from 'winreg';
+import { parentPort, workerData } from 'worker_threads';
+import { IRegistryValue } from './windowsRegistry';
+
+const WinReg = require('winreg');
+
+const regKey = new WinReg(workerData);
+
+function copyRegistryValues(values: IRegistryValue[]): IRegistryValue[] {
+    // Use the map function to create a new array with copies of the specified properties.
+    return values.map((value) => ({
+        hive: value.hive,
+        arch: value.arch,
+        key: value.key,
+        name: value.name,
+        type: value.type,
+        value: value.value,
+    }));
+}
+
+regKey.values((err: Error, res: RegistryItem[]) => {
+    if (!parentPort) {
+        throw new Error('Not in a worker thread');
+    }
+    const messageRes = copyRegistryValues(res);
+    parentPort.postMessage({ err, res: messageRes });
+});

src/client/pythonEnvironments/common/windowsRegistry.ts

@@ -1,8 +1,11 @@
+/* eslint-disable @typescript-eslint/no-explicit-any */
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
 import { HKCU, HKLM, Options, REG_SZ, Registry, RegistryItem } from 'winreg';
+import * as path from 'path';
 import { createDeferred } from '../../common/utils/async';
+import { executeWorkerFile } from './externalDependencies';
 
 export { HKCU, HKLM, REG_SZ, Options };
 
@@ -22,30 +25,36 @@ export interface IRegistryValue {
     value: string;
 }
 
-export async function readRegistryValues(options: Options): Promise<IRegistryValue[]> {
-    // eslint-disable-next-line global-require
-    const WinReg = require('winreg');
-    const regKey = new WinReg(options);
-    const deferred = createDeferred<RegistryItem[]>();
-    regKey.values((err: Error, res: RegistryItem[]) => {
-        if (err) {
-            deferred.reject(err);
-        }
-        deferred.resolve(res);
-    });
-    return deferred.promise;
+export async function readRegistryValues(options: Options, useWorkerThreads: boolean): Promise<IRegistryValue[]> {
+    if (!useWorkerThreads) {
+        // eslint-disable-next-line global-require
+        const WinReg = require('winreg');
+        const regKey = new WinReg(options);
+        const deferred = createDeferred<RegistryItem[]>();
+        regKey.values((err: Error, res: RegistryItem[]) => {
+            if (err) {
+                deferred.reject(err);
+            }
+            deferred.resolve(res);
+        });
+        return deferred.promise;
+    }
+    return executeWorkerFile(path.join(__dirname, 'registryValuesWorker.js'), options);
 }
 
-export async function readRegistryKeys(options: Options): Promise<IRegistryKey[]> {
-    // eslint-disable-next-line global-require
-    const WinReg = require('winreg');
-    const regKey = new WinReg(options);
-    const deferred = createDeferred<Registry[]>();
-    regKey.keys((err: Error, res: Registry[]) => {
-        if (err) {
-            deferred.reject(err);
-        }
-        deferred.resolve(res);
-    });
-    return deferred.promise;
+export async function readRegistryKeys(options: Options, useWorkerThreads: boolean): Promise<IRegistryKey[]> {
+    if (!useWorkerThreads) {
+        // eslint-disable-next-line global-require
+        const WinReg = require('winreg');
+        const regKey = new WinReg(options);
+        const deferred = createDeferred<Registry[]>();
+        regKey.keys((err: Error, res: Registry[]) => {
+            if (err) {
+                deferred.reject(err);
+            }
+            deferred.resolve(res);
+        });
+        return deferred.promise;
+    }
+    return executeWorkerFile(path.join(__dirname, 'registryKeysWorker.js'), options);
 }