[mlir-vscode] Refactor server creation to be lazy

[rriddle]

Summary

We currently proactively create language clients for every workspace folder,
and every language. This makes startup time more costly, and also emits errors
for missing language servers in contexts that the user currently isn't in. For example,
if a user opens a .mlir file we don't want to emit errors about .pdll files. We also don't
want to emit errors for missing servers in workspace folders that don't even utilize
MLIR.

This commit refactors client creation to lazy-load when a document that requires the
server is opened.

[jpienaar]

I might even suggest the error emission from not finding LSP should be logged (flagged in some way) rather than popup given the extension is the combination of LSP & syntax highlighting. And I'm reading these files more than I traverse via LSP assistance, so in fresh client/machine without lsp built the extension is still useful and the missing LSP warning interrupts my flow before I even tried to use any LSP functionality (and I may not need any at all). But that's just a nit.

[rriddle]

In D123184#3433072, @jpienaar wrote:

I might even suggest the error emission from not finding LSP should be logged (flagged in some way) rather than popup given the extension is the combination of LSP & syntax highlighting. And I'm reading these files more than I traverse via LSP assistance, so in fresh client/machine without lsp built the extension is still useful and the missing LSP warning interrupts my flow before I even tried to use any LSP functionality (and I may not need any at all). But that's just a nit.

