Improve time taken to trigger language server startup once extension activation is triggered (#22514)

For #22146 

Improves time taken to trigger language server startup once extension
activation is triggered

- Do not block discovery on windows registry
- Do not blocking auto-selection on validation of all interpreters
- Make Windows Path locator faster

Author: Kartik Raj

src/client/activation/activationManager.ts

@@ -11,6 +11,7 @@ import { PYTHON_LANGUAGE } from '../common/constants';
 import { IFileSystem } from '../common/platform/types';
 import { IDisposable, IInterpreterPathService, Resource } from '../common/types';
 import { Deferred } from '../common/utils/async';
+import { StopWatch } from '../common/utils/stopWatch';
 import { IInterpreterAutoSelectionService } from '../interpreter/autoSelection/types';
 import { traceDecoratorError } from '../logging';
 import { sendActivationTelemetry } from '../telemetry/envFileTelemetry';
@@ -69,20 +70,20 @@ export class ExtensionActivationManager implements IExtensionActivationManager {
         }
     }
 
-    public async activate(): Promise<void> {
+    public async activate(startupStopWatch: StopWatch): Promise<void> {
         this.filterServices();
         await this.initialize();
 
         // Activate all activation services together.
 
         await Promise.all([
             ...this.singleActivationServices.map((item) => item.activate()),
-            this.activateWorkspace(this.activeResourceService.getActiveResource()),
+            this.activateWorkspace(this.activeResourceService.getActiveResource(), startupStopWatch),
         ]);
     }
 
     @traceDecoratorError('Failed to activate a workspace')
-    public async activateWorkspace(resource: Resource): Promise<void> {
+    public async activateWorkspace(resource: Resource, startupStopWatch?: StopWatch): Promise<void> {
         const folder = this.workspaceService.getWorkspaceFolder(resource);
         resource = folder ? folder.uri : undefined;
         const key = this.getWorkspaceKey(resource);
@@ -97,7 +98,7 @@ export class ExtensionActivationManager implements IExtensionActivationManager {
             await this.interpreterPathService.copyOldInterpreterStorageValuesToNew(resource);
         }
         await sendActivationTelemetry(this.fileSystem, this.workspaceService, resource);
-        await Promise.all(this.activationServices.map((item) => item.activate(resource)));
+        await Promise.all(this.activationServices.map((item) => item.activate(resource, startupStopWatch)));
         await this.appDiagnostics.performPreStartupHealthCheck(resource);
     }
 

src/client/activation/types.ts

@@ -6,6 +6,7 @@
 import { Event } from 'vscode';
 import { LanguageClient, LanguageClientOptions } from 'vscode-languageclient/node';
 import type { IDisposable, ILogOutputChannel, Resource } from '../common/types';
+import { StopWatch } from '../common/utils/stopWatch';
 import { PythonEnvironment } from '../pythonEnvironments/info';
 
 export const IExtensionActivationManager = Symbol('IExtensionActivationManager');
@@ -23,7 +24,7 @@ export interface IExtensionActivationManager extends IDisposable {
      * @returns {Promise<void>}
      * @memberof IExtensionActivationManager
      */
-    activate(): Promise<void>;
+    activate(startupStopWatch: StopWatch): Promise<void>;
     /**
      * Method invoked when a workspace is loaded.
      * This is where we place initialization scripts for each workspace.
@@ -47,7 +48,7 @@ export const IExtensionActivationService = Symbol('IExtensionActivationService')
  */
 export interface IExtensionActivationService {
     supportedWorkspaceTypes: { untrustedWorkspace: boolean; virtualWorkspace: boolean };
-    activate(resource: Resource): Promise<void>;
+    activate(resource: Resource, startupStopWatch?: StopWatch): Promise<void>;
 }
 
 export enum LanguageServerType {

src/client/extension.ts

@@ -110,6 +110,7 @@ async function activateUnsafe(
     const activationDeferred = createDeferred<void>();
     displayProgress(activationDeferred.promise);
     startupDurations.startActivateTime = startupStopWatch.elapsedTime;
+    const activationStopWatch = new StopWatch();
 
     //===============================================
     // activation starts here
@@ -127,7 +128,7 @@ async function activateUnsafe(
     const components = await initializeComponents(ext);
 
     // Then we finish activating.
-    const componentsActivated = await activateComponents(ext, components);
+    const componentsActivated = await activateComponents(ext, components, activationStopWatch);
     activateFeatures(ext, components);
 
     const nonBlocking = componentsActivated.map((r) => r.fullyReady);

src/client/extensionActivation.ts

@@ -51,11 +51,13 @@ import { registerCreateEnvironmentTriggers } from './pythonEnvironments/creation
 import { initializePersistentStateForTriggers } from './common/persistentState';
 import { logAndNotifyOnLegacySettings } from './logging/settingLogs';
 import { DebuggerTypeName } from './debugger/constants';
+import { StopWatch } from './common/utils/stopWatch';
 
 export async function activateComponents(
     // `ext` is passed to any extra activation funcs.
     ext: ExtensionState,
     components: Components,
+    startupStopWatch: StopWatch,
 ): Promise<ActivationResult[]> {
     // Note that each activation returns a promise that resolves
     // when that activation completes.  However, it might have started
@@ -73,7 +75,7 @@ export async function activateComponents(
     // activate them in parallel with the other components.
     // https://github.com/microsoft/vscode-python/issues/15380
     // These will go away eventually once everything is refactored into components.
-    const legacyActivationResult = await activateLegacy(ext);
+    const legacyActivationResult = await activateLegacy(ext, startupStopWatch);
     const workspaceService = new WorkspaceService();
     if (!workspaceService.isTrusted) {
         return [legacyActivationResult];
@@ -105,7 +107,7 @@ export function activateFeatures(ext: ExtensionState, _components: Components):
 // init and activation: move them to activateComponents().
 // See https://github.com/microsoft/vscode-python/issues/10454.
 
-async function activateLegacy(ext: ExtensionState): Promise<ActivationResult> {
+async function activateLegacy(ext: ExtensionState, startupStopWatch: StopWatch): Promise<ActivationResult> {
     const { legacyIOC } = ext;
     const { serviceManager, serviceContainer } = legacyIOC;
 
@@ -183,7 +185,7 @@ async function activateLegacy(ext: ExtensionState): Promise<ActivationResult> {
     const manager = serviceContainer.get<IExtensionActivationManager>(IExtensionActivationManager);
     disposables.push(manager);
 
-    const activationPromise = manager.activate();
+    const activationPromise = manager.activate(startupStopWatch);
 
     return { fullyReady: activationPromise };
 }

src/client/interpreter/autoSelection/index.ts

@@ -6,11 +6,13 @@
 import { inject, injectable } from 'inversify';
 import { Event, EventEmitter, Uri } from 'vscode';
 import { IWorkspaceService } from '../../common/application/types';
+import { DiscoveryUsingWorkers } from '../../common/experiments/groups';
 import '../../common/extensions';
 import { IFileSystem } from '../../common/platform/types';
-import { IPersistentState, IPersistentStateFactory, Resource } from '../../common/types';
+import { IExperimentService, IPersistentState, IPersistentStateFactory, Resource } from '../../common/types';
 import { createDeferred, Deferred } from '../../common/utils/async';
 import { compareSemVerLikeVersions } from '../../pythonEnvironments/base/info/pythonVersion';
+import { ProgressReportStage } from '../../pythonEnvironments/base/locator';
 import { PythonEnvironment } from '../../pythonEnvironments/info';
 import { sendTelemetryEvent } from '../../telemetry';
 import { EventName } from '../../telemetry/constants';
@@ -44,6 +46,7 @@ export class InterpreterAutoSelectionService implements IInterpreterAutoSelectio
         @inject(IInterpreterComparer) private readonly envTypeComparer: IInterpreterComparer,
         @inject(IInterpreterAutoSelectionProxyService) proxy: IInterpreterAutoSelectionProxyService,
         @inject(IInterpreterHelper) private readonly interpreterHelper: IInterpreterHelper,
+        @inject(IExperimentService) private readonly experimentService: IExperimentService,
     ) {
         proxy.registerInstance!(this);
     }
@@ -183,7 +186,7 @@ export class InterpreterAutoSelectionService implements IInterpreterAutoSelectio
 
     private getAutoSelectionQueriedOnceState(): IPersistentState<boolean | undefined> {
         const key = `autoSelectionInterpretersQueriedOnce`;
-        return this.stateFactory.createWorkspacePersistentState(key, undefined);
+        return this.stateFactory.createGlobalPersistentState(key, undefined);
     }
 
     /**
@@ -199,22 +202,44 @@ export class InterpreterAutoSelectionService implements IInterpreterAutoSelectio
     private async autoselectInterpreterWithLocators(resource: Resource): Promise<void> {
         // Do not perform a full interpreter search if we already have cached interpreters for this workspace.
         const queriedState = this.getAutoSelectionInterpretersQueryState(resource);
-        if (queriedState.value !== true && resource) {
+        const globalQueriedState = this.getAutoSelectionQueriedOnceState();
+        if (globalQueriedState.value && queriedState.value !== true && resource) {
             await this.interpreterService.triggerRefresh({
                 searchLocations: { roots: [resource], doNotIncludeNonRooted: true },
             });
         }
 
-        const globalQueriedState = this.getAutoSelectionQueriedOnceState();
-        if (!globalQueriedState.value) {
-            // Global interpreters are loaded the first time an extension loads, after which we don't need to
-            // wait on global interpreter promise refresh.
-            await this.interpreterService.refreshPromise;
-        }
-        const interpreters = this.interpreterService.getInterpreters(resource);
+        const inExperiment = this.experimentService.inExperimentSync(DiscoveryUsingWorkers.experiment);
         const workspaceUri = this.interpreterHelper.getActiveWorkspaceUri(resource);
+        let recommendedInterpreter: PythonEnvironment | undefined;
+        if (inExperiment) {
+            if (!globalQueriedState.value) {
+                // Global interpreters are loaded the first time an extension loads, after which we don't need to
+                // wait on global interpreter promise refresh.
+                // Do not wait for validation of all interpreters to finish, we only need to validate the recommended interpreter.
+                await this.interpreterService.getRefreshPromise({ stage: ProgressReportStage.allPathsDiscovered });
+            }
+            let interpreters = this.interpreterService.getInterpreters(resource);
+
+            recommendedInterpreter = this.envTypeComparer.getRecommended(interpreters, workspaceUri?.folderUri);
+            const details = recommendedInterpreter
+                ? await this.interpreterService.getInterpreterDetails(recommendedInterpreter.path)
+                : undefined;
+            if (!details || !recommendedInterpreter) {
+                await this.interpreterService.refreshPromise; // Interpreter is invalid, wait for all of validation to finish.
+                interpreters = this.interpreterService.getInterpreters(resource);
+                recommendedInterpreter = this.envTypeComparer.getRecommended(interpreters, workspaceUri?.folderUri);
+            }
+        } else {
+            if (!globalQueriedState.value) {
+                // Global interpreters are loaded the first time an extension loads, after which we don't need to
+                // wait on global interpreter promise refresh.
+                await this.interpreterService.refreshPromise;
+            }
+            const interpreters = this.interpreterService.getInterpreters(resource);
 
-        const recommendedInterpreter = this.envTypeComparer.getRecommended(interpreters, workspaceUri?.folderUri);
+            recommendedInterpreter = this.envTypeComparer.getRecommended(interpreters, workspaceUri?.folderUri);
+        }
         if (!recommendedInterpreter) {
             return;
         }

src/client/interpreter/contracts.ts

@@ -4,6 +4,7 @@ import { FileChangeType } from '../common/platform/fileSystemWatcher';
 import { Resource } from '../common/types';
 import { PythonEnvSource } from '../pythonEnvironments/base/info';
 import {
+    GetRefreshEnvironmentsOptions,
     ProgressNotificationEvent,
     PythonLocatorQuery,
     TriggerRefreshOptions,
@@ -22,7 +23,7 @@ export const IComponentAdapter = Symbol('IComponentAdapter');
 export interface IComponentAdapter {
     readonly onProgress: Event<ProgressNotificationEvent>;
     triggerRefresh(query?: PythonLocatorQuery, options?: TriggerRefreshOptions): Promise<void>;
-    getRefreshPromise(): Promise<void> | undefined;
+    getRefreshPromise(options?: GetRefreshEnvironmentsOptions): Promise<void> | undefined;
     readonly onChanged: Event<PythonEnvironmentsChangedEvent>;
     // VirtualEnvPrompt
     onDidCreate(resource: Resource, callback: () => void): Disposable;
@@ -74,6 +75,7 @@ export const IInterpreterService = Symbol('IInterpreterService');
 export interface IInterpreterService {
     triggerRefresh(query?: PythonLocatorQuery, options?: TriggerRefreshOptions): Promise<void>;
     readonly refreshPromise: Promise<void> | undefined;
+    getRefreshPromise(options?: GetRefreshEnvironmentsOptions): Promise<void> | undefined;
     readonly onDidChangeInterpreters: Event<PythonEnvironmentsChangedEvent>;
     onDidChangeInterpreterConfiguration: Event<Uri | undefined>;
     onDidChangeInterpreter: Event<Uri | undefined>;

src/client/interpreter/interpreterService.ts

@@ -38,7 +38,11 @@ import { Interpreters } from '../common/utils/localize';
 import { sendTelemetryEvent } from '../telemetry';
 import { EventName } from '../telemetry/constants';
 import { cache } from '../common/utils/decorators';
-import { PythonLocatorQuery, TriggerRefreshOptions } from '../pythonEnvironments/base/locator';
+import {
+    GetRefreshEnvironmentsOptions,
+    PythonLocatorQuery,
+    TriggerRefreshOptions,
+} from '../pythonEnvironments/base/locator';
 import { sleep } from '../common/utils/async';
 
 type StoredPythonEnvironment = PythonEnvironment & { store?: boolean };
@@ -59,6 +63,10 @@ export class InterpreterService implements Disposable, IInterpreterService {
         return this.pyenvs.getRefreshPromise();
     }
 
+    public getRefreshPromise(options?: GetRefreshEnvironmentsOptions): Promise<void> | undefined {
+        return this.pyenvs.getRefreshPromise(options);
+    }
+
     public get onDidChangeInterpreter(): Event<Uri | undefined> {
         return this.didChangeInterpreterEmitter.event;
     }

src/client/languageServer/watcher.ts

@@ -29,6 +29,7 @@ import { PylanceLSExtensionManager } from './pylanceLSExtensionManager';
 import { ILanguageServerExtensionManager, ILanguageServerWatcher } from './types';
 import { sendTelemetryEvent } from '../telemetry';
 import { EventName } from '../telemetry/constants';
+import { StopWatch } from '../common/utils/stopWatch';
 
 @injectable()
 /**
@@ -73,14 +74,18 @@ export class LanguageServerWatcher implements IExtensionActivationService, ILang
 
     // IExtensionActivationService
 
-    public async activate(resource?: Resource): Promise<void> {
+    public async activate(resource?: Resource, startupStopWatch?: StopWatch): Promise<void> {
         this.register();
-        await this.startLanguageServer(this.languageServerType, resource);
+        await this.startLanguageServer(this.languageServerType, resource, startupStopWatch);
     }
 
     // ILanguageServerWatcher
-    public async startLanguageServer(languageServerType: LanguageServerType, resource?: Resource): Promise<void> {
-        await this.startAndGetLanguageServer(languageServerType, resource);
+    public async startLanguageServer(
+        languageServerType: LanguageServerType,
+        resource?: Resource,
+        startupStopWatch?: StopWatch,
+    ): Promise<void> {
+        await this.startAndGetLanguageServer(languageServerType, resource, startupStopWatch);
     }
 
     public register(): void {
@@ -124,6 +129,7 @@ export class LanguageServerWatcher implements IExtensionActivationService, ILang
     private async startAndGetLanguageServer(
         languageServerType: LanguageServerType,
         resource?: Resource,
+        startupStopWatch?: StopWatch,
     ): Promise<ILanguageServerExtensionManager> {
         const lsResource = this.getWorkspaceUri(resource);
         const currentInterpreter = this.workspaceInterpreters.get(lsResource.fsPath);
@@ -170,6 +176,10 @@ export class LanguageServerWatcher implements IExtensionActivationService, ILang
 
         if (languageServerExtensionManager.canStartLanguageServer(interpreter)) {
             // Start the language server.
+            if (startupStopWatch) {
+                // It means that startup is triggering this code, track time it takes since startup to activate this code.
+                sendTelemetryEvent(EventName.LANGUAGE_SERVER_TRIGGER_DURATION, startupStopWatch.elapsedTime);
+            }
             await languageServerExtensionManager.startLanguageServer(lsResource, interpreter);
 
             logStartup(languageServerType, lsResource);

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.
      */

src/client/pythonEnvironments/base/locatorUtils.ts

@@ -85,7 +85,7 @@ export async function getEnvs<I = PythonEnvInfo>(iterator: IPythonEnvsIterator<I
                 }
                 updatesDone.resolve();
                 listener.dispose();
-            } else {
+            } else if (event.index !== undefined) {
                 const { index, update } = event;
                 if (envs[index] === undefined) {
                     const json = JSON.stringify(update);
@@ -95,6 +95,8 @@ 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,9 +43,18 @@ export abstract class LazyResourceBasedLocator extends Locator<BasicEnvInfo> imp
         await this.disposables.dispose();
     }
 
-    public async *iterEnvs(query?: PythonLocatorQuery): IPythonEnvsIterator<BasicEnvInfo> {
-        await this.activate();
+    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> {
+        await this.activate();
         if (query?.envPath) {
             let result = await iterator.next();
             while (!result.done) {