vscode: fix memory leak on server restart centralize disposing

The memory leak was because on the server restrart the array of extensionContext.substiptions was not cleared

Author: Veetaha

editors/code/src/ctx.ts

@@ -3,16 +3,19 @@ import * as lc from 'vscode-languageclient';
 
 import { Config } from './config';
 import { createClient } from './client';
-import { isRustEditor, RustEditor } from './util';
+import { isRustEditor, RustEditor, dispose, Disposable } from './util';
 
 export class Ctx {
     private constructor(
         readonly config: Config,
         private readonly extCtx: vscode.ExtensionContext,
         readonly client: lc.LanguageClient,
         readonly serverPath: string,
-    ) {
+    ) { }
 
+    async dispose() {
+        dispose(this.extCtx.subscriptions);
+        await this.client.stop();
     }
 
     static async create(config: Config, extCtx: vscode.ExtensionContext, serverPath: string): Promise<Ctx> {
@@ -54,7 +57,4 @@ export class Ctx {
     }
 }
 
-export interface Disposable {
-    dispose(): void;
-}
 export type Cmd = (...args: any[]) => unknown;

editors/code/src/highlighting.ts

@@ -4,7 +4,7 @@ import * as ra from './rust-analyzer-api';
 import { ColorTheme, TextMateRuleSettings } from './color_theme';
 
 import { Ctx } from './ctx';
-import { sendRequestWithRetry, isRustDocument } from './util';
+import { sendRequestWithRetry, isRustDocument, dispose } from './util';
 
 export function activateHighlighting(ctx: Ctx) {
     const highlighter = new Highlighter(ctx);
@@ -83,10 +83,7 @@ class Highlighter {
         }
 
         // Decorations are removed when the object is disposed
-        for (const decoration of this.decorations.values()) {
-            decoration.dispose();
-        }
-
+        dispose(this.decorations);
         this.decorations = null;
     }
 

editors/code/src/inlay_hints.ts

@@ -2,26 +2,26 @@ import * as lc from "vscode-languageclient";
 import * as vscode from 'vscode';
 import * as ra from './rust-analyzer-api';
 
-import { Ctx, Disposable } from './ctx';
-import { sendRequestWithRetry, isRustDocument, RustDocument, RustEditor } from './util';
+import { Ctx } from './ctx';
+import { sendRequestWithRetry, isRustDocument, RustDocument, RustEditor, dispose, Disposable } from './util';
 
 
 export function activateInlayHints(ctx: Ctx) {
     const maybeUpdater = {
-        updater: null as null | HintsUpdater,
+        updater: undefined as undefined | HintsUpdater,
         onConfigChange() {
             if (
                 !ctx.config.inlayHints.typeHints &&
                 !ctx.config.inlayHints.parameterHints &&
                 !ctx.config.inlayHints.chainingHints
             ) {
-                return this.dispose();
+                return dispose(this);
             }
             if (!this.updater) this.updater = new HintsUpdater(ctx);
         },
         dispose() {
-            this.updater?.dispose();
-            this.updater = null;
+            dispose(this.updater);
+            this.updater = undefined;
         }
     };
 
@@ -115,8 +115,10 @@ class HintsUpdater implements Disposable {
 
     dispose() {
         this.sourceFiles.forEach(file => file.inlaysRequest?.cancel());
-        this.ctx.visibleRustEditors.forEach(editor => this.renderDecorations(editor, { param: [], type: [], chaining: [] }));
-        this.disposables.forEach(d => d.dispose());
+        this.ctx.visibleRustEditors.forEach(
+            editor => this.renderDecorations(editor, { param: [], type: [], chaining: [] })
+        );
+        dispose(this.disposables);
     }
 
     onDidChangeTextDocument({ contentChanges, document }: vscode.TextDocumentChangeEvent) {

editors/code/src/main.ts

@@ -9,7 +9,7 @@ import { activateStatusDisplay } from './status_display';
 import { Ctx } from './ctx';
 import { activateHighlighting } from './highlighting';
 import { Config, NIGHTLY_TAG } from './config';
-import { log, assert } from './util';
+import { log, assert, dispose } from './util';
 import { PersistentState } from './persistent_state';
 import { fetchRelease, download } from './net';
 import { spawnSync } from 'child_process';
@@ -48,21 +48,12 @@ export async function activate(context: vscode.ExtensionContext) {
     ctx = await Ctx.create(config, context, serverPath);
 
     // Commands which invokes manually via command palette, shortcut, etc.
-    ctx.registerCommand('reload', (ctx) => {
-        return async () => {
-            vscode.window.showInformationMessage('Reloading rust-analyzer...');
-            // @DanTup maneuver
-            // https://github.com/microsoft/vscode/issues/45774#issuecomment-373423895
-            await deactivate();
-            for (const sub of ctx.subscriptions) {
-                try {
-                    sub.dispose();
-                } catch (e) {
-                    log.error(e);
-                }
-            }
-            await activate(context);
-        };
+
+    // Reloading is inspired by @DanTup maneuver: https://github.com/microsoft/vscode/issues/45774#issuecomment-373423895
+    ctx.registerCommand('reload', _ => () => {
+        void vscode.window.showInformationMessage('Reloading rust-analyzer...');
+
+        deactivate().then(() => activate(context)).catch(log.error);
     });
 
     ctx.registerCommand('analyzerStatus', commands.analyzerStatus);
@@ -74,7 +65,7 @@ export async function activate(context: vscode.ExtensionContext) {
     ctx.registerCommand('expandMacro', commands.expandMacro);
     ctx.registerCommand('run', commands.run);
 
-    defaultOnEnter.dispose();
+    dispose(defaultOnEnter);
     ctx.registerCommand('onEnter', commands.onEnter);
 
     ctx.registerCommand('ssr', commands.ssr);
@@ -96,7 +87,7 @@ export async function activate(context: vscode.ExtensionContext) {
 }
 
 export async function deactivate() {
-    await ctx?.client?.stop();
+    await ctx?.dispose();
     ctx = undefined;
 }
 

editors/code/src/status_display.ts

@@ -3,6 +3,7 @@ import * as vscode from 'vscode';
 import { WorkDoneProgress, WorkDoneProgressBegin, WorkDoneProgressReport, WorkDoneProgressEnd, Disposable } from 'vscode-languageclient';
 
 import { Ctx } from './ctx';
+import { dispose } from './util';
 
 const spinnerFrames = ['⠋', '⠙', '⠹', '⠸', '⠼', '⠴', '⠦', '⠧', '⠇', '⠏'];
 
@@ -64,7 +65,7 @@ class StatusDisplay implements Disposable {
             this.timer = undefined;
         }
 
-        this.statusBarItem.dispose();
+        dispose(this.statusBarItem);
     }
 
     refreshLabel() {

editors/code/src/util.ts

@@ -82,3 +82,32 @@ export function isRustDocument(document: vscode.TextDocument): document is RustD
 export function isRustEditor(editor: vscode.TextEditor): editor is RustEditor {
     return isRustDocument(editor.document);
 }
+
+
+export interface Disposable {
+    dispose(): void;
+}
+
+
+export function dispose(disposable: undefined | Disposable | Disposable[] | Map<unknown, Disposable>) {
+    if (disposable === undefined) return;
+
+    if (disposable instanceof Array) {
+        while (disposable.length > 0) {
+            dispose(disposable.pop());
+        }
+        return;
+    }
+
+    if (disposable instanceof Map) {
+        disposable.forEach(dispose);
+        disposable.clear();
+        return;
+    }
+
+    try {
+        disposable.dispose();
+    } catch (err) {
+        log.error("Error was thrown while disposing an object", err);
+    }
+}