From e9b158e8f5ec436793a03ff8fdc384912c3019aa Mon Sep 17 00:00:00 2001
From: Bojan Serafimov <bojan.serafimov7@gmail.com>
Date: Thu, 15 Sep 2022 21:14:07 -0400
Subject: [PATCH] Add draw binary

---
 pageserver/src/bin/draw_trace.rs             | 52 ++++++++++++++++++++
 test_runner/fixtures/neon_fixtures.py        |  8 +++
 test_runner/performance/test_trace_replay.py |  5 +-
 3 files changed, 63 insertions(+), 2 deletions(-)
 create mode 100644 pageserver/src/bin/draw_trace.rs

diff --git a/pageserver/src/bin/draw_trace.rs b/pageserver/src/bin/draw_trace.rs
new file mode 100644
index 0000000000..3ea5b27a90
--- /dev/null
+++ b/pageserver/src/bin/draw_trace.rs
@@ -0,0 +1,52 @@
+use clap::{App, Arg};
+use pageserver::page_service::PagestreamFeMessage;
+use std::{
+    fs::{read_dir, File},
+    io::BufReader,
+    path::PathBuf,
+    str::FromStr,
+};
+use utils::{
+    id::{TenantId, TimelineId},
+    pq_proto::{BeMessage, FeMessage},
+};
+
+fn main() -> anyhow::Result<()> {
+    // TODO upgrade to struct macro arg parsing
+    let arg_matches = App::new("Pageserver trace visualization tool")
+        .about("Makes a svg file that displays the read pattern")
+        .arg(
+            Arg::new("traces_dir")
+                .takes_value(true)
+                .help("Directory where the read traces are stored"),
+        )
+        .get_matches();
+
+    let traces_dir = PathBuf::from(arg_matches.value_of("traces_dir").unwrap());
+    for tenant_dir in read_dir(traces_dir)? {
+        let entry = tenant_dir?;
+        let path = entry.path();
+        let _tenant_id = TenantId::from_str(path.file_name().unwrap().to_str().unwrap())?;
+
+        for timeline_dir in read_dir(path)? {
+            let entry = timeline_dir?;
+            let path = entry.path();
+            let _timeline_id = TimelineId::from_str(path.file_name().unwrap().to_str().unwrap())?;
+
+            for trace_dir in read_dir(path)? {
+                let entry = trace_dir?;
+                let path = entry.path();
+                let _conn_id = TimelineId::from_str(path.file_name().unwrap().to_str().unwrap())?;
+
+                let file = File::open(path.clone())?;
+                let mut reader = BufReader::new(file);
+                while let Ok(msg) = PagestreamFeMessage::parse(&mut reader) {
+                    println!("Parsed message {:?}", msg);
+                    // TODO add to svg
+                }
+            }
+        }
+    }
+
+    Ok(())
+}
diff --git a/test_runner/fixtures/neon_fixtures.py b/test_runner/fixtures/neon_fixtures.py
index f51013706e..084fa52d8f 100644
--- a/test_runner/fixtures/neon_fixtures.py
+++ b/test_runner/fixtures/neon_fixtures.py
@@ -1651,6 +1651,14 @@ class ReplayBin:
         # return subprocess.run(args, capture_output=True).stdout.decode("UTF-8").strip()
         subprocess.run(args)
 
+    def draw_all(self):
+        draw_binpath = os.path.join(str(neon_binpath), "draw_trace")
+        args = [
+            draw_binpath,
+            self.traces_dir,
+        ]
+        subprocess.run(args)
+
 
 @pytest.fixture(scope="function")
 def replay_bin(test_output_dir):
diff --git a/test_runner/performance/test_trace_replay.py b/test_runner/performance/test_trace_replay.py
index b20b11b640..45da31a044 100644
--- a/test_runner/performance/test_trace_replay.py
+++ b/test_runner/performance/test_trace_replay.py
@@ -47,10 +47,11 @@ def test_trace_replay(
     # trace_path = env.repo_dir / "traces" / str(tenant) / str(timeline) / str(timeline)
     # assert trace_path.exists()
 
+    replay_bin.draw_all()
+    return
+
     print("replaying")
     ps_connstr = env.pageserver.connstr()
-
-    # ps_connstr = "host=localhost port=15004 dbname=postgres user=neon_admin"
     with zenbenchmark.record_duration("replay"):
         output = replay_bin.replay_all(ps_connstr)
     print(output)