From 173aad609027048a0e0a3b5860ff12e8014fca10 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cl=C3=A9ment=20Renault?= <clement@meilisearch.com>
Date: Thu, 8 Feb 2024 15:20:05 +0100
Subject: [PATCH] Introduce a tracing tool to filter logs associated to high
 memory usage

---
 Cargo.lock                            | 229 +++++++++++++++++++++++++-
 tracing-trace/Cargo.toml              |   7 +-
 tracing-trace/src/bin/filter-trace.rs |  89 ++++++++++
 tracing-trace/src/entry.rs            |  14 ++
 tracing-trace/src/processor/fmt.rs    |   4 +-
 5 files changed, 331 insertions(+), 12 deletions(-)
 create mode 100644 tracing-trace/src/bin/filter-trace.rs

diff --git a/Cargo.lock b/Cargo.lock
index b402ca2cd..fbdc0b241 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -382,6 +382,12 @@ dependencies = [
  "derive_arbitrary",
 ]
 
+[[package]]
+name = "arrayvec"
+version = "0.7.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "96d30a06541fbafbc7f82ed10c06164cfbd2c401138f6addd8404629c4b16711"
+
 [[package]]
 name = "arroy"
 version = "0.2.0"
@@ -585,6 +591,18 @@ dependencies = [
  "serde",
 ]
 
+[[package]]
+name = "bitvec"
+version = "1.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1bc2832c24239b0141d5674bb9174f9d68a8b5b3f2753311927c172ca46f7e9c"
+dependencies = [
+ "funty",
+ "radium",
+ "tap",
+ "wyz",
+]
+
 [[package]]
 name = "block-buffer"
 version = "0.10.4"
@@ -594,6 +612,30 @@ dependencies = [
  "generic-array",
 ]
 
+[[package]]
+name = "borsh"
+version = "1.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f58b559fd6448c6e2fd0adb5720cd98a2506594cafa4737ff98c396f3e82f667"
+dependencies = [
+ "borsh-derive",
+ "cfg_aliases",
+]
+
+[[package]]
+name = "borsh-derive"
+version = "1.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7aadb5b6ccbd078890f6d7003694e33816e6b784358f18e15e7e6d9f065a57cd"
+dependencies = [
+ "once_cell",
+ "proc-macro-crate",
+ "proc-macro2",
+ "quote",
+ "syn 2.0.48",
+ "syn_derive",
+]
+
 [[package]]
 name = "brotli"
 version = "3.4.0"
@@ -642,6 +684,39 @@ dependencies = [
  "utf8-width",
 ]
 
+[[package]]
+name = "byte-unit"
+version = "5.1.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "33ac19bdf0b2665407c39d82dbc937e951e7e2001609f0fb32edd0af45a2d63e"
+dependencies = [
+ "rust_decimal",
+ "serde",
+ "utf8-width",
+]
+
+[[package]]
+name = "bytecheck"
+version = "0.6.12"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "23cdc57ce23ac53c931e88a43d06d070a6fd142f2617be5855eb75efc9beb1c2"
+dependencies = [
+ "bytecheck_derive",
+ "ptr_meta",
+ "simdutf8",
+]
+
+[[package]]
+name = "bytecheck_derive"
+version = "0.6.12"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3db406d29fbcd95542e92559bed4d8ad92636d1ca8b3b72ede10b4bcc010e659"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 1.0.109",
+]
+
 [[package]]
 name = "bytecount"
 version = "0.6.3"
@@ -866,6 +941,12 @@ version = "1.0.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"
 
+[[package]]
+name = "cfg_aliases"
+version = "0.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fd16c4719339c4530435d38e511904438d07cce7950afa3718a84ac36c10e89e"
+
 [[package]]
 name = "change-detection"
 version = "1.2.0"
