Handle multiple downloads + workspaces more robustly

Author: lukel97

src/extension.ts

@@ -25,17 +25,22 @@ import { DocsBrowser } from './docsBrowser';
 import { downloadHaskellLanguageServer } from './hlsBinaries';
 import { executableExists } from './utils';
 
-const clients: Map<string, LanguageClient> = new Map();
+// The current map of documents & folders to language servers.
+// It may be null to indicate that we are in the process of launching a server,
+// in which case don't try to launch another one for that uri
+const clients: Map<string, LanguageClient | null> = new Map();
 
 // This is the entrypoint to our extension
 export async function activate(context: ExtensionContext) {
-  // Register HIE to check every time a text document gets opened, to
-  // support multi-root workspaces.
-
-  workspace.onDidOpenTextDocument(async (document: TextDocument) => await activateHie(context, document));
-  workspace.textDocuments.forEach(async (document: TextDocument) => await activateHie(context, document));
-
-  // Stop HIE from any workspace folders that are removed.
+  // (Possibly) launch the language server every time a document is opened, so
+  // it works across multiple workspace folders. Eventually, haskell-lsp should
+  // just support
+  // https://microsoft.github.io/language-server-protocol/specifications/specification-3-15/#workspace_workspaceFolders
+  // and then we can just launch one server
+  workspace.onDidOpenTextDocument(async (document: TextDocument) => await activeServer(context, document));
+  workspace.textDocuments.forEach(async (document: TextDocument) => await activeServer(context, document));
+
+  // Stop the server from any workspace folders that are removed.
   workspace.onDidChangeWorkspaceFolders((event) => {
     for (const folder of event.removed) {
       const client = clients.get(folder.uri.toString());
@@ -49,8 +54,8 @@ export async function activate(context: ExtensionContext) {
   // Register editor commands for HIE, but only register the commands once at activation.
   const restartCmd = commands.registerCommand(CommandNames.RestartHieCommandName, async () => {
     for (const langClient of clients.values()) {
-      await langClient.stop();
-      langClient.start();
+      await langClient?.stop();
+      langClient?.start();
     }
   });
   context.subscriptions.push(restartCmd);
@@ -107,7 +112,7 @@ function findLocalServer(context: ExtensionContext, uri: Uri, folder?: Workspace
   return null;
 }
 
-async function activateHie(context: ExtensionContext, document: TextDocument) {
+async function activeServer(context: ExtensionContext, document: TextDocument) {
   // We are only interested in Haskell files.
   if (
     (document.languageId !== 'haskell' &&
@@ -121,15 +126,20 @@ async function activateHie(context: ExtensionContext, document: TextDocument) {
   const uri = document.uri;
   const folder = workspace.getWorkspaceFolder(uri);
 
-  // If the client already has an LSP server for this folder, then don't start a new one.
-  if (folder && clients.has(folder.uri.toString())) {
+  activateServerForFolder(context, uri, folder);
+}
+
+async function activateServerForFolder(context: ExtensionContext, uri: Uri, folder?: WorkspaceFolder) {
+  const clientsKey = folder ? folder.uri.toString() : uri.toString();
+
+  // If the client already has an LSP server for this uri/folder, then don't start a new one.
+  if (clients.has(clientsKey)) {
     return;
   }
-  activateHieNoCheck(context, uri, folder);
-}
+  // Set the key to null to prevent multiple servers being launched at once
+  clients.set(clientsKey, null);
 
-async function activateHieNoCheck(context: ExtensionContext, uri: Uri, folder?: WorkspaceFolder) {
-  // Stop right here, if HIE is disabled in the resource/workspace folder.
+  // Stop right here, if Haskell is disabled in the resource/workspace folder.
   const enable = workspace.getConfiguration('haskell', uri).enable;
   if (!enable) {
     return;
@@ -223,35 +233,27 @@ async function activateHieNoCheck(context: ExtensionContext, uri: Uri, folder?:
   };
 
   // Create the LSP client.
-  const langClient = new LanguageClient(langName, langName, serverOptions, clientOptions, true);
+  const langClient = new LanguageClient(langName, langName, serverOptions, clientOptions);
 
   // Register ClientCapabilities for stuff like window/progress
   langClient.registerProposedFeatures();
 
-  // If the client already has an LSP server, then don't start a new one.
-  // We check this again, as there may be multiple parallel requests.
-  if (folder && clients.has(folder.uri.toString())) {
-    return;
-  }
-
   // Finally start the client and add it to the list of clients.
   langClient.start();
-  if (folder) {
-    clients.set(folder.uri.toString(), langClient);
-  } else {
-    clients.set(uri.toString(), langClient);
-  }
+  clients.set(clientsKey, langClient);
 }
 
 /*
  * Deactivate each of the LSP servers.
  */
-export function deactivate(): Thenable<void> {
+export async function deactivate() {
   const promises: Array<Thenable<void>> = [];
   for (const client of clients.values()) {
-    promises.push(client.stop());
+    if (client) {
+      promises.push(client.stop());
+    }
   }
-  return Promise.all(promises).then(() => undefined);
+  await Promise.all(promises);
 }
 
 function showNotInstalledErrorMessage(uri: Uri) {

src/hlsBinaries.ts

@@ -2,8 +2,7 @@ import * as child_process from 'child_process';
 import * as fs from 'fs';
 import * as https from 'https';
 import * as path from 'path';
-import * as url from 'url';
-import { env, ExtensionContext, ProgressLocation, Uri, window, workspace, WorkspaceFolder } from 'vscode';
+import { env, ExtensionContext, ProgressLocation, Uri, window, WorkspaceFolder } from 'vscode';
 import { downloadFile, executableExists, userAgentHeader } from './utils';
 
 /** GitHub API release */
@@ -120,8 +119,11 @@ async function getProjectGhcVersion(context: ExtensionContext, dir: string, rele
     throw Error(`Couldn't find any ${assetName} binaries for release ${release.tag_name}`);
   }
 
-  const wrapperUrl = url.parse(wrapperAsset.browser_download_url);
-  await downloadFile('Downloading haskell-language-server-wrapper', wrapperUrl, downloadedWrapper);
+  await downloadFile(
+    'Downloading haskell-language-server-wrapper',
+    wrapperAsset.browser_download_url,
+    downloadedWrapper
+  );
 
   return callWrapper(downloadedWrapper);
 }
@@ -199,14 +201,13 @@ export async function downloadHaskellLanguageServer(
     );
     return null;
   }
-  const binaryURL = url.parse(asset.browser_download_url);
 
   const serverName = `haskell-language-server-${release.tag_name}-${process.platform}-${ghcVersion}${exeExtension}`;
   const binaryDest = path.join(context.globalStoragePath, serverName);
 
   const title = `Downloading haskell-language-server ${release.tag_name} for GHC ${ghcVersion}`;
   try {
-    await downloadFile(title, binaryURL, binaryDest);
+    await downloadFile(title, asset.browser_download_url, binaryDest);
     return binaryDest;
   } catch (e) {
     if (e instanceof Error) {

src/utils.ts

@@ -14,7 +14,25 @@ import { createGunzip } from 'zlib';
  */
 export const userAgentHeader = { 'User-Agent': 'vscode-hie-server' };
 
-export async function downloadFile(titleMsg: string, srcUrl: url.UrlWithStringQuery, dest: string): Promise<void> {
+/** downloadFile may get called twice on the same src and destination:
+ * When this happens, we should only download the file once but return two
+ * promises that wait on the same download. This map keeps track of which
+ * files are currently being downloaded and we short circuit any calls to
+ * downloadFile which have a hit in this map by returning the promise stored
+ * here.
+ * Note that we have to use a double nested map since array/pointer/object
+ * equality is by reference, not value in Map. And we are using a tuple of
+ * [src, dest] as the key.
+ */
+const inFlightDownloads = new Map<string, Map<string, Thenable<void>>>();
+
+export async function downloadFile(titleMsg: string, src: string, dest: string): Promise<void> {
+  // Check to see if we're already in the process of downloading the same thing
+  const inFlightDownload = inFlightDownloads.get(src)?.get(dest);
+  if (inFlightDownload) {
+    return inFlightDownload;
+  }
+
   // If it already is downloaded just use that
   if (fs.existsSync(dest)) {
     return;
@@ -28,17 +46,20 @@ export async function downloadFile(titleMsg: string, srcUrl: url.UrlWithStringQu
     fs.unlinkSync(downloadDest);
   }
 
-  const downloadHieTask = window.withProgress(
+  const downloadTask = window.withProgress(
     {
       location: ProgressLocation.Notification,
       title: titleMsg,
       cancellable: false,
     },
     async (progress) => {
       const p = new Promise<void>((resolve, reject) => {
+        const srcUrl = url.parse(src);
         const opts: https.RequestOptions = {
           host: srcUrl.host,
           path: srcUrl.path,
+          protocol: srcUrl.protocol,
+          port: srcUrl.port,
           headers: userAgentHeader,
         };
         getWithRedirects(opts, (res) => {
@@ -67,17 +88,27 @@ export async function downloadFile(titleMsg: string, srcUrl: url.UrlWithStringQu
           fileStream.on('close', resolve);
         }).on('error', reject);
       });
-      await p;
-      // Finally rename it to the actual dest
-      fs.renameSync(downloadDest, dest);
+      try {
+        await p;
+        // Finally rename it to the actual dest
+        fs.renameSync(downloadDest, dest);
+      } finally {
+        // And remember to remove it from the list of current downloads
+        inFlightDownloads.get(src)?.delete(dest);
+      }
     }
   );
 
   try {
-    return downloadHieTask;
+    if (inFlightDownloads.has(src)) {
+      inFlightDownloads.get(src)?.set(dest, downloadTask);
+    } else {
+      inFlightDownloads.set(src, new Map([[dest, downloadTask]]));
+    }
+    return downloadTask;
   } catch (e) {
     fs.unlinkSync(downloadDest);
-    throw new Error(`Failed to download ${url.format(srcUrl)}`);
+    throw new Error(`Failed to download ${url}`);
   }
 }