Add query cancellation capabilities in pgbench init phase

This can be useful to stop data generation happening on the server for
long-running queries caused by large scale factors.  This cannot happen
by default as data is generated on the client, but it is possible to
control the initialization steps of pgbench to do that.

Reported-by: Fujii Masao
Author: Fabien Coelho
Discussion: https://postgr.es/m/alpine.DEB.2.21.1910311939430.27369@lancre
Discussion: https://postgr.es/m/CAHGQGwHWEyTXxZh46qgFY8a2bDF_EYeUdp3+_Hy=qLZSzwVPKg@mail.gmail.com

1 files changed, 8 insertions, 0 deletions

diff --git a/src/bin/pgbench/pgbench.c b/src/bin/pgbench/pgbench.c
index 4a7ac1f8213..5129aea5160 100644
--- a/src/bin/pgbench/pgbench.c
+++ b/src/bin/pgbench/pgbench.c
@@ -59,6 +59,7 @@
 
 #include "common/int.h"
 #include "common/logging.h"
+#include "fe_utils/cancel.h"
 #include "fe_utils/conditional.h"
 #include "getopt_long.h"
 #include "libpq-fe.h"
@@ -3894,6 +3895,9 @@ initGenerateDataClientSide(PGconn *con)
 			exit(1);
 		}
 
+		if (CancelRequested)
+			break;
+
 		/*
 		 * If we want to stick with the original logging, print a message each
 		 * 100k inserted rows.
@@ -4109,6 +4113,9 @@ runInitSteps(const char *initialize_steps)
 	if ((con = doConnect()) == NULL)
 		exit(1);
 
+	setup_cancel_handler(NULL);
+	SetCancelConn(con);
+
 	for (step = initialize_steps; *step != '\0'; step++)
 	{
 		instr_time	start;
@@ -4176,6 +4183,7 @@ runInitSteps(const char *initialize_steps)
 	}
 
 	fprintf(stderr, "done in %.2f s (%s).\n", run_time, stats.data);
+	ResetCancelConn();
 	PQfinish(con);
 	termPQExpBuffer(&stats);
 }