feat: CLI for optd-perftest (#106)

**Summary**: An extensible CLI built with `clap` for the `optd-perftest`
package

**Demo**:


https://github.com/cmu-db/optd/assets/20631215/05315dcc-6a91-4c4d-a9ea-fdb8d15ba675


**Details**:
* The CLI currently only exposes the `cardtest` option (cardinality
testing) but can be extended to support more types of performance tests
* The `cardtest` option takes in scale factor and seed as options, but
both have defaults as well
* Fixed a bug where we forgot to reconnect after reinitializing pgdata

Author: wangpatrick57

Cargo.lock

@@ -0,0 +1,19 @@
+#!/bin/bash
+# a sketchy way to kill "orphaned" postgres processes (i.e. processes whose pgdata dir has already been deleted)
+# I named it mac cuz I have no idea if it'll work on linux
+# use this script to avoid having to manually do "ps aux" and "kill [pid]" every time cardtest_integration (or something similar) fails
+
+# the "main" process has bin/postgres in it
+# we kill it in a loop because sometimes it's required for some reason I don't understand
+while pid=$(ps aux | grep bin/postgres | grep -v grep | head -n1 | awk '{print $2}'); do
+  if [ -z "$pid" ]; then
+    break
+  else
+    # we'll sometimes kill pids that don't exist. don't show the error in these cases
+    kill $pid &>/dev/null
+    sleep 1 # sleep so we don't loop too fast
+  fi
+done
+
+# even after it's gone, wait for it to completely shut down
+sleep 1

mac_kill_postgres.sh

@@ -20,5 +20,8 @@ tokio = { version = "1.24", features = [
 shlex = "1.3"
 tokio-postgres = "0.7"
 regex = "1.10"
+clap = { version = "4.5", features = [
+    "derive",
+] }
 log = "0.4"
 env_logger = "0.11"

optd-perftest/Cargo.toml

@@ -25,12 +25,12 @@ impl CardtestRunner {
     ///   of the entire query, not of a subtree of the query. This detail is specified in Section 7.1 of
     ///   [Yang 2020](https://arxiv.org/pdf/2006.08109.pdf)
     pub async fn eval_benchmark_qerrors_alldbs(
-        &self,
+        &mut self,
         benchmark: &Benchmark,
     ) -> anyhow::Result<HashMap<String, Vec<f64>>> {
         let mut qerrors_alldbs = HashMap::new();
 
-        for database in &self.databases {
+        for database in &mut self.databases {
             let estcards = database.eval_benchmark_estcards(benchmark).await?;
             let truecards = database.eval_benchmark_truecards(benchmark).await?;
             assert!(truecards.len() == estcards.len());
@@ -76,7 +76,16 @@ pub trait CardtestRunnerDBHelper {
     // get_name() has &self so that we're able to do Box<dyn CardtestRunnerDBHelper>
     fn get_name(&self) -> &str;
 
-    // the order of queries has to be the same between these two functions
-    async fn eval_benchmark_estcards(&self, benchmark: &Benchmark) -> anyhow::Result<Vec<usize>>;
-    async fn eval_benchmark_truecards(&self, benchmark: &Benchmark) -> anyhow::Result<Vec<usize>>;
+    // The order of queries has to be the same between these two functions.
+    // They take mutable references because evaluation sometimes involves mutating self.
+    //   One example of this is in PostgresDb where we may need to reconnect to the database,
+    //   which requires modifying the PostgresDb object.
+    async fn eval_benchmark_estcards(
+        &mut self,
+        benchmark: &Benchmark,
+    ) -> anyhow::Result<Vec<usize>>;
+    async fn eval_benchmark_truecards(
+        &mut self,
+        benchmark: &Benchmark,
+    ) -> anyhow::Result<Vec<usize>>;
 }

optd-perftest/src/cardtest.rs

@@ -8,11 +8,17 @@ impl CardtestRunnerDBHelper for DatafusionDb {
         "DataFusion"
     }
 
-    async fn eval_benchmark_truecards(&self, _benchmark: &Benchmark) -> anyhow::Result<Vec<usize>> {
+    async fn eval_benchmark_truecards(
+        &mut self,
+        _benchmark: &Benchmark,
+    ) -> anyhow::Result<Vec<usize>> {
         Ok(vec![])
     }
 
-    async fn eval_benchmark_estcards(&self, _benchmark: &Benchmark) -> anyhow::Result<Vec<usize>> {
+    async fn eval_benchmark_estcards(
+        &mut self,
+        _benchmark: &Benchmark,
+    ) -> anyhow::Result<Vec<usize>> {
         Ok(vec![])
     }
 }

optd-perftest/src/datafusion_db_cardtest.rs

@@ -1,5 +1,5 @@
-use anyhow::Result;
 use cardtest::{CardtestRunner, CardtestRunnerDBHelper};
+use clap::{Parser, Subcommand};
 use postgres_db::PostgresDb;
 
 use crate::{
@@ -14,18 +14,47 @@ mod postgres_db;
 mod shell;
 mod tpch;
 
+#[derive(Parser)]
+struct Cli {
+    #[command(subcommand)]
+    command: Commands,
+}
+
+#[derive(Subcommand)]
+enum Commands {
+    Cardtest {
+        #[arg(long)]
+        #[clap(default_value = "0.01")]
+        scale_factor: f64,
+        #[arg(long)]
+        #[clap(default_value = "15721")]
+        seed: i32,
+    },
+}
+
 #[tokio::main]
-async fn main() -> Result<()> {
+async fn main() -> anyhow::Result<()> {
     env_logger::init();
+    let cli = Cli::parse();
+
+    match &cli.command {
+        Commands::Cardtest { scale_factor, seed } => {
+            let tpch_config = TpchConfig {
+                database: String::from(TPCH_KIT_POSTGRES),
+                scale_factor: *scale_factor,
+                seed: *seed,
+            };
+            cardtest(tpch_config).await
+        }
+    }
+}
+
+async fn cardtest(tpch_config: TpchConfig) -> anyhow::Result<()> {
     let pg_db = PostgresDb::build().await?;
     let databases: Vec<Box<dyn CardtestRunnerDBHelper>> = vec![Box::new(pg_db)];
-    let tpch_config = TpchConfig {
-        database: String::from(TPCH_KIT_POSTGRES),
-        scale_factor: 1,
-        seed: 15721,
-    };
+
     let tpch_benchmark = Benchmark::Tpch(tpch_config.clone());
-    let cardtest_runner = CardtestRunner::new(databases).await?;
+    let mut cardtest_runner = CardtestRunner::new(databases).await?;
     let qerrors = cardtest_runner
         .eval_benchmark_qerrors_alldbs(&tpch_benchmark)
         .await?;

optd-perftest/src/main.rs

@@ -33,6 +33,8 @@ pub struct PostgresDb {
 ///   - Setup should be done in build() unless it requires more information (like benchmark)
 impl PostgresDb {
     pub async fn build() -> anyhow::Result<Self> {
+        log::debug!("[start] building PostgresDb");
+
         // build paths, sometimes creating them if they don't exist
         let curr_dpath = env::current_dir()?;
         let postgres_db_dpath = Path::new(file!())
@@ -61,6 +63,7 @@ impl PostgresDb {
         db.start_postgres().await?;
         db.connect_to_postgres().await?;
 
+        log::debug!("[end] building PostgresDb");
         Ok(db)
     }
 
@@ -178,7 +181,7 @@ impl PostgresDb {
         Ok(())
     }
 
-    async fn load_benchmark_data(&self, benchmark: &Benchmark) -> anyhow::Result<()> {
+    async fn load_benchmark_data(&mut self, benchmark: &Benchmark) -> anyhow::Result<()> {
         let benchmark_stringid = benchmark.get_stringid();
         if benchmark.is_readonly() {
             let done_fname = format!("{}_done", benchmark_stringid);
@@ -200,7 +203,7 @@ impl PostgresDb {
     }
 
     /// Load the benchmark data without worrying about caching
-    async fn load_benchmark_data_raw(&self, benchmark: &Benchmark) -> anyhow::Result<()> {
+    async fn load_benchmark_data_raw(&mut self, benchmark: &Benchmark) -> anyhow::Result<()> {
         match benchmark {
             Benchmark::Tpch(tpch_config) => self.load_tpch_data_raw(tpch_config).await?,
             _ => unimplemented!(),
@@ -209,13 +212,15 @@ impl PostgresDb {
     }
 
     /// Load the TPC-H data without worrying about caching
-    async fn load_tpch_data_raw(&self, tpch_config: &TpchConfig) -> anyhow::Result<()> {
+    async fn load_tpch_data_raw(&mut self, tpch_config: &TpchConfig) -> anyhow::Result<()> {
         // start from a clean slate
         self.remove_pgdata().await?;
         // since we deleted pgdata we'll need to re-init it
         self.init_pgdata().await?;
         // postgres must be started again since remove_pgdata() stops it
         self.start_postgres().await?;
+        // deleting pgdata would also delete the old connection so we have to reconnect
+        self.connect_to_postgres().await?;
         // load the schema
         let tpch_kit = TpchKit::build()?;
         shell::run_command_with_status_check(&format!(
@@ -248,15 +253,21 @@ impl CardtestRunnerDBHelper for PostgresDb {
         "Postgres"
     }
 
-    async fn eval_benchmark_estcards(&self, benchmark: &Benchmark) -> anyhow::Result<Vec<usize>> {
+    async fn eval_benchmark_estcards(
+        &mut self,
+        benchmark: &Benchmark,
+    ) -> anyhow::Result<Vec<usize>> {
         self.load_benchmark_data(benchmark).await?;
         match benchmark {
             Benchmark::Test => unimplemented!(),
             Benchmark::Tpch(tpch_config) => self.eval_tpch_estcards(tpch_config).await,
         }
     }
 
-    async fn eval_benchmark_truecards(&self, benchmark: &Benchmark) -> anyhow::Result<Vec<usize>> {
+    async fn eval_benchmark_truecards(
+        &mut self,
+        benchmark: &Benchmark,
+    ) -> anyhow::Result<Vec<usize>> {
         self.load_benchmark_data(benchmark).await?;
         match benchmark {
             Benchmark::Test => unimplemented!(),

optd-perftest/src/postgres_db.rs

@@ -11,10 +11,10 @@ const TPCH_KIT_REPO_URL: &str = "[email protected]:lmwnshn/tpch-kit.git";
 pub const TPCH_KIT_POSTGRES: &str = "POSTGRESQL";
 const NUM_TPCH_QUERIES: usize = 22;
 
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct TpchConfig {
     pub database: String,
-    pub scale_factor: i32,
+    pub scale_factor: f64,
     pub seed: i32,
 }
 
@@ -44,6 +44,8 @@ pub struct TpchKit {
 /// I keep the same conventions for these methods as I do for PostgresDb
 impl TpchKit {
     pub fn build() -> io::Result<Self> {
+        log::debug!("[start] building TpchKit");
+
         // build paths, sometimes creating them if they don't exist
         // we assume that this is being run in the base optd repo dir
         let optd_repo_dpath = env::current_dir()?;
@@ -88,6 +90,7 @@ impl TpchKit {
         // do setup after creating kit
         kit.clonepull_tpch_kit_repo()?;
 
+        log::debug!("[end] building TpchKit");
         Ok(kit)
     }