Fix duplicated environments in interpreter list (#22964)

Failing tests are due to
#22965, this can still
be reviewed.

Closes #22918 closes
#22146

Author: Kartik Raj

src/client/pythonEnvironments/base/locator.ts

@@ -20,7 +20,7 @@ export type PythonEnvUpdatedEvent<I = PythonEnvInfo> = {
     /**
      * The iteration index of The env info that was previously provided.
      */
-    index?: number;
+    index: number;
     /**
      * The env info that was previously provided.
      */
@@ -243,7 +243,7 @@ export interface IDiscoveryAPI {
     resolveEnv(path: string): Promise<PythonEnvInfo | undefined>;
 }
 
-interface IEmitter<E> {
+export interface IEmitter<E> {
     fire(e: E): void;
 }
 

src/client/pythonEnvironments/base/locatorUtils.ts

@@ -95,8 +95,6 @@ export async function getEnvs<I = PythonEnvInfo>(iterator: IPythonEnvsIterator<I
                 }
                 // We don't worry about if envs[index] is set already.
                 envs[index] = update;
-            } else if (event.update) {
-                envs.push(event.update);
             }
         });
     }

src/client/pythonEnvironments/base/locators/common/resourceBasedLocator.ts

@@ -43,18 +43,9 @@ export abstract class LazyResourceBasedLocator extends Locator<BasicEnvInfo> imp
         await this.disposables.dispose();
     }
 
