[mlir-vscode] Refactor server creation to be lazy

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.

Differential Revision: https://reviews.llvm.org/D123184

mlir/utils/vscode/src/configWatcher.ts

@@ -43,46 +43,39 @@ async function promptRestart(settingName: string, promptMessage: string) {
  */
 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}`;
+    if (event.affectsConfiguration(expandedSetting, workspaceFolder)) {
+      promptRestart(
+          'onSettingsChanged',
+          `setting '${
+              expandedSetting}' has changed. Do you want to reload the server?`);
     }
   }));
 
-  // 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;
-    }
-
-    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(); }));
+  // If the server path actually exists, track it in case it changes. Check that
+  // the path actually exists.
+  if (serverPath === '') {
+    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

@@ -9,14 +9,13 @@ 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();
 }
 
 /**
@@ -25,46 +24,85 @@ class WorkspaceFolderContext {
  */
 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;
       }
-    }
-    this.workspaceFolders.push(await this.activateWorkspaceFolder(
-        null, outputChannel, warnOnEmptyServerPath));
+      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());
+            }
+          }
+        }));
   }
 
   /**
-   *  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;
   }
 
   /**
@@ -165,7 +203,7 @@ export class MLIRContext implements vscode.Disposable {
     // 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 ];
   }
 
@@ -225,6 +263,7 @@ export class MLIRContext implements vscode.Disposable {
   dispose() {
     this.subscriptions.forEach((d) => { d.dispose(); });
     this.subscriptions = [];
-    this.workspaceFolders = [];
+    this.workspaceFolders.forEach((d) => { d.dispose(); });
+    this.workspaceFolders.clear();
   }
 }