From 5d915057f93f49d9bbabfce33df12ac18404814a Mon Sep 17 00:00:00 2001 From: Kartik Raj Date: Wed, 15 Nov 2023 13:49:10 -0800 Subject: [PATCH 1/6] Use worker threads for fetching Windows Registry interpreters --- src/client/pythonEnvironments/base/locator.ts | 2 +- .../base/locators/composite/resolverUtils.ts | 2 +- .../lowLevel/windowsRegistryLocator.ts | 4 +- .../common/environmentManagers/conda.ts | 2 +- .../common/registryKeysWorker.ts | 24 +++++ .../common/registryValuesWorker.ts | 27 +++++ .../common/windowsRegistry.ts | 98 ++++++++++++++----- .../pythonEnvironments/common/windowsUtils.ts | 26 +++-- .../windowsRegistryLocator.testvirtualenvs.ts | 21 ++++ 9 files changed, 169 insertions(+), 37 deletions(-) create mode 100644 src/client/pythonEnvironments/common/registryKeysWorker.ts create mode 100644 src/client/pythonEnvironments/common/registryValuesWorker.ts create mode 100644 src/test/pythonEnvironments/base/locators/lowLevel/windowsRegistryLocator.testvirtualenvs.ts diff --git a/src/client/pythonEnvironments/base/locator.ts b/src/client/pythonEnvironments/base/locator.ts index ab3b17629bc5..8fefe8e2d0dd 100644 --- a/src/client/pythonEnvironments/base/locator.ts +++ b/src/client/pythonEnvironments/base/locator.ts @@ -178,7 +178,7 @@ export interface ILocator 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): IPythonEnvsIterator; + iterEnvs(query?: QueryForEvent, useWorkerThreads?: boolean): IPythonEnvsIterator; } export type ICompositeLocator = Omit, 'providerId'>; diff --git a/src/client/pythonEnvironments/base/locators/composite/resolverUtils.ts b/src/client/pythonEnvironments/base/locators/composite/resolverUtils.ts index 0cca49e2b4c5..fd45b4909aff 100644 --- a/src/client/pythonEnvironments/base/locators/composite/resolverUtils.ts +++ b/src/client/pythonEnvironments/base/locators/composite/resolverUtils.ts @@ -109,7 +109,7 @@ async function updateEnvUsingRegistry(env: PythonEnvInfo): Promise { 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) { diff --git a/src/client/pythonEnvironments/base/locators/lowLevel/windowsRegistryLocator.ts b/src/client/pythonEnvironments/base/locators/lowLevel/windowsRegistryLocator.ts index b52e9c35779f..947d90ee9cb5 100644 --- a/src/client/pythonEnvironments/base/locators/lowLevel/windowsRegistryLocator.ts +++ b/src/client/pythonEnvironments/base/locators/lowLevel/windowsRegistryLocator.ts @@ -12,10 +12,10 @@ export class WindowsRegistryLocator extends Locator { public readonly providerId: string = 'windows-registry'; // eslint-disable-next-line class-methods-use-this - public iterEnvs(): IPythonEnvsIterator { + public iterEnvs(_?: unknown, useWorkerThreads = true): IPythonEnvsIterator { 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. diff --git a/src/client/pythonEnvironments/common/environmentManagers/conda.ts b/src/client/pythonEnvironments/common/environmentManagers/conda.ts index 8f048ddd0676..842f81b99a4e 100644 --- a/src/client/pythonEnvironments/common/environmentManagers/conda.ts +++ b/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)); diff --git a/src/client/pythonEnvironments/common/registryKeysWorker.ts b/src/client/pythonEnvironments/common/registryKeysWorker.ts new file mode 100644 index 000000000000..05996d057f11 --- /dev/null +++ b/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 }); +}); diff --git a/src/client/pythonEnvironments/common/registryValuesWorker.ts b/src/client/pythonEnvironments/common/registryValuesWorker.ts new file mode 100644 index 000000000000..eaef7cbd58a7 --- /dev/null +++ b/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 }); +}); diff --git a/src/client/pythonEnvironments/common/windowsRegistry.ts b/src/client/pythonEnvironments/common/windowsRegistry.ts index 30047e1c9907..74b63f82d469 100644 --- a/src/client/pythonEnvironments/common/windowsRegistry.ts +++ b/src/client/pythonEnvironments/common/windowsRegistry.ts @@ -2,6 +2,8 @@ // Licensed under the MIT License. import { HKCU, HKLM, Options, REG_SZ, Registry, RegistryItem } from 'winreg'; +import * as path from 'path'; +import { Worker } from 'worker_threads'; import { createDeferred } from '../../common/utils/async'; export { HKCU, HKLM, REG_SZ, Options }; @@ -22,30 +24,82 @@ export interface IRegistryValue { value: string; } -export async function readRegistryValues(options: Options): Promise { - // eslint-disable-next-line global-require - const WinReg = require('winreg'); - const regKey = new WinReg(options); - const deferred = createDeferred(); - regKey.values((err: Error, res: RegistryItem[]) => { - if (err) { - deferred.reject(err); - } - deferred.resolve(res); +export async function readRegistryValues(options: Options, useWorkerThreads: boolean): Promise { + if (!useWorkerThreads) { + // eslint-disable-next-line global-require + const WinReg = require('winreg'); + const regKey = new WinReg(options); + const deferred = createDeferred(); + regKey.values((err: Error, res: RegistryItem[]) => { + if (err) { + deferred.reject(err); + } + deferred.resolve(res); + }); + return deferred.promise; + } + return new Promise((resolve, reject) => { + // Create a worker thread + const worker = new Worker(path.join(__dirname, 'registryValuesWorker.js'), { workerData: options }); + + // Listen for messages from the worker + worker.on('message', (res: { err: Error; res: IRegistryValue[] }) => { + if (res.err) { + reject(res.err); + } + resolve(res.res); + }); + + // Listen for any errors that occur + worker.on('error', (ex: Error) => { + reject(ex); + }); + + // Listen for exit if the worker thread exits prematurely + worker.on('exit', (code) => { + if (code !== 0) { + reject(new Error(`Worker stopped with exit code ${code}`)); + } + }); }); - return deferred.promise; } -export async function readRegistryKeys(options: Options): Promise { - // eslint-disable-next-line global-require - const WinReg = require('winreg'); - const regKey = new WinReg(options); - const deferred = createDeferred(); - regKey.keys((err: Error, res: Registry[]) => { - if (err) { - deferred.reject(err); - } - deferred.resolve(res); +export async function readRegistryKeys(options: Options, useWorkerThreads: boolean): Promise { + if (!useWorkerThreads) { + // eslint-disable-next-line global-require + const WinReg = require('winreg'); + const regKey = new WinReg(options); + const deferred = createDeferred(); + regKey.keys((err: Error, res: Registry[]) => { + if (err) { + deferred.reject(err); + } + deferred.resolve(res); + }); + return deferred.promise; + } + return new Promise((resolve, reject) => { + // Create a worker thread + const worker = new Worker(path.join(__dirname, 'registryKeysWorker.js'), { workerData: options }); + + // Listen for messages from the worker + worker.on('message', (res: { err: Error; res: Registry[] }) => { + if (res.err) { + reject(res.err); + } + resolve(res.res); + }); + + // Listen for any errors that occur + worker.on('error', (ex: Error) => { + reject(ex); + }); + + // Listen for exit if the worker thread exits prematurely + worker.on('exit', (code) => { + if (code !== 0) { + reject(new Error(`Worker stopped with exit code ${code}`)); + } + }); }); - return deferred.promise; } diff --git a/src/client/pythonEnvironments/common/windowsUtils.ts b/src/client/pythonEnvironments/common/windowsUtils.ts index a47025ceef6f..e9210371be05 100644 --- a/src/client/pythonEnvironments/common/windowsUtils.ts +++ b/src/client/pythonEnvironments/common/windowsUtils.ts @@ -54,13 +54,14 @@ export interface IRegistryInterpreterData { async function getInterpreterDataFromKey( { arch, hive, key }: IRegistryKey, distroOrgName: string, + useWorkerThreads: boolean, ): Promise { const result: IRegistryInterpreterData = { interpreterPath: '', distroOrgName, }; - const values: IRegistryValue[] = await readRegistryValues({ arch, hive, key }); + const values: IRegistryValue[] = await readRegistryValues({ arch, hive, key }, useWorkerThreads); for (const value of values) { switch (value.name) { case 'SysArchitecture': @@ -80,10 +81,10 @@ async function getInterpreterDataFromKey( } } - const subKeys: IRegistryKey[] = await readRegistryKeys({ arch, hive, key }); + const subKeys: IRegistryKey[] = await readRegistryKeys({ arch, hive, key }, useWorkerThreads); const subKey = subKeys.map((s) => s.key).find((s) => s.endsWith('InstallPath')); if (subKey) { - const subKeyValues: IRegistryValue[] = await readRegistryValues({ arch, hive, key: subKey }); + const subKeyValues: IRegistryValue[] = await readRegistryValues({ arch, hive, key: subKey }, useWorkerThreads); const value = subKeyValues.find((v) => v.name === 'ExecutablePath'); if (value) { result.interpreterPath = value.value; @@ -103,10 +104,13 @@ export async function getInterpreterDataFromRegistry( arch: string, hive: string, key: string, + useWorkerThreads: boolean, ): Promise { - const subKeys = await readRegistryKeys({ arch, hive, key }); + const subKeys = await readRegistryKeys({ arch, hive, key }, useWorkerThreads); const distroOrgName = key.substr(key.lastIndexOf('\\') + 1); - const allData = await Promise.all(subKeys.map((subKey) => getInterpreterDataFromKey(subKey, distroOrgName))); + const allData = await Promise.all( + subKeys.map((subKey) => getInterpreterDataFromKey(subKey, distroOrgName, useWorkerThreads)), + ); return (allData.filter((data) => data !== undefined) || []) as IRegistryInterpreterData[]; } @@ -122,15 +126,15 @@ export function getRegistryInterpretersSync(): IRegistryInterpreterData[] | unde let registryInterpretersPromise: Promise | undefined; -export async function getRegistryInterpreters(): Promise { +export async function getRegistryInterpreters(useWorkerThreads: boolean): Promise { if (!isTestExecution() && registryInterpretersPromise !== undefined) { return registryInterpretersPromise; } - registryInterpretersPromise = getRegistryInterpretersImpl(); + registryInterpretersPromise = getRegistryInterpretersImpl(useWorkerThreads); return registryInterpretersPromise; } -async function getRegistryInterpretersImpl(): Promise { +async function getRegistryInterpretersImpl(useWorkerThreads: boolean): Promise { let registryData: IRegistryInterpreterData[] = []; for (const arch of ['x64', 'x86']) { @@ -138,13 +142,15 @@ async function getRegistryInterpretersImpl(): Promise k.key); + keys = (await readRegistryKeys({ arch, hive, key: root }, useWorkerThreads)).map((k) => k.key); } catch (ex) { traceError(`Failed to access Registry: ${arch}\\${hive}\\${root}`, ex); } for (const key of keys) { - registryData = registryData.concat(await getInterpreterDataFromRegistry(arch, hive, key)); + registryData = registryData.concat( + await getInterpreterDataFromRegistry(arch, hive, key, useWorkerThreads), + ); } } } diff --git a/src/test/pythonEnvironments/base/locators/lowLevel/windowsRegistryLocator.testvirtualenvs.ts b/src/test/pythonEnvironments/base/locators/lowLevel/windowsRegistryLocator.testvirtualenvs.ts new file mode 100644 index 000000000000..2f0e050e439b --- /dev/null +++ b/src/test/pythonEnvironments/base/locators/lowLevel/windowsRegistryLocator.testvirtualenvs.ts @@ -0,0 +1,21 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +import { getEnvs } from '../../../../../client/pythonEnvironments/base/locatorUtils'; +import { WindowsRegistryLocator } from '../../../../../client/pythonEnvironments/base/locators/lowLevel/windowsRegistryLocator'; +import { assertBasicEnvsEqual } from '../envTestUtils'; +import { TEST_TIMEOUT } from '../../../../constants'; + +suite('Windows Registry Locator', async () => { + let locator: WindowsRegistryLocator; + + setup(() => { + locator = new WindowsRegistryLocator(); + }); + + test('Make sure worker thread to fetch environments is working', async () => { + const items = await getEnvs(locator.iterEnvs(undefined, false)); + const workerItems = await getEnvs(locator.iterEnvs(undefined, true)); + assertBasicEnvsEqual(items, workerItems); + }).timeout(TEST_TIMEOUT * 2); +}); From cc48a0dc818dcc3b813435ba2986623434d8abe8 Mon Sep 17 00:00:00 2001 From: Kartik Raj Date: Wed, 15 Nov 2023 14:00:58 -0800 Subject: [PATCH 2/6] Only run on windows --- .../lowLevel/windowsRegistryLocator.testvirtualenvs.ts | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/test/pythonEnvironments/base/locators/lowLevel/windowsRegistryLocator.testvirtualenvs.ts b/src/test/pythonEnvironments/base/locators/lowLevel/windowsRegistryLocator.testvirtualenvs.ts index 2f0e050e439b..c89af1b45027 100644 --- a/src/test/pythonEnvironments/base/locators/lowLevel/windowsRegistryLocator.testvirtualenvs.ts +++ b/src/test/pythonEnvironments/base/locators/lowLevel/windowsRegistryLocator.testvirtualenvs.ts @@ -5,17 +5,23 @@ import { getEnvs } from '../../../../../client/pythonEnvironments/base/locatorUt import { WindowsRegistryLocator } from '../../../../../client/pythonEnvironments/base/locators/lowLevel/windowsRegistryLocator'; import { assertBasicEnvsEqual } from '../envTestUtils'; import { TEST_TIMEOUT } from '../../../../constants'; +import { getOSType, OSType } from '../../../../../client/common/utils/platform'; suite('Windows Registry Locator', async () => { let locator: WindowsRegistryLocator; - setup(() => { + setup(function () { + if (getOSType() !== OSType.Windows) { + return this.skip(); + } locator = new WindowsRegistryLocator(); + return undefined; }); test('Make sure worker thread to fetch environments is working', async () => { const items = await getEnvs(locator.iterEnvs(undefined, false)); const workerItems = await getEnvs(locator.iterEnvs(undefined, true)); + // Make sure items returned when using worker threads v/s not are the same. assertBasicEnvsEqual(items, workerItems); }).timeout(TEST_TIMEOUT * 2); }); From cd6d37dfbb8b74115f3ce83d723040a7c8cfb578 Mon Sep 17 00:00:00 2001 From: Kartik Raj Date: Wed, 15 Nov 2023 14:31:17 -0800 Subject: [PATCH 3/6] Enable venv tests on Windows --- .github/workflows/pr-check.yml | 2 +- .../locators/lowLevel/windowsRegistryLocator.testvirtualenvs.ts | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/workflows/pr-check.yml b/.github/workflows/pr-check.yml index 4ade4fd2af1e..3283d3281e5b 100644 --- a/.github/workflows/pr-check.yml +++ b/.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: diff --git a/src/test/pythonEnvironments/base/locators/lowLevel/windowsRegistryLocator.testvirtualenvs.ts b/src/test/pythonEnvironments/base/locators/lowLevel/windowsRegistryLocator.testvirtualenvs.ts index c89af1b45027..4e4f4f629e85 100644 --- a/src/test/pythonEnvironments/base/locators/lowLevel/windowsRegistryLocator.testvirtualenvs.ts +++ b/src/test/pythonEnvironments/base/locators/lowLevel/windowsRegistryLocator.testvirtualenvs.ts @@ -21,6 +21,7 @@ suite('Windows Registry Locator', async () => { test('Make sure worker thread to fetch environments is working', async () => { const items = await getEnvs(locator.iterEnvs(undefined, false)); const workerItems = await getEnvs(locator.iterEnvs(undefined, true)); + console.log('Number of items Windows registry locator returned', items.length); // Make sure items returned when using worker threads v/s not are the same. assertBasicEnvsEqual(items, workerItems); }).timeout(TEST_TIMEOUT * 2); From b57161592aba9606dbfe2f8f0af715d07c8e8861 Mon Sep 17 00:00:00 2001 From: Kartik Raj Date: Wed, 15 Nov 2023 15:00:41 -0800 Subject: [PATCH 4/6] Add description --- package.json | 4 ++++ package.nls.json | 1 + src/client/common/experiments/groups.ts | 4 ++++ .../pythonEnvironments/base/locators/wrappers.ts | 11 ++++++++++- .../pythonEnvironments/common/externalDependencies.ts | 7 ++++++- .../windowsRegistryLocator.testvirtualenvs.ts | 2 +- 6 files changed, 26 insertions(+), 3 deletions(-) diff --git a/package.json b/package.json index 5c25b2b84832..068f04114cae 100644 --- a/package.json +++ b/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%" ] diff --git a/package.nls.json b/package.nls.json index 71b146b5fac2..ab1eec6d3065 100644 --- a/package.nls.json +++ b/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.", diff --git a/src/client/common/experiments/groups.ts b/src/client/common/experiments/groups.ts index 8f8ecc631caf..5e302c926ccb 100644 --- a/src/client/common/experiments/groups.ts +++ b/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', diff --git a/src/client/pythonEnvironments/base/locators/wrappers.ts b/src/client/pythonEnvironments/base/locators/wrappers.ts index bfaede584f6f..1334c60fa2e0 100644 --- a/src/client/pythonEnvironments/base/locators/wrappers.ts +++ b/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 extends Locators { + private readonly useWorkerThreads: boolean; + constructor( // These are expected to be low-level locators (e.g. system). private readonly nonWorkspace: ILocator[], @@ -25,6 +30,10 @@ export class ExtensionLocators extends Locators { private readonly workspace: ILocator, ) { super([...nonWorkspace, workspace]); + this.useWorkerThreads = inExperiment(DiscoveryUsingWorkers.experiment); + if (this.useWorkerThreads) { + traceLog('Using worker threads for discovery...'); + } } public iterEnvs(query?: PythonLocatorQuery): IPythonEnvsIterator { @@ -33,7 +42,7 @@ export class ExtensionLocators extends Locators { 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); } diff --git a/src/client/pythonEnvironments/common/externalDependencies.ts b/src/client/pythonEnvironments/common/externalDependencies.ts index b64d47f42269..6eb525b92e31 100644 --- a/src/client/pythonEnvironments/common/externalDependencies.ts +++ b/src/client/pythonEnvironments/common/externalDependencies.ts @@ -6,7 +6,7 @@ import * as path from 'path'; 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 +29,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); + return service.inExperimentSync(experimentName); +} + // Workspace export function isVirtualWorkspace(): boolean { diff --git a/src/test/pythonEnvironments/base/locators/lowLevel/windowsRegistryLocator.testvirtualenvs.ts b/src/test/pythonEnvironments/base/locators/lowLevel/windowsRegistryLocator.testvirtualenvs.ts index 4e4f4f629e85..6f0ff0e1bb5e 100644 --- a/src/test/pythonEnvironments/base/locators/lowLevel/windowsRegistryLocator.testvirtualenvs.ts +++ b/src/test/pythonEnvironments/base/locators/lowLevel/windowsRegistryLocator.testvirtualenvs.ts @@ -21,7 +21,7 @@ suite('Windows Registry Locator', async () => { test('Make sure worker thread to fetch environments is working', async () => { const items = await getEnvs(locator.iterEnvs(undefined, false)); const workerItems = await getEnvs(locator.iterEnvs(undefined, true)); - console.log('Number of items Windows registry locator returned', items.length); + console.log('Number of items Windows registry locator returned:', items.length); // Make sure items returned when using worker threads v/s not are the same. assertBasicEnvsEqual(items, workerItems); }).timeout(TEST_TIMEOUT * 2); From ba02c03ed8e2229511a79fd85f9ecdb95f5d485b Mon Sep 17 00:00:00 2001 From: Kartik Raj Date: Wed, 15 Nov 2023 15:12:52 -0800 Subject: [PATCH 5/6] Clean up --- .../common/windowsRegistry.ts | 45 +++++-------------- 1 file changed, 11 insertions(+), 34 deletions(-) diff --git a/src/client/pythonEnvironments/common/windowsRegistry.ts b/src/client/pythonEnvironments/common/windowsRegistry.ts index 74b63f82d469..0fc5f078192a 100644 --- a/src/client/pythonEnvironments/common/windowsRegistry.ts +++ b/src/client/pythonEnvironments/common/windowsRegistry.ts @@ -1,3 +1,4 @@ +/* eslint-disable @typescript-eslint/no-explicit-any */ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. @@ -5,6 +6,7 @@ import { HKCU, HKLM, Options, REG_SZ, Registry, RegistryItem } from 'winreg'; import * as path from 'path'; import { Worker } from 'worker_threads'; import { createDeferred } from '../../common/utils/async'; +import { traceError } from '../../logging'; export { HKCU, HKLM, REG_SZ, Options }; @@ -38,30 +40,7 @@ export async function readRegistryValues(options: Options, useWorkerThreads: boo }); return deferred.promise; } - return new Promise((resolve, reject) => { - // Create a worker thread - const worker = new Worker(path.join(__dirname, 'registryValuesWorker.js'), { workerData: options }); - - // Listen for messages from the worker - worker.on('message', (res: { err: Error; res: IRegistryValue[] }) => { - if (res.err) { - reject(res.err); - } - resolve(res.res); - }); - - // Listen for any errors that occur - worker.on('error', (ex: Error) => { - reject(ex); - }); - - // Listen for exit if the worker thread exits prematurely - worker.on('exit', (code) => { - if (code !== 0) { - reject(new Error(`Worker stopped with exit code ${code}`)); - } - }); - }); + return executeWorkerFile('registryValuesWorker.js', options); } export async function readRegistryKeys(options: Options, useWorkerThreads: boolean): Promise { @@ -78,27 +57,25 @@ export async function readRegistryKeys(options: Options, useWorkerThreads: boole }); return deferred.promise; } - return new Promise((resolve, reject) => { - // Create a worker thread - const worker = new Worker(path.join(__dirname, 'registryKeysWorker.js'), { workerData: options }); + return executeWorkerFile('registryKeysWorker.js', options); +} - // Listen for messages from the worker - worker.on('message', (res: { err: Error; res: Registry[] }) => { +async function executeWorkerFile(workerFileName: string, workerData: any): Promise { + return new Promise((resolve, reject) => { + const worker = new Worker(path.join(__dirname, workerFileName), { workerData }); + worker.on('message', (res: { err: Error; res: unknown }) => { if (res.err) { reject(res.err); } resolve(res.res); }); - - // Listen for any errors that occur worker.on('error', (ex: Error) => { + traceError(`Error in worker ${workerFileName}`, ex); reject(ex); }); - - // Listen for exit if the worker thread exits prematurely worker.on('exit', (code) => { if (code !== 0) { - reject(new Error(`Worker stopped with exit code ${code}`)); + reject(new Error(`Worker ${workerFileName} stopped with exit code ${code}`)); } }); }); From e56aa22934e61a8a9673abb2e47f083d74d73792 Mon Sep 17 00:00:00 2001 From: Kartik Raj Date: Wed, 15 Nov 2023 15:32:56 -0800 Subject: [PATCH 6/6] Cleanup --- .../common/externalDependencies.ts | 23 +++++++++++++++ .../common/windowsRegistry.ts | 28 ++----------------- 2 files changed, 26 insertions(+), 25 deletions(-) diff --git a/src/client/pythonEnvironments/common/externalDependencies.ts b/src/client/pythonEnvironments/common/externalDependencies.ts index 6eb525b92e31..357967236c9e 100644 --- a/src/client/pythonEnvironments/common/externalDependencies.ts +++ b/src/client/pythonEnvironments/common/externalDependencies.ts @@ -3,6 +3,7 @@ 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'; @@ -200,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 { + 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}`)); + } + }); + }); +} diff --git a/src/client/pythonEnvironments/common/windowsRegistry.ts b/src/client/pythonEnvironments/common/windowsRegistry.ts index 0fc5f078192a..efac5bb3209f 100644 --- a/src/client/pythonEnvironments/common/windowsRegistry.ts +++ b/src/client/pythonEnvironments/common/windowsRegistry.ts @@ -4,9 +4,8 @@ import { HKCU, HKLM, Options, REG_SZ, Registry, RegistryItem } from 'winreg'; import * as path from 'path'; -import { Worker } from 'worker_threads'; import { createDeferred } from '../../common/utils/async'; -import { traceError } from '../../logging'; +import { executeWorkerFile } from './externalDependencies'; export { HKCU, HKLM, REG_SZ, Options }; @@ -40,7 +39,7 @@ export async function readRegistryValues(options: Options, useWorkerThreads: boo }); return deferred.promise; } - return executeWorkerFile('registryValuesWorker.js', options); + return executeWorkerFile(path.join(__dirname, 'registryValuesWorker.js'), options); } export async function readRegistryKeys(options: Options, useWorkerThreads: boolean): Promise { @@ -57,26 +56,5 @@ export async function readRegistryKeys(options: Options, useWorkerThreads: boole }); return deferred.promise; } - return executeWorkerFile('registryKeysWorker.js', options); -} - -async function executeWorkerFile(workerFileName: string, workerData: any): Promise { - return new Promise((resolve, reject) => { - const worker = new Worker(path.join(__dirname, 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}`)); - } - }); - }); + return executeWorkerFile(path.join(__dirname, 'registryKeysWorker.js'), options); }