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

Author: michaelpq

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);
 }