Environment info cache class (#14065)

* Add persistent storage external deps
* PythonEnvInfoCache class + tests
* Instantiate & initialize cache class in createAPI
* Add extra test for flush() and initialize()
* Env cache fixes: storage key + find() result check
* Update src/client/pythonEnvironments/common/externalDependencies.ts
Co-authored-by: Kartik Raj <karraj@microsoft.com>
* Use areSameEnvironment in getEnv
* Don't ping persistent storage for every initialize
* No need to export CompleteEnvInfoFunction
* PythonEnvInfoCache doc comment
* Rename createGlobalPersistentStoreStub to get...
* Preemptively drop id key (#14051)
* Return deep copies
* IPersistentStore wrapper around IPersistentState
* Use correct areSameEnvironment + fix stub
* Remove obsolete comment
* getEnv -> filterEnvs
* Remove stubbing of areSameEnvironment
* Update areSameEnv
* Move IPersistentStateFactory registration to registerForIOC
* Revert "Move IPersistentStateFactory registration to registerForIOC"
This reverts commit edc6ce5.
* Don't instantiate nor initialize cache for now
Co-authored-by: Kartik Raj <karraj@microsoft.com>

src/client/pythonEnvironments/base/envsCache.ts

@@ -0,0 +1,100 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import { cloneDeep } from 'lodash';
+import { getGlobalPersistentStore, IPersistentStore } from '../common/externalDependencies';
+import { PythonEnvInfo } from './info';
+import { areSameEnv } from './info/env';
+
+/**
+ * Represents the environment info cache to be used by the cache locator.
+ */
+export interface IEnvsCache {
+    /**
+     * Initialization logic to be done outside of the constructor, for example reading from persistent storage.
+     */
+    initialize(): void;
+
+    /**
+     * Return all environment info currently in memory for this session.
+     *
+     * @return An array of cached environment info, or `undefined` if there are none.
+     */
+    getAllEnvs(): PythonEnvInfo[] | undefined;
+
+    /**
+     * Replace all environment info currently in memory for this session.
+     *
+     * @param envs The array of environment info to store in the in-memory cache.
+     */
+    setAllEnvs(envs: PythonEnvInfo[]): void;
+
+    /**
+     * If the cache has been initialized, return environmnent info objects that match a query object.
+     * If none of the environments in the cache match the query data, return an empty array.
+     * If the in-memory cache has not been initialized prior to calling `filterEnvs`, return `undefined`.
+     *
+     * @param env The environment info data that will be used to look for
+     * environment info objects in the cache, or a unique environment key.
+     * If passing an environment info object, it may contain incomplete environment info.
+     * @return The environment info objects matching the `env` param,
+     * or `undefined` if the in-memory cache is not initialized.
+     */
+    filterEnvs(env: PythonEnvInfo | string): PythonEnvInfo[] | undefined;
+
+    /**
+     * Writes the content of the in-memory cache to persistent storage.
+     */
+    flush(): Promise<void>;
+}
+
+type CompleteEnvInfoFunction = (envInfo: PythonEnvInfo) => boolean;
+
+/**
+ * Environment info cache using persistent storage to save and retrieve pre-cached env info.
+ */
+export class PythonEnvInfoCache implements IEnvsCache {
+    private initialized = false;
+
+    private envsList: PythonEnvInfo[] | undefined;
+
+    private persistentStorage: IPersistentStore<PythonEnvInfo[]> | undefined;
+
+    constructor(private readonly isComplete: CompleteEnvInfoFunction) {}
+
+    public initialize(): void {
+        if (this.initialized) {
+            return;
+        }
+
+        this.initialized = true;
+        this.persistentStorage = getGlobalPersistentStore<PythonEnvInfo[]>('PYTHON_ENV_INFO_CACHE');
+        this.envsList = this.persistentStorage?.get();
+    }
+
+    public getAllEnvs(): PythonEnvInfo[] | undefined {
+        return cloneDeep(this.envsList);
+    }
+
+    public setAllEnvs(envs: PythonEnvInfo[]): void {
+        this.envsList = cloneDeep(envs);
+    }
+
+    public filterEnvs(env: PythonEnvInfo | string): PythonEnvInfo[] | undefined {
+        const result = this.envsList?.filter((info) => areSameEnv(info, env));
+
+        if (result) {
+            return cloneDeep(result);
+        }
+
+        return undefined;
+    }
+
+    public async flush(): Promise<void> {
+        const completeEnvs = this.envsList?.filter(this.isComplete);
+
+        if (completeEnvs?.length) {
+            await this.persistentStorage?.set(completeEnvs);
+        }
+    }
+}

src/client/pythonEnvironments/common/externalDependencies.ts

@@ -4,6 +4,7 @@
 import * as fsapi from 'fs-extra';
 import * as path from 'path';
 import { ExecutionResult, IProcessServiceFactory } from '../../common/process/types';
+import { IPersistentStateFactory } from '../../common/types';
 import { getOSType, OSType } from '../../common/utils/platform';
 import { IServiceContainer } from '../../ioc/types';
 
@@ -37,3 +38,22 @@ export function arePathsSame(path1: string, path2: string): boolean {
     }
     return path1 === path2;
 }
+
+function getPersistentStateFactory(): IPersistentStateFactory {
+    return internalServiceContainer.get<IPersistentStateFactory>(IPersistentStateFactory);
+}
+
+export interface IPersistentStore<T> {
+    get(): T | undefined;
+    set(value: T): Promise<void>;
+}
+
+export function getGlobalPersistentStore<T>(key: string): IPersistentStore<T> {
+    const factory = getPersistentStateFactory();
+    const state = factory.createGlobalPersistentState<T>(key, undefined);
+
+    return {
+        get() { return state.value; },
+        set(value: T) { return state.updateValue(value); },
+    };
+}

src/client/pythonEnvironments/index.ts

@@ -55,7 +55,7 @@ export function createAPI(): [PythonEnvironments, () => void] {
         () => {
             activateLocators();
             // Any other activation needed for the API will go here later.
-        }
+        },
     ];
 }
 

