g3bench: add option to ignore fatal error

bytedance · Oct 20, 2023 · e21bc6c · e21bc6c
1 parent 8c545b3
commit e21bc6c
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 3 deletions.
diff --git a/g3bench/src/opts.rs b/g3bench/src/opts.rs
@@ -45,6 +45,7 @@ const GLOBAL_ARG_TIME_LIMIT: &str = "time-limit";
 const GLOBAL_ARG_REQUESTS: &str = "requests";
 const GLOBAL_ARG_RESOLVE: &str = "resolve";
 const GLOBAL_ARG_LOG_ERROR: &str = "log-error";
+const GLOBAL_ARG_IGNORE_FATAL_ERROR: &str = "ignore-fatal-error";
 const GLOBAL_ARG_EMIT_METRICS: &str = "emit-metrics";
 const GLOBAL_ARG_STATSD_TARGET_UDP: &str = "statsd-target-udp";
 const GLOBAL_ARG_STATSD_TARGET_UNIX: &str = "statsd-target-unix";
@@ -62,6 +63,7 @@ pub struct ProcArgs {
     pub(super) requests: Option<usize>,
     pub(super) time_limit: Option<Duration>,
     pub(super) log_error_count: usize,
+    pub(super) ignore_fatal_error: bool,
     pub(super) task_unconstrained: bool,
     resolver: AHashMap<UpstreamAddr, IpAddr>,
     pub(super) use_unaided_worker: bool,
@@ -84,6 +86,7 @@ impl Default for ProcArgs {
             requests: None,
             time_limit: None,
             log_error_count: 0,
+            ignore_fatal_error: false,
             task_unconstrained: false,
             resolver: AHashMap::new(),
             use_unaided_worker: false,
@@ -322,6 +325,13 @@ pub fn add_global_args(app: Command) -> Command {
             .num_args(1)
             .value_parser(value_parser!(usize)),
     )
+    .arg(
+        Arg::new(GLOBAL_ARG_IGNORE_FATAL_ERROR)
+            .help("Continue even if fatal error occurred")
+            .long(GLOBAL_ARG_IGNORE_FATAL_ERROR)
+            .global(true)
+            .action(ArgAction::SetTrue),
+    )
     .arg(
         Arg::new(GLOBAL_ARG_EMIT_METRICS)
             .help("Set if we need to emit metrics to statsd")
@@ -435,6 +445,9 @@ pub fn parse_global_args(args: &ArgMatches) -> anyhow::Result<ProcArgs> {
     if let Some(n) = args.get_one::<usize>(GLOBAL_ARG_LOG_ERROR) {
         proc_args.log_error_count = *n;
     }
+    if args.get_flag(GLOBAL_ARG_IGNORE_FATAL_ERROR) {
+        proc_args.ignore_fatal_error = true;
+    }
 
     if args.get_flag(GLOBAL_ARG_EMIT_METRICS) {
         let mut config =

diff --git a/g3bench/src/target/mod.rs b/g3bench/src/target/mod.rs
@@ -179,6 +179,7 @@ where
 
         let task_unconstrained = proc_args.task_unconstrained;
         let latency = proc_args.latency;
+        let ignore_fatal_error = proc_args.ignore_fatal_error;
         let rt = super::worker::select_handle(i).unwrap_or_else(tokio::runtime::Handle::current);
         rt.spawn(async move {
             sem.add_permits(1);
@@ -217,15 +218,21 @@ where
                     Err(BenchError::Fatal(e)) => {
                         context.mark_task_failed();
                         global_state.add_failed();
-                        eprintln!("!! Fatal error with task context {i}: {e:?}");
-                        break;
+                        if ignore_fatal_error {
+                            if global_state.check_log_error() {
+                                eprintln!("! request {task_id} failed: {e:?}\n");
+                            }
+                        } else {
+                            eprintln!("!! Fatal error with task context {i}: {e:?}");
+                            break;
+                        }
                     }
                     Err(BenchError::Task(e)) => {
                         context.mark_task_failed();
+                        global_state.add_failed();
                         if global_state.check_log_error() {
                             eprintln!("! request {task_id} failed: {e:?}\n");
                         }
-                        global_state.add_failed();
                     }
                 }
                 req_count += 1;