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

diff --git a/src/bin/pgbench/pgbench.c b/src/bin/pgbench/pgbench.c
index 4a7ac1f82137580ae5953fe5384c65a1d559aef9..5129aea51603fe241288abe624084cb487c8d275 100644 (file)
--- 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);
 }