Support memory profiling with dhat

Unfortunately, this requires a custom build of r-a, and it's quite slow.

Author: ChayimFriedman2

src/tools/rust-analyzer/Cargo.lock

@@ -418,6 +418,22 @@ dependencies = [
  "syn",
 ]
 
+[[package]]
+name = "dhat"
+version = "0.3.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "98cd11d84628e233de0ce467de10b8633f4ddaecafadefc86e13b84b8739b827"
+dependencies = [
+ "backtrace",
+ "lazy_static",
+ "mintex",
+ "parking_lot",
+ "rustc-hash 1.1.0",
+ "serde",
+ "serde_json",
+ "thousands",
+]
+
 [[package]]
 name = "dirs"
 version = "6.0.0"
@@ -1383,6 +1399,12 @@ dependencies = [
  "adler2",
 ]
 
+[[package]]
+name = "mintex"
+version = "0.1.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c505b3e17ed6b70a7ed2e67fbb2c560ee327353556120d6e72f5232b6880d536"
+
 [[package]]
 name = "mio"
 version = "1.1.0"
@@ -1452,7 +1474,7 @@ version = "0.50.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5"
 dependencies = [
- "windows-sys 0.60.2",
+ "windows-sys 0.61.0",
 ]
 
 [[package]]
@@ -2011,6 +2033,7 @@ dependencies = [
  "cargo_metadata 0.21.0",
  "cfg",
  "crossbeam-channel",
+ "dhat",
  "dirs",
  "dissimilar",
  "expect-test",
@@ -2528,6 +2551,12 @@ dependencies = [
  "syn",
 ]
 
+[[package]]
+name = "thousands"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3bf63baf9f5039dadc247375c29eb13706706cfde997d0330d05aa63a77d8820"
+
 [[package]]
 name = "thread_local"
 version = "1.1.9"

src/tools/rust-analyzer/crates/rust-analyzer/Cargo.toml

@@ -53,6 +53,7 @@ semver.workspace = true
 memchr = "2.7.5"
 cargo_metadata.workspace = true
 process-wrap.workspace = true
+dhat = { version = "0.3.3", optional = true }
 
 cfg.workspace = true
 hir-def.workspace = true
@@ -105,6 +106,7 @@ in-rust-tree = [
   "hir-ty/in-rust-tree",
   "load-cargo/in-rust-tree",
 ]
+dhat = ["dep:dhat"]
 
 [lints]
 workspace = true

src/tools/rust-analyzer/crates/rust-analyzer/src/config.rs

@@ -378,6 +378,12 @@ config_data! {
         /// Internal config, path to proc-macro server executable.
         procMacro_server: Option<Utf8PathBuf> = None,
 
+        /// The path where to save memory profiling output.
+        ///
+        /// **Note:** Memory profiling is not enabled by default in rust-analyzer builds, you need to build
+        /// from source for it.
+        profiling_memoryProfile: Option<Utf8PathBuf> = None,
+
         /// Exclude imports from find-all-references.
         references_excludeImports: bool = false,
 
@@ -2165,6 +2171,11 @@ impl Config {
         Some(AbsPathBuf::try_from(path).unwrap_or_else(|path| self.root_path.join(path)))
     }
 
+    pub fn dhat_output_file(&self) -> Option<AbsPathBuf> {
+        let path = self.profiling_memoryProfile().clone()?;
+        Some(AbsPathBuf::try_from(path).unwrap_or_else(|path| self.root_path.join(path)))
+    }
+
     pub fn ignored_proc_macros(
         &self,
         source_root: Option<SourceRootId>,

src/tools/rust-analyzer/crates/rust-analyzer/src/handlers/request.rs

@@ -126,17 +126,35 @@ pub(crate) fn handle_analyzer_status(
     Ok(buf)
 }
 
-pub(crate) fn handle_memory_usage(state: &mut GlobalState, _: ()) -> anyhow::Result<String> {
+pub(crate) fn handle_memory_usage(_state: &mut GlobalState, _: ()) -> anyhow::Result<String> {
     let _p = tracing::info_span!("handle_memory_usage").entered();
-    let mem = state.analysis_host.per_query_memory_usage();
 
-    let mut out = String::new();
-    for (name, bytes, entries) in mem {
-        format_to!(out, "{:>8} {:>6} {}\n", bytes, entries, name);
+    #[cfg(not(feature = "dhat"))]
+    {
+        Err(anyhow::anyhow!(
+            "Memory profiling is not enabled for this build of rust-analyzer.\n\n\
+            To build rust-analyzer with profiling support, pass `--features dhat --profile dev-rel` to `cargo build`
+            when building from source, or pass `--enable-profiling` to `cargo xtask`."
+        ))
+    }
+    #[cfg(feature = "dhat")]
+    {
+        if let Some(dhat_output_file) = _state.config.dhat_output_file() {
+            let mutprofiler = crate::DHAT_PROFILER.lock().unwrap();
+            let old_profiler = profiler.take();
+            // Need to drop the old profiler before creating a new one.
+            drop(old_profiler);
+            *profiler = Some(dhat::Profiler::builder().file_name(&dhat_output_file).build());
+            Ok(format!(
+                "Memory profile was saved successfully to {dhat_output_file}.\n\n\
+                See https://docs.rs/dhat/latest/dhat/#viewing for how to inspect the profile."
+            ))
+        } else {
+            Err(anyhow::anyhow!(
+                "Please set `rust-analyzer.profiling.memoryProfile` to the path where you want to save the profile."
+            ))
+        }
     }
-    format_to!(out, "{:>8}        Remaining\n", profile::memory_usage().allocated);
-
-    Ok(out)
 }
 
 pub(crate) fn handle_view_syntax_tree(

src/tools/rust-analyzer/crates/rust-analyzer/src/lib.rs

@@ -82,3 +82,10 @@ macro_rules! try_default_ {
     };
 }
 pub(crate) use try_default_ as try_default;
+
+#[cfg(feature = "dhat")]
+#[global_allocator]
+static ALLOC: dhat::Alloc = dhat::Alloc;
+
+#[cfg(feature = "dhat")]
+static DHAT_PROFILER: std::sync::Mutex<Option<dhat::Profiler>> = std::sync::Mutex::new(None);

src/tools/rust-analyzer/crates/rust-analyzer/src/main_loop.rs

@@ -60,6 +60,14 @@ pub fn main_loop(config: Config, connection: Connection) -> anyhow::Result<()> {
         SetThreadPriority(thread, thread_priority_above_normal);
     }
 
+    #[cfg(feature = "dhat")]
+    {
+        if let Some(dhat_output_file) = config.dhat_output_file() {
+            *crate::DHAT_PROFILER.lock().unwrap() =
+                Some(dhat::Profiler::builder().file_name(&dhat_output_file).build());
+        }
+    }
+
     GlobalState::new(connection.sender, config).run(connection.receiver)
 }
 

src/tools/rust-analyzer/docs/book/src/configuration_generated.md

@@ -1289,6 +1289,16 @@ Default: `null`
 Internal config, path to proc-macro server executable.
 
 
+## rust-analyzer.profiling.memoryProfile {#profiling.memoryProfile}
+
+Default: `null`
+
+The path where to save memory profiling output.
+
+**Note:** Memory profiling is not enabled by default in rust-analyzer builds, you need to build
+from source for it.
+
+
 ## rust-analyzer.references.excludeImports {#references.excludeImports}
 
 Default: `false`

src/tools/rust-analyzer/editors/code/package.json

@@ -2749,6 +2749,19 @@
                     }
                 }
             },
+            {
+                "title": "Profiling",
+                "properties": {
+                    "rust-analyzer.profiling.memoryProfile": {
+                        "markdownDescription": "The path where to save memory profiling output.\n\n**Note:** Memory profiling is not enabled by default in rust-analyzer builds, you need to build\nfrom source for it.",
+                        "default": null,
+                        "type": [
+                            "null",
+                            "string"
+                        ]
+                    }
+                }
+            },
             {
                 "title": "References",
                 "properties": {

src/tools/rust-analyzer/editors/code/src/commands.ts

@@ -71,32 +71,9 @@ export function analyzerStatus(ctx: CtxInit): Cmd {
 }
 
 export function memoryUsage(ctx: CtxInit): Cmd {
-    const tdcp = new (class implements vscode.TextDocumentContentProvider {
-        readonly uri = vscode.Uri.parse("rust-analyzer-memory://memory");
-        readonly eventEmitter = new vscode.EventEmitter<vscode.Uri>();
-
-        provideTextDocumentContent(_uri: vscode.Uri): vscode.ProviderResult<string> {
-            if (!vscode.window.activeTextEditor) return "";
-
-            // eslint-disable-next-line @typescript-eslint/no-explicit-any
-            return ctx.client.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> {
-            return this.eventEmitter.event;
-        }
-    })();
-
-    ctx.pushExtCleanup(
-        vscode.workspace.registerTextDocumentContentProvider("rust-analyzer-memory", tdcp),
-    );
-
     return async () => {
-        tdcp.eventEmitter.fire(tdcp.uri);
-        const document = await vscode.workspace.openTextDocument(tdcp.uri);
-        return vscode.window.showTextDocument(document, vscode.ViewColumn.Two, true);
+        const response = await ctx.client.sendRequest(ra.memoryUsage);
+        vscode.window.showInformationMessage(response);
     };
 }
 

src/tools/rust-analyzer/editors/code/src/snippets.ts

@@ -24,7 +24,9 @@ export async function applySnippetWorkspaceEdit(
                 for (const indel of edits) {
                     assert(
                         !(indel instanceof vscode.SnippetTextEdit),
-                        `bad ws edit: snippet received with multiple edits: ${JSON.stringify(edit)}`,
+                        `bad ws edit: snippet received with multiple edits: ${JSON.stringify(
+                            edit,
+                        )}`,
                     );
                     builder.replace(indel.range, indel.newText);
                 }