Refactor server lifecycle

12 files changed, 218 insertions, 201 deletions

diff --git a/editors/code/src/client.ts b/editors/code/src/client.ts
new file mode 100644
index 000000000..94948b10f
--- /dev/null
+++ b/editors/code/src/client.ts
@@ -0,0 +1,90 @@
+import { homedir } from 'os';
+import * as lc from 'vscode-languageclient';
+
+import { window, workspace } from 'vscode';
+import { Config } from './config';
+
+export function createClient(config: Config): lc.LanguageClient {
+    // '.' Is the fallback if no folder is open
+    // TODO?: Workspace folders support Uri's (eg: file://test.txt). It might be a good idea to test if the uri points to a file.
+    let folder: string = '.';
+    if (workspace.workspaceFolders !== undefined) {
+        folder = workspace.workspaceFolders[0].uri.fsPath.toString();
+    }
+
+    const command = expandPathResolving(config.raLspServerPath);
+    const run: lc.Executable = {
+        command,
+        options: { cwd: folder },
+    };
+    const serverOptions: lc.ServerOptions = {
+        run,
+        debug: run,
+    };
+    const traceOutputChannel = window.createOutputChannel(
+        'Rust Analyzer Language Server Trace',
+    );
+    const clientOptions: lc.LanguageClientOptions = {
+        documentSelector: [{ scheme: 'file', language: 'rust' }],
+        initializationOptions: {
+            publishDecorations: true,
+            lruCapacity: config.lruCapacity,
+            maxInlayHintLength: config.maxInlayHintLength,
+            cargoWatchEnable: config.cargoWatchOptions.enable,
+            cargoWatchArgs: config.cargoWatchOptions.arguments,
+            cargoWatchCommand: config.cargoWatchOptions.command,
+            cargoWatchAllTargets:
+                config.cargoWatchOptions.allTargets,
+            excludeGlobs: config.excludeGlobs,
+            useClientWatching: config.useClientWatching,
+            featureFlags: config.featureFlags,
+            withSysroot: config.withSysroot,
+            cargoFeatures: config.cargoFeatures,
+        },
+        traceOutputChannel,
+    };
+
+    const res = new lc.LanguageClient(
+        'rust-analyzer',
+        'Rust Analyzer Language Server',
+        serverOptions,
+        clientOptions,
+    );
+
+    // HACK: This is an awful way of filtering out the decorations notifications
+    // However, pending proper support, this is the most effecitve approach
+    // Proper support for this would entail a change to vscode-languageclient to allow not notifying on certain messages
+    // Or the ability to disable the serverside component of highlighting (but this means that to do tracing we need to disable hihlighting)
+    // This also requires considering our settings strategy, which is work which needs doing
+    // @ts-ignore The tracer is private to vscode-languageclient, but we need access to it to not log publishDecorations requests
+    res._tracer = {
+        log: (messageOrDataObject: string | any, data?: string) => {
+            if (typeof messageOrDataObject === 'string') {
+                if (
+                    messageOrDataObject.includes(
+                        'rust-analyzer/publishDecorations',
+                    ) ||
+                    messageOrDataObject.includes(
+                        'rust-analyzer/decorationsRequest',
+                    )
+                ) {
+                    // Don't log publish decorations requests
+                } else {
+                    // @ts-ignore This is just a utility function
+                    res.logTrace(messageOrDataObject, data);
+                }
+            } else {
+                // @ts-ignore
+                res.logObjectTrace(messageOrDataObject);
+            }
+        },
+    };
+    res.registerProposedFeatures()
+    return res;
+}
+function expandPathResolving(path: string) {
+    if (path.startsWith('~/')) {
+        return path.replace('~', homedir());
+    }
+    return path;
+}
diff --git a/editors/code/src/commands/analyzer_status.ts b/editors/code/src/commands/analyzer_status.ts
index 2c8362286..cf37dc6f0 100644
--- a/editors/code/src/commands/analyzer_status.ts
+++ b/editors/code/src/commands/analyzer_status.ts
@@ -49,9 +49,10 @@ class TextDocumentContentProvider
         _uri: vscode.Uri,
     ): vscode.ProviderResult<string> {
         const editor = vscode.window.activeTextEditor;
-        if (editor == null) return '';
+        const client = this.ctx.client
+        if (!editor || !client) return '';
 
-        return this.ctx.client.sendRequest<string>(
+        return client.sendRequest<string>(
             'rust-analyzer/analyzerStatus',
             null,
         );
diff --git a/editors/code/src/commands/expand_macro.ts b/editors/code/src/commands/expand_macro.ts
index da208257a..472f43b8d 100644
--- a/editors/code/src/commands/expand_macro.ts
+++ b/editors/code/src/commands/expand_macro.ts
@@ -52,14 +52,15 @@ class TextDocumentContentProvider
 
     async provideTextDocumentContent(_uri: vscode.Uri): Promise<string> {
         const editor = vscode.window.activeTextEditor;
-        if (editor == null) return '';
+        const client = this.ctx.client
+        if (!editor || !client) return '';
 
         const position = editor.selection.active;
         const request: lc.TextDocumentPositionParams = {
             textDocument: { uri: editor.document.uri.toString() },
             position,
         };
-        const expanded = await this.ctx.client.sendRequest<ExpandedMacro>(
+        const expanded = await client.sendRequest<ExpandedMacro>(
             'rust-analyzer/expandMacro',
             request,
         );
diff --git a/editors/code/src/commands/index.ts b/editors/code/src/commands/index.ts
index c28709c8a..4431fdcf6 100644
--- a/editors/code/src/commands/index.ts
+++ b/editors/code/src/commands/index.ts
@@ -15,18 +15,21 @@ import { run, runSingle } from './runnables';
 
 function collectGarbage(ctx: Ctx): Cmd {
     return async () => {
-        ctx.client.sendRequest<null>('rust-analyzer/collectGarbage', null);
+        ctx.client?.sendRequest<null>('rust-analyzer/collectGarbage', null);
     };
 }
 
 function showReferences(ctx: Ctx): Cmd {
     return (uri: string, position: lc.Position, locations: lc.Location[]) => {
-        vscode.commands.executeCommand(
-            'editor.action.showReferences',
-            vscode.Uri.parse(uri),
-            ctx.client.protocol2CodeConverter.asPosition(position),
-            locations.map(ctx.client.protocol2CodeConverter.asLocation),
-        );
+        let client = ctx.client;
+        if (client) {
+            vscode.commands.executeCommand(
+                'editor.action.showReferences',
+                vscode.Uri.parse(uri),
+                client.protocol2CodeConverter.asPosition(position),
+                locations.map(client.protocol2CodeConverter.asLocation),
+            );
+        }
     };
 }
 
@@ -36,6 +39,13 @@ function applySourceChange(ctx: Ctx): Cmd {
     }
 }
 
+function reload(ctx: Ctx): Cmd {
+    return async () => {
+        vscode.window.showInformationMessage('Reloading rust-analyzer...');
+        await ctx.restartServer();
+    }
+}
+
 export {
     analyzerStatus,
     expandMacro,
@@ -49,4 +59,5 @@ export {
     runSingle,
     showReferences,
     applySourceChange,
+    reload
 };
diff --git a/editors/code/src/commands/join_lines.ts b/editors/code/src/commands/join_lines.ts
index f4f902cf9..7b08c3255 100644
--- a/editors/code/src/commands/join_lines.ts
+++ b/editors/code/src/commands/join_lines.ts
@@ -6,13 +6,14 @@ import { applySourceChange, SourceChange } from '../source_change';
 export function joinLines(ctx: Ctx): Cmd {
     return async () => {
         const editor = ctx.activeRustEditor;
-        if (!editor) return;
+        const client = ctx.client;
+        if (!editor || !client) return;
 
         const request: JoinLinesParams = {
-            range: ctx.client.code2ProtocolConverter.asRange(editor.selection),
+            range: client.code2ProtocolConverter.asRange(editor.selection),
             textDocument: { uri: editor.document.uri.toString() },
         };
-        const change = await ctx.client.sendRequest<SourceChange>(
+        const change = await client.sendRequest<SourceChange>(
             'rust-analyzer/joinLines',
             request,
         );
diff --git a/editors/code/src/ctx.ts b/editors/code/src/ctx.ts
index 393d6a602..13988056a 100644
--- a/editors/code/src/ctx.ts
+++ b/editors/code/src/ctx.ts
@@ -1,19 +1,38 @@
 import * as vscode from 'vscode';
 import * as lc from 'vscode-languageclient';
-import { Server } from './server';
 import { Config } from './config';
+import { createClient } from './client'
 
 export class Ctx {
     readonly config: Config;
+    // Because we have "reload server" action, various listeners **will** face a
+    // situation where the client is not ready yet, and should be prepared to
+    // deal with it.
+    //
+    // Ideally, this should be replaced with async getter though.
+    client: lc.LanguageClient | null = null
     private extCtx: vscode.ExtensionContext;
+    private onDidRestartHooks: Array<(client: lc.LanguageClient) => void> = [];
 
     constructor(extCtx: vscode.ExtensionContext) {
         this.config = new Config(extCtx)
         this.extCtx = extCtx;
     }
 
-    get client(): lc.LanguageClient {
-        return Server.client;
+    async restartServer() {
+        let old = this.client;
+        if (old) {
+            await old.stop()
+        }
+        this.client = null;
+        const client = createClient(this.config);
+        this.pushCleanup(client.start());
+        await client.onReady();
+
+        this.client = client
+        for (const hook of this.onDidRestartHooks) {
+            hook(client)
+        }
     }
 
     get activeRustEditor(): vscode.TextEditor | undefined {
@@ -60,35 +79,34 @@ export class Ctx {
         this.extCtx.subscriptions.push(d);
     }
 
-    async sendRequestWithRetry<R>(
-        method: string,
-        param: any,
-        token?: vscode.CancellationToken,
-    ): Promise<R> {
-        await this.client.onReady();
-        for (const delay of [2, 4, 6, 8, 10, null]) {
-            try {
-                return await (token ? this.client.sendRequest(method, param, token) : this.client.sendRequest(method, param));
-            } catch (e) {
-                if (
-                    e.code === lc.ErrorCodes.ContentModified &&
-                    delay !== null
-                ) {
-                    await sleep(10 * (1 << delay));
-                    continue;
-                }
-                throw e;
-            }
-        }
-        throw 'unreachable';
-    }
-
-    onNotification(method: string, handler: lc.GenericNotificationHandler) {
-        this.client.onReady()
-            .then(() => this.client.onNotification(method, handler))
+    onDidRestart(hook: (client: lc.LanguageClient) => void) {
+        this.onDidRestartHooks.push(hook)
     }
 }
 
 export type Cmd = (...args: any[]) => any;
 
+export async function sendRequestWithRetry<R>(
+    client: lc.LanguageClient,
+    method: string,
+    param: any,
+    token?: vscode.CancellationToken,
+): Promise<R> {
+    for (const delay of [2, 4, 6, 8, 10, null]) {
+        try {
+            return await (token ? client.sendRequest(method, param, token) : client.sendRequest(method, param));
+        } catch (e) {
+            if (
+                e.code === lc.ErrorCodes.ContentModified &&
+                delay !== null
+            ) {
+                await sleep(10 * (1 << delay));
+                continue;
+            }
+            throw e;
+        }
+    }
+    throw 'unreachable';
+}
+
 const sleep = (ms: number) => new Promise(resolve => setTimeout(resolve, ms));
diff --git a/editors/code/src/highlighting.ts b/editors/code/src/highlighting.ts
index d4e961b5b..f9d2e9d90 100644
--- a/editors/code/src/highlighting.ts
+++ b/editors/code/src/highlighting.ts
@@ -5,31 +5,32 @@ const seedrandom = seedrandom_; // https://github.com/jvandemo/generator-angular
 
 import { ColorTheme, TextMateRuleSettings } from './color_theme';
 
-import { Ctx } from './ctx';
+import { Ctx, sendRequestWithRetry } from './ctx';
 
 export function activateHighlighting(ctx: Ctx) {
     const highlighter = new Highlighter(ctx);
-
-    ctx.onNotification(
-        'rust-analyzer/publishDecorations',
-        (params: PublishDecorationsParams) => {
-            if (!ctx.config.highlightingOn) return;
-
-            const targetEditor = vscode.window.visibleTextEditors.find(
-                editor => {
-                    const unescapedUri = unescape(
-                        editor.document.uri.toString(),
-                    );
-                    // Unescaped URI looks like:
-                    // file:///c:/Workspace/ra-test/src/main.rs
-                    return unescapedUri === params.uri;
-                },
-            );
-            if (!targetEditor) return;
-
-            highlighter.setHighlights(targetEditor, params.decorations);
-        },
-    );
+    ctx.onDidRestart(client => {
+        client.onNotification(
+            'rust-analyzer/publishDecorations',
+            (params: PublishDecorationsParams) => {
+                if (!ctx.config.highlightingOn) return;
+
+                const targetEditor = vscode.window.visibleTextEditors.find(
+                    editor => {
+                        const unescapedUri = unescape(
+                            editor.document.uri.toString(),
+                        );
+                        // Unescaped URI looks like:
+                        // file:///c:/Workspace/ra-test/src/main.rs
+                        return unescapedUri === params.uri;
+                    },
+                );
+                if (!targetEditor) return;
+
+                highlighter.setHighlights(targetEditor, params.decorations);
+            },
+        );
+    })
 
     vscode.workspace.onDidChangeConfiguration(
         _ => highlighter.removeHighlights(),
@@ -40,11 +41,14 @@ export function activateHighlighting(ctx: Ctx) {
         async (editor: vscode.TextEditor | undefined) => {
             if (!editor || editor.document.languageId !== 'rust') return;
             if (!ctx.config.highlightingOn) return;
+            let client = ctx.client;
+            if (!client) return;
 
             const params: lc.TextDocumentIdentifier = {
                 uri: editor.document.uri.toString(),
             };
-            const decorations = await ctx.sendRequestWithRetry<Decoration[]>(
+            const decorations = await sendRequestWithRetry<Decoration[]>(
+                client,
                 'rust-analyzer/decorationsRequest',
                 params,
             );
@@ -103,6 +107,8 @@ class Highlighter {
     }
 
     public setHighlights(editor: vscode.TextEditor, highlights: Decoration[]) {
+        let client = this.ctx.client;
+        if (!client) return;
         // Initialize decorations if necessary
         //
         // Note: decoration objects need to be kept around so we can dispose them
@@ -135,13 +141,13 @@ class Highlighter {
                 colorfulIdents
                     .get(d.bindingHash)![0]
                     .push(
-                        this.ctx.client.protocol2CodeConverter.asRange(d.range),
+                        client.protocol2CodeConverter.asRange(d.range),
                     );
             } else {
                 byTag
                     .get(d.tag)!
                     .push(
-                        this.ctx.client.protocol2CodeConverter.asRange(d.range),
+                        client.protocol2CodeConverter.asRange(d.range),
                     );
             }
         }
diff --git a/editors/code/src/inlay_hints.ts b/editors/code/src/inlay_hints.ts
index b6eb70168..e74d6996f 100644
--- a/editors/code/src/inlay_hints.ts
+++ b/editors/code/src/inlay_hints.ts
@@ -1,7 +1,7 @@
 import * as vscode from 'vscode';
 import * as lc from 'vscode-languageclient';
 
-import { Ctx } from './ctx';
+import { Ctx, sendRequestWithRetry } from './ctx';
 
 export function activateInlayHints(ctx: Ctx) {
     const hintsUpdater = new HintsUpdater(ctx);
@@ -19,9 +19,7 @@ export function activateInlayHints(ctx: Ctx) {
         hintsUpdater.setEnabled(ctx.config.displayInlayHints);
     }, ctx.subscriptions);
 
-    // XXX: don't await here;
-    // Who knows what happens if an exception is thrown here...
-    hintsUpdater.refresh();
+    ctx.onDidRestart(_ => hintsUpdater.setEnabled(ctx.config.displayInlayHints))
 }
 
 interface InlayHintsParams {
@@ -97,6 +95,8 @@ class HintsUpdater {
     }
 
     private async queryHints(documentUri: string): Promise<InlayHint[] | null> {
+        let client = this.ctx.client;
+        if (!client) return null
         const request: InlayHintsParams = {
             textDocument: { uri: documentUri },
         };
@@ -105,7 +105,8 @@ class HintsUpdater {
         if (prev) prev.cancel();
         this.pending.set(documentUri, tokenSource);
         try {
-            return await this.ctx.sendRequestWithRetry<InlayHint[] | null>(
+            return await sendRequestWithRetry<InlayHint[] | null>(
+                client,
                 'rust-analyzer/inlayHints',
                 request,
                 tokenSource.token,
diff --git a/editors/code/src/main.ts b/editors/code/src/main.ts
index 3d9107927..22450060b 100644
--- a/editors/code/src/main.ts
+++ b/editors/code/src/main.ts
@@ -3,7 +3,6 @@ import * as vscode from 'vscode';
 import * as commands from './commands';
 import { activateInlayHints } from './inlay_hints';
 import { activateStatusDisplay } from './status_display';
-import { Server } from './server';
 import { Ctx } from './ctx';
 import { activateHighlighting } from './highlighting';
 
@@ -21,6 +20,7 @@ export async function activate(context: vscode.ExtensionContext) {
     ctx.registerCommand('syntaxTree', commands.syntaxTree);
     ctx.registerCommand('expandMacro', commands.expandMacro);
     ctx.registerCommand('run', commands.run);
+    ctx.registerCommand('reload', commands.reload);
 
     // Internal commands which are invoked by the server.
     ctx.registerCommand('runSingle', commands.runSingle);
@@ -30,38 +30,17 @@ export async function activate(context: vscode.ExtensionContext) {
     if (ctx.config.enableEnhancedTyping) {
         ctx.overrideCommand('type', commands.onEnter);
     }
-
-    const startServer = () => Server.start(ctx.config);
-    const reloadCommand = () => reloadServer(startServer);
-
-    vscode.commands.registerCommand('rust-analyzer.reload', reloadCommand);
-
+    activateStatusDisplay(ctx);
+    activateHighlighting(ctx);
+    activateInlayHints(ctx);
     // Start the language server, finally!
     try {
-        await startServer();
+        await ctx.restartServer();
     } catch (e) {
         vscode.window.showErrorMessage(e.message);
     }
-
-    activateStatusDisplay(ctx);
-    activateHighlighting(ctx);
-
-    if (ctx.config.displayInlayHints) {
-        activateInlayHints(ctx);
-    }
 }
 
-export function deactivate(): Thenable<void> {
-    if (!Server.client) {
-        return Promise.resolve();
-    }
-    return Server.client.stop();
-}
-
-async function reloadServer(startServer: () => Promise<void>) {
-    if (Server.client != null) {
-        vscode.window.showInformationMessage('Reloading rust-analyzer...');
-        await Server.client.stop();
-        await startServer();
-    }
+export async function deactivate() {
+    await ctx?.client?.stop();
 }
diff --git a/editors/code/src/server.ts b/editors/code/src/server.ts
deleted file mode 100644
index ab9f3bfa6..000000000
--- a/editors/code/src/server.ts
+++ /dev/null
@@ -1,96 +0,0 @@
-import { homedir } from 'os';
-import * as lc from 'vscode-languageclient';
-
-import { window, workspace } from 'vscode';
-import { Config } from './config';
-
-function expandPathResolving(path: string) {
-    if (path.startsWith('~/')) {
-        return path.replace('~', homedir());
-    }
-    return path;
-}
-
-export class Server {
-    static config: Config;
-    public static client: lc.LanguageClient;
-
-    public static async start(config: Config) {
-        // '.' Is the fallback if no folder is open
-        // TODO?: Workspace folders support Uri's (eg: file://test.txt). It might be a good idea to test if the uri points to a file.
-        let folder: string = '.';
-        if (workspace.workspaceFolders !== undefined) {
-            folder = workspace.workspaceFolders[0].uri.fsPath.toString();
-        }
-
-        this.config = config;
-        const command = expandPathResolving(this.config.raLspServerPath);
-        const run: lc.Executable = {
-            command,
-            options: { cwd: folder },
-        };
-        const serverOptions: lc.ServerOptions = {
-            run,
-            debug: run,
-        };
-        const traceOutputChannel = window.createOutputChannel(
-            'Rust Analyzer Language Server Trace',
-        );
-        const clientOptions: lc.LanguageClientOptions = {
-            documentSelector: [{ scheme: 'file', language: 'rust' }],
-            initializationOptions: {
-                publishDecorations: true,
-                lruCapacity: Server.config.lruCapacity,
-                maxInlayHintLength: Server.config.maxInlayHintLength,
-                cargoWatchEnable: Server.config.cargoWatchOptions.enable,
-                cargoWatchArgs: Server.config.cargoWatchOptions.arguments,
-                cargoWatchCommand: Server.config.cargoWatchOptions.command,
-                cargoWatchAllTargets:
-                    Server.config.cargoWatchOptions.allTargets,
-                excludeGlobs: Server.config.excludeGlobs,
-                useClientWatching: Server.config.useClientWatching,
-                featureFlags: Server.config.featureFlags,
-                withSysroot: Server.config.withSysroot,
-                cargoFeatures: Server.config.cargoFeatures,
-            },
-            traceOutputChannel,
-        };
-
-        Server.client = new lc.LanguageClient(
-            'rust-analyzer',
-            'Rust Analyzer Language Server',
-            serverOptions,
-            clientOptions,
-        );
-        // HACK: This is an awful way of filtering out the decorations notifications
-        // However, pending proper support, this is the most effecitve approach
-        // Proper support for this would entail a change to vscode-languageclient to allow not notifying on certain messages
-        // Or the ability to disable the serverside component of highlighting (but this means that to do tracing we need to disable hihlighting)
-        // This also requires considering our settings strategy, which is work which needs doing
-        // @ts-ignore The tracer is private to vscode-languageclient, but we need access to it to not log publishDecorations requests
-        Server.client._tracer = {
-            log: (messageOrDataObject: string | any, data?: string) => {
-                if (typeof messageOrDataObject === 'string') {
-                    if (
-                        messageOrDataObject.includes(
-                            'rust-analyzer/publishDecorations',
-                        ) ||
-                        messageOrDataObject.includes(
-                            'rust-analyzer/decorationsRequest',
-                        )
-                    ) {
-                        // Don't log publish decorations requests
-                    } else {
-                        // @ts-ignore This is just a utility function
-                        Server.client.logTrace(messageOrDataObject, data);
-                    }
-                } else {
-                    // @ts-ignore
-                    Server.client.logObjectTrace(messageOrDataObject);
-                }
-            },
-        };
-        Server.client.registerProposedFeatures();
-        Server.client.start();
-    }
-}
diff --git a/editors/code/src/source_change.ts b/editors/code/src/source_change.ts
index a4f9068b2..887191d9e 100644
--- a/editors/code/src/source_change.ts
+++ b/editors/code/src/source_change.ts
@@ -10,7 +10,10 @@ export interface SourceChange {
 }
 
 export async function applySourceChange(ctx: Ctx, change: SourceChange) {
-    const wsEdit = ctx.client.protocol2CodeConverter.asWorkspaceEdit(
+    const client = ctx.client;
+    if (!client) return
+
+    const wsEdit = client.protocol2CodeConverter.asWorkspaceEdit(
         change.workspaceEdit,
     );
     let created;
@@ -32,10 +35,10 @@ export async function applySourceChange(ctx: Ctx, change: SourceChange) {
         const doc = await vscode.workspace.openTextDocument(toOpenUri);
         await vscode.window.showTextDocument(doc);
     } else if (toReveal) {
-        const uri = ctx.client.protocol2CodeConverter.asUri(
+        const uri = client.protocol2CodeConverter.asUri(
             toReveal.textDocument.uri,
         );
-        const position = ctx.client.protocol2CodeConverter.asPosition(
+        const position = client.protocol2CodeConverter.asPosition(
             toReveal.position,
         );
         const editor = vscode.window.activeTextEditor;
diff --git a/editors/code/src/status_display.ts b/editors/code/src/status_display.ts
index e3719075b..1454bf8b0 100644
--- a/editors/code/src/status_display.ts
+++ b/editors/code/src/status_display.ts
@@ -7,7 +7,9 @@ const spinnerFrames = ['⠋', '⠙', '⠹', '⠸', '⠼', '⠴', '⠦', '⠧', '
 export function activateStatusDisplay(ctx: Ctx) {
     const statusDisplay = new StatusDisplay(ctx.config.cargoWatchOptions.command);
     ctx.pushCleanup(statusDisplay);
-    ctx.onNotification('$/progress', params => statusDisplay.handleProgressNotification(params));
+    ctx.onDidRestart(client => {
+        client.onNotification('$/progress', params => statusDisplay.handleProgressNotification(params));
+    })
 }
 
 class StatusDisplay implements vscode.Disposable {