Support memory profiling with dhat

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

Author: ChayimFriedman2

Cargo.lock

@@ -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

crates/rust-analyzer/Cargo.toml

@@ -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>,

crates/rust-analyzer/src/config.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(

crates/rust-analyzer/src/handlers/request.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);

crates/rust-analyzer/src/lib.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)
 }
 

crates/rust-analyzer/src/main_loop.rs

@@ -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);
     };
 }
 

editors/code/src/commands.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);
                 }

editors/code/src/snippets.ts

@@ -45,11 +45,22 @@ impl flags::Dist {
                 allocator,
                 self.zig,
                 self.pgo,
+                // Profiling requires debug information.
+                self.enable_profiling,
             )?;
             let release_tag = if stable { date_iso(sh)? } else { "nightly".to_owned() };
             dist_client(sh, &version, &release_tag, &target)?;
         } else {
-            dist_server(sh, "0.0.0-standalone", &target, allocator, self.zig, self.pgo)?;
+            dist_server(
+                sh,
+                "0.0.0-standalone",
+                &target,
+                allocator,
+                self.zig,
+                self.pgo,
+                // Profiling requires debug information.
+                self.enable_profiling,
+            )?;
         }
         Ok(())
     }
@@ -92,9 +103,11 @@ fn dist_server(
     allocator: Malloc,
     zig: bool,
     pgo: Option<PgoTrainingCrate>,
+    dev_rel: bool,
 ) -> anyhow::Result<()> {
     let _e = sh.push_env("CFG_RELEASE", release);
     let _e = sh.push_env("CARGO_PROFILE_RELEASE_LTO", "thin");
+    let _e = sh.push_env("CARGO_PROFILE_DEV_REL_LTO", "thin");
 
     // Uncomment to enable debug info for releases. Note that:
     //   * debug info is split on windows and macs, so it does nothing for those platforms,
@@ -120,7 +133,7 @@ fn dist_server(
         None
     };
 
-    let mut cmd = build_command(sh, command, &target_name, features);
+    let mut cmd = build_command(sh, command, &target_name, features, dev_rel);
     if let Some(profile) = pgo_profile {
         cmd = cmd.env("RUSTFLAGS", format!("-Cprofile-use={}", profile.to_str().unwrap()));
     }
@@ -141,10 +154,12 @@ fn build_command<'a>(
     command: &str,
     target_name: &str,
     features: &[&str],
+    dev_rel: bool,
 ) -> Cmd<'a> {
+    let profile = if dev_rel { "dev-rel" } else { "release" };
     cmd!(
         sh,
-        "cargo {command} --manifest-path ./crates/rust-analyzer/Cargo.toml --bin rust-analyzer --target {target_name} {features...} --release"
+        "cargo {command} --manifest-path ./crates/rust-analyzer/Cargo.toml --bin rust-analyzer --target {target_name} {features...} --profile {profile}"
     )
 }
 

xtask/src/dist.rs

@@ -42,6 +42,10 @@ xflags::xflags! {
             optional --mimalloc
             /// Use jemalloc allocator for server.
             optional --jemalloc
+            // Enable memory profiling support.
+            //
+            // **Warning:** This will produce a slower build of rust-analyzer, use only for profiling.
+            optional --enable-profiling
 
             /// Install the proc-macro server.
             optional --proc-macro-server
@@ -67,6 +71,10 @@ xflags::xflags! {
             optional --mimalloc
             /// Use jemalloc allocator for server
             optional --jemalloc
+            // Enable memory profiling support.
+            //
+            // **Warning:** This will produce a slower build of rust-analyzer, use only for profiling.
+            optional --enable-profiling
             optional --client-patch-version version: String
             /// Use cargo-zigbuild
             optional --zig
@@ -125,6 +133,7 @@ pub struct Install {
     pub server: bool,
     pub mimalloc: bool,
     pub jemalloc: bool,
+    pub enable_profiling: bool,
     pub proc_macro_server: bool,
     pub dev_rel: bool,
     pub force_always_assert: bool,
@@ -143,6 +152,7 @@ pub struct Release {
 pub struct Dist {
     pub mimalloc: bool,
     pub jemalloc: bool,
+    pub enable_profiling: bool,
     pub client_patch_version: Option<String>,
     pub zig: bool,
     pub pgo: Option<PgoTrainingCrate>,
@@ -280,6 +290,7 @@ pub(crate) enum Malloc {
     System,
     Mimalloc,
     Jemalloc,
+    Dhat,
 }
 
 impl Malloc {
@@ -288,6 +299,7 @@ impl Malloc {
             Malloc::System => &[][..],
             Malloc::Mimalloc => &["--features", "mimalloc"],
             Malloc::Jemalloc => &["--features", "jemalloc"],
+            Malloc::Dhat => &["--features", "dhat"],
         }
     }
 }
@@ -301,12 +313,15 @@ impl Install {
             Malloc::Mimalloc
         } else if self.jemalloc {
             Malloc::Jemalloc
+        } else if self.enable_profiling {
+            Malloc::Dhat
         } else {
             Malloc::System
         };
         Some(ServerOpt {
             malloc,
-            dev_rel: self.dev_rel,
+            // Profiling requires debug information.
+            dev_rel: self.dev_rel || self.enable_profiling,
             pgo: self.pgo.clone(),
             force_always_assert: self.force_always_assert,
         })
@@ -331,6 +346,8 @@ impl Dist {
             Malloc::Mimalloc
         } else if self.jemalloc {
             Malloc::Jemalloc
+        } else if self.enable_profiling {
+            Malloc::Dhat
         } else {
             Malloc::System
         }