-    public iterEnvs(query?: PythonLocatorQuery): IPythonEnvsIterator<BasicEnvInfo> {
-        const iterator = this.doIterEnvs(query);
-        const it = this._iterEnvs(iterator, query);
-        it.onUpdated = iterator.onUpdated;
-        return it;
-    }
-
-    private async *_iterEnvs(
-        iterator: IPythonEnvsIterator<BasicEnvInfo>,
-        query?: PythonLocatorQuery,
-    ): IPythonEnvsIterator<BasicEnvInfo> {
+    public async *iterEnvs(query?: PythonLocatorQuery): IPythonEnvsIterator<BasicEnvInfo> {
         await this.activate();
+        const iterator = this.doIterEnvs(query);
         if (query?.envPath) {
             let result = await iterator.next();
             while (!result.done) {

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

@@ -141,12 +141,12 @@ export class EnvsCollectionService extends PythonEnvsWatcher<PythonEnvCollection
         const updatesDone = createDeferred<void>();
 
         if (iterator.onUpdated !== undefined) {
-            iterator.onUpdated(async (event) => {
+            const listener = iterator.onUpdated(async (event) => {
                 if (isProgressEvent(event)) {
                     switch (event.stage) {
                         case ProgressReportStage.discoveryFinished:
                             state.done = true;
-                            // listener.dispose();
+                            listener.dispose();
                             break;
                         case ProgressReportStage.allPathsDiscovered:
                             if (!query) {
@@ -164,10 +164,6 @@ export class EnvsCollectionService extends PythonEnvsWatcher<PythonEnvCollection
                         seen[event.index] = event.update;
                     }
                     state.pending -= 1;
-                } else if (event.update) {
-                    // New env, add it to cache.
-                    seen.push(event.update);
-                    this.cache.addEnv(event.update);
                 }
                 if (state.done && state.pending === 0) {
                     updatesDone.resolve();

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

@@ -49,15 +49,12 @@ async function* iterEnvsIterator(
     };
     const seen: BasicEnvInfo[] = [];
 
-    didUpdate.fire({ stage: ProgressReportStage.discoveryStarted });
     if (iterator.onUpdated !== undefined) {
-        iterator.onUpdated((event) => {
+        const listener = iterator.onUpdated((event) => {
             if (isProgressEvent(event)) {
                 if (event.stage === ProgressReportStage.discoveryFinished) {
                     state.done = true;
-                    // For super slow locators such as Windows registry, we expect updates even after discovery
-                    // is "officially" finished, hence do not dispose listeners.
-                    // listener.dispose();
+                    listener.dispose();
                 } else {
                     didUpdate.fire(event);
                 }
@@ -69,11 +66,15 @@ async function* iterEnvsIterator(
                 const oldEnv = seen[event.index];
                 seen[event.index] = event.update;
                 didUpdate.fire({ index: event.index, old: oldEnv, update: event.update });
-            } else if (event.update) {
-                didUpdate.fire({ update: event.update });
+            } else {
+                // This implies a problem in a downstream locator
+                traceVerbose(`Expected already iterated env, got ${event.old} (#${event.index})`);
             }
+            state.pending -= 1;
             checkIfFinishedAndNotify(state, didUpdate);
         });
+    } else {
+        didUpdate.fire({ stage: ProgressReportStage.discoveryStarted });
     }
 
     let result = await iterator.next();
@@ -89,8 +90,10 @@ async function* iterEnvsIterator(
         }
         result = await iterator.next();
     }
-    state.done = true;
-    checkIfFinishedAndNotify(state, didUpdate);
+    if (iterator.onUpdated === undefined) {
+        state.done = true;
+        checkIfFinishedAndNotify(state, didUpdate);
+    }
 }
 
 async function resolveDifferencesInBackground(
@@ -124,8 +127,8 @@ function checkIfFinishedAndNotify(
 ) {
     if (state.done && state.pending === 0) {
         didUpdate.fire({ stage: ProgressReportStage.discoveryFinished });
+        didUpdate.dispose();
         traceVerbose(`Finished with environment reducer`);
-        state.done = false; // No need to notify again.
     }
 }
 

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

@@ -81,30 +81,29 @@ export class PythonEnvsResolver implements IResolvingLocator {
         const seen: PythonEnvInfo[] = [];
 
         if (iterator.onUpdated !== undefined) {
-            iterator.onUpdated(async (event) => {
+            const listener = iterator.onUpdated(async (event) => {
                 state.pending += 1;
                 if (isProgressEvent(event)) {
                     if (event.stage === ProgressReportStage.discoveryFinished) {
                         didUpdate.fire({ stage: ProgressReportStage.allPathsDiscovered });
                         state.done = true;
-                        // For super slow locators such as Windows registry, we expect updates even after discovery
-                        // is "officially" finished, hence do not dispose listeners.
-                        // listener.dispose();
+                        listener.dispose();
                     } else {
                         didUpdate.fire(event);
                     }
                 } else if (event.update === undefined) {
                     throw new Error(
                         'Unsupported behavior: `undefined` environment updates are not supported from downstream locators in resolver',
                     );
-                } else if (event.index && seen[event.index] !== undefined) {
+                } else if (event.index !== undefined && seen[event.index] !== undefined) {
                     const old = seen[event.index];
                     await setKind(event.update, environmentKinds);
                     seen[event.index] = await resolveBasicEnv(event.update);
                     didUpdate.fire({ old, index: event.index, update: seen[event.index] });
                     this.resolveInBackground(event.index, state, didUpdate, seen).ignoreErrors();
                 } else {
-                    didUpdate.fire({ update: await this.resolveEnv(event.update.executablePath) });
+                    // This implies a problem in a downstream locator
+                    traceVerbose(`Expected already iterated env, got ${event.old} (#${event.index})`);
                 }
                 state.pending -= 1;
                 checkIfFinishedAndNotify(state, didUpdate);
@@ -174,6 +173,7 @@ function checkIfFinishedAndNotify(
 ) {
     if (state.done && state.pending === 0) {
         didUpdate.fire({ stage: ProgressReportStage.discoveryFinished });
+        didUpdate.dispose();
         traceVerbose(`Finished with environment resolver`);
     }
 }

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

@@ -3,76 +3,53 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-import { EventEmitter } from 'vscode';
 import { PythonEnvKind, PythonEnvSource } from '../../info';
-import {
-    BasicEnvInfo,
-    IPythonEnvsIterator,
-    Locator,
-    ProgressNotificationEvent,
-    ProgressReportStage,
-    PythonEnvUpdatedEvent,
-} from '../../locator';
+import { BasicEnvInfo, IPythonEnvsIterator, Locator, PythonLocatorQuery, IEmitter } from '../../locator';
 import { getRegistryInterpreters } from '../../../common/windowsUtils';
 import { traceError, traceVerbose } from '../../../../logging';
 import { isMicrosoftStoreDir } from '../../../common/environmentManagers/microsoftStoreEnv';
 import { inExperiment } from '../../../common/externalDependencies';
 import { DiscoveryUsingWorkers } from '../../../../common/experiments/groups';
+import { PythonEnvsChangedEvent } from '../../watcher';
+
+export const WINDOWS_REG_PROVIDER_ID = 'windows-registry';
 
 export class WindowsRegistryLocator extends Locator<BasicEnvInfo> {
-    public readonly providerId: string = 'windows-registry';
+    public readonly providerId: string = WINDOWS_REG_PROVIDER_ID;
 
     // eslint-disable-next-line class-methods-use-this
     public iterEnvs(
-        _?: unknown,
+        query?: PythonLocatorQuery,
         useWorkerThreads = inExperiment(DiscoveryUsingWorkers.experiment),
     ): IPythonEnvsIterator<BasicEnvInfo> {
-        const didUpdate = new EventEmitter<PythonEnvUpdatedEvent<BasicEnvInfo> | ProgressNotificationEvent>();
-        const iterator = useWorkerThreads ? iterEnvsIterator(didUpdate) : oldIterEnvsIterator();
         if (useWorkerThreads) {
-            iterator.onUpdated = didUpdate.event;
+            /**
+             * 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 iterateEnvsLazily(this.emitter);
         }
-        return iterator;
+        return iterateEnvs(false);
     }
 }
 
-/**
- * Windows registry is slow and often not necessary, so notify completion immediately, while still updating lazily as we find stuff.
- * To accomplish this, use an empty iterator while lazily firing environments using updates.
- */
-async function* iterEnvsIterator(
-    didUpdate: EventEmitter<PythonEnvUpdatedEvent<BasicEnvInfo> | ProgressNotificationEvent>,
-): IPythonEnvsIterator<BasicEnvInfo> {
-    updateLazily(didUpdate).ignoreErrors();
+async function* iterateEnvsLazily(changed: IEmitter<PythonEnvsChangedEvent>): IPythonEnvsIterator<BasicEnvInfo> {
+    loadAllEnvs(changed).ignoreErrors();
 }
 
-async function updateLazily(didUpdate: EventEmitter<PythonEnvUpdatedEvent<BasicEnvInfo> | ProgressNotificationEvent>) {
+async function loadAllEnvs(changed: IEmitter<PythonEnvsChangedEvent>) {
     traceVerbose('Searching for windows registry interpreters');
-    const interpreters = await getRegistryInterpreters(true);
-    for (const interpreter of interpreters) {
-        try {
-            // Filter out Microsoft Store app directories. We have a store app locator that handles this.
-            // The python.exe available in these directories might not be python. It can be a store install
-            // shortcut that takes you to microsoft store.
-            if (isMicrosoftStoreDir(interpreter.interpreterPath)) {
-                continue;
-            }
-            const env: BasicEnvInfo = {
-                kind: PythonEnvKind.OtherGlobal,
-                executablePath: interpreter.interpreterPath,
-                source: [PythonEnvSource.WindowsRegistry],
-            };
-            didUpdate.fire({ update: env });
-        } catch (ex) {
-            traceError(`Failed to process environment: ${interpreter}`, ex);
-        }
-    }
-    didUpdate.fire({ stage: ProgressReportStage.discoveryFinished });
+    await getRegistryInterpreters(true);
+    changed.fire({ providerId: WINDOWS_REG_PROVIDER_ID });
     traceVerbose('Finished searching for windows registry interpreters');
 }
 
-async function* oldIterEnvsIterator(): IPythonEnvsIterator<BasicEnvInfo> {
-    const interpreters = await getRegistryInterpreters(false);
+async function* iterateEnvs(useWorkerThreads: boolean): IPythonEnvsIterator<BasicEnvInfo> {
+    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/test/pythonEnvironments/base/locators/lowLevel/windowsRegistryLocator.testvirtualenvs.ts

@@ -2,7 +2,10 @@
 // Licensed under the MIT License.
 
 import { getEnvs } from '../../../../../client/pythonEnvironments/base/locatorUtils';
-import { WindowsRegistryLocator } from '../../../../../client/pythonEnvironments/base/locators/lowLevel/windowsRegistryLocator';
+import {
+    WindowsRegistryLocator,
+    WINDOWS_REG_PROVIDER_ID,
+} from '../../../../../client/pythonEnvironments/base/locators/lowLevel/windowsRegistryLocator';
 import { assertBasicEnvsEqual } from '../envTestUtils';
 import { TEST_TIMEOUT } from '../../../../constants';
 import { getOSType, OSType } from '../../../../../client/common/utils/platform';
@@ -19,8 +22,8 @@ suite('Windows Registry Locator', async () => {
     });
 
     test('Worker thread to fetch registry interpreters is working', async () => {
-        const items = await getEnvs(locator.iterEnvs(undefined, false));
-        const workerItems = await getEnvs(locator.iterEnvs(undefined, true));
+        const items = await getEnvs(locator.iterEnvs({ providerId: WINDOWS_REG_PROVIDER_ID }, false));
+        const workerItems = await getEnvs(locator.iterEnvs({ providerId: WINDOWS_REG_PROVIDER_ID }, 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);

src/test/pythonEnvironments/base/locators/lowLevel/windowsRegistryLocator.unit.test.ts

@@ -4,10 +4,14 @@
 import * as assert from 'assert';
 import * as path from 'path';
 import * as sinon from 'sinon';
+import { expect } from 'chai';
 import { PythonEnvKind, PythonEnvSource } from '../../../../../client/pythonEnvironments/base/info';
 import { getEnvs } from '../../../../../client/pythonEnvironments/base/locatorUtils';
 import * as winreg from '../../../../../client/pythonEnvironments/common/windowsRegistry';
-import { WindowsRegistryLocator } from '../../../../../client/pythonEnvironments/base/locators/lowLevel/windowsRegistryLocator';
+import {
+    WindowsRegistryLocator,
+    WINDOWS_REG_PROVIDER_ID,
+} from '../../../../../client/pythonEnvironments/base/locators/lowLevel/windowsRegistryLocator';
 import { createBasicEnv } from '../../common';
 import { TEST_LAYOUT_ROOT } from '../../../common/commonTestConstants';
 import { assertBasicEnvsEqual } from '../envTestUtils';
@@ -201,7 +205,7 @@ suite('Windows Registry', () => {
     }
 
     setup(async () => {
-        sinon.stub(externalDependencies, 'inExperiment').returns(false);
+        sinon.stub(externalDependencies, 'inExperiment').returns(true);
         stubReadRegistryValues = sinon.stub(winreg, 'readRegistryValues');
         stubReadRegistryKeys = sinon.stub(winreg, 'readRegistryKeys');
         stubReadRegistryValues.callsFake(fakeRegistryValues);
@@ -222,18 +226,29 @@ suite('Windows Registry', () => {
             createBasicEnv(PythonEnvKind.OtherGlobal, path.join(regTestRoot, 'python38', 'python.exe')),
         ].map((e) => ({ ...e, source: [PythonEnvSource.WindowsRegistry] }));
 
-        const iterator = locator.iterEnvs(undefined, true);
+        const lazyIterator = locator.iterEnvs(undefined, true);
+        const envs = await getEnvs(lazyIterator);
+        expect(envs.length).to.equal(0);
+
+        const iterator = locator.iterEnvs({ providerId: WINDOWS_REG_PROVIDER_ID }, true);
         const actualEnvs = await getEnvs(iterator);
 
         assertBasicEnvsEqual(actualEnvs, expectedEnvs);
     });
 
+    test('iterEnvs(): query is undefined', async () => {
+        // Iterate no envs when query is `undefined`, i.e notify completion immediately.
+        const lazyIterator = locator.iterEnvs(undefined, true);
+        const envs = await getEnvs(lazyIterator);
+        expect(envs.length).to.equal(0);
+    });
+
     test('iterEnvs(): no registry permission', async () => {
         stubReadRegistryKeys.callsFake(() => {
             throw Error();
         });
 
-        const iterator = locator.iterEnvs(undefined, true);
+        const iterator = locator.iterEnvs({ providerId: WINDOWS_REG_PROVIDER_ID }, true);
         const actualEnvs = await getEnvs(iterator);
 
         assert.deepStrictEqual(actualEnvs, []);
@@ -252,7 +267,7 @@ suite('Windows Registry', () => {
             createBasicEnv(PythonEnvKind.OtherGlobal, path.join(regTestRoot, 'python38', 'python.exe')),
         ].map((e) => ({ ...e, source: [PythonEnvSource.WindowsRegistry] }));
 
-        const iterator = locator.iterEnvs(undefined, true);
+        const iterator = locator.iterEnvs({ providerId: WINDOWS_REG_PROVIDER_ID }, true);
         const actualEnvs = await getEnvs(iterator);
 
         assertBasicEnvsEqual(actualEnvs, expectedEnvs);