From 0d6d8fefd30c6b7eaf9aa9593629ddba3925c655 Mon Sep 17 00:00:00 2001
From: Bojan Serafimov <bojan.serafimov7@gmail.com>
Date: Mon, 12 Sep 2022 16:10:19 -0400
Subject: [PATCH] Add replay binary

---
 pageserver/src/bin/replay.rs                 | 51 ++++++++++++++++++++
 test_runner/fixtures/neon_fixtures.py        | 21 ++++++++
 test_runner/performance/test_trace_replay.py |  7 ++-
 3 files changed, 77 insertions(+), 2 deletions(-)
 create mode 100644 pageserver/src/bin/replay.rs

diff --git a/pageserver/src/bin/replay.rs b/pageserver/src/bin/replay.rs
new file mode 100644
index 0000000000..83e9da018c
--- /dev/null
+++ b/pageserver/src/bin/replay.rs
@@ -0,0 +1,51 @@
+use std::{
+    fs::{read_dir, File},
+    io::BufReader,
+    path::PathBuf,
+    str::FromStr,
+};
+
+use clap::{App, Arg};
+use utils::zid::{ZTenantId, ZTimelineId};
+
+async fn replay_trace<R: std::io::Read>(reader: R) -> anyhow::Result<()> {
+    Ok(())
+}
+
+#[tokio::main]
+async fn main() -> anyhow::Result<()> {
+    // TODO upgrade to struct macro arg parsing
+    let arg_matches = App::new("Pageserver trace replay tool")
+        .about("Replays wal or read traces to test pageserver performance")
+        .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 = ZTenantId::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 = ZTimelineId::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 = ZTimelineId::from_str(path.file_name().unwrap().to_str().unwrap())?;
+
+                let file = File::open(path)?;
+                let reader = BufReader::new(file);
+                replay_trace(reader).await?;
+            }
+        }
+    }
+
+    Ok(())
+}
diff --git a/test_runner/fixtures/neon_fixtures.py b/test_runner/fixtures/neon_fixtures.py
index b47e560325..041378d741 100644
--- a/test_runner/fixtures/neon_fixtures.py
+++ b/test_runner/fixtures/neon_fixtures.py
@@ -1636,6 +1636,27 @@ def pg_bin(test_output_dir: Path) -> PgBin:
     return PgBin(test_output_dir)
 
 
+@dataclass
+class ReplayBin:
+    """A helper class for replaying pageserver wal and read traces."""
+
+    traces_dir: str
+
+    def replay_all(self) -> str:
+        replay_binpath = os.path.join(str(neon_binpath), "replay")
+        args = [
+            replay_binpath,
+            self.traces_dir,
+        ]
+        return subprocess.run(args, capture_output=True).stdout.decode("UTF-8").strip()
+
+
+@pytest.fixture(scope="function")
+def replay_bin(test_output_dir):
+    traces_dir = os.path.join(test_output_dir, "repo", "traces")
+    return ReplayBin(traces_dir)
+
+
 class VanillaPostgres(PgProtocol):
     def __init__(self, pgdatadir: Path, pg_bin: PgBin, port: int, init=True):
         super().__init__(host="localhost", port=port, dbname="postgres")
diff --git a/test_runner/performance/test_trace_replay.py b/test_runner/performance/test_trace_replay.py
index cb52fb4e5e..1209422af9 100644
--- a/test_runner/performance/test_trace_replay.py
+++ b/test_runner/performance/test_trace_replay.py
@@ -1,7 +1,7 @@
-from fixtures.neon_fixtures import NeonEnvBuilder
+from fixtures.neon_fixtures import NeonEnvBuilder, ReplayBin
 
 
-def test_trace_replay(neon_env_builder: NeonEnvBuilder):
+def test_trace_replay(neon_env_builder: NeonEnvBuilder, replay_bin: ReplayBin):
     neon_env_builder.num_safekeepers = 1
     env = neon_env_builder.init_start()
 
@@ -17,3 +17,6 @@ def test_trace_replay(neon_env_builder: NeonEnvBuilder):
 
     trace_path = env.repo_dir / "traces" / str(tenant) / str(timeline) / str(timeline)
     assert trace_path.exists()
+
+    output = replay_bin.replay_all()
+    print(output)