Make environments.known API faster (microsoft/vscode-python#23010)

Closes microsoft/vscode-python#22994 closes
microsoft/vscode-python#22993

Temporarily build our own known cache from which we return envs instead.

Author: Kartik Raj

extensions/positron-python/src/client/environmentApi.ts

@@ -32,6 +32,7 @@ import {
     Resource,
 } from './api/types';
 import { buildEnvironmentCreationApi } from './pythonEnvironments/creation/createEnvApi';
+import { EnvironmentKnownCache } from './environmentKnownCache';
 
 type ActiveEnvironmentChangeEvent = {
     resource: WorkspaceFolder | undefined;
@@ -120,6 +121,15 @@ export function buildEnvironmentApi(
     const disposables = serviceContainer.get<IDisposableRegistry>(IDisposableRegistry);
     const extensions = serviceContainer.get<IExtensions>(IExtensions);
     const envVarsProvider = serviceContainer.get<IEnvironmentVariablesProvider>(IEnvironmentVariablesProvider);
+    let knownCache: EnvironmentKnownCache;
+
+    function initKnownCache() {
+        const knownEnvs = discoveryApi
+            .getEnvs()
+            .filter((e) => filterUsingVSCodeContext(e))
+            .map((e) => updateReference(e));
+        return new EnvironmentKnownCache(knownEnvs);
+    }
     function sendApiTelemetry(apiName: string, args?: unknown) {
         extensions
             .determineExtensionFromCallStack()
@@ -139,19 +149,26 @@ export function buildEnvironmentApi(
                 // Filter out environments that are not in the current workspace.
                 return;
             }
+            if (!knownCache) {
+                knownCache = initKnownCache();
+            }
             if (e.old) {
                 if (e.new) {
+                    const newEnv = updateReference(e.new);
+                    knownCache.updateEnv(convertEnvInfo(e.old), newEnv);
                     traceVerbose('Python API env change detected', env.id, 'update');
-                    onEnvironmentsChanged.fire({ type: 'update', env: convertEnvInfoAndGetReference(e.new) });
+                    onEnvironmentsChanged.fire({ type: 'update', env: newEnv });
                     reportInterpretersChanged([
                         {
                             path: getEnvPath(e.new.executable.filename, e.new.location).path,
                             type: 'update',
                         },
                     ]);
                 } else {
+                    const oldEnv = updateReference(e.old);
+                    knownCache.updateEnv(oldEnv, undefined);
                     traceVerbose('Python API env change detected', env.id, 'remove');
-                    onEnvironmentsChanged.fire({ type: 'remove', env: convertEnvInfoAndGetReference(e.old) });
+                    onEnvironmentsChanged.fire({ type: 'remove', env: oldEnv });
                     reportInterpretersChanged([
                         {
                             path: getEnvPath(e.old.executable.filename, e.old.location).path,
@@ -160,8 +177,10 @@ export function buildEnvironmentApi(
                     ]);
                 }
             } else if (e.new) {
+                const newEnv = updateReference(e.new);
+                knownCache.addEnv(newEnv);
                 traceVerbose('Python API env change detected', env.id, 'add');
-                onEnvironmentsChanged.fire({ type: 'add', env: convertEnvInfoAndGetReference(e.new) });
+                onEnvironmentsChanged.fire({ type: 'add', env: newEnv });
                 reportInterpretersChanged([
                     {
                         path: getEnvPath(e.new.executable.filename, e.new.location).path,
@@ -179,6 +198,9 @@ export function buildEnvironmentApi(
         onEnvironmentsChanged,
         onEnvironmentVariablesChanged,
     );
+    if (!knownCache!) {
+        knownCache = initKnownCache();
+    }
 
     const environmentApi: PythonExtension['environments'] = {
         getEnvironmentVariables: (resource?: Resource) => {
@@ -234,11 +256,9 @@ export function buildEnvironmentApi(
             return resolveEnvironment(path, discoveryApi);
         },
         get known(): Environment[] {
-            sendApiTelemetry('known');
-            return discoveryApi
-                .getEnvs()
-                .filter((e) => filterUsingVSCodeContext(e))
-                .map((e) => convertEnvInfoAndGetReference(e));
+            // Do not send telemetry for "known", as this may be called 1000s of times so it can significant:
+            // sendApiTelemetry('known');
+            return knownCache.envs;
         },
         async refreshEnvironments(options?: RefreshOptions) {
             if (!workspace.isTrusted) {
@@ -351,7 +371,7 @@ export function convertEnvInfo(env: PythonEnvInfo): Environment {
     return convertedEnv as Environment;
 }
 
-function convertEnvInfoAndGetReference(env: PythonEnvInfo): Environment {
+function updateReference(env: PythonEnvInfo): Environment {
     return getEnvReference(convertEnvInfo(env));
 }
 

extensions/positron-python/src/client/environmentKnownCache.ts

@@ -0,0 +1,37 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import { Environment } from './api/types';
+
+/**
+ * Workaround temp cache until types are consolidated.
+ */
+export class EnvironmentKnownCache {
+    private _envs: Environment[] = [];
+
+    constructor(envs: Environment[]) {
+        this._envs = envs;
+    }
+
+    public get envs(): Environment[] {
+        return this._envs;
+    }
+
+    public addEnv(env: Environment): void {
+        const found = this._envs.find((e) => env.id === e.id);
+        if (!found) {
+            this._envs.push(env);
+        }
+    }
+
+    public updateEnv(oldValue: Environment, newValue: Environment | undefined): void {
+        const index = this._envs.findIndex((e) => oldValue.id === e.id);
+        if (index !== -1) {
+            if (newValue === undefined) {
+                this._envs.splice(index, 1);
+            } else {
+                this._envs[index] = newValue;
+            }
+        }
+    }
+}

extensions/positron-python/src/test/api.functional.test.ts

@@ -37,6 +37,7 @@ suite('Extension API', () => {
         interpreterService = mock(InterpreterService);
         environmentVariablesProvider = mock<IEnvironmentVariablesProvider>();
         discoverAPI = mock<IDiscoveryAPI>();
+        when(discoverAPI.getEnvs()).thenReturn([]);
 
         when(serviceContainer.get<IConfigurationService>(IConfigurationService)).thenReturn(
             instance(configurationService),

extensions/positron-python/src/test/environmentApi.unit.test.ts

@@ -74,8 +74,7 @@ suite('Python Environment API', () => {
         envVarsProvider = typemoq.Mock.ofType<IEnvironmentVariablesProvider>();
         extensions
             .setup((e) => e.determineExtensionFromCallStack())
-            .returns(() => Promise.resolve({ extensionId: 'id', displayName: 'displayName', apiName: 'apiName' }))
-            .verifiable(typemoq.Times.atLeastOnce());
+            .returns(() => Promise.resolve({ extensionId: 'id', displayName: 'displayName', apiName: 'apiName' }));
         interpreterPathService = typemoq.Mock.ofType<IInterpreterPathService>();
         configService = typemoq.Mock.ofType<IConfigurationService>();
         onDidChangeRefreshState = new EventEmitter();
@@ -94,13 +93,12 @@ suite('Python Environment API', () => {
 
         discoverAPI.setup((d) => d.onProgress).returns(() => onDidChangeRefreshState.event);
         discoverAPI.setup((d) => d.onChanged).returns(() => onDidChangeEnvironments.event);
+        discoverAPI.setup((d) => d.getEnvs()).returns(() => []);
 
         environmentApi = buildEnvironmentApi(discoverAPI.object, serviceContainer.object);
     });
 
     teardown(() => {
-        // Verify each API method sends telemetry regarding who called the API.
-        extensions.verifyAll();
         sinon.restore();
     });
 
@@ -325,6 +323,7 @@ suite('Python Environment API', () => {
             },
         ];
         discoverAPI.setup((d) => d.getEnvs()).returns(() => envs);
+        environmentApi = buildEnvironmentApi(discoverAPI.object, serviceContainer.object);
         const actual = environmentApi.known;
         const actualEnvs = actual?.map((a) => (a as EnvironmentReference).internal);
         assert.deepEqual(
@@ -409,10 +408,10 @@ suite('Python Environment API', () => {
         // Update events
         events = [];
         expectedEvents = [];
-        const updatedEnv = cloneDeep(envs[0]);
-        updatedEnv.arch = Architecture.x86;
-        onDidChangeEnvironments.fire({ old: envs[0], new: updatedEnv });
-        expectedEvents.push({ env: convertEnvInfo(updatedEnv), type: 'update' });
+        const updatedEnv0 = cloneDeep(envs[0]);
+        updatedEnv0.arch = Architecture.x86;
+        onDidChangeEnvironments.fire({ old: envs[0], new: updatedEnv0 });
+        expectedEvents.push({ env: convertEnvInfo(updatedEnv0), type: 'update' });
         eventValues = events.map((e) => ({ env: (e.env as EnvironmentReference).internal, type: e.type }));
         assert.deepEqual(eventValues, expectedEvents);
 
@@ -423,6 +422,11 @@ suite('Python Environment API', () => {
         expectedEvents.push({ env: convertEnvInfo(envs[2]), type: 'remove' });
         eventValues = events.map((e) => ({ env: (e.env as EnvironmentReference).internal, type: e.type }));
         assert.deepEqual(eventValues, expectedEvents);
+
+        const expectedEnvs = [convertEnvInfo(updatedEnv0), convertEnvInfo(envs[1])].sort();
+        const knownEnvs = environmentApi.known.map((e) => (e as EnvironmentReference).internal).sort();
+
+        assert.deepEqual(expectedEnvs, knownEnvs);
     });
 
     test('updateActiveEnvironmentPath: no resource', async () => {