Make poetry discovery faster (#15860)

* Add caching to command output in discovery component

* Add doc comments

* Fix tests

* Cache for lesser expiry duration after startup

* Get python executables using `getInterpreterPathFromDir` instead of finding interpreters in all directories

* Update doc comment

Author: Kartik Raj

src/client/common/utils/decorators.ts

@@ -4,6 +4,7 @@ import { traceError, traceVerbose } from '../logger';
 import { createDeferred, Deferred } from './async';
 import { getCacheKeyFromFunctionArgs, getGlobalCacheStore } from './cacheUtils';
 import { TraceInfo, tracing } from './misc';
+import { StopWatch } from './stopWatch';
 
 const _debounce = require('lodash/debounce') as typeof import('lodash/debounce');
 
@@ -121,7 +122,25 @@ export function makeDebounceAsyncDecorator(wait?: number) {
 
 type PromiseFunctionWithAnyArgs = (...any: any) => Promise<any>;
 const cacheStoreForMethods = getGlobalCacheStore();
-export function cache(expiryDurationMs: number) {
+
+/**
+ * Extension start up time is considered the duration until extension is likely to keep running commands in background.
+ * It is observed on CI it can take upto 3 minutes, so this is an intelligent guess.
+ */
+const extensionStartUpTime = 200_000;
+/**
+ * Tracks the time since the module was loaded. For caching purposes, we consider this time to approximately signify
+ * how long extension has been active.
+ */
+const moduleLoadWatch = new StopWatch();
+/**
+ * Caches function value until a specific duration.
+ * @param expiryDurationMs Duration to cache the result for. If set as '-1', the cache will never expire for the session.
+ * @param cachePromise If true, cache the promise instead of the promise result.
+ * @param expiryDurationAfterStartUpMs If specified, this is the duration to cache the result for after extension startup (until extension is likely to
+ * keep running commands in background)
+ */
+export function cache(expiryDurationMs: number, cachePromise = false, expiryDurationAfterStartUpMs?: number) {
     return function (
         target: Object,
         propertyName: string,
@@ -136,16 +155,22 @@ export function cache(expiryDurationMs: number) {
             }
             const key = getCacheKeyFromFunctionArgs(keyPrefix, args);
             const cachedItem = cacheStoreForMethods.get(key);
-            if (cachedItem && cachedItem.expiry > Date.now()) {
+            if (cachedItem && (cachedItem.expiry > Date.now() || expiryDurationMs === -1)) {
                 traceVerbose(`Cached data exists ${key}`);
                 return Promise.resolve(cachedItem.data);
             }
+            const expiryMs =
+                expiryDurationAfterStartUpMs && moduleLoadWatch.elapsedTime > extensionStartUpTime
+                    ? expiryDurationAfterStartUpMs
+                    : expiryDurationMs;
             const promise = originalMethod.apply(this, args) as Promise<any>;
-            promise
-                .then((result) =>
-                    cacheStoreForMethods.set(key, { data: result, expiry: Date.now() + expiryDurationMs }),
-                )
-                .ignoreErrors();
+            if (cachePromise) {
+                cacheStoreForMethods.set(key, { data: promise, expiry: Date.now() + expiryMs });
+            } else {
+                promise
+                    .then((result) => cacheStoreForMethods.set(key, { data: result, expiry: Date.now() + expiryMs }))
+                    .ignoreErrors();
+            }
             return promise;
         };
     };

src/client/pythonEnvironments/common/externalDependencies.ts

@@ -61,10 +61,18 @@ export function pathExists(absPath: string): Promise<boolean> {
     return fsapi.pathExists(absPath);
 }
 
+export function pathExistsSync(absPath: string): boolean {
+    return fsapi.pathExistsSync(absPath);
+}
+
 export function readFile(filePath: string): Promise<string> {
     return fsapi.readFile(filePath, 'utf-8');
 }
 
+export function readFileSync(filePath: string): string {
+    return fsapi.readFileSync(filePath, 'utf-8');
+}
+
 /**
  * Returns true if given file path exists within the given parent directory, false otherwise.
  * @param filePath File path to check for

src/client/pythonEnvironments/discovery/locators/services/conda.ts

@@ -10,6 +10,7 @@ import { parseVersion } from '../../../base/info/pythonVersion';
 import { getRegistryInterpreters } from '../../../common/windowsUtils';
 import { EnvironmentType, PythonEnvironment } from '../../../info';
 import { IDisposable } from '../../../../common/types';
+import { cache } from '../../../../common/utils/decorators';
 
 export const AnacondaCompanyNames = ['Anaconda, Inc.', 'Continuum Analytics, Inc.'];
 
@@ -347,6 +348,7 @@ export class Conda {
      * Retrieves list of Python environments known to this conda.
      * Corresponds to "conda env list --json", but also computes environment names.
      */
+    @cache(30_000, true, 10_000)
     public async getEnvList(): Promise<CondaEnvInfo[]> {
         const info = await this.getInfo();
         const { envs } = info;

src/client/pythonEnvironments/discovery/locators/services/poetry.ts

@@ -9,13 +9,14 @@ import { getOSType, getUserHomeDir, OSType } from '../../../../common/utils/plat
 import {
     getPythonSetting,
     isParentPath,
-    pathExists,
-    readFile,
+    pathExistsSync,
+    readFileSync,
     shellExecute,
 } from '../../../common/externalDependencies';
 import { getEnvironmentDirFromPath } from '../../../common/commonUtils';
 import { isVirtualenvEnvironment } from './virtualEnvironmentIdentifier';
 import { StopWatch } from '../../../../common/utils/stopWatch';
+import { cache } from '../../../../common/utils/decorators';
 
 /**
  * Global virtual env dir for a project is named as:
@@ -59,7 +60,7 @@ async function isLocalPoetryEnvironment(interpreterPath: string): Promise<boolea
         return false;
     }
     const project = path.dirname(envDir);
-    if (!(await hasValidPyprojectToml(project))) {
+    if (!hasValidPyprojectToml(project)) {
         return false;
     }
     // The assumption is that we need to be able to run poetry CLI for an environment in order to mark it as poetry.
@@ -111,8 +112,16 @@ export class Poetry {
         this.fixCwd();
     }
 
+    /**
+     * Returns a Poetry instance corresponding to the binary which can be used to run commands for the cwd.
+     *
+     * Poetry commands can be slow and so can be bottleneck to overall discovery time. So trigger command
+     * execution as soon as possible. To do that we need to ensure the operations before the command are
+     * performed synchronously.
+     */
     public static async getPoetry(cwd: string): Promise<Poetry | undefined> {
-        if (!(await hasValidPyprojectToml(cwd))) {
+        // Following check should be performed synchronously so we trigger poetry execution as soon as possible.
+        if (!hasValidPyprojectToml(cwd)) {
             // This check is not expensive and may change during a session, so we need not cache it.
             return undefined;
         }
@@ -123,12 +132,12 @@ export class Poetry {
         return Poetry._poetryPromise.get(cwd);
     }
 
-    /**
-     * Returns a Poetry instance corresponding to the binary.
-     */
     private static async locate(cwd: string): Promise<Poetry | undefined> {
+        // First thing this method awaits on should be poetry command execution, hence perform all operations
+        // before that synchronously.
+
         // Produce a list of candidate binaries to be probed by exec'ing them.
-        async function* getCandidates() {
+        function* getCandidates() {
             const customPoetryPath = getPythonSetting<string>('poetryPath');
             if (customPoetryPath && customPoetryPath !== 'poetry') {
                 // If user has specified a custom poetry path, use it first.
@@ -139,26 +148,21 @@ export class Poetry {
             const home = getUserHomeDir();
             if (home) {
                 const defaultPoetryPath = path.join(home, '.poetry', 'bin', 'poetry');
-                if (await pathExists(defaultPoetryPath)) {
+                if (pathExistsSync(defaultPoetryPath)) {
                     yield defaultPoetryPath;
                 }
             }
         }
 
         // Probe the candidates, and pick the first one that exists and does what we need.
-        for await (const poetryPath of getCandidates()) {
+        for (const poetryPath of getCandidates()) {
             traceVerbose(`Probing poetry binary for ${cwd}: ${poetryPath}`);
             const poetry = new Poetry(poetryPath, cwd);
-            const stopWatch = new StopWatch();
             const virtualenvs = await poetry.getEnvList();
             if (virtualenvs !== undefined) {
                 traceVerbose(`Found poetry via filesystem probing for ${cwd}: ${poetryPath}`);
                 return poetry;
             }
-            traceVerbose(
-                `Time taken to run ${poetryPath} env list --full-path in ms for ${cwd}`,
-                stopWatch.elapsedTime,
-            );
         }
 
         // Didn't find anything.
@@ -176,6 +180,15 @@ export class Poetry {
      * Corresponds to "poetry env list --full-path". Swallows errors if any.
      */
     public async getEnvList(): Promise<string[] | undefined> {
+        return this.getEnvListCached(this.cwd);
+    }
+
+    /**
+     * Method created to faciliate caching. The caching decorator uses function arguments as cache key,
+     * so pass in cwd on which we need to cache.
+     */
+    @cache(30_000, true, 10_000)
+    private async getEnvListCached(_cwd: string): Promise<string[] | undefined> {
         const result = await this.safeShellExecute(`${this._command} env list --full-path`);
         if (!result) {
             return undefined;
@@ -203,6 +216,15 @@ export class Poetry {
      * Corresponds to "poetry env info -p". Swallows errors if any.
      */
     public async getActiveEnvPath(): Promise<string | undefined> {
+        return this.getActiveEnvPathCached(this.cwd);
+    }
+
+    /**
+     * Method created to faciliate caching. The caching decorator uses function arguments as cache key,
+     * so pass in cwd on which we need to cache.
+     */
+    @cache(20_000, true, 10_000)
+    private async getActiveEnvPathCached(_cwd: string): Promise<string | undefined> {
         const result = await this.safeShellExecute(`${this._command} env info -p`, true);
         if (!result) {
             return undefined;
@@ -288,12 +310,12 @@ export async function isPoetryEnvironmentRelatedToFolder(
  *
  * @param folder Folder to look for pyproject.toml file in.
  */
-async function hasValidPyprojectToml(folder: string): Promise<boolean> {
+function hasValidPyprojectToml(folder: string): boolean {
     const pyprojectToml = path.join(folder, 'pyproject.toml');
-    if (!(await pathExists(pyprojectToml))) {
+    if (!pathExistsSync(pyprojectToml)) {
         return false;
     }
-    const content = await readFile(pyprojectToml);
+    const content = readFileSync(pyprojectToml);
     if (!content.includes('[tool.poetry]')) {
         return false;
     }

src/client/pythonEnvironments/discovery/locators/services/poetryLocator.ts

@@ -13,10 +13,9 @@ import { buildEnvInfo } from '../../../base/info/env';
 import { IPythonEnvsIterator } from '../../../base/locator';
 import { FSWatchingLocator } from '../../../base/locators/lowLevel/fsWatchingLocator';
 import {
-    findInterpretersInDir,
     getEnvironmentDirFromPath,
+    getInterpreterPathFromDir,
     getPythonVersionFromPath,
-    looksLikeBasicVirtualPython,
 } from '../../../common/commonUtils';
 import { getFileInfo, isParentPath, pathExists } from '../../../common/externalDependencies';
 import { isPoetryEnvironment, localPoetryEnvDirName, Poetry } from './poetry';
@@ -114,29 +113,20 @@ export class PoetryLocator extends FSWatchingLocator {
                     traceVerbose(`Searching for poetry virtual envs in: ${envDir}`);
 
                     const isLocal = isParentPath(envDir, root);
-                    const executables = findInterpretersInDir(envDir, 1);
-
-                    for await (const entry of executables) {
-                        const { filename } = entry;
-                        // We only care about python.exe (on windows) and python (on linux/mac)
-                        // Other version like python3.exe or python3.8 are often symlinks to
-                        // python.exe or python in the same directory in the case of virtual
-                        // environments.
-                        if (await looksLikeBasicVirtualPython(entry)) {
+                    const filename = await getInterpreterPathFromDir(envDir);
+                    if (filename !== undefined) {
+                        const kind = PythonEnvKind.Poetry;
+                        if (isLocal && !(await isPoetryEnvironment(filename))) {
+                            // This is not a poetry env.
+                            traceVerbose(
+                                `Poetry Virtual Environment: [skipped] ${filename} (reason: Not poetry environment)`,
+                            );
+                        } else {
                             // We should extract the kind here to avoid doing is*Environment()
                             // check multiple times. Those checks are file system heavy and
                             // we can use the kind to determine this anyway.
-                            yield buildVirtualEnvInfo(
-                                filename,
-                                // Global environments are fetched using 'poetry env list' so we already
-                                // know they're poetry environments, no need to get kind for them.
-                                isLocal ? await getVirtualEnvKind(filename) : PythonEnvKind.Poetry,
-                                undefined,
-                                isLocal,
-                            );
+                            yield buildVirtualEnvInfo(filename, kind, undefined, isLocal);
                             traceVerbose(`Poetry Virtual Environment: [added] ${filename}`);
-                        } else {
-                            traceVerbose(`Poetry Virtual Environment: [skipped] ${filename}`);
                         }
                     }
                 }

src/test/pythonEnvironments/discovery/locators/lowLevel/poetry.unit.test.ts

@@ -94,7 +94,6 @@ suite('isPoetryEnvironment Tests', () => {
 suite('Poetry binary is located correctly', async () => {
     let shellExecute: sinon.SinonStub;
     let getPythonSetting: sinon.SinonStub;
-    let pathExists: sinon.SinonStub;
 
     suiteSetup(() => {
         Poetry._poetryPromise = new Map();
@@ -175,9 +174,9 @@ suite('Poetry binary is located correctly', async () => {
             return;
         }
         const defaultPoetry = path.join(home, '.poetry', 'bin', 'poetry');
-        pathExists = sinon.stub(externalDependencies, 'pathExists');
-        pathExists.withArgs(defaultPoetry).resolves(true);
-        pathExists.callThrough();
+        const pathExistsSync = sinon.stub(externalDependencies, 'pathExistsSync');
+        pathExistsSync.withArgs(defaultPoetry).returns(true);
+        pathExistsSync.callThrough();
         getPythonSetting.returns('poetry');
         shellExecute.callsFake((command: string, options: ShellOptions) => {
             if (