ztracing: Enable memory profiling and callstack sampling (#47052)

While at it, annotate more functions that are potentially related to
language parsing in buffers.

Also, on macOS, in order to actually have callstack frames properly
recorded by Tracy, you need to manually run `dsymutil` on the binary.

Release Notes:

- N/A

Cargo.lock

@@ -2373,9 +2373,11 @@ dependencies = [
  "settings",
  "sum_tree",
  "text",
+ "tracing",
  "unindent",
  "util",
  "zlog",
+ "ztracing",
 ]
 
 [[package]]
@@ -5972,10 +5974,12 @@ dependencies = [
  "task",
  "tempfile",
  "toml 0.8.23",
+ "tracing",
  "url",
  "util",
  "wasm-encoder 0.221.3",
  "wasmparser 0.221.3",
+ "ztracing",
 ]
 
 [[package]]
@@ -6044,12 +6048,14 @@ dependencies = [
  "theme",
  "theme_extension",
  "toml 0.8.23",
+ "tracing",
  "url",
  "util",
  "wasmparser 0.221.3",
  "wasmtime",
  "wasmtime-wasi",
  "zlog",
+ "ztracing",
 ]
 
 [[package]]
@@ -8929,6 +8935,7 @@ dependencies = [
  "text",
  "theme",
  "toml 0.8.23",
+ "tracing",
  "tree-sitter",
  "tree-sitter-elixir",
  "tree-sitter-embedded-template",
@@ -8945,6 +8952,7 @@ dependencies = [
  "util",
  "watch",
  "zlog",
+ "ztracing",
 ]
 
 [[package]]
@@ -8964,7 +8972,9 @@ dependencies = [
  "project",
  "serde",
  "serde_json",
+ "tracing",
  "util",
+ "ztracing",
 ]
 
 [[package]]
@@ -20419,8 +20429,10 @@ dependencies = [
  "smol",
  "sum_tree",
  "text",
+ "tracing",
  "util",
  "zlog",
+ "ztracing",
 ]
 
 [[package]]
@@ -21351,6 +21363,7 @@ dependencies = [
  "tracing",
  "tracing-subscriber",
  "tracing-tracy",
+ "tracy-client",
  "zlog",
  "ztracing_macro",
 ]

crates/buffer_diff/Cargo.toml

@@ -26,7 +26,9 @@ rope.workspace = true
 settings = { workspace = true, optional = true }
 sum_tree.workspace = true
 text.workspace = true
+tracing.workspace = true
 util.workspace = true
+ztracing.workspace = true
 
 [dev-dependencies]
 ctor.workspace = true
@@ -37,3 +39,6 @@ settings.workspace = true
 text = { workspace = true, features = ["test-support"] }
 unindent.workspace = true
 zlog.workspace = true
+
+[package.metadata.cargo-machete]
+ignored = ["tracing"]

crates/buffer_diff/src/buffer_diff.rs

@@ -216,6 +216,7 @@ impl BufferDiffSnapshot {
         self.secondary_diff.as_deref()
     }
 
+    #[ztracing::instrument(skip_all)]
     pub fn hunks_intersecting_range<'a>(
         &'a self,
         range: Range<Anchor>,
@@ -1279,6 +1280,7 @@ impl BufferDiff {
         })
     }
 
