From 902d700a24f58669dfd20249d70a265ed5288315 Mon Sep 17 00:00:00 2001
From: Louis Dureuil <louis@meilisearch.com>
Date: Wed, 7 Feb 2024 15:48:07 +0100
Subject: [PATCH] Tracing trace: toggle the profiling of memory at runtime

---
 tracing-trace/src/layer.rs | 14 ++++++++++----
 tracing-trace/src/main.rs  |  2 +-
 2 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/tracing-trace/src/layer.rs b/tracing-trace/src/layer.rs
index b30e1ad38..f2f1d64ae 100644
--- a/tracing-trace/src/layer.rs
+++ b/tracing-trace/src/layer.rs
@@ -20,24 +20,26 @@ pub struct TraceLayer {
     sender: tokio::sync::mpsc::UnboundedSender<Entry>,
     callsites: RwLock<HashMap<OpaqueIdentifier, ResourceId>>,
     start_time: std::time::Instant,
+    profile_memory: bool,
 }
 
 impl Trace {
-    pub fn new() -> (Self, TraceLayer) {
+    pub fn new(profile_memory: bool) -> (Self, TraceLayer) {
         let (sender, receiver) = tokio::sync::mpsc::unbounded_channel();
         let trace = Trace { receiver };
         let layer = TraceLayer {
             sender,
             callsites: Default::default(),
             start_time: std::time::Instant::now(),
+            profile_memory,
         };
         (trace, layer)
     }
 }
 
 impl<W: Write> TraceWriter<W> {
-    pub fn new(writer: W) -> (Self, TraceLayer) {
-        let (trace, layer) = Trace::new();
+    pub fn new(writer: W, profile_memory: bool) -> (Self, TraceLayer) {
+        let (trace, layer) = Trace::new(profile_memory);
         (trace.into_writer(writer), layer)
     }
 
@@ -97,7 +99,11 @@ impl TraceLayer {
     }
 
     fn memory_stats(&self) -> Option<MemoryStats> {
-        MemoryStats::fetch()
+        if self.profile_memory {
+            MemoryStats::fetch()
+        } else {
+            None
+        }
     }
 
     fn send(&self, entry: Entry) {
diff --git a/tracing-trace/src/main.rs b/tracing-trace/src/main.rs
index c74bf1fb7..b9202727d 100644
--- a/tracing-trace/src/main.rs
+++ b/tracing-trace/src/main.rs
@@ -74,7 +74,7 @@ fn on_panic(info: &std::panic::PanicInfo) {
 
 fn main() {
     let (mut trace, profiling_layer) =
-        tracing_trace::TraceWriter::new(std::fs::File::create("trace.json").unwrap());
+        tracing_trace::TraceWriter::new(std::fs::File::create("trace.json").unwrap(), true);
 
     let subscriber = tracing_subscriber::registry()
         // any number of other subscriber layers may be added before or