Add an optional onUpdated event to the iterator returned by ILocator.iterEnvs(). (#13950)

In order to ensure that ILocator.iterEnvs() can finish as fast as possible, we add a side-channel event for each iteration. This event fires any time an already-yielded env info object is updated (e.g. reduced/merged or resolved/completed). The update only relates to operations triggered by that particular iteration. ILocator.onChanged remains separate and only relates to when the locators finds a new env, notices one was removed, or that one was otherwise fundamentally changed.

There are two similar approaches we could take for this update event. Either we added the event as a property of IPythonEnvsIterator or we change the return value of ILocator.iterEnvs() to be a 2-tuple ([PythonEnvsIterator, Event<PythonEnvUpdatedEvent>]). We took the property approach since most of the time callers won't need to worry about the update events.

Author: ericsnowcurrently

src/client/pythonEnvironments/base/locator.ts

@@ -4,17 +4,68 @@
 import { Event, Uri } from 'vscode';
 import { iterEmpty } from '../../common/utils/async';
 import { PythonEnvInfo, PythonEnvKind } from './info';
-import { BasicPythonEnvsChangedEvent, IPythonEnvsWatcher, PythonEnvsChangedEvent, PythonEnvsWatcher } from './watcher';
+import {
+    BasicPythonEnvsChangedEvent,
+    IPythonEnvsWatcher,
+    PythonEnvsChangedEvent,
+    PythonEnvsWatcher
+} from './watcher';
 
 /**
- * An async iterator of `PythonEnvInfo`.
+ * A single update to a previously provided Python env object.
  */
-export type PythonEnvsIterator = AsyncIterator<PythonEnvInfo, void>;
+export type PythonEnvUpdatedEvent = {
+    /**
+     * The env info that was previously provided.
+     *
+     * If the event comes from `IPythonEnvsIterator.onUpdated` then
+     * `old` was previously yielded during iteration.
+     */
+    old: PythonEnvInfo;
+    /**
+     * The env info that replaces the old info.
+     */
+    new: PythonEnvInfo;
+};
+
+/**
+ * A fast async iterator of Python envs, which may have incomplete info.
+ *
+ * Each object yielded by the iterator represents a unique Python
+ * environment.
+ *
+ * The iterator is not required to have provide all info about
+ * an environment.  However, each yielded item will at least
+ * include all the `PythonEnvBaseInfo` data.
+ *
+ * During iteration the information for an already
+ * yielded object may be updated.  Rather than updating the yielded
+ * object or yielding it again with updated info, the update is
+ * emitted by the iterator's `onUpdated` (event) property. Once there are no more updates, the event emits
+ * `null`.
+ *
+ * If the iterator does not have `onUpdated` then it means the
+ * provider does not support updates.
+ *
+ * Callers can usually ignore the update event entirely and rely on
+ * the locator to provide sufficiently complete information.
+ */
+export interface IPythonEnvsIterator extends AsyncIterator<PythonEnvInfo, void> {
+    /**
+     * Provides possible updates for already-iterated envs.
+     *
+     * Once there are no more updates, `null` is emitted.
+     *
+     * If this property is not provided then it means the iterator does
+     * not support updates.
+     */
+    onUpdated?: Event<PythonEnvUpdatedEvent | null>;
+}
 
 /**
  * An empty Python envs iterator.
  */
-export const NOOP_ITERATOR: PythonEnvsIterator = iterEmpty<PythonEnvInfo>();
+export const NOOP_ITERATOR: IPythonEnvsIterator = iterEmpty<PythonEnvInfo>();
 
 /**
  * The most basic info to send to a locator when requesting environments.
@@ -64,11 +115,18 @@ export interface ILocator<E extends BasicPythonEnvsChangedEvent = PythonEnvsChan
      *
      * Locators are not required to have provide all info about
      * an environment.  However, each yielded item will at least
-     * include all the `PythonEnvBaseInfo` data.
+     * include all the `PythonEnvBaseInfo` data.  To ensure all
+     * possible information is filled in, call `ILocator.resolveEnv()`.
+     *
+     * Updates to yielded objects may be provided via the optional
+     * `onUpdated` property of the iterator.  However, callers can
+     * usually ignore the update event entirely and rely on the
+     * locator to provide sufficiently complete information.
      *
      * @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>): PythonEnvsIterator;
+    iterEnvs(query?: QueryForEvent<E>): IPythonEnvsIterator;
 
     /**
      * Find the given Python environment and fill in as much missing info as possible.
@@ -111,7 +169,7 @@ export abstract class LocatorBase<E extends BasicPythonEnvsChangedEvent = Python
         this.onChanged = watcher.onChanged;
     }
 
-    public abstract iterEnvs(query?: QueryForEvent<E>): PythonEnvsIterator;
+    public abstract iterEnvs(query?: QueryForEvent<E>): IPythonEnvsIterator;
 
     public async resolveEnv(_env: string | PythonEnvInfo): Promise<PythonEnvInfo | undefined> {
         return undefined;

src/client/pythonEnvironments/base/locators.ts

@@ -1,11 +1,48 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+import { EventEmitter } from 'vscode';
 import { chain } from '../../common/utils/async';
 import { PythonEnvInfo } from './info';
-import { ILocator, NOOP_ITERATOR, PythonEnvsIterator, PythonLocatorQuery } from './locator';
+import {
+    ILocator,
+    IPythonEnvsIterator,
+    NOOP_ITERATOR,
+    PythonEnvUpdatedEvent,
+    PythonLocatorQuery
+} from './locator';
 import { DisableableEnvsWatcher, PythonEnvsWatchers } from './watchers';
 
+/**
+ * Combine the `onUpdated` event of the given iterators into a single event.
+ */
+export function combineIterators(iterators: IPythonEnvsIterator[]): IPythonEnvsIterator {
+    const result: IPythonEnvsIterator = chain(iterators);
+    const events = iterators.map((it) => it.onUpdated).filter((v) => v);
+    if (!events || events.length === 0) {
+        // There are no sub-events, so we leave `onUpdated` undefined.
+        return result;
+    }
+
+    const emitter = new EventEmitter<PythonEnvUpdatedEvent | null>();
+    let numActive = events.length;
+    events.forEach((event) => {
+        event!((e: PythonEnvUpdatedEvent | null) => {
+            if (e === null) {
+                numActive -= 1;
+                if (numActive === 0) {
+                    // All the sub-events are done so we're done.
+                    emitter.fire(null);
+                }
+            } else {
+                emitter.fire(e);
+            }
+        });
+    });
+    result.onUpdated = emitter.event;
+    return result;
+}
+
 /**
  * A wrapper around a set of locators, exposing them as a single locator.
  *
@@ -19,9 +56,9 @@ export class Locators extends PythonEnvsWatchers implements ILocator {
         super(locators);
     }
 
-    public iterEnvs(query?: PythonLocatorQuery): PythonEnvsIterator {
+    public iterEnvs(query?: PythonLocatorQuery): IPythonEnvsIterator {
         const iterators = this.locators.map((loc) => loc.iterEnvs(query));
-        return chain(iterators);
+        return combineIterators(iterators);
     }
 
     public async resolveEnv(env: string | PythonEnvInfo): Promise<PythonEnvInfo | undefined> {
@@ -50,7 +87,7 @@ export class DisableableLocator extends DisableableEnvsWatcher implements ILocat
         super(locator);
     }
 
-    public iterEnvs(query?: PythonLocatorQuery): PythonEnvsIterator {
+    public iterEnvs(query?: PythonLocatorQuery): IPythonEnvsIterator {
         if (!this.enabled) {
             return NOOP_ITERATOR;
         }

src/client/pythonEnvironments/discovery/locators/index.ts

@@ -5,7 +5,7 @@ import {
 import { traceDecorators } from '../../../common/logger';
 import { IPlatformService } from '../../../common/platform/types';
 import { IDisposableRegistry } from '../../../common/types';
-import { chain, createDeferred, Deferred } from '../../../common/utils/async';
+import { createDeferred, Deferred } from '../../../common/utils/async';
 import { OSType } from '../../../common/utils/platform';
 import {
     CONDA_ENV_FILE_SERVICE,
@@ -21,8 +21,18 @@ import {
 } from '../../../interpreter/contracts';
 import { IServiceContainer } from '../../../ioc/types';
 import { PythonEnvInfo } from '../../base/info';
-import { ILocator, Locator, NOOP_ITERATOR, PythonEnvsIterator, PythonLocatorQuery } from '../../base/locator';
-import { DisableableLocator, Locators } from '../../base/locators';
+import {
+    ILocator,
+    IPythonEnvsIterator,
+    Locator,
+    NOOP_ITERATOR,
+    PythonLocatorQuery,
+} from '../../base/locator';
+import {
+    combineIterators,
+    DisableableLocator,
+    Locators,
+} from '../../base/locators';
 import { PythonEnvironment } from '../../info';
 import { isHiddenInterpreter } from './services/interpreterFilter';
 import { GetInterpreterLocatorOptions } from './types';
@@ -83,7 +93,7 @@ export class WorkspaceLocators extends Locator {
         folders.onRemoved((root: Uri) => this.removeRoot(root));
     }
 
-    public iterEnvs(query?: PythonLocatorQuery): PythonEnvsIterator {
+    public iterEnvs(query?: PythonLocatorQuery): IPythonEnvsIterator {
         const iterators = Object.keys(this.locators).map((key) => {
             if (query?.searchLocations) {
                 const root = this.roots[key];
@@ -95,7 +105,7 @@ export class WorkspaceLocators extends Locator {
             const locator = this.locators[key];
             return locator.iterEnvs(query);
         });
-        return chain(iterators);
+        return combineIterators(iterators);
     }
 
     public async resolveEnv(env: string | PythonEnvInfo): Promise<PythonEnvInfo | undefined> {

src/client/pythonEnvironments/index.ts

@@ -4,7 +4,7 @@
 import * as vscode from 'vscode';
 import { IServiceContainer, IServiceManager } from '../ioc/types';
 import { PythonEnvInfo } from './base/info';
-import { ILocator, PythonEnvsIterator, PythonLocatorQuery } from './base/locator';
+import { ILocator, IPythonEnvsIterator, PythonLocatorQuery } from './base/locator';
 import { PythonEnvsChangedEvent } from './base/watcher';
 import { ExtensionLocators, WorkspaceLocators } from './discovery/locators';
 import { registerForIOC } from './legacyIOC';
@@ -33,7 +33,7 @@ export class PythonEnvironments implements ILocator {
         return this.locators.onChanged;
     }
 
-    public iterEnvs(query?: PythonLocatorQuery): PythonEnvsIterator {
+    public iterEnvs(query?: PythonLocatorQuery): IPythonEnvsIterator {
         return this.locators.iterEnvs(query);
     }
 

src/test/pythonEnvironments/base/common.ts

@@ -10,7 +10,7 @@ import {
     PythonReleaseLevel,
     PythonVersion
 } from '../../../client/pythonEnvironments/base/info';
-import { Locator, PythonEnvsIterator, PythonLocatorQuery } from '../../../client/pythonEnvironments/base/locator';
+import { IPythonEnvsIterator, Locator, PythonLocatorQuery } from '../../../client/pythonEnvironments/base/locator';
 import { PythonEnvsChangedEvent } from '../../../client/pythonEnvironments/base/watcher';
 
 export function createEnv(
@@ -111,7 +111,7 @@ export class SimpleLocator extends Locator {
     public fire(event: PythonEnvsChangedEvent) {
         this.emitter.fire(event);
     }
-    public iterEnvs(query?: PythonLocatorQuery): PythonEnvsIterator {
+    public iterEnvs(query?: PythonLocatorQuery): IPythonEnvsIterator {
         const deferred = this.deferred;
         const callbacks = this.callbacks;
         let envs = this.envs;
@@ -161,6 +161,6 @@ export class SimpleLocator extends Locator {
     }
 }
 
-export async function getEnvs(iterator: PythonEnvsIterator): Promise<PythonEnvInfo[]> {
+export async function getEnvs(iterator: IPythonEnvsIterator): Promise<PythonEnvInfo[]> {
     return flattenIterator(iterator);
 }