src/test/pythonEnvironments/base/envsCache.unit.test.ts

@@ -0,0 +1,175 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import * as assert from 'assert';
+import * as sinon from 'sinon';
+import { PythonEnvInfoCache } from '../../../client/pythonEnvironments/base/envsCache';
+import { PythonEnvInfo, PythonEnvKind } from '../../../client/pythonEnvironments/base/info';
+import * as externalDependencies from '../../../client/pythonEnvironments/common/externalDependencies';
+
+suite('Environment Info cache', () => {
+    let getGlobalPersistentStoreStub: sinon.SinonStub;
+    let updatedValues: PythonEnvInfo[] | undefined;
+
+    const allEnvsComplete = () => true;
+    const envInfoArray = [
+        {
+            kind: PythonEnvKind.Conda, executable: { filename: 'my-conda-env' },
+        },
+        {
+            kind: PythonEnvKind.Venv, executable: { filename: 'my-venv-env' },
+        },
+        {
+            kind: PythonEnvKind.Pyenv, executable: { filename: 'my-pyenv-env' },
+        },
+    ] as PythonEnvInfo[];
+
+    setup(() => {
+        getGlobalPersistentStoreStub = sinon.stub(externalDependencies, 'getGlobalPersistentStore');
+        getGlobalPersistentStoreStub.returns({
+            get() { return envInfoArray; },
+            set(envs: PythonEnvInfo[]) {
+                updatedValues = envs;
+                return Promise.resolve();
+            },
+        });
+    });
+
+    teardown(() => {
+        getGlobalPersistentStoreStub.restore();
+        updatedValues = undefined;
+    });
+
+    test('`initialize` reads from persistent storage', () => {
+        const envsCache = new PythonEnvInfoCache(allEnvsComplete);
+
+        envsCache.initialize();
+
+        assert.ok(getGlobalPersistentStoreStub.calledOnce);
+    });
+
+    test('The in-memory env info array is undefined if there is no value in persistent storage when initializing the cache', () => {
+        const envsCache = new PythonEnvInfoCache(allEnvsComplete);
+
+        getGlobalPersistentStoreStub.returns({ get() { return undefined; } });
+        envsCache.initialize();
+        const result = envsCache.getAllEnvs();
+
+        assert.strictEqual(result, undefined);
+    });
+
+    test('`getAllEnvs` should return a deep copy of the environments currently in memory', () => {
+        const envsCache = new PythonEnvInfoCache(allEnvsComplete);
+
+        envsCache.initialize();
+        const envs = envsCache.getAllEnvs()!;
+
+        envs[0].name = 'some-other-name';
+
+        assert.ok(envs[0] !== envInfoArray[0]);
+    });
+
+    test('`getAllEnvs` should return undefined if nothing has been set', () => {
+        const envsCache = new PythonEnvInfoCache(allEnvsComplete);
+
+        const envs = envsCache.getAllEnvs();
+
+        assert.deepStrictEqual(envs, undefined);
+    });
+
+    test('`setAllEnvs` should clone the environment info array passed as a parameter', () => {
+        const envsCache = new PythonEnvInfoCache(allEnvsComplete);
+
+        envsCache.setAllEnvs(envInfoArray);
+        const envs = envsCache.getAllEnvs();
+
+        assert.deepStrictEqual(envs, envInfoArray);
+        assert.strictEqual(envs === envInfoArray, false);
+    });
+
+    test('`filterEnvs` should return environments that match its argument using areSameEnvironmnet', () => {
+        const env:PythonEnvInfo = { executable: { filename: 'my-venv-env' } } as unknown as PythonEnvInfo;
+        const envsCache = new PythonEnvInfoCache(allEnvsComplete);
+
+        envsCache.initialize();
+
+        const result = envsCache.filterEnvs(env);
+
+        assert.deepStrictEqual(result, [{
+            kind: PythonEnvKind.Venv, executable: { filename: 'my-venv-env' },
+        }]);
+    });
+
+    test('`filterEnvs` should return a deep copy of the matched environments', () => {
+        const envToFind = {
+            kind: PythonEnvKind.System, executable: { filename: 'my-system-env' },
+        } as unknown as PythonEnvInfo;
+        const env:PythonEnvInfo = { executable: { filename: 'my-system-env' } } as unknown as PythonEnvInfo;
+        const envsCache = new PythonEnvInfoCache(allEnvsComplete);
+
+        envsCache.setAllEnvs([...envInfoArray, envToFind]);
+
+        const result = envsCache.filterEnvs(env)!;
+        result[0].name = 'some-other-name';
+
+        assert.notDeepStrictEqual(result[0], envToFind);
+    });
+
+    test('`filterEnvs` should return an empty array if no environment matches the properties of its argument', () => {
+        const env:PythonEnvInfo = { executable: { filename: 'my-nonexistent-env' } } as unknown as PythonEnvInfo;
+        const envsCache = new PythonEnvInfoCache(allEnvsComplete);
+
+        envsCache.initialize();
+
+        const result = envsCache.filterEnvs(env);
+
+        assert.deepStrictEqual(result, []);
+    });
+
+    test('`filterEnvs` should return undefined if the cache hasn\'t been initialized', () => {
+        const env:PythonEnvInfo = { executable: { filename: 'my-nonexistent-env' } } as unknown as PythonEnvInfo;
+        const envsCache = new PythonEnvInfoCache(allEnvsComplete);
+
+        const result = envsCache.filterEnvs(env);
+
+        assert.strictEqual(result, undefined);
+    });
+
+    test('`flush` should write complete environment info objects to persistent storage', async () => {
+        const otherEnv = {
+            kind: PythonEnvKind.OtherGlobal,
+            executable: { filename: 'my-other-env' },
+            defaultDisplayName: 'other-env',
+        };
+        const updatedEnvInfoArray = [
+            otherEnv, { kind: PythonEnvKind.System, executable: { filename: 'my-system-env' } },
+        ] as PythonEnvInfo[];
+        const expected = [
+            otherEnv,
+        ];
+        const envsCache = new PythonEnvInfoCache((env) => env.defaultDisplayName !== undefined);
+
+        envsCache.initialize();
+        envsCache.setAllEnvs(updatedEnvInfoArray);
+        await envsCache.flush();
+
+        assert.deepStrictEqual(updatedValues, expected);
+    });
+
+    test('`flush` should not write to persistent storage if there are no environment info objects in-memory', async () => {
+        const envsCache = new PythonEnvInfoCache((env) => env.kind === PythonEnvKind.MacDefault);
+
+        await envsCache.flush();
+
+        assert.strictEqual(updatedValues, undefined);
+    });
+
+    test('`flush` should not write to persistent storage if there are no complete environment info objects', async () => {
+        const envsCache = new PythonEnvInfoCache((env) => env.kind === PythonEnvKind.MacDefault);
+
+        envsCache.initialize();
+        await envsCache.flush();
+
+        assert.strictEqual(updatedValues, undefined);
+    });
+});