Refactor language client handling

1 files changed, 89 insertions, 94 deletions

diff --git a/editors/code/src/commands.ts b/editors/code/src/commands.ts
index b9ad525e361..cbdeb28c99f 100644
--- a/editors/code/src/commands.ts
+++ b/editors/code/src/commands.ts
@@ -21,16 +21,16 @@ export function analyzerStatus(ctx: Ctx): Cmd {
         readonly uri = vscode.Uri.parse("rust-analyzer-status://status");
         readonly eventEmitter = new vscode.EventEmitter<vscode.Uri>();
 
-        provideTextDocumentContent(_uri: vscode.Uri): vscode.ProviderResult<string> {
+        async provideTextDocumentContent(_uri: vscode.Uri): Promise<string> {
             if (!vscode.window.activeTextEditor) return "";
+            const client = await ctx.getClient();
 
             const params: ra.AnalyzerStatusParams = {};
             const doc = ctx.activeRustEditor?.document;
             if (doc != null) {
-                params.textDocument =
-                    ctx.client.code2ProtocolConverter.asTextDocumentIdentifier(doc);
+                params.textDocument = client.code2ProtocolConverter.asTextDocumentIdentifier(doc);
             }
-            return ctx.client.sendRequest(ra.analyzerStatus, params);
+            return await client.sendRequest(ra.analyzerStatus, params);
         }
 
         get onDidChange(): vscode.Event<vscode.Uri> {
@@ -38,7 +38,7 @@ export function analyzerStatus(ctx: Ctx): Cmd {
         }
     })();
 
-    ctx.pushCleanup(
+    ctx.pushExtCleanup(
         vscode.workspace.registerTextDocumentContentProvider("rust-analyzer-status", tdcp)
     );
 
@@ -60,9 +60,14 @@ export function memoryUsage(ctx: Ctx): Cmd {
         provideTextDocumentContent(_uri: vscode.Uri): vscode.ProviderResult<string> {
             if (!vscode.window.activeTextEditor) return "";
 
-            return ctx.client.sendRequest(ra.memoryUsage).then((mem: any) => {
-                return "Per-query memory usage:\n" + mem + "\n(note: database has been cleared)";
-            });
+            return ctx
+                .getClient()
+                .then((it) => it.sendRequest(ra.memoryUsage))
+                .then((mem: any) => {
+                    return (
+                        "Per-query memory usage:\n" + mem + "\n(note: database has been cleared)"
+                    );
+                });
         }
 
         get onDidChange(): vscode.Event<vscode.Uri> {
@@ -70,7 +75,7 @@ export function memoryUsage(ctx: Ctx): Cmd {
         }
     })();
 
-    ctx.pushCleanup(
+    ctx.pushExtCleanup(
         vscode.workspace.registerTextDocumentContentProvider("rust-analyzer-memory", tdcp)
     );
 
@@ -83,23 +88,19 @@ export function memoryUsage(ctx: Ctx): Cmd {
 
 export function shuffleCrateGraph(ctx: Ctx): Cmd {
     return async () => {
-        const client = ctx.client;
-        if (!client) return;
-
-        await client.sendRequest(ra.shuffleCrateGraph);
+        return ctx.getClient().then((it) => it.sendRequest(ra.shuffleCrateGraph));
     };
 }
 
 export function matchingBrace(ctx: Ctx): Cmd {
     return async () => {
         const editor = ctx.activeRustEditor;
-        const client = ctx.client;
-        if (!editor || !client) return;
+        if (!editor) return;
+
+        const client = await ctx.getClient();
 
         const response = await client.sendRequest(ra.matchingBrace, {
-            textDocument: ctx.client.code2ProtocolConverter.asTextDocumentIdentifier(
-                editor.document
-            ),
+            textDocument: client.code2ProtocolConverter.asTextDocumentIdentifier(editor.document),
             positions: editor.selections.map((s) =>
                 client.code2ProtocolConverter.asPosition(s.active)
             ),
@@ -116,14 +117,13 @@ export function matchingBrace(ctx: Ctx): Cmd {
 export function joinLines(ctx: Ctx): Cmd {
     return async () => {
         const editor = ctx.activeRustEditor;
-        const client = ctx.client;
-        if (!editor || !client) return;
+        if (!editor) return;
+
+        const client = await ctx.getClient();
 
         const items: lc.TextEdit[] = await client.sendRequest(ra.joinLines, {
             ranges: editor.selections.map((it) => client.code2ProtocolConverter.asRange(it)),
-            textDocument: ctx.client.code2ProtocolConverter.asTextDocumentIdentifier(
-                editor.document
-            ),
+            textDocument: client.code2ProtocolConverter.asTextDocumentIdentifier(editor.document),
         });
         const textEdits = await client.protocol2CodeConverter.asTextEdits(items);
         await editor.edit((builder) => {
@@ -145,14 +145,12 @@ export function moveItemDown(ctx: Ctx): Cmd {
 export function moveItem(ctx: Ctx, direction: ra.Direction): Cmd {
     return async () => {
         const editor = ctx.activeRustEditor;
-        const client = ctx.client;
-        if (!editor || !client) return;
+        if (!editor) return;
+        const client = await ctx.getClient();
 
         const lcEdits = await client.sendRequest(ra.moveItem, {
             range: client.code2ProtocolConverter.asRange(editor.selection),
-            textDocument: ctx.client.code2ProtocolConverter.asTextDocumentIdentifier(
-                editor.document
-            ),
+            textDocument: client.code2ProtocolConverter.asTextDocumentIdentifier(editor.document),
             direction,
         });
 
@@ -166,13 +164,13 @@ export function moveItem(ctx: Ctx, direction: ra.Direction): Cmd {
 export function onEnter(ctx: Ctx): Cmd {
     async function handleKeypress() {
         const editor = ctx.activeRustEditor;
-        const client = ctx.client;
 
-        if (!editor || !client) return false;
+        if (!editor) return false;
 
+        const client = await ctx.getClient();
         const lcEdits = await client
             .sendRequest(ra.onEnter, {
-                textDocument: ctx.client.code2ProtocolConverter.asTextDocumentIdentifier(
+                textDocument: client.code2ProtocolConverter.asTextDocumentIdentifier(
                     editor.document
                 ),
                 position: client.code2ProtocolConverter.asPosition(editor.selection.active),
@@ -198,14 +196,13 @@ export function onEnter(ctx: Ctx): Cmd {
 export function parentModule(ctx: Ctx): Cmd {
     return async () => {
         const editor = vscode.window.activeTextEditor;
-        const client = ctx.client;
-        if (!editor || !client) return;
+        if (!editor) return;
         if (!(isRustDocument(editor.document) || isCargoTomlDocument(editor.document))) return;
 
+        const client = await ctx.getClient();
+
         const locations = await client.sendRequest(ra.parentModule, {
-            textDocument: ctx.client.code2ProtocolConverter.asTextDocumentIdentifier(
-                editor.document
-            ),
+            textDocument: client.code2ProtocolConverter.asTextDocumentIdentifier(editor.document),
             position: client.code2ProtocolConverter.asPosition(editor.selection.active),
         });
         if (!locations) return;
@@ -236,13 +233,11 @@ export function parentModule(ctx: Ctx): Cmd {
 export function openCargoToml(ctx: Ctx): Cmd {
     return async () => {
         const editor = ctx.activeRustEditor;
-        const client = ctx.client;
-        if (!editor || !client) return;
+        if (!editor) return;
 
+        const client = await ctx.getClient();
         const response = await client.sendRequest(ra.openCargoToml, {
-            textDocument: ctx.client.code2ProtocolConverter.asTextDocumentIdentifier(
-                editor.document
-            ),
+            textDocument: client.code2ProtocolConverter.asTextDocumentIdentifier(editor.document),
         });
         if (!response) return;
 
@@ -259,12 +254,13 @@ export function openCargoToml(ctx: Ctx): Cmd {
 export function ssr(ctx: Ctx): Cmd {
     return async () => {
         const editor = vscode.window.activeTextEditor;
-        const client = ctx.client;
-        if (!editor || !client) return;
+        if (!editor) return;
+
+        const client = await ctx.getClient();
 
         const position = editor.selection.active;
         const selections = editor.selections;
-        const textDocument = ctx.client.code2ProtocolConverter.asTextDocumentIdentifier(
+        const textDocument = client.code2ProtocolConverter.asTextDocumentIdentifier(
             editor.document
         );
 
@@ -354,21 +350,22 @@ export function syntaxTree(ctx: Ctx): Cmd {
             }
         }
 
-        provideTextDocumentContent(
+        async provideTextDocumentContent(
             uri: vscode.Uri,
             ct: vscode.CancellationToken
-        ): vscode.ProviderResult<string> {
+        ): Promise<string> {
             const rustEditor = ctx.activeRustEditor;
             if (!rustEditor) return "";
+            const client = await ctx.getClient();
 
             // When the range based query is enabled we take the range of the selection
             const range =
                 uri.query === "range=true" && !rustEditor.selection.isEmpty
-                    ? ctx.client.code2ProtocolConverter.asRange(rustEditor.selection)
+                    ? client.code2ProtocolConverter.asRange(rustEditor.selection)
                     : null;
 
             const params = { textDocument: { uri: rustEditor.document.uri.toString() }, range };
-            return ctx.client.sendRequest(ra.syntaxTree, params, ct);
+            return client.sendRequest(ra.syntaxTree, params, ct);
         }
 
         get onDidChange(): vscode.Event<vscode.Uri> {
@@ -378,10 +375,10 @@ export function syntaxTree(ctx: Ctx): Cmd {
 
     void new AstInspector(ctx);
 
-    ctx.pushCleanup(
+    ctx.pushExtCleanup(
         vscode.workspace.registerTextDocumentContentProvider("rust-analyzer-syntax-tree", tdcp)
     );
-    ctx.pushCleanup(
+    ctx.pushExtCleanup(
         vscode.languages.setLanguageConfiguration("ra_syntax_tree", {
             brackets: [["[", ")"]],
         })
@@ -437,14 +434,14 @@ export function viewHir(ctx: Ctx): Cmd {
             }
         }
 
-        provideTextDocumentContent(
+        async provideTextDocumentContent(
             _uri: vscode.Uri,
             ct: vscode.CancellationToken
-        ): vscode.ProviderResult<string> {
+        ): Promise<string> {
             const rustEditor = ctx.activeRustEditor;
-            const client = ctx.client;
-            if (!rustEditor || !client) return "";
+            if (!rustEditor) return "";
 
+            const client = await ctx.getClient();
             const params = {
                 textDocument: client.code2ProtocolConverter.asTextDocumentIdentifier(
                     rustEditor.document
@@ -459,7 +456,7 @@ export function viewHir(ctx: Ctx): Cmd {
         }
     })();
 
-    ctx.pushCleanup(
+    ctx.pushExtCleanup(
         vscode.workspace.registerTextDocumentContentProvider("rust-analyzer-hir", tdcp)
     );
 
@@ -503,13 +500,13 @@ export function viewFileText(ctx: Ctx): Cmd {
             }
         }
 
-        provideTextDocumentContent(
+        async provideTextDocumentContent(
             _uri: vscode.Uri,
             ct: vscode.CancellationToken
-        ): vscode.ProviderResult<string> {
+        ): Promise<string> {
             const rustEditor = ctx.activeRustEditor;
-            const client = ctx.client;
-            if (!rustEditor || !client) return "";
+            if (!rustEditor) return "";
+            const client = await ctx.getClient();
 
             const params = client.code2ProtocolConverter.asTextDocumentIdentifier(
                 rustEditor.document
@@ -522,7 +519,7 @@ export function viewFileText(ctx: Ctx): Cmd {
         }
     })();
 
-    ctx.pushCleanup(
+    ctx.pushExtCleanup(
         vscode.workspace.registerTextDocumentContentProvider("rust-analyzer-file-text", tdcp)
     );
 
@@ -566,13 +563,13 @@ export function viewItemTree(ctx: Ctx): Cmd {
             }
         }
 
-        provideTextDocumentContent(
+        async provideTextDocumentContent(
             _uri: vscode.Uri,
             ct: vscode.CancellationToken
-        ): vscode.ProviderResult<string> {
+        ): Promise<string> {
             const rustEditor = ctx.activeRustEditor;
-            const client = ctx.client;
-            if (!rustEditor || !client) return "";
+            if (!rustEditor) return "";
+            const client = await ctx.getClient();
 
             const params = {
                 textDocument: client.code2ProtocolConverter.asTextDocumentIdentifier(
@@ -587,7 +584,7 @@ export function viewItemTree(ctx: Ctx): Cmd {
         }
     })();
 
-    ctx.pushCleanup(
+    ctx.pushExtCleanup(
         vscode.workspace.registerTextDocumentContentProvider("rust-analyzer-item-tree", tdcp)
     );
 
@@ -618,8 +615,8 @@ function crateGraph(ctx: Ctx, full: boolean): Cmd {
         const params = {
             full: full,
         };
-
-        const dot = await ctx.client.sendRequest(ra.viewCrateGraph, params);
+        const client = await ctx.getClient();
+        const dot = await client.sendRequest(ra.viewCrateGraph, params);
         const uri = panel.webview.asWebviewUri(nodeModulesPath);
 
         const html = `
@@ -690,13 +687,13 @@ export function expandMacro(ctx: Ctx): Cmd {
         eventEmitter = new vscode.EventEmitter<vscode.Uri>();
         async provideTextDocumentContent(_uri: vscode.Uri): Promise<string> {
             const editor = vscode.window.activeTextEditor;
-            const client = ctx.client;
-            if (!editor || !client) return "";
+            if (!editor) return "";
+            const client = await ctx.getClient();
 
             const position = editor.selection.active;
 
             const expanded = await client.sendRequest(ra.expandMacro, {
-                textDocument: ctx.client.code2ProtocolConverter.asTextDocumentIdentifier(
+                textDocument: client.code2ProtocolConverter.asTextDocumentIdentifier(
                     editor.document
                 ),
                 position,
@@ -712,7 +709,7 @@ export function expandMacro(ctx: Ctx): Cmd {
         }
     })();
 
-    ctx.pushCleanup(
+    ctx.pushExtCleanup(
         vscode.workspace.registerTextDocumentContentProvider("rust-analyzer-expand-macro", tdcp)
     );
 
@@ -724,11 +721,11 @@ export function expandMacro(ctx: Ctx): Cmd {
 }
 
 export function reloadWorkspace(ctx: Ctx): Cmd {
-    return async () => ctx.client.sendRequest(ra.reloadWorkspace);
+    return async () => (await ctx.getClient()).sendRequest(ra.reloadWorkspace);
 }
 
 async function showReferencesImpl(
-    client: LanguageClient,
+    client: LanguageClient | undefined,
     uri: string,
     position: lc.Position,
     locations: lc.Location[]
@@ -745,7 +742,7 @@ async function showReferencesImpl(
 
 export function showReferences(ctx: Ctx): Cmd {
     return async (uri: string, position: lc.Position, locations: lc.Location[]) => {
-        await showReferencesImpl(ctx.client, uri, position, locations);
+        await showReferencesImpl(await ctx.getClient(), uri, position, locations);
     };
 }
 
@@ -762,25 +759,23 @@ export function applyActionGroup(_ctx: Ctx): Cmd {
 
 export function gotoLocation(ctx: Ctx): Cmd {
     return async (locationLink: lc.LocationLink) => {
-        const client = ctx.client;
-        if (client) {
-            const uri = client.protocol2CodeConverter.asUri(locationLink.targetUri);
-            let range = client.protocol2CodeConverter.asRange(locationLink.targetSelectionRange);
-            // collapse the range to a cursor position
-            range = range.with({ end: range.start });
+        const client = await ctx.getClient();
+        const uri = client.protocol2CodeConverter.asUri(locationLink.targetUri);
+        let range = client.protocol2CodeConverter.asRange(locationLink.targetSelectionRange);
+        // collapse the range to a cursor position
+        range = range.with({ end: range.start });
 
-            await vscode.window.showTextDocument(uri, { selection: range });
-        }
+        await vscode.window.showTextDocument(uri, { selection: range });
     };
 }
 
 export function openDocs(ctx: Ctx): Cmd {
     return async () => {
-        const client = ctx.client;
         const editor = vscode.window.activeTextEditor;
-        if (!editor || !client) {
+        if (!editor) {
             return;
         }
+        const client = await ctx.getClient();
 
         const position = editor.selection.active;
         const textDocument = { uri: editor.document.uri.toString() };
@@ -795,20 +790,21 @@ export function openDocs(ctx: Ctx): Cmd {
 
 export function cancelFlycheck(ctx: Ctx): Cmd {
     return async () => {
-        await ctx.client.sendRequest(ra.cancelFlycheck);
+        const client = await ctx.getClient();
+        await client.sendRequest(ra.cancelFlycheck);
     };
 }
 
 export function resolveCodeAction(ctx: Ctx): Cmd {
-    const client = ctx.client;
     return async (params: lc.CodeAction) => {
+        const client = await ctx.getClient();
         params.command = undefined;
-        const item = await client.sendRequest(lc.CodeActionResolveRequest.type, params);
-        if (!item.edit) {
+        const item = await client?.sendRequest(lc.CodeActionResolveRequest.type, params);
+        if (!item?.edit) {
             return;
         }
         const itemEdit = item.edit;
-        const edit = await client.protocol2CodeConverter.asWorkspaceEdit(itemEdit);
+        const edit = await client?.protocol2CodeConverter.asWorkspaceEdit(itemEdit);
         // filter out all text edits and recreate the WorkspaceEdit without them so we can apply
         // snippet edits on our own
         const lcFileSystemEdit = {
@@ -847,11 +843,10 @@ export function run(ctx: Ctx): Cmd {
 }
 
 export function peekTests(ctx: Ctx): Cmd {
-    const client = ctx.client;
-
     return async () => {
         const editor = ctx.activeRustEditor;
-        if (!editor || !client) return;
+        if (!editor) return;
+        const client = await ctx.getClient();
 
         await vscode.window.withProgress(
             {
@@ -937,10 +932,10 @@ export function newDebugConfig(ctx: Ctx): Cmd {
     };
 }
 
-export function linkToCommand(ctx: Ctx): Cmd {
+export function linkToCommand(_: Ctx): Cmd {
     return async (commandId: string) => {
         const link = LINKED_COMMANDS.get(commandId);
-        if (ctx.client && link) {
+        if (link) {
             const { command, arguments: args = [] } = link;
             await vscode.commands.executeCommand(command, ...args);
         }