+    #[ztracing::instrument(skip_all)]
     pub fn language_changed(
         &mut self,
         language: Option<Arc<Language>>,

crates/extension/Cargo.toml

@@ -31,10 +31,12 @@ serde.workspace = true
 serde_json.workspace = true
 task.workspace = true
 toml.workspace = true
+tracing.workspace = true
 url.workspace = true
 util.workspace = true
 wasm-encoder.workspace = true
 wasmparser.workspace = true
+ztracing.workspace = true
 
 [dev-dependencies]
 fs = { workspace = true, "features" = ["test-support"] }
@@ -42,3 +44,6 @@ gpui = { workspace = true, "features" = ["test-support"] }
 indoc.workspace = true
 pretty_assertions.workspace = true
 tempfile.workspace = true
+
+[package.metadata.cargo-machete]
+ignored = ["tracing"]

crates/extension/src/extension_host_proxy.rs

@@ -223,6 +223,7 @@ pub trait ExtensionGrammarProxy: Send + Sync + 'static {
 }
 
 impl ExtensionGrammarProxy for ExtensionHostProxy {
+    #[ztracing::instrument(skip_all)]
     fn register_grammars(&self, grammars: Vec<(Arc<str>, PathBuf)>) {
         let Some(proxy) = self.grammar_proxy.read().clone() else {
             return;
@@ -250,6 +251,7 @@ pub trait ExtensionLanguageProxy: Send + Sync + 'static {
 }
 
 impl ExtensionLanguageProxy for ExtensionHostProxy {
+    #[ztracing::instrument(skip_all, fields(lang = language.0.as_str()))]
     fn register_language(
         &self,
         language: LanguageName,

crates/extension_host/Cargo.toml

@@ -47,11 +47,13 @@ task.workspace = true
 telemetry.workspace = true
 tempfile.workspace = true
 toml.workspace = true
+tracing.workspace = true
 url.workspace = true
 util.workspace = true
 wasmparser.workspace = true
 wasmtime-wasi.workspace = true
 wasmtime.workspace = true
+ztracing.workspace = true
 
 [dev-dependencies]
 criterion.workspace = true
@@ -71,3 +73,6 @@ zlog.workspace = true
 [[bench]]
 name = "extension_compilation_benchmark"
 harness = false
+
+[package.metadata.cargo-machete]
+ignored = ["tracing"]

crates/extension_host/src/extension_host.rs

@@ -1074,6 +1074,7 @@ impl ExtensionStore {
     /// no longer in the manifest, or whose files have changed on disk.
     /// Then it loads any themes, languages, or grammars that are newly
     /// added to the manifest, or whose files have changed on disk.
+    #[ztracing::instrument(skip_all)]
     fn extensions_updated(
         &mut self,
         mut new_index: ExtensionIndex,

crates/git_ui/src/project_diff.rs

@@ -690,6 +690,7 @@ impl ProjectDiff {
         }
     }
 
+    #[instrument(skip_all)]
     pub async fn refresh(
         this: WeakEntity<Self>,
         reason: RefreshReason,

crates/language/Cargo.toml

@@ -62,6 +62,7 @@ sum_tree.workspace = true
 task.workspace = true
 text.workspace = true
 theme.workspace = true
+tracing.workspace = true
 tree-sitter-md = { workspace = true, optional = true }
 tree-sitter-python = { workspace = true, optional = true }
 tree-sitter-rust = { workspace = true, optional = true }
@@ -71,6 +72,7 @@ unicase = "2.6"
 util.workspace = true
 watch.workspace = true
 zlog.workspace = true
+ztracing.workspace = true
 diffy = "0.4.2"
 
 [dev-dependencies]
@@ -98,3 +100,6 @@ toml.workspace = true
 unindent.workspace = true
 util = { workspace = true, features = ["test-support"] }
 zlog.workspace = true
+
+[package.metadata.cargo-machete]
+ignored = ["tracing"]

crates/language/src/buffer.rs

@@ -1087,6 +1087,7 @@ impl Buffer {
     }
 
     /// Assign a language to the buffer, blocking for up to 1ms to reparse the buffer, returning the buffer.
+    #[ztracing::instrument(skip_all, fields(lang = language.config.name.0.as_str()))]
     pub fn with_language(mut self, language: Arc<Language>, cx: &mut Context<Self>) -> Self {
         self.set_language(Some(language), cx);
         self
@@ -1414,6 +1415,7 @@ impl Buffer {
 
     /// Retrieve a snapshot of the buffer's raw text, without any
     /// language-related state like the syntax tree or diagnostics.
+    #[ztracing::instrument(skip_all)]
     pub fn text_snapshot(&self) -> text::BufferSnapshot {
         self.text.snapshot()
     }
@@ -1463,6 +1465,7 @@ impl Buffer {
         self.set_language_(language, true, cx);
     }
 
+    #[ztracing::instrument(skip_all)]
     fn set_language_(
         &mut self,
         language: Option<Arc<Language>>,
@@ -1745,6 +1748,7 @@ impl Buffer {
     /// initiate an additional reparse recursively. To avoid concurrent parses
     /// for the same buffer, we only initiate a new parse if we are not already
     /// parsing in the background.
+    #[ztracing::instrument(skip_all)]
     pub fn reparse(&mut self, cx: &mut Context<Self>, may_block: bool) {
         if self.text.version() != *self.tree_sitter_data.version() {
             self.invalidate_tree_sitter_data(self.text.snapshot());
@@ -3590,6 +3594,7 @@ impl BufferSnapshot {
         None
     }
 
+    #[ztracing::instrument(skip_all)]
     fn get_highlights(&self, range: Range<usize>) -> (SyntaxMapCaptures<'_>, Vec<HighlightMap>) {
         let captures = self.syntax.captures(range, &self.text, |grammar| {
             grammar
@@ -3609,6 +3614,7 @@ impl BufferSnapshot {
     /// in an arbitrary way due to being stored in a [`Rope`](text::Rope). The text is also
     /// returned in chunks where each chunk has a single syntax highlighting style and
     /// diagnostic status.
+    #[ztracing::instrument(skip_all)]
     pub fn chunks<T: ToOffset>(&self, range: Range<T>, language_aware: bool) -> BufferChunks<'_> {
         let range = range.start.to_offset(self)..range.end.to_offset(self);
 

crates/language/src/language_registry.rs

@@ -922,6 +922,7 @@ impl LanguageRegistry {
         available_language
     }
 
+    #[ztracing::instrument(skip_all)]
     pub fn load_language(
         self: &Arc<Self>,
         language: &AvailableLanguage,
@@ -1009,6 +1010,7 @@ impl LanguageRegistry {
         rx
     }
 
+    #[ztracing::instrument(skip_all)]
     fn get_or_load_language(
         self: &Arc<Self>,
         callback: impl Fn(
@@ -1030,6 +1032,8 @@ impl LanguageRegistry {
         self: &Arc<Self>,
         name: Arc<str>,
     ) -> impl Future<Output = Result<tree_sitter::Language>> {
+        let span = ztracing::debug_span!("get_or_load_grammar", name = &*name.clone());
+        let _enter = span.enter();
         let (tx, rx) = oneshot::channel();
         let mut state = self.state.write();
 

crates/language/src/syntax_map.rs

@@ -21,7 +21,7 @@ use sum_tree::{Bias, Dimensions, SeekTarget, SumTree};
 use text::{Anchor, BufferSnapshot, OffsetRangeExt, Point, Rope, ToOffset, ToPoint};
 use tree_sitter::{
     Node, Query, QueryCapture, QueryCaptures, QueryCursor, QueryMatch, QueryMatches,
-    QueryPredicateArg, Tree,
+    QueryPredicateArg,
 };
 
 pub const MAX_BYTES_TO_QUERY: usize = 16 * 1024;
@@ -117,7 +117,7 @@ impl SyntaxLayerContent {
         }
     }
 
-    fn tree(&self) -> Option<&Tree> {
+    fn tree(&self) -> Option<&tree_sitter::Tree> {
         match self {
             SyntaxLayerContent::Parsed { tree, .. } => Some(tree),
             SyntaxLayerContent::Pending { .. } => None,
@@ -133,7 +133,7 @@ pub struct SyntaxLayer<'a> {
     pub language: &'a Arc<Language>,
     pub included_sub_ranges: Option<&'a [Range<Anchor>]>,
     pub(crate) depth: usize,
-    tree: &'a Tree,
+    tree: &'a tree_sitter::Tree,
     pub(crate) offset: (usize, tree_sitter::Point),
 }
 
@@ -296,6 +296,7 @@ impl SyntaxSnapshot {
         self.update_count
     }
 
+    #[ztracing::instrument(skip_all)]
     pub fn interpolate(&mut self, text: &BufferSnapshot) {
         let edits = text
             .anchored_edits_since::<Dimensions<usize, Point>>(&self.interpolated_version)
@@ -434,6 +435,7 @@ impl SyntaxSnapshot {
         self.reparse_(text, registry, root_language, Some(budget))
     }
 
+    #[ztracing::instrument(skip_all, fields(lang = root_language.config.name.0.as_str()))]
     fn reparse_(
         &mut self,
         text: &BufferSnapshot,
@@ -497,6 +499,7 @@ impl SyntaxSnapshot {
         Ok(())
     }
 
+    #[ztracing::instrument(skip_all)]
     fn reparse_with_ranges(
         &mut self,
         text: &BufferSnapshot,
@@ -885,7 +888,7 @@ impl SyntaxSnapshot {
     pub fn single_tree_captures<'a>(
         range: Range<usize>,
         text: &'a Rope,
-        tree: &'a Tree,
+        tree: &'a tree_sitter::Tree,
         language: &'a Arc<Language>,
         query: fn(&Grammar) -> Option<&Query>,
     ) -> SyntaxMapCaptures<'a> {
@@ -1422,14 +1425,15 @@ impl std::fmt::Display for ParseTimeout {
     }
 }
 
+#[ztracing::instrument(skip_all)]
 fn parse_text(
     grammar: &Grammar,
     text: &Rope,
     start_byte: usize,
     ranges: &[tree_sitter::Range],
-    old_tree: Option<Tree>,
+    old_tree: Option<tree_sitter::Tree>,
     parse_budget: &mut Option<Duration>,
-) -> anyhow::Result<Tree> {
+) -> anyhow::Result<tree_sitter::Tree> {
     with_parser(|parser| {
         let mut timed_out = false;
         let now = Instant::now();
@@ -1474,6 +1478,7 @@ fn parse_text(
     })
 }
 
+#[ztracing::instrument(skip_all)]
 fn get_injections(
     config: &InjectionConfig,
     text: &BufferSnapshot,
@@ -1705,6 +1710,7 @@ pub(crate) fn splice_included_ranges(
 /// different lines. For performance, only iterate through the given range of
 /// indices. All of the ranges in the array are relative to a given start byte
 /// and point.
+#[ztracing::instrument(skip_all)]
 fn insert_newlines_between_ranges(
     indices: Range<usize>,
     ranges: &mut Vec<tree_sitter::Range>,

crates/language_extension/Cargo.toml

@@ -25,4 +25,9 @@ lsp.workspace = true
 project.workspace = true
 serde.workspace = true
 serde_json.workspace = true
+tracing.workspace = true
 util.workspace = true
+ztracing.workspace = true
+
+[package.metadata.cargo-machete]
+ignored = ["tracing"]

crates/language_extension/src/language_extension.rs

@@ -37,6 +37,7 @@ struct LanguageServerRegistryProxy {
 }
 
 impl ExtensionGrammarProxy for LanguageServerRegistryProxy {
+    #[ztracing::instrument(skip_all)]
     fn register_grammars(&self, grammars: Vec<(Arc<str>, PathBuf)>) {
         self.language_registry.register_wasm_grammars(grammars)
     }

crates/multi_buffer/src/multi_buffer.rs

@@ -7928,6 +7928,7 @@ impl<'a> MultiBufferChunks<'a> {
         }
     }
 
+    #[ztracing::instrument(skip_all)]
     fn next_excerpt_chunk(&mut self) -> Option<Chunk<'a>> {
         loop {
             if self.excerpt_offset_range.is_empty() {
@@ -7975,6 +7976,7 @@ impl<'a> Iterator for ReversedMultiBufferChunks<'a> {
 impl<'a> Iterator for MultiBufferChunks<'a> {
     type Item = Chunk<'a>;
 
+    #[ztracing::instrument(skip_all)]
     fn next(&mut self) -> Option<Chunk<'a>> {
         if self.range.start >= self.range.end {
             return None;

crates/project/src/buffer_store.rs

@@ -86,6 +86,7 @@ enum OpenBuffer {
 
 pub enum BufferStoreEvent {
     BufferAdded(Entity<Buffer>),
+    // TODO(jk): this event seems unused
     BufferOpened {
         buffer: Entity<Buffer>,
         project_path: ProjectPath,
@@ -634,6 +635,7 @@ impl LocalBufferStore {
         self.save_local_buffer(buffer, worktree, path.path, true, cx)
     }
 
+    #[ztracing::instrument(skip_all)]
     fn open_buffer(
         &self,
         path: Arc<RelPath>,
@@ -844,6 +846,7 @@ impl BufferStore {
         }
     }
 
+    #[ztracing::instrument(skip_all)]
     pub fn open_buffer(
         &mut self,
         project_path: ProjectPath,

crates/project/src/git_store.rs

@@ -738,6 +738,7 @@ impl GitStore {
         })
     }
 
+    #[ztracing::instrument(skip_all)]
     pub fn open_uncommitted_diff(
         &mut self,
         buffer: Entity<Buffer>,
@@ -790,6 +791,7 @@ impl GitStore {
         cx.background_spawn(async move { task.await.map_err(|e| anyhow!("{e}")) })
     }
 
+    #[ztracing::instrument(skip_all)]
     async fn open_diff_internal(
         this: WeakEntity<Self>,
         kind: DiffKind,
@@ -2924,6 +2926,7 @@ impl BufferGitState {
         }
     }
 
+    #[ztracing::instrument(skip_all)]
     fn buffer_language_changed(&mut self, buffer: Entity<Buffer>, cx: &mut Context<Self>) {
         self.language = buffer.read(cx).language().cloned();
         self.language_changed = true;
@@ -3080,6 +3083,7 @@ impl BufferGitState {
         self.recalculate_diffs(buffer, cx)
     }
 
+    #[ztracing::instrument(skip_all)]
     fn recalculate_diffs(&mut self, buffer: text::BufferSnapshot, cx: &mut Context<Self>) {
         *self.recalculating_tx.borrow_mut() = true;
 

crates/project/src/project.rs

@@ -2820,6 +2820,7 @@ impl Project {
         }
     }
 
+    #[ztracing::instrument(skip_all)]
     pub fn open_buffer(
         &mut self,
         path: impl Into<ProjectPath>,
@@ -2872,6 +2873,7 @@ impl Project {
             .update(cx, |git_store, cx| git_store.open_unstaged_diff(buffer, cx))
     }
 
+    #[ztracing::instrument(skip_all)]
     pub fn open_uncommitted_diff(
         &mut self,
         buffer: Entity<Buffer>,

crates/worktree/Cargo.toml

@@ -48,7 +48,9 @@ smallvec.workspace = true
 smol.workspace = true
 sum_tree.workspace = true
 text.workspace = true
+tracing.workspace = true
 util.workspace = true
+ztracing.workspace = true
 
 [dev-dependencies]
 clock = { workspace = true, features = ["test-support"] }
@@ -63,3 +65,6 @@ rpc = { workspace = true, features = ["test-support"] }
 settings = { workspace = true, features = ["test-support"] }
 util = { workspace = true, features = ["test-support"] }
 zlog.workspace = true
+
+[package.metadata.cargo-machete]
+ignored = ["tracing"]

crates/worktree/src/worktree.rs

@@ -1342,6 +1342,7 @@ impl LocalWorktree {
         })
     }
 
+    #[ztracing::instrument(skip_all)]
     fn load_file(&self, path: &RelPath, cx: &Context<Worktree>) -> Task<Result<LoadedFile>> {
         let path = Arc::from(path);
         let abs_path = self.absolutize(&path);

crates/ztracing/Cargo.toml

@@ -9,7 +9,7 @@ license = "GPL-3.0-or-later"
 workspace = true
 
 [features]
-tracy = ["tracing-tracy"]
+tracy = ["tracing-tracy", "tracy-client"]
 
 [dependencies]
 zlog.workspace = true
@@ -17,5 +17,6 @@ tracing.workspace = true
 
 tracing-subscriber = "0.3.22"
 tracing-tracy = { version = "0.11.4", optional = true, features = ["enable", "ondemand"] }
+tracy-client = { version = "0.18.2", optional = true, features = ["enable", "ondemand"] }
 
 ztracing_macro.workspace = true

crates/ztracing/build.rs

@@ -2,8 +2,13 @@ use std::env;
 
 fn main() {
     if env::var_os("ZTRACING").is_some() {
-        println!(r"cargo::rustc-cfg=ztracing");
+        println!("cargo::rustc-cfg=ztracing");
+    }
+    if env::var_os("ZTRACING_WITH_MEMORY").is_some() {
+        println!("cargo::rustc-cfg=ztracing");
+        println!("cargo::rustc-cfg=ztracing_with_memory");
     }
     println!("cargo::rerun-if-changed=build.rs");
     println!("cargo::rerun-if-env-changed=ZTRACING");
+    println!("cargo::rerun-if-env-changed=ZTRACING_WITH_MEMORY");
 }

crates/ztracing/src/lib.rs

@@ -7,6 +7,11 @@ pub use tracing::{
 #[cfg(not(ztracing))]
 pub use ztracing_macro::instrument;
 
+#[cfg(all(ztracing, ztracing_with_memory))]
+#[global_allocator]
+static GLOBAL: tracy_client::ProfiledAllocator<std::alloc::System> =
+    tracy_client::ProfiledAllocator::new(std::alloc::System, 100);
+
 #[cfg(not(ztracing))]
 pub use __consume_all_tokens as trace_span;
 #[cfg(not(ztracing))]
@@ -46,10 +51,38 @@ impl Span {
 
 #[cfg(ztracing)]
 pub fn init() {
-    zlog::info!("Starting tracy subscriber, you can now connect the profiler");
+    use tracing_subscriber::fmt::format::DefaultFields;
     use tracing_subscriber::prelude::*;
+
+    #[derive(Default)]
+    struct TracyLayerConfig {
+        fmt: DefaultFields,
+    }
+
+    impl tracing_tracy::Config for TracyLayerConfig {
+        type Formatter = DefaultFields;
+
+        fn formatter(&self) -> &Self::Formatter {
+            &self.fmt
+        }
+
+        fn stack_depth(&self, _: &tracing::Metadata) -> u16 {
+            8
+        }
+
+        fn format_fields_in_zone_name(&self) -> bool {
+            true
+        }
+
+        fn on_error(&self, client: &tracy_client::Client, error: &'static str) {
+            client.color_message(error, 0xFF000000, 0);
+        }
+    }
+
+    zlog::info!("Starting tracy subscriber, you can now connect the profiler");
     tracing::subscriber::set_global_default(
-        tracing_subscriber::registry().with(tracing_tracy::TracyLayer::default()),
+        tracing_subscriber::registry()
+            .with(tracing_tracy::TracyLayer::new(TracyLayerConfig::default())),
     )
     .expect("setup tracy layer");
 }