Yeah, we could do that. We can just ignore missing servers if the options are empty (it's already behind a flag, but I can make it always true). I'll prepare a followup and then bump the version afterwards.

[rriddle]

In D123184#3433535, @rriddle wrote:

In D123184#3433072, @jpienaar wrote:

I might even suggest the error emission from not finding LSP should be logged (flagged in some way) rather than popup given the extension is the combination of LSP & syntax highlighting. And I'm reading these files more than I traverse via LSP assistance, so in fresh client/machine without lsp built the extension is still useful and the missing LSP warning interrupts my flow before I even tried to use any LSP functionality (and I may not need any at all). But that's just a nit.

Yeah, we could do that. We can just ignore missing servers if the options are empty (it's already behind a flag, but I can make it always true). I'll prepare a followup and then bump the version afterwards.

Sent out https://reviews.llvm.org/D123240 for that.

mlir/utils/vscode/src/configWatcher.ts

 }
 
 /**
  *  Activate watchers that track configuration changes for the given workspace
  *  folder, or null if the workspace is top-level.
  */
 export async function activate(mlirContext: MLIRContext,
                                workspaceFolder: vscode.WorkspaceFolder,
-                               serverPathsToWatch: string[]) {
+                               serverSetting: string, serverPath: string) {
   // When a configuration change happens, check to see if we should restart the
   // server.
   mlirContext.subscriptions.push(vscode.workspace.onDidChangeConfiguration(event => {
-    const settings: string[] = [ 'server_path', 'pdll_server_path' ];
-    for (const setting of settings) {
-      const expandedSetting = `mlir.${setting}`;
-      if (event.affectsConfiguration(expandedSetting, workspaceFolder)) {
-        promptRestart(
-            'onSettingsChanged',
-            `setting '${
-                expandedSetting}' has changed. Do you want to reload the server?`);
-        break;
-      }
+    const expandedSetting = `mlir.${serverSetting}`;
     }
   }));
 
-  // Track the server file in case it changes. We use `fs` here because the
-  // server may not be in a workspace directory.
-  for (const serverPath of serverPathsToWatch) {
-    // Check that the path actually exists.
-    if (serverPath === '') {
-      continue;
-    }
+  // If the server path actually exists, track it in case it changes. Check that
+  // the path actually exists.
+    return;
 
-    const fileWatcherConfig = {
-      disableGlobbing : true,
-      followSymlinks : true,
-      ignoreInitial : true,
-      awaitWriteFinish : true,
-    };
-    const fileWatcher = chokidar.watch(serverPath, fileWatcherConfig);
-    fileWatcher.on('all', (event, _filename, _details) => {
-      if (event != 'unlink') {
-        promptRestart(
-            'onSettingsChanged',
-            'MLIR language server binary has changed. Do you want to reload the server?');
-      }
-    });
-    mlirContext.subscriptions.push(
-        new vscode.Disposable(() => { fileWatcher.close(); }));
-  }
-}

mlir/utils/vscode/src/mlirContext.ts

 import * as fs from 'fs';
 import * as path from 'path';
 import * as vscode from 'vscode';
 import * as vscodelc from 'vscode-languageclient';
 
 import * as config from './config';
 import * as configWatcher from './configWatcher';
 
 /**
  *  This class represents the context of a specific workspace folder.
  */
-class WorkspaceFolderContext {
-  constructor(mlirServer: vscodelc.LanguageClient,
-              pdllServer: vscodelc.LanguageClient) {
-    this.mlirServer = mlirServer;
-    this.pdllServer = pdllServer;
+class WorkspaceFolderContext implements vscode.Disposable {
+  dispose() {
+    this.clients.forEach(client => client.stop());
+    this.clients.clear();
   }
-  mlirServer!: vscodelc.LanguageClient;
-  pdllServer!: vscodelc.LanguageClient;
+
+  clients: Map<string, vscodelc.LanguageClient> = new Map();
 }
 
 /**
  *  This class manages all of the MLIR extension state,
  *  including the language client.
  */
 export class MLIRContext implements vscode.Disposable {
   subscriptions: vscode.Disposable[] = [];
-  workspaceFolders: WorkspaceFolderContext[] = [];
+  workspaceFolders: Map<string, WorkspaceFolderContext> = new Map();
 
   /**
    *  Activate the MLIR context, and start the language clients.
    */
   async activate(outputChannel: vscode.OutputChannel,
                  warnOnEmptyServerPath: boolean) {
-    // Start clients for each workspace folder.
-    if (vscode.workspace.workspaceFolders &&
-        vscode.workspace.workspaceFolders.length > 0) {
-      for (const workspaceFolder of vscode.workspace.workspaceFolders) {
-        this.workspaceFolders.push(await this.activateWorkspaceFolder(
-            workspaceFolder, outputChannel, warnOnEmptyServerPath));
+    // This lambda is used to lazily start language clients for the given
+    // document. It removes the need to pro-actively start language clients for
+    // every folder within the workspace and every language type we provide.
+    const startClientOnOpenDocument = async (document: vscode.TextDocument) => {
+      if (document.uri.scheme !== 'file') {
+        return;
+      }
+      let serverSettingName: string;
+      if (document.languageId === 'mlir') {
+        serverSettingName = 'server_path';
+      } else if (document.languageId === 'pdll') {
+        serverSettingName = 'pdll_server_path';
+      } else {
+        return;
+      }
+
+      // Resolve the workspace folder if this document is in one. We use the
+      // workspace folder when determining if a server needs to be started.
+      const uri = document.uri;
+      let workspaceFolder = vscode.workspace.getWorkspaceFolder(uri);
+      let workspaceFolderStr =
+          workspaceFolder ? workspaceFolder.uri.toString() : "";
+
+      // Get or create a client context for this folder.
+      let folderContext = this.workspaceFolders.get(workspaceFolderStr);
+      if (!folderContext) {
+        folderContext = new WorkspaceFolderContext();
+        this.workspaceFolders.set(workspaceFolderStr, folderContext);
+      }
+      // Start the client for this language if necessary.
+      if (!folderContext.clients.has(document.languageId)) {
+        let client = await this.activateWorkspaceFolder(
+            workspaceFolder, serverSettingName, document.languageId,
+            outputChannel, warnOnEmptyServerPath);
+        folderContext.clients.set(document.languageId, client);
+      }
+    };
+    // Process any existing documents.
+    vscode.workspace.textDocuments.forEach(startClientOnOpenDocument);
+
+    // Watch any new documents to spawn servers when necessary.
+    this.subscriptions.push(
+        vscode.workspace.onDidOpenTextDocument(startClientOnOpenDocument));
+    this.subscriptions.push(
+        vscode.workspace.onDidChangeWorkspaceFolders((event) => {
+          for (const folder of event.removed) {
+            const client = this.workspaceFolders.get(folder.uri.toString());
+            if (client) {
+              client.dispose();
+              this.workspaceFolders.delete(folder.uri.toString());
             }
           }
-    this.workspaceFolders.push(await this.activateWorkspaceFolder(
-        null, outputChannel, warnOnEmptyServerPath));
+        }));
   }
 
   /**
-   *  Activate the context for the given workspace folder, and start the
-   *  language clients.
+   *  Activate the language client for the given language in the given workspace
+   *  folder.
    */
   async activateWorkspaceFolder(workspaceFolder: vscode.WorkspaceFolder,
+                                serverSettingName: string, languageName: string,
                                 outputChannel: vscode.OutputChannel,
                                 warnOnEmptyServerPath: boolean):
-      Promise<WorkspaceFolderContext> {
-    // Create the language clients for mlir and pdll.
-    const [mlirServer, mlirServerPath] = await this.startLanguageClient(
-        workspaceFolder, outputChannel, warnOnEmptyServerPath, 'server_path',
-        'mlir');
-    const [pdllServer, pdllServerPath] = await this.startLanguageClient(
+      Promise<vscodelc.LanguageClient> {
+    const [server, serverPath] = await this.startLanguageClient(
         workspaceFolder, outputChannel, warnOnEmptyServerPath,
-        'pdll_server_path', 'pdll');
+        serverSettingName, languageName);
 
     // Watch for configuration changes on this folder.
-    const serverPathsToWatch = [ mlirServerPath, pdllServerPath ];
-    await configWatcher.activate(this, workspaceFolder, serverPathsToWatch);
-
-    return new WorkspaceFolderContext(mlirServer, pdllServer);
+    await configWatcher.activate(this, workspaceFolder, serverSettingName,
+                                 serverPath);
+    return server;
   }
 
   /**
    *  Start a new language client for the given language. Returns an array
    *  containing the opened server, or null if the server could not be started,
    *  and the resolved server path.
    */
   async startLanguageClient(workspaceFolder: vscode.WorkspaceFolder,
const clientOptions: vscodelc.LanguageClientOptions = {
       outputChannel : outputChannel,
       workspaceFolder : workspaceFolder,
       middleware : middleware
     };
 
     // Create the language client and start the client.
     let languageClient = new vscodelc.LanguageClient(
         languageName + '-lsp', clientTitle, serverOptions, clientOptions);
-    this.subscriptions.push(languageClient.start());
+    languageClient.start();
     return [ languageClient, serverPath ];
   }
 
   /**
    * Given a server setting, return the default server path.
    */
   static getDefaultServerFilename(serverSettingName: string): string {
     if (serverSettingName === 'pdll_server_path') {
async resolveServerPath(serverSettingName: string,
     }
     // Otherwise, return the resolved path.
     return foundUris[0].fsPath;
   }
 
   dispose() {
     this.subscriptions.forEach((d) => { d.dispose(); });
     this.subscriptions = [];
-    this.workspaceFolders = [];
+    this.workspaceFolders.forEach((d) => { d.dispose(); });
+    this.workspaceFolders.clear();
   }
 }