@@ -957,9 +1038,9 @@ dependencies = [
 
 [[package]]
 name = "clap"
-version = "4.4.17"
+version = "4.4.18"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "80932e03c33999b9235edb8655bc9df3204adc9887c2f95b50cb1deb9fd54253"
+checksum = "1e578d6ec4194633722ccf9544794b71b1385c3c027efe0c55db226fc880865c"
 dependencies = [
  "clap_builder",
  "clap_derive",
@@ -967,9 +1048,9 @@ dependencies = [
 
 [[package]]
 name = "clap_builder"
-version = "4.4.17"
+version = "4.4.18"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d6c0db58c659eef1c73e444d298c27322a1b52f6927d2ad470c0c0f96fa7b8fa"
+checksum = "4df4df40ec50c46000231c914968278b1eb05098cf8f1b3a518a95030e71d1c7"
 dependencies = [
  "anstream",
  "anstyle",
@@ -1861,6 +1942,12 @@ version = "0.4.7"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "7ab85b9b05e3978cc9a9cf8fea7f01b494e1a09ed3037e16ba39edc7a29eb61a"
 
+[[package]]
+name = "funty"
+version = "2.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e6d5a32815ae3f33302d95fdcb2ce17862f8c65363dcfd29360480ba1001fc9c"
+
 [[package]]
 name = "futures"
 version = "0.3.30"
@@ -3678,7 +3765,7 @@ dependencies = [
  "async-trait",
  "brotli",
  "bstr",
- "byte-unit",
+ "byte-unit 4.0.19",
  "bytes",
  "cargo_toml",
  "clap",
@@ -4484,6 +4571,15 @@ version = "0.2.17"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "5b40af805b3121feab8a3c29f04d8ad262fa8e0561883e7653e024ae4479e6de"
 
+[[package]]
+name = "proc-macro-crate"
+version = "3.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6d37c51ca738a55da99dc0c4a34860fd675453b8b36209178c2249bb13651284"
+dependencies = [
+ "toml_edit",
+]
+
 [[package]]
 name = "proc-macro-error"
 version = "1.0.4"
@@ -4553,6 +4649,26 @@ version = "2.28.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "106dd99e98437432fed6519dedecfade6a06a73bb7b2a1e019fdd2bee5778d94"
 
+[[package]]
+name = "ptr_meta"
+version = "0.1.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0738ccf7ea06b608c10564b31debd4f5bc5e197fc8bfe088f68ae5ce81e7a4f1"
+dependencies = [
+ "ptr_meta_derive",
+]
+
+[[package]]
+name = "ptr_meta_derive"
+version = "0.1.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "16b845dbfca988fa33db069c0e230574d15a3088f147a87b64c7589eb662c9ac"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 1.0.109",
+]
+
 [[package]]
 name = "puffin"
 version = "0.16.0"
@@ -4591,6 +4707,12 @@ dependencies = [
  "proc-macro2",
 ]
 
+[[package]]
+name = "radium"
+version = "0.7.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "dc33ff2d4973d518d823d61aa239014831e521c75da58e3df4840d3f47749d09"
+
 [[package]]
 name = "rand"
 version = "0.8.5"
@@ -4769,6 +4891,15 @@ version = "0.8.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "c08c74e62047bb2de4ff487b251e4a92e24f48745648451635cec7d591162d9f"
 
+[[package]]
+name = "rend"
+version = "0.4.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "71fe3824f5629716b1589be05dacd749f6aa084c87e00e016714a8cdfccc997c"
+dependencies = [
+ "bytecheck",
+]
+
 [[package]]
 name = "reqwest"
 version = "0.11.23"
@@ -4850,6 +4981,35 @@ dependencies = [
  "windows-sys 0.48.0",
 ]
 
+[[package]]
+name = "rkyv"
+version = "0.7.44"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5cba464629b3394fc4dbc6f940ff8f5b4ff5c7aef40f29166fd4ad12acbc99c0"
+dependencies = [
+ "bitvec",
+ "bytecheck",
+ "bytes",
+ "hashbrown 0.12.3",
+ "ptr_meta",
+ "rend",
+ "rkyv_derive",
+ "seahash",
+ "tinyvec",
+ "uuid",
+]
+
+[[package]]
+name = "rkyv_derive"
+version = "0.7.44"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a7dddfff8de25e6f62b9d64e6e432bf1c6736c57d20323e15ee10435fbda7c65"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn 1.0.109",
+]
+
 [[package]]
 name = "roaring"
 version = "0.10.2"
@@ -4874,6 +5034,22 @@ dependencies = [
  "smallvec",
 ]
 
+[[package]]
+name = "rust_decimal"
+version = "1.34.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "755392e1a2f77afd95580d3f0d0e94ac83eeeb7167552c9b5bca549e61a94d83"
+dependencies = [
+ "arrayvec",
+ "borsh",
+ "bytes",
+ "num-traits",
+ "rand",
+ "rkyv",
+ "serde",
+ "serde_json",
+]
+
 [[package]]
 name = "rustc-demangle"
 version = "0.1.23"
@@ -5012,6 +5188,12 @@ dependencies = [
  "untrusted 0.7.1",
 ]
 
+[[package]]
+name = "seahash"
+version = "4.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1c107b6f4780854c8b126e228ea8869f4d7b71260f962fefb57b996b8959ba6b"
+
 [[package]]
 name = "segment"
 version = "0.2.3"
@@ -5169,6 +5351,12 @@ dependencies = [
  "libc",
 ]
 
+[[package]]
+name = "simdutf8"
+version = "0.1.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f27f6278552951f1f2b8cf9da965d10969b2efdea95a6ec47987ab46edfe263a"
+
 [[package]]
 name = "similar"
 version = "2.2.1"
@@ -5359,6 +5547,18 @@ dependencies = [
  "unicode-ident",
 ]
 
+[[package]]
+name = "syn_derive"
+version = "0.1.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1329189c02ff984e9736652b1631330da25eaa6bc639089ed4915d25446cbe7b"
+dependencies = [
+ "proc-macro-error",
+ "proc-macro2",
+ "quote",
+ "syn 2.0.48",
+]
+
 [[package]]
 name = "synchronoise"
 version = "1.0.1"
@@ -5430,6 +5630,12 @@ dependencies = [
  "libc",
 ]
 
+[[package]]
+name = "tap"
+version = "1.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "55937e1799185b12863d447f42597ed69d9928686b8d88a1df17376a097d8369"
+
 [[package]]
 name = "tar"
 version = "0.4.40"
@@ -5813,7 +6019,9 @@ dependencies = [
 name = "tracing-trace"
 version = "0.1.0"
 dependencies = [
- "byte-unit",
+ "anyhow",
+ "byte-unit 5.1.4",
+ "clap",
  "color-spantrace",
  "fxprof-processed-profile",
  "libproc",
@@ -6467,6 +6675,15 @@ version = "0.5.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "c0af0c3d13faebf8dda0b5256fa7096a2d5ccb662f7b9f54a40fe201077ab1c2"
 
+[[package]]
+name = "wyz"
+version = "0.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "05f360fc0b24296329c78fda852a1e9ae82de9cf7b27dae4b7f62f118f77b9ed"
+dependencies = [
+ "tap",
+]
+
 [[package]]
 name = "xattr"
 version = "1.0.1"
diff --git a/tracing-trace/Cargo.toml b/tracing-trace/Cargo.toml
index 64848bff9..7f169b345 100644
--- a/tracing-trace/Cargo.toml
+++ b/tracing-trace/Cargo.toml
@@ -13,11 +13,10 @@ serde_json = "1.0.111"
 tracing = "0.1.40"
 tracing-error = "0.2.0"
 tracing-subscriber = "0.3.18"
-byte-unit = { version = "4.0.19", default-features = false, features = [
-    "std",
-    "serde",
-] }
 tokio = { version = "1.35.1", features = ["sync"] }
+clap = { version = "4.4.18", features = ["derive"] }
+anyhow = "1.0.79"
+byte-unit = { version = "5.1.4", features = ["byte"] }
 
 [target.'cfg(any(target_os = "linux", target_os = "macos"))'.dependencies]
 libproc = "0.14.2"
diff --git a/tracing-trace/src/bin/filter-trace.rs b/tracing-trace/src/bin/filter-trace.rs
new file mode 100644
index 000000000..b8b0b5506
--- /dev/null
+++ b/tracing-trace/src/bin/filter-trace.rs
@@ -0,0 +1,89 @@
+use std::collections::vec_deque::Drain;
+use std::collections::VecDeque;
+use std::io::{self, BufReader, Write};
+use std::mem;
+
+use anyhow::Context;
+use byte_unit::Byte;
+use clap::Parser;
+use tracing_trace::entry::Entry;
+
+/// A program that filters trace logs to only keeps
+/// the logs related to memory usage above the given threshold.
+#[derive(Parser, Debug)]
+#[command(author, version, about, long_about = None)]
+struct Args {
+    /// The threshold that a log must have to be returned by this program.
+    #[arg(short, long)]
+    memory_threshold: Byte,
+
+    /// Number of context lines to keep around high memory log lines.
+    #[arg(long, default_value_t = 10)]
+    context: usize,
+}
+
+fn main() -> anyhow::Result<()> {
+    let Args { memory_threshold, context } = Args::parse();
+
+    let mut context = EntryContext::new(context);
+    let mut currently_in_threshold = false;
+
+    let input = BufReader::new(io::stdin());
+    let mut output = io::BufWriter::new(io::stdout());
+    for result in tracing_trace::TraceReader::new(input) {
+        let entry = result?;
+        if entry.memory().map_or(true, |m| m.resident < memory_threshold.as_u64()) {
+            if mem::replace(&mut currently_in_threshold, false) {
+                for entry in context.drain() {
+                    serde_json::to_writer(&mut output, &entry)
+                        .context("while serializing and writing to stdout")?;
+                }
+            }
+            context.push(entry);
+        } else {
+            currently_in_threshold = true;
+            for entry in context.drain() {
+                serde_json::to_writer(&mut output, &entry)
+                    .context("while serializing and writing to stdout")?;
+            }
+            serde_json::to_writer(&mut output, &entry)
+                .context("while serializing and writing to stdout")?;
+        }
+    }
+
+    for entry in context.drain() {
+        serde_json::to_writer(&mut output, &entry)
+            .context("while serializing and writing to stdout")?;
+    }
+
+    output.flush().context("flushing stdout")?;
+    Ok(())
+}
+
+/// Keeps only the last `size` element in memory.
+/// It's basically a sliding window.
+pub struct EntryContext {
+    size: usize,
+    queue: VecDeque<Entry>,
+}
+
+impl EntryContext {
+    pub fn new(size: usize) -> EntryContext {
+        EntryContext { size, queue: VecDeque::with_capacity(size) }
+    }
+
+    pub fn is_full(&self) -> bool {
+        self.size >= self.queue.len()
+    }
+
+    pub fn push(&mut self, entry: Entry) {
+        if self.queue.len() == self.size {
+            self.queue.pop_front();
+        }
+        self.queue.push_back(entry);
+    }
+
+    pub fn drain(&mut self) -> Drain<Entry> {
+        self.queue.drain(..)
+    }
+}
diff --git a/tracing-trace/src/entry.rs b/tracing-trace/src/entry.rs
index 26e543ba0..46a83f8e8 100644
--- a/tracing-trace/src/entry.rs
+++ b/tracing-trace/src/entry.rs
@@ -38,6 +38,20 @@ pub enum Entry {
     Event(Event),
 }
 
+impl Entry {
+    pub fn memory(&self) -> Option<MemoryStats> {
+        match self {
+            Entry::NewCallsite(_)
+            | Entry::NewThread(_)
+            | Entry::NewSpan(_)
+            | Entry::SpanClose(_) => None,
+            Entry::SpanEnter(event) => event.memory,
+            Entry::SpanExit(event) => event.memory,
+            Entry::Event(event) => event.memory,
+        }
+    }
+}
+
 #[derive(Clone, Copy, Debug, Serialize, Deserialize, PartialEq, Eq, Hash)]
 pub struct SpanId(u64);
 
diff --git a/tracing-trace/src/processor/fmt.rs b/tracing-trace/src/processor/fmt.rs
index 68e95c00c..767415da9 100644
--- a/tracing-trace/src/processor/fmt.rs
+++ b/tracing-trace/src/processor/fmt.rs
@@ -189,7 +189,7 @@ fn print_duration(duration: std::time::Duration) -> String {
 
 /// Format only the allocated bytes, deallocated bytes and reallocated bytes in GiB, MiB, KiB, Bytes.
 fn print_memory(MemoryStats { resident }: MemoryStats) -> String {
-    use byte_unit::Byte;
-    let rss_bytes = Byte::from_bytes(resident).get_appropriate_unit(true);
+    use byte_unit::{Byte, UnitType};
+    let rss_bytes = Byte::from_u64(resident).get_appropriate_unit(UnitType::Binary);
     format!("RSS {rss_bytes:.2}")
 }