From 498b1d2a903de92a92be6e4a0577900c8a403e19 Mon Sep 17 00:00:00 2001
From: Florian Loitsch <florian@toit.io>
Date: Thu, 23 Jun 2022 15:54:13 +0200
Subject: [PATCH] Fix race for the LSP server. (#200)

Sometimes we would see multiple LSP servers because of a race condition.
Now we store the promise instead of the actual LSP client. This way we
don't accidentally start two clients.
---
 vscode/src/lspClient.ts | 33 +++++++++++++++++++++++----------
 1 file changed, 23 insertions(+), 10 deletions(-)

diff --git a/vscode/src/lspClient.ts b/vscode/src/lspClient.ts
index 0bf289d..40f8d52 100644
--- a/vscode/src/lspClient.ts
+++ b/vscode/src/lspClient.ts
@@ -10,7 +10,7 @@ import { DocumentSelector, LanguageClient, LanguageClientOptions, ServerOptions
 
 // Untitled documents, or documents outside all workspaces go to a default client.
 let nonFileClient: LanguageClient;
-const clients: Map<string, LanguageClient> = new Map();
+const clients: Map<string, Promise<LanguageClient>> = new Map();
 const clientCounts: Map<string, number> = new Map();
 let _workspaceFolders: Set<string>|undefined;
 
@@ -191,13 +191,26 @@ export function activateLsp(context: ExtensionContext, lspCommand: Array<string>
     }
     const workingDir = config.workingDir!;
     if (!clients.has(workingDir)) {
-      const client = await startToitLsp(context, lspCommand, outputChannel, config);
-      clients.set(workingDir, client);
+      const clientPromise = startToitLsp(context, lspCommand, outputChannel, config);
+      clients.set(workingDir, clientPromise);
       clientCounts.set(workingDir, 1);
     } else {
       const oldCount = clientCounts.get(workingDir)!;
       clientCounts.set(workingDir, oldCount + 1);
     }
+    const clientPromise = clients.get(workingDir);
+    try {
+      await clientPromise;
+    } catch (e) {
+      // Delete the current client promise from the map.
+      // It's not completely clear if the check for the clientPromise is necessary, but
+      // it can't hurt.
+      if (clients.has(workingDir) && clientCounts.get(workingDir) === clientPromise) {
+        clients.delete(workingDir);
+        clientCounts.delete(workingDir);
+      }
+      throw e;
+    }
   }
 
   async function didOpenTextDocument(document: TextDocument): Promise<void> {
@@ -224,10 +237,10 @@ export function activateLsp(context: ExtensionContext, lspCommand: Array<string>
     if (clients.has(workingDir)) {
       const oldCount = clientCounts.get(workingDir)!;
       if (oldCount === 1) {
-        const client = clients.get(workingDir)!;
+        const clientPromise = clients.get(workingDir)!;
         clients.delete(workingDir);
         clientCounts.delete(workingDir);
-        client.stop();
+        clientPromise.then((client) => client.stop());
       } else {
         clientCounts.set(workingDir, oldCount - 1);
       }
@@ -239,10 +252,10 @@ export function activateLsp(context: ExtensionContext, lspCommand: Array<string>
   Workspace.onDidChangeWorkspaceFolders(event => {
     for (const folder of event.removed) {
       const uriString = folder.uri.toString();
-      const client = clients.get(uriString);
-      if (client) {
+      const clientPromise = clients.get(uriString);
+      if (clientPromise) {
         clients.delete(uriString);
-        client.stop();
+        clientPromise.then((client) => client.stop());
       }
     }
   });
@@ -254,8 +267,8 @@ export function deactivateLsp(): Thenable<void> {
   if (nonFileClient) {
     promises.push(nonFileClient.stop());
   }
-  for (const client of clients.values()) {
-    promises.push(client.stop());
+  for (const clientPromise of clients.values()) {
+    promises.push(clientPromise.then((client) => client.stop()));
   }
   return Promise.all(promises).then(() => undefined);
 }