Allow pgbench to exit immediately when any client is aborted.

Previously when client was aborted due to some error during
benchmarking, other clients continued their run until certain number
of transactions specified -t was reached or the time specified by -T
was expired. At the end, the results are printed with caution: "Run
was aborted; the above results are incomplete" shows.

New option "--exit-on-abort" allows pgbench to exit immediately in
this case so that users could quickly fix the cause of the failure and
try again another round of benchmarking.

Author: Yugo Nagata
Reviewed-by: Fabien COELHO, Tatsuo Ishii
Discussion: https://postgr.es/m/flat/20230804130325.df32e60879c38c92bca64207%40sraoss.co.jp

2 files changed, 58 insertions, 1 deletions

diff --git a/src/bin/pgbench/pgbench.c b/src/bin/pgbench/pgbench.c
index 763c4b946a4..617b2023b18 100644
--- a/src/bin/pgbench/pgbench.c
+++ b/src/bin/pgbench/pgbench.c
@@ -765,6 +765,8 @@ static int64 total_weight = 0;
 
 static bool verbose_errors = false; /* print verbose messages of all errors */
 
+static bool exit_on_abort = false;	/* exit when any client is aborted */
+
 /* Builtin test scripts */
 typedef struct BuiltinScript
 {
@@ -911,6 +913,7 @@ usage(void)
 		   "  -T, --time=NUM           duration of benchmark test in seconds\n"
 		   "  -v, --vacuum-all         vacuum all four standard tables before tests\n"
 		   "  --aggregate-interval=NUM aggregate data over NUM seconds\n"
+		   "  --exit-on-abort          exit when any client is aborted\n"
 		   "  --failures-detailed      report the failures grouped by basic types\n"
 		   "  --log-prefix=PREFIX      prefix for transaction time log file\n"
 		   "                           (default: \"pgbench_log\")\n"
@@ -6617,6 +6620,7 @@ main(int argc, char **argv)
 		{"failures-detailed", no_argument, NULL, 13},
 		{"max-tries", required_argument, NULL, 14},
 		{"verbose-errors", no_argument, NULL, 15},
+		{"exit-on-abort", no_argument, NULL, 16},
 		{NULL, 0, NULL, 0}
 	};
 
@@ -6950,6 +6954,10 @@ main(int argc, char **argv)
 				benchmarking_option_set = true;
 				verbose_errors = true;
 				break;
+			case 16:			/* exit-on-abort */
+				benchmarking_option_set = true;
+				exit_on_abort = true;
+				break;
 			default:
 				/* getopt_long already emitted a complaint */
 				pg_log_error_hint("Try \"%s --help\" for more information.", progname);
@@ -7559,10 +7567,17 @@ threadRun(void *arg)
 			advanceConnectionState(thread, st, &aggs);
 
 			/*
+			 * If --exit-on-abort is used, the program is going to exit
+			 * when any client is aborted.
+			 */
+			if (exit_on_abort && st->state == CSTATE_ABORTED)
+				goto done;
+			/*
 			 * If advanceConnectionState changed client to finished state,
 			 * that's one fewer client that remains.
 			 */
-			if (st->state == CSTATE_FINISHED || st->state == CSTATE_ABORTED)
+			else if (st->state == CSTATE_FINISHED ||
+					 st->state == CSTATE_ABORTED)
 				remains--;
 		}
 
@@ -7595,6 +7610,22 @@ threadRun(void *arg)
 	}
 
 done:
+	if (exit_on_abort)
+	{
+		/*
+		 * Abort if any client is not finished, meaning some error occurred.
+		 */
+		for (int i = 0; i < nstate; i++)
+		{
+			if (state[i].state != CSTATE_FINISHED)
+			{
+				pg_log_error("Run was aborted due to an error in thread %d",
+							 thread->tid);
+				exit(2);
+			}
+		}
+	}
+
 	disconnect_all(state, nstate);
 
 	if (thread->logfile)
diff --git a/src/bin/pgbench/t/001_pgbench_with_server.pl b/src/bin/pgbench/t/001_pgbench_with_server.pl
index 142f966300a..96be529d6bb 100644
--- a/src/bin/pgbench/t/001_pgbench_with_server.pl
+++ b/src/bin/pgbench/t/001_pgbench_with_server.pl
@@ -1475,6 +1475,32 @@ SELECT pg_advisory_unlock_all();
 # Clean up
 $node->safe_psql('postgres', 'DROP TABLE first_client_table, xy;');
 
+# Test --exit-on-abort
+$node->safe_psql('postgres',
+	'CREATE TABLE counter(i int); '.
+	'INSERT INTO counter VALUES (0);'
+);
+
+$node->pgbench(
+	'-t 10 -c 2 -j 2 --exit-on-abort',
+	2,
+	[],
+	[
+		qr{division by zero},
+		qr{Run was aborted due to an error in thread}
+	],
+	'test --exit-on-abort',
+	{
+		'001_exit_on_abort' => q{
+update counter set i = i+1 returning i \gset
+\if :i = 5
+\set y 1/0
+\endif
+}
+	});
+
+# Clean up
+$node->safe_psql('postgres', 'DROP TABLE counter;');
 
 # done
 $node->safe_psql('postgres', 'DROP TABLESPACE regress_pgbench_tap_1_ts');