From b72146eee409e3dfdc282a327a649670e67fc31b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cl=C3=A9ment=20Renault?= <clement@meilisearch.com>
Date: Mon, 22 Jan 2024 14:32:57 +0100
Subject: [PATCH] Give the allocator to the tracer when necessary

---
 meilisearch/src/main.rs    | 4 +++-
 tracing-trace/src/layer.rs | 5 +++--
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/meilisearch/src/main.rs b/meilisearch/src/main.rs
index fd3879b36..373280e58 100644
--- a/meilisearch/src/main.rs
+++ b/meilisearch/src/main.rs
@@ -35,8 +35,10 @@ fn setup(opt: &Opt) -> anyhow::Result<()> {
 
     let file = std::fs::File::create(&trace_file)
         .with_context(|| format!("could not create trace file at '{}'", trace_file))?;
-    // TODO kero: Pass the allocator stats to Trace here
+    #[cfg(not(feature = "stats_alloc"))]
     let (mut trace, layer) = tracing_trace::Trace::new(file);
+    #[cfg(feature = "stats_alloc")]
+    let (mut trace, layer) = tracing_trace::Trace::with_stats_alloc(file, &ALLOC);
 
     let subscriber = tracing_subscriber::registry()
         .with(
diff --git a/tracing-trace/src/layer.rs b/tracing-trace/src/layer.rs
index f5464b699..66d54c674 100644
--- a/tracing-trace/src/layer.rs
+++ b/tracing-trace/src/layer.rs
@@ -83,7 +83,7 @@ enum OpaqueIdentifier {
     Call(tracing::callsite::Identifier),
 }
 
-impl TraceLayer {
+impl<A: GlobalAlloc> TraceLayer<A> {
     fn resource_id(&self, opaque: OpaqueIdentifier) -> Option<ResourceId> {
         self.callsites.read().unwrap().get(&opaque).copied()
     }
@@ -132,9 +132,10 @@ impl TraceLayer {
     }
 }
 
-impl<S> Layer<S> for TraceLayer
+impl<S, A> Layer<S> for TraceLayer<A>
 where
     S: Subscriber,
+    A: GlobalAlloc,
 {
     fn on_new_span(&self, attrs: &Attributes<'_>, id: &TracingId, _ctx: Context<'_, S>) {
         let call_id = self