From 1d468b9ad81b9139b4a0b16b416c3597925af4b0 Mon Sep 17 00:00:00 2001
From: Michael Paquier <michael@paquier.xyz>
Date: Mon, 2 Dec 2019 11:42:28 +0900
Subject: [PATCH] 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
---
 src/bin/pgbench/pgbench.c | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/src/bin/pgbench/pgbench.c b/src/bin/pgbench/pgbench.c
index 4a7ac1f821..5129aea516 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);
 }