From a61d8c59ff198c6ad6081e5f584a90efb01be0b0 Mon Sep 17 00:00:00 2001
From: Louis Dureuil <louis@meilisearch.com>
Date: Wed, 24 Jan 2024 09:50:51 +0100
Subject: [PATCH] Add span stats processor

---
 tracing-trace/src/bin/trace-to-callstats.rs | 16 +++++
 tracing-trace/src/processor/mod.rs          |  1 +
 tracing-trace/src/processor/span_stats.rs   | 79 +++++++++++++++++++++
 3 files changed, 96 insertions(+)
 create mode 100644 tracing-trace/src/bin/trace-to-callstats.rs
 create mode 100644 tracing-trace/src/processor/span_stats.rs

diff --git a/tracing-trace/src/bin/trace-to-callstats.rs b/tracing-trace/src/bin/trace-to-callstats.rs
new file mode 100644
index 000000000..3644b7bff
--- /dev/null
+++ b/tracing-trace/src/bin/trace-to-callstats.rs
@@ -0,0 +1,16 @@
+use std::ffi::OsString;
+use std::io::Write;
+
+fn main() {
+    let input_file = std::env::args_os().nth(1).expect("missing <INPUT> file");
+    let input =
+        std::io::BufReader::new(std::fs::File::open(&input_file).expect("could not open <INPUT>"));
+    let trace = tracing_trace::TraceReader::new(input);
+    let profile = tracing_trace::processor::span_stats::to_call_stats(trace).unwrap();
+    let mut output_file = OsString::new();
+    output_file.push("callstats-");
+    output_file.push(input_file);
+    let mut output_file = std::io::BufWriter::new(std::fs::File::create(output_file).unwrap());
+    serde_json::to_writer(&mut output_file, &profile).unwrap();
+    output_file.flush().unwrap();
+}
diff --git a/tracing-trace/src/processor/mod.rs b/tracing-trace/src/processor/mod.rs
index a84cb3b63..ea445b0a5 100644
--- a/tracing-trace/src/processor/mod.rs
+++ b/tracing-trace/src/processor/mod.rs
@@ -1,2 +1,3 @@
 pub mod firefox_profiler;
 pub mod fmt;
+pub mod span_stats;
diff --git a/tracing-trace/src/processor/span_stats.rs b/tracing-trace/src/processor/span_stats.rs
new file mode 100644
index 000000000..63b6ae5c1
--- /dev/null
+++ b/tracing-trace/src/processor/span_stats.rs
@@ -0,0 +1,79 @@
+use std::collections::{BTreeMap, HashMap};
+use std::time::Duration;
+
+use serde::Serialize;
+
+use crate::entry::{Entry, NewCallsite, SpanClose, SpanEnter, SpanExit};
+use crate::{Error, TraceReader};
+
+#[derive(Debug, Clone, Copy)]
+enum SpanStatus {
+    Outside,
+    Inside(std::time::Duration),
+}
+
+#[derive(Serialize)]
+pub struct CallStats {
+    nb: usize,
+    ns: u64,
+}
+
+pub fn to_call_stats<R: std::io::Read>(
+    trace: TraceReader<R>,
+) -> Result<BTreeMap<String, CallStats>, Error> {
+    let mut calls = HashMap::new();
+    let mut spans = HashMap::new();
+    for entry in trace {
+        let entry = entry?;
+        match entry {
+            Entry::NewCallsite(callsite) => {
+                calls.insert(callsite.call_id, (callsite, vec![]));
+            }
+            Entry::NewThread(_) => {}
+            Entry::NewSpan(span) => {
+                spans.insert(span.id, (span, SpanStatus::Outside));
+            }
+            Entry::SpanEnter(SpanEnter { id, time, memory: _ }) => {
+                let (_, status) = spans.get_mut(&id).unwrap();
+
+                let SpanStatus::Outside = status else {
+                    continue;
+                };
+
+                *status = SpanStatus::Inside(time);
+            }
+            Entry::SpanExit(SpanExit { id, time: end, memory: _ }) => {
+                let (span, status) = spans.get_mut(&id).unwrap();
+
+                let SpanStatus::Inside(begin) = status else {
+                    continue;
+                };
+                let begin = *begin;
+
+                *status = SpanStatus::Outside;
+
+                let span = *span;
+                let (_, call_list) = calls.get_mut(&span.call_id).unwrap();
+                call_list.push(end - begin);
+            }
+            Entry::SpanClose(SpanClose { id, time: _ }) => {
+                spans.remove(&id);
+            }
+            Entry::Event(_) => {}
+        }
+    }
+
+    Ok(calls
+        .into_iter()
+        .map(|(_, (call_site, calls))| (site_to_string(call_site), calls_to_stats(calls)))
+        .collect())
+}
+
+fn site_to_string(call_site: NewCallsite) -> String {
+    format!("{}::{}", call_site.target, call_site.name)
+}
+fn calls_to_stats(calls: Vec<Duration>) -> CallStats {
+    let nb = calls.len();
+    let sum: Duration = calls.iter().sum();
+    CallStats { nb, ns: sum.as_nanos() as u64 }
+}