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/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/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/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/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/externalDependencies.ts b/src/client/pythonEnvironments/common/externalDependencies.ts index b64d47f42269..357967236c9e 100644 --- a/src/client/pythonEnvironments/common/externalDependencies.ts +++ b/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); + 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 { + 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/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..efac5bb3209f 100644 --- a/src/client/pythonEnvironments/common/windowsRegistry.ts +++ b/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 { - // 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; +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 executeWorkerFile(path.join(__dirname, 'registryValuesWorker.js'), options); } -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); - }); - return deferred.promise; +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 executeWorkerFile(path.join(__dirname, 'registryKeysWorker.js'), options); } 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..6f0ff0e1bb5e --- /dev/null +++ b/src/test/pythonEnvironments/base/locators/lowLevel/windowsRegistryLocator.testvirtualenvs.ts @@ -0,0 +1,28 @@ +// 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'; +import { getOSType, OSType } from '../../../../../client/common/utils/platform'; + +suite('Windows Registry Locator', async () => { + let locator: WindowsRegistryLocator; + + 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)); + 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); +});