From 497d15685c5875a90a655f4187f38dba58c72372 Mon Sep 17 00:00:00 2001
From: Louis Dureuil <louis@meilisearch.com>
Date: Mon, 17 Jun 2024 16:37:50 +0200
Subject: [PATCH] Use inspector allocator

---
 Cargo.lock                                  | 10 ++++++++++
 Cargo.toml                                  |  2 +-
 meilisearch/Cargo.toml                      |  2 ++
 meilisearch/src/main.rs                     |  8 +++-----
 meilisearch/src/routes/indexes/documents.rs | 12 ++++++++++++
 5 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index fad60e8da..2c052bdab 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2506,6 +2506,14 @@ dependencies = [
  "generic-array",
 ]
 
+[[package]]
+name = "inspecting-allocator"
+version = "1.8.0"
+dependencies = [
+ "tracing",
+ "tracing-error",
+]
+
 [[package]]
 name = "insta"
 version = "1.34.0"
@@ -3316,6 +3324,7 @@ dependencies = [
  "http 0.2.11",
  "index-scheduler",
  "indexmap",
+ "inspecting-allocator",
  "insta",
  "is-terminal",
  "itertools 0.11.0",
@@ -3365,6 +3374,7 @@ dependencies = [
  "toml",
  "tracing",
  "tracing-actix-web",
+ "tracing-error",
  "tracing-subscriber",
  "tracing-trace",
  "url",
diff --git a/Cargo.toml b/Cargo.toml
index 10996065a..7da2c8d98 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -18,7 +18,7 @@ members = [
     "fuzzers",
     "tracing-trace",
     "xtask",
-    "build-info",
+    "build-info", "inspecting-allocator",
 ]
 
 [workspace.package]
diff --git a/meilisearch/Cargo.toml b/meilisearch/Cargo.toml
index 4a2b11b21..d7a278996 100644
--- a/meilisearch/Cargo.toml
+++ b/meilisearch/Cargo.toml
@@ -108,6 +108,8 @@ tracing-subscriber = { version = "0.3.18", features = ["json"] }
 tracing-trace = { version = "0.1.0", path = "../tracing-trace" }
 tracing-actix-web = "0.7.9"
 build-info = { version = "1.7.0", path = "../build-info" }
+inspecting-allocator = { version = "1.8.0", path = "../inspecting-allocator" }
+tracing-error = { version = "0.2.0", default-features = false }
 
 [dev-dependencies]
 actix-rt = "2.9.0"
diff --git a/meilisearch/src/main.rs b/meilisearch/src/main.rs
index af02f58e1..c2dd95b9d 100644
--- a/meilisearch/src/main.rs
+++ b/meilisearch/src/main.rs
@@ -16,15 +16,11 @@ use meilisearch::{
     LogStderrType, Opt, SubscriberForSecondLayer,
 };
 use meilisearch_auth::{generate_master_key, AuthController, MASTER_KEY_MIN_SIZE};
-use mimalloc::MiMalloc;
 use termcolor::{Color, ColorChoice, ColorSpec, StandardStream, WriteColor};
 use tracing::level_filters::LevelFilter;
 use tracing_subscriber::layer::SubscriberExt as _;
 use tracing_subscriber::Layer;
 
-#[global_allocator]
-static ALLOC: MiMalloc = MiMalloc;
-
 fn default_log_route_layer() -> LogRouteType {
     None.with_filter(tracing_subscriber::filter::Targets::new().with_target("", LevelFilter::OFF))
 }
@@ -56,8 +52,10 @@ fn setup(opt: &Opt) -> anyhow::Result<(LogRouteHandle, LogStderrHandle)> {
     let (stderr_layer, stderr_layer_handle) =
         tracing_subscriber::reload::Layer::new(default_log_stderr_layer(opt));
     let route_layer: tracing_subscriber::reload::Layer<_, _> = route_layer;
+    let error_layer = tracing_error::ErrorLayer::default();
 
-    let subscriber = tracing_subscriber::registry().with(route_layer).with(stderr_layer);
+    let subscriber =
+        tracing_subscriber::registry().with(route_layer).with(stderr_layer).with(error_layer);
 
     // set the subscriber as the default for the application
     tracing::subscriber::set_global_default(subscriber).unwrap();
diff --git a/meilisearch/src/routes/indexes/documents.rs b/meilisearch/src/routes/indexes/documents.rs
index 43fab1dae..4d1374265 100644
--- a/meilisearch/src/routes/indexes/documents.rs
+++ b/meilisearch/src/routes/indexes/documents.rs
@@ -8,6 +8,7 @@ use deserr::actix_web::{AwebJson, AwebQueryParameter};
 use deserr::Deserr;
 use futures::StreamExt;
 use index_scheduler::{IndexScheduler, TaskId};
+use inspecting_allocator::InspectingAllocator;
 use meilisearch_types::deserr::query_params::Param;
 use meilisearch_types::deserr::{DeserrJsonError, DeserrQueryParamError};
 use meilisearch_types::document_formats::{read_csv, read_json, read_ndjson, PayloadType};
@@ -20,6 +21,7 @@ use meilisearch_types::milli::DocumentId;
 use meilisearch_types::star_or::OptionStarOrList;
 use meilisearch_types::tasks::KindWithContent;
 use meilisearch_types::{milli, Document, Index};
+use mimalloc::MiMalloc;
 use mime::Mime;
 use once_cell::sync::Lazy;
 use serde::Deserialize;
@@ -46,6 +48,9 @@ static ACCEPTED_CONTENT_TYPE: Lazy<Vec<String>> = Lazy::new(|| {
     vec!["application/json".to_string(), "application/x-ndjson".to_string(), "text/csv".to_string()]
 });
 
+#[global_allocator]
+static ALLOC: InspectingAllocator<MiMalloc> = InspectingAllocator::wrap(MiMalloc);
+
 /// Extracts the mime type from the content type and return
 /// a meilisearch error if anything bad happen.
 fn extract_mime_type(req: &HttpRequest) -> Result<Option<Mime>, MeilisearchHttpError> {
@@ -468,6 +473,13 @@ async fn document_addition(
     };
 
     let scheduler = index_scheduler.clone();
+    for (address, entry) in ALLOC.find_older_generations(5) {
+        println!(
+            "Found allocation older than 5 generations: {address:p} in generation {}. Span trace",
+            entry.generation()
+        );
+        println!("{entry}")
+    }
     let task = match tokio::task::spawn_blocking(move || scheduler.register(task, task_id, dry_run))
         .await?
     {