let user specify status addr (#3)

Signed-off-by: ekexium <ekexium@fastmail.com>
PingCAP-QE · Jun 8, 2022 · 7e587b6 · 7e587b6
1 parent b1fa5ce
commit 7e587b6
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 5 deletions.
diff --git a/src/config.rs b/src/config.rs
@@ -9,6 +9,7 @@ pub struct Config {
     pub assertion: String,
     pub limit: u32,
     pub uri: String,
+    pub status_addr: String,
     pub log_path: String,
 }
 
@@ -64,6 +65,14 @@ pub fn init_app() -> Config {
                 .required(false)
                 .default_value("corrupttest.log"),
         )
+        .arg(
+            Arg::new("status_addr")
+                .short('s')
+                .long("status_addr")
+                .takes_value(true)
+                .required(false)
+                .default_value("127.0.0.1:10080"),
+        )
         .get_matches();
     let config = Config {
         workload_name: matches
@@ -78,6 +87,7 @@ pub fn init_app() -> Config {
             .parse::<u32>()
             .expect("limit must be a non-negative number"),
         uri: matches.value_of("uri").unwrap().to_owned(),
+        status_addr: matches.value_of("status_addr").unwrap().to_owned(),
         log_path: matches.value_of("log_path").unwrap().to_owned(),
     };
     config

diff --git a/src/failpoint.rs b/src/failpoint.rs
@@ -6,18 +6,20 @@ use slog::Logger;
 use std::sync::atomic::Ordering;
 use tokio::time;
 
-const STATUS_ADDRESS: &str = "127.0.0.1:10080";
-
 pub async fn enable_failpoint(
     log: &Logger,
     client: &reqwest::Client,
+    status_addr: impl Into<String>,
     name: impl Into<String>,
     value: impl Into<String>,
 ) -> Result<()> {
     let start = time::Instant::now();
-    let name = name.into();
     let res = client
-        .put(format!("http://{}/fail/{}", STATUS_ADDRESS, &name))
+        .put(format!(
+            "http://{}/fail/{}",
+            status_addr.into(),
+            name.into()
+        ))
         .body(value.into())
         .send()
         .await?;
@@ -35,11 +37,16 @@ pub async fn enable_failpoint(
 pub async fn disable_failpoint(
     log: &Logger,
     client: &reqwest::Client,
+    status_addr: impl Into<String>,
     name: impl Into<String>,
 ) -> Result<()> {
     let start = time::Instant::now();
     let res = client
-        .delete(format!("http://{}/fail/{}", STATUS_ADDRESS, name.into()))
+        .delete(format!(
+            "http://{}/fail/{}",
+            status_addr.into(),
+            name.into()
+        ))
         .send()
         .await?;
     let duration = start.elapsed();

diff --git a/src/workload.rs b/src/workload.rs
@@ -132,6 +132,7 @@ impl Workload for SingleInsertion {
             enable_failpoint(
                 &log,
                 client,
+                config.status_addr.clone(),
                 "github.com/pingcap/tidb/table/tables/corruptMutations",
                 format!("return(\"{}\")", injection),
             )
@@ -148,6 +149,7 @@ impl Workload for SingleInsertion {
             disable_failpoint(
                 &log,
                 client,
+                config.status_addr.clone(),
                 "github.com/pingcap/tidb/table/tables/corruptMutations",
             )
             .await?;
@@ -184,6 +186,7 @@ impl Workload for DoubleInsertion {
             enable_failpoint(
                 &log,
                 client,
+                config.status_addr.clone(),
                 "github.com/pingcap/tidb/table/tables/corruptMutations",
                 format!("1*return(\"{}\")", injection),
             )
@@ -213,6 +216,7 @@ impl Workload for DoubleInsertion {
             disable_failpoint(
                 &log,
                 client,
+                config.status_addr.clone(),
                 "github.com/pingcap/tidb/table/tables/corruptMutations",
             )
             .await?;
@@ -252,6 +256,7 @@ impl Workload for T2 {
             enable_failpoint(
                 &log,
                 client,
+                config.status_addr.clone(),
                 "github.com/pingcap/tidb/table/tables/corruptMutations",
                 format!("1*return(\"{}\")", injection),
             )
@@ -287,6 +292,7 @@ impl Workload for T2 {
             disable_failpoint(
                 &log,
                 client,
+                config.status_addr.clone(),
                 "github.com/pingcap/tidb/table/tables/corruptMutations",
             )
             .await?;
@@ -326,6 +332,7 @@ impl Workload for T3 {
             enable_failpoint(
                 &log,
                 client,
+                config.status_addr.clone(),
                 "github.com/pingcap/tidb/table/tables/corruptMutations",
                 format!("1*return(\"{}\")", injection),
             )
@@ -368,6 +375,7 @@ impl Workload for T3 {
             disable_failpoint(
                 &log,
                 client,
+                config.status_addr.clone(),
                 "github.com/pingcap/tidb/table/tables/corruptMutations",
             )
             .await?;
@@ -423,6 +431,7 @@ impl Workload for T4 {
                 enable_failpoint(
                     &log,
                     client,
+                    config.status_addr.clone(),
                     "github.com/pingcap/tidb/table/tables/corruptMutations",
                     format!("1*return(\"{}\")", injection),
                 )
@@ -442,6 +451,7 @@ impl Workload for T4 {
             disable_failpoint(
                 &log,
                 client,
+                config.status_addr.clone(),
                 "github.com/pingcap/tidb/table/tables/corruptMutations",
             )
             .await?;