Make pgbench exit on SIGINT more reliably

Started by Tristan Partinover 2 years ago13 messages
#1Tristan Partin
tristan@neon.tech
2 attachment(s)

Hello,

The way that pgbench handled SIGINT changed in
1d468b9ad81b9139b4a0b16b416c3597925af4b0. Unfortunately this had a
couple of unintended consequences, at least from what I can tell[1]/messages/by-id/alpine.DEB.2.21.1910311939430.27369@lancre -- Tristan Partin Neon (https://neon.tech).

- CTRL-C no longer stops the program unless the right point in pgbench
execution is hit
- pgbench no longer exits with a non-zero exit code

An easy reproduction of these problems is to run with a large scale
factor like: pgbench -i -s 500000. Then try to CTRL-C the program.

The attached set of patches fixes this problem by allowing callers of
setup_cancel_handler() to attach a post-PQcancel callback. In this case,
we just call _exit(2). In addition, I noticed that psql had an EXIT_USER
constant, so I moved the common exit codes from src/bin/psql/settings.h
to src/include/fe_utils/exit_codes.h and made pgbench exit with
EXIT_USER.

[1]: /messages/by-id/alpine.DEB.2.21.1910311939430.27369@lancre -- Tristan Partin Neon (https://neon.tech)
--
Tristan Partin
Neon (https://neon.tech)

Attachments:

v1-0001-Move-exit-code-definitions-to-fe_utils.patchtext/x-patch; charset=utf-8; name=v1-0001-Move-exit-code-definitions-to-fe_utils.patchDownload
From 3cff6aec8f154eb8a47524efc15f16a6b7e95f37 Mon Sep 17 00:00:00 2001
From: Tristan Partin <tristan@neon.tech>
Date: Mon, 22 May 2023 08:07:37 -0500
Subject: [PATCH postgres v1 1/2] Move exit code definitions to fe_utils

This makes sharing exit code definitions much simpler.
---
 src/bin/psql/common.c             |  3 ++-
 src/bin/psql/mainloop.c           |  1 +
 src/bin/psql/settings.h           | 13 -------------
 src/bin/psql/startup.c            |  1 +
 src/include/fe_utils/exit_codes.h | 24 ++++++++++++++++++++++++
 5 files changed, 28 insertions(+), 14 deletions(-)
 create mode 100644 src/include/fe_utils/exit_codes.h

diff --git a/src/bin/psql/common.c b/src/bin/psql/common.c
index c0e6e8e6ed..a36e188944 100644
--- a/src/bin/psql/common.c
+++ b/src/bin/psql/common.c
@@ -25,6 +25,7 @@
 #include "copy.h"
 #include "crosstabview.h"
 #include "fe_utils/cancel.h"
+#include "fe_utils/exit_codes.h"
 #include "fe_utils/mbprint.h"
 #include "fe_utils/string_utils.h"
 #include "portability/instr_time.h"
@@ -273,7 +274,7 @@ psql_cancel_callback(void)
 void
 psql_setup_cancel_handler(void)
 {
-	setup_cancel_handler(psql_cancel_callback);
+	setup_cancel_handler(psql_cancel_callback, NULL);
 }
 
 
diff --git a/src/bin/psql/mainloop.c b/src/bin/psql/mainloop.c
index 692c6db34c..98df0e0a97 100644
--- a/src/bin/psql/mainloop.c
+++ b/src/bin/psql/mainloop.c
@@ -10,6 +10,7 @@
 #include "command.h"
 #include "common.h"
 #include "common/logging.h"
+#include "fe_utils/exit_codes.h"
 #include "input.h"
 #include "mainloop.h"
 #include "mb/pg_wchar.h"
diff --git a/src/bin/psql/settings.h b/src/bin/psql/settings.h
index 73d4b393bc..506d6db0a4 100644
--- a/src/bin/psql/settings.h
+++ b/src/bin/psql/settings.h
@@ -157,17 +157,4 @@ typedef struct _psqlSettings
 
 extern PsqlSettings pset;
 
-
-#ifndef EXIT_SUCCESS
-#define EXIT_SUCCESS 0
-#endif
-
-#ifndef EXIT_FAILURE
-#define EXIT_FAILURE 1
-#endif
-
-#define EXIT_BADCONN 2
-
-#define EXIT_USER 3
-
 #endif
diff --git a/src/bin/psql/startup.c b/src/bin/psql/startup.c
index 5a28b6f713..877408f65b 100644
--- a/src/bin/psql/startup.c
+++ b/src/bin/psql/startup.c
@@ -19,6 +19,7 @@
 #include "common/logging.h"
 #include "common/string.h"
 #include "describe.h"
+#include "fe_utils/exit_codes.h"
 #include "fe_utils/print.h"
 #include "getopt_long.h"
 #include "help.h"
diff --git a/src/include/fe_utils/exit_codes.h b/src/include/fe_utils/exit_codes.h
new file mode 100644
index 0000000000..d7136f8b50
--- /dev/null
+++ b/src/include/fe_utils/exit_codes.h
@@ -0,0 +1,24 @@
+/*
+ * psql - the PostgreSQL interactive terminal
+ *
+ * Copyright (c) 2000-2023, PostgreSQL Global Development Group
+ *
+ * src/include/fe_utils/exit_codes.h
+ */
+
+#ifndef EXIT_CODES_H
+#define EXIT_CODES_H
+
+#ifndef EXIT_SUCCESS
+#define EXIT_SUCCESS 0
+#endif
+
+#ifndef EXIT_FAILURE
+#define EXIT_FAILURE 1
+#endif
+
+#define EXIT_BADCONN 2
+
+#define EXIT_USER 3
+
+#endif
-- 
-- 
Tristan Partin
Neon (https://neon.tech)

v1-0002-Add-a-post-PQcancel-hook-to-setup_cancel_handler.patchtext/x-patch; charset=utf-8; name=v1-0002-Add-a-post-PQcancel-hook-to-setup_cancel_handler.patchDownload
From 4aa771e49095730cd1c5751c0517d147a2db6eeb Mon Sep 17 00:00:00 2001
From: Tristan Partin <tristan@neon.tech>
Date: Mon, 22 May 2023 08:19:27 -0500
Subject: [PATCH postgres v1 2/2] Add a post-PQcancel hook to
 setup_cancel_handler

This is a follow-up to 1d468b9ad81b9139b4a0b16b416c3597925af4b0. This
patch allowed pgbench to cancel server-side operations, but kept pgbench
from exiting until the only CancelRequested check was evaluated. In
addition, pgbench would not exit with a non-zero exit value given a
SIGINT.
---
 src/bin/pg_amcheck/pg_amcheck.c |  2 +-
 src/bin/pgbench/pgbench.c       | 17 +++++++++++++----
 src/bin/scripts/clusterdb.c     |  2 +-
 src/bin/scripts/reindexdb.c     |  2 +-
 src/bin/scripts/vacuumdb.c      |  2 +-
 src/fe_utils/cancel.c           | 27 ++++++++++++++++++---------
 src/include/fe_utils/cancel.h   |  5 ++++-
 7 files changed, 39 insertions(+), 18 deletions(-)

diff --git a/src/bin/pg_amcheck/pg_amcheck.c b/src/bin/pg_amcheck/pg_amcheck.c
index 68f8180c19..f3f31cde0f 100644
--- a/src/bin/pg_amcheck/pg_amcheck.c
+++ b/src/bin/pg_amcheck/pg_amcheck.c
@@ -470,7 +470,7 @@ main(int argc, char *argv[])
 	cparams.dbname = NULL;
 	cparams.override_dbname = NULL;
 
-	setup_cancel_handler(NULL);
+	setup_cancel_handler(NULL, NULL);
 
 	/* choose the database for our initial connection */
 	if (opts.alldb)
diff --git a/src/bin/pgbench/pgbench.c b/src/bin/pgbench/pgbench.c
index 70ed034e70..b8568f47de 100644
--- a/src/bin/pgbench/pgbench.c
+++ b/src/bin/pgbench/pgbench.c
@@ -39,6 +39,7 @@
 #include <math.h>
 #include <signal.h>
 #include <time.h>
+#include <unistd.h>
 #include <sys/time.h>
 #include <sys/resource.h>		/* for getrlimit */
 
@@ -60,6 +61,7 @@
 #include "common/username.h"
 #include "fe_utils/cancel.h"
 #include "fe_utils/conditional.h"
+#include "fe_utils/exit_codes.h"
 #include "fe_utils/option_utils.h"
 #include "fe_utils/string_utils.h"
 #include "getopt_long.h"
@@ -4944,9 +4946,6 @@ initGenerateDataClientSide(PGconn *con)
 		if (PQputline(con, sql.data))
 			pg_fatal("PQputline failed");
 
-		if (CancelRequested)
-			break;
-
 		/*
 		 * If we want to stick with the original logging, print a message each
 		 * 100k inserted rows.
@@ -5139,6 +5138,16 @@ checkInitSteps(const char *initialize_steps)
 	}
 }
 
+/*
+ * This function is called within a signal handler. Only use signal-safe
+ * functions. See signal-safety(7) for more information.
+ */
+static void
+post_cancel_callback(void)
+{
+	_exit(EXIT_USER);
+}
+
 /*
  * Invoke each initialization step in the given string
  */
@@ -5156,7 +5165,7 @@ runInitSteps(const char *initialize_steps)
 	if ((con = doConnect()) == NULL)
 		pg_fatal("could not create connection for initialization");
 
-	setup_cancel_handler(NULL);
+	setup_cancel_handler(NULL, post_cancel_callback);
 	SetCancelConn(con);
 
 	for (step = initialize_steps; *step != '\0'; step++)
diff --git a/src/bin/scripts/clusterdb.c b/src/bin/scripts/clusterdb.c
index e3585b3272..f7a4785ec5 100644
--- a/src/bin/scripts/clusterdb.c
+++ b/src/bin/scripts/clusterdb.c
@@ -140,7 +140,7 @@ main(int argc, char *argv[])
 	cparams.prompt_password = prompt_password;
 	cparams.override_dbname = NULL;
 
-	setup_cancel_handler(NULL);
+	setup_cancel_handler(NULL, NULL);
 
 	if (alldb)
 	{
diff --git a/src/bin/scripts/reindexdb.c b/src/bin/scripts/reindexdb.c
index 3e8f6aca40..9a6af80f9e 100644
--- a/src/bin/scripts/reindexdb.c
+++ b/src/bin/scripts/reindexdb.c
@@ -201,7 +201,7 @@ main(int argc, char *argv[])
 	cparams.prompt_password = prompt_password;
 	cparams.override_dbname = NULL;
 
-	setup_cancel_handler(NULL);
+	setup_cancel_handler(NULL, NULL);
 
 	if (alldb)
 	{
diff --git a/src/bin/scripts/vacuumdb.c b/src/bin/scripts/vacuumdb.c
index 687af9c1f3..4a49f83489 100644
--- a/src/bin/scripts/vacuumdb.c
+++ b/src/bin/scripts/vacuumdb.c
@@ -364,7 +364,7 @@ main(int argc, char *argv[])
 	cparams.prompt_password = prompt_password;
 	cparams.override_dbname = NULL;
 
-	setup_cancel_handler(NULL);
+	setup_cancel_handler(NULL, NULL);
 
 	/* Avoid opening extra connections. */
 	if (tbl_count && (concurrentCons > tbl_count))
diff --git a/src/fe_utils/cancel.c b/src/fe_utils/cancel.c
index 10c0cd4554..06cdaa27cc 100644
--- a/src/fe_utils/cancel.c
+++ b/src/fe_utils/cancel.c
@@ -63,10 +63,14 @@ static CRITICAL_SECTION cancelConnLock;
 #endif
 
 /*
- * Additional callback for cancellations.
+ * Additional callback for cancellations which is called prior to PQcancel.
  */
-static void (*cancel_callback) (void) = NULL;
+static cancel_callback *pre_cancel_callback = NULL;
 
+/*
+ * Additional callback for cancellations which is called after PQcancel.
+ */
+static cancel_callback *post_cancel_callback = NULL;
 
 /*
  * SetCancelConn
@@ -157,8 +161,8 @@ handle_sigint(SIGNAL_ARGS)
 
 	CancelRequested = true;
 
-	if (cancel_callback != NULL)
-		cancel_callback();
+	if (pre_cancel_callback != NULL)
+		pre_cancel_callback();
 
 	/* Send QueryCancel if we are processing a database query */
 	if (cancelConn != NULL)
@@ -174,6 +178,9 @@ handle_sigint(SIGNAL_ARGS)
 		}
 	}
 
+	if (post_cancel_callback != NULL)
+		post_cancel_callback();
+
 	errno = save_errno;			/* just in case the write changed it */
 }
 
@@ -183,9 +190,11 @@ handle_sigint(SIGNAL_ARGS)
  * Register query cancellation callback for SIGINT.
  */
 void
-setup_cancel_handler(void (*query_cancel_callback) (void))
+setup_cancel_handler(cancel_callback *query_pre_cancel_callback,
+					 cancel_callback *query_post_cancel_callback)
 {
-	cancel_callback = query_cancel_callback;
+	pre_cancel_callback = query_pre_cancel_callback;
+	post_cancel_callback = query_post_cancel_callback;
 	cancel_sent_msg = _("Cancel request sent\n");
 	cancel_not_sent_msg = _("Could not send cancel request: ");
 
@@ -204,8 +213,8 @@ consoleHandler(DWORD dwCtrlType)
 	{
 		CancelRequested = true;
 
-		if (cancel_callback != NULL)
-			cancel_callback();
+		if (pre_cancel_callback != NULL)
+			pre_cancel_callback();
 
 		/* Send QueryCancel if we are processing a database query */
 		EnterCriticalSection(&cancelConnLock);
@@ -234,7 +243,7 @@ consoleHandler(DWORD dwCtrlType)
 void
 setup_cancel_handler(void (*callback) (void))
 {
-	cancel_callback = callback;
+	pre_cancel_callback = callback;
 	cancel_sent_msg = _("Cancel request sent\n");
 	cancel_not_sent_msg = _("Could not send cancel request: ");
 
diff --git a/src/include/fe_utils/cancel.h b/src/include/fe_utils/cancel.h
index 6e8b9ddfc4..0597186bed 100644
--- a/src/include/fe_utils/cancel.h
+++ b/src/include/fe_utils/cancel.h
@@ -18,6 +18,8 @@
 
 #include "libpq-fe.h"
 
+typedef void cancel_callback(void);
+
 extern PGDLLIMPORT volatile sig_atomic_t CancelRequested;
 
 extern void SetCancelConn(PGconn *conn);
@@ -27,6 +29,7 @@ extern void ResetCancelConn(void);
  * A callback can be optionally set up to be called at cancellation
  * time.
  */
-extern void setup_cancel_handler(void (*query_cancel_callback) (void));
+extern void setup_cancel_handler(cancel_callback *query_pre_cancel_callback,
+								 cancel_callback *query_post_cancel_callback);
 
 #endif							/* CANCEL_H */
-- 
-- 
Tristan Partin
Neon (https://neon.tech)

#2Tristan Partin
tristan@neon.tech
In reply to: Tristan Partin (#1)
2 attachment(s)
Re: Make pgbench exit on SIGINT more reliably

Here is a v2 that handles the Windows case that I seemingly missed in my
first readthrough of this code.

--
Tristan Partin
Neon (https://neon.tech)

Attachments:

v2-0001-Move-exit-code-definitions-to-fe_utils.patchtext/x-patch; charset=utf-8; name=v2-0001-Move-exit-code-definitions-to-fe_utils.patchDownload
From 3cff6aec8f154eb8a47524efc15f16a6b7e95f37 Mon Sep 17 00:00:00 2001
From: Tristan Partin <tristan@neon.tech>
Date: Mon, 22 May 2023 08:07:37 -0500
Subject: [PATCH postgres v2 1/2] Move exit code definitions to fe_utils

This makes sharing exit code definitions much simpler.
---
 src/bin/psql/common.c             |  3 ++-
 src/bin/psql/mainloop.c           |  1 +
 src/bin/psql/settings.h           | 13 -------------
 src/bin/psql/startup.c            |  1 +
 src/include/fe_utils/exit_codes.h | 24 ++++++++++++++++++++++++
 5 files changed, 28 insertions(+), 14 deletions(-)
 create mode 100644 src/include/fe_utils/exit_codes.h

diff --git a/src/bin/psql/common.c b/src/bin/psql/common.c
index c0e6e8e6ed..a36e188944 100644
--- a/src/bin/psql/common.c
+++ b/src/bin/psql/common.c
@@ -25,6 +25,7 @@
 #include "copy.h"
 #include "crosstabview.h"
 #include "fe_utils/cancel.h"
+#include "fe_utils/exit_codes.h"
 #include "fe_utils/mbprint.h"
 #include "fe_utils/string_utils.h"
 #include "portability/instr_time.h"
@@ -273,7 +274,7 @@ psql_cancel_callback(void)
 void
 psql_setup_cancel_handler(void)
 {
-	setup_cancel_handler(psql_cancel_callback);
+	setup_cancel_handler(psql_cancel_callback, NULL);
 }
 
 
diff --git a/src/bin/psql/mainloop.c b/src/bin/psql/mainloop.c
index 692c6db34c..98df0e0a97 100644
--- a/src/bin/psql/mainloop.c
+++ b/src/bin/psql/mainloop.c
@@ -10,6 +10,7 @@
 #include "command.h"
 #include "common.h"
 #include "common/logging.h"
+#include "fe_utils/exit_codes.h"
 #include "input.h"
 #include "mainloop.h"
 #include "mb/pg_wchar.h"
diff --git a/src/bin/psql/settings.h b/src/bin/psql/settings.h
index 73d4b393bc..506d6db0a4 100644
--- a/src/bin/psql/settings.h
+++ b/src/bin/psql/settings.h
@@ -157,17 +157,4 @@ typedef struct _psqlSettings
 
 extern PsqlSettings pset;
 
-
-#ifndef EXIT_SUCCESS
-#define EXIT_SUCCESS 0
-#endif
-
-#ifndef EXIT_FAILURE
-#define EXIT_FAILURE 1
-#endif
-
-#define EXIT_BADCONN 2
-
-#define EXIT_USER 3
-
 #endif
diff --git a/src/bin/psql/startup.c b/src/bin/psql/startup.c
index 5a28b6f713..877408f65b 100644
--- a/src/bin/psql/startup.c
+++ b/src/bin/psql/startup.c
@@ -19,6 +19,7 @@
 #include "common/logging.h"
 #include "common/string.h"
 #include "describe.h"
+#include "fe_utils/exit_codes.h"
 #include "fe_utils/print.h"
 #include "getopt_long.h"
 #include "help.h"
diff --git a/src/include/fe_utils/exit_codes.h b/src/include/fe_utils/exit_codes.h
new file mode 100644
index 0000000000..d7136f8b50
--- /dev/null
+++ b/src/include/fe_utils/exit_codes.h
@@ -0,0 +1,24 @@
+/*
+ * psql - the PostgreSQL interactive terminal
+ *
+ * Copyright (c) 2000-2023, PostgreSQL Global Development Group
+ *
+ * src/include/fe_utils/exit_codes.h
+ */
+
+#ifndef EXIT_CODES_H
+#define EXIT_CODES_H
+
+#ifndef EXIT_SUCCESS
+#define EXIT_SUCCESS 0
+#endif
+
+#ifndef EXIT_FAILURE
+#define EXIT_FAILURE 1
+#endif
+
+#define EXIT_BADCONN 2
+
+#define EXIT_USER 3
+
+#endif
-- 
-- 
Tristan Partin
Neon (https://neon.tech)

v2-0002-Add-a-post-PQcancel-hook-to-setup_cancel_handler.patchtext/x-patch; charset=utf-8; name=v2-0002-Add-a-post-PQcancel-hook-to-setup_cancel_handler.patchDownload
From b90da559dc27d671095ebafa9c805f1562029508 Mon Sep 17 00:00:00 2001
From: Tristan Partin <tristan@neon.tech>
Date: Mon, 22 May 2023 08:19:27 -0500
Subject: [PATCH postgres v2 2/2] Add a post-PQcancel hook to
 setup_cancel_handler

This is a follow-up to 1d468b9ad81b9139b4a0b16b416c3597925af4b0. This
patch allowed pgbench to cancel server-side operations, but kept pgbench
from exiting until the only CancelRequested check was evaluated. In
addition, pgbench would not exit with a non-zero exit value given a
SIGINT.
---
 src/bin/pg_amcheck/pg_amcheck.c |  2 +-
 src/bin/pgbench/pgbench.c       | 17 +++++++++++++----
 src/bin/scripts/clusterdb.c     |  2 +-
 src/bin/scripts/reindexdb.c     |  2 +-
 src/bin/scripts/vacuumdb.c      |  2 +-
 src/fe_utils/cancel.c           | 34 +++++++++++++++++++++++----------
 src/include/fe_utils/cancel.h   |  5 ++++-
 7 files changed, 45 insertions(+), 19 deletions(-)

diff --git a/src/bin/pg_amcheck/pg_amcheck.c b/src/bin/pg_amcheck/pg_amcheck.c
index 68f8180c19..f3f31cde0f 100644
--- a/src/bin/pg_amcheck/pg_amcheck.c
+++ b/src/bin/pg_amcheck/pg_amcheck.c
@@ -470,7 +470,7 @@ main(int argc, char *argv[])
 	cparams.dbname = NULL;
 	cparams.override_dbname = NULL;
 
-	setup_cancel_handler(NULL);
+	setup_cancel_handler(NULL, NULL);
 
 	/* choose the database for our initial connection */
 	if (opts.alldb)
diff --git a/src/bin/pgbench/pgbench.c b/src/bin/pgbench/pgbench.c
index 70ed034e70..b8568f47de 100644
--- a/src/bin/pgbench/pgbench.c
+++ b/src/bin/pgbench/pgbench.c
@@ -39,6 +39,7 @@
 #include <math.h>
 #include <signal.h>
 #include <time.h>
+#include <unistd.h>
 #include <sys/time.h>
 #include <sys/resource.h>		/* for getrlimit */
 
@@ -60,6 +61,7 @@
 #include "common/username.h"
 #include "fe_utils/cancel.h"
 #include "fe_utils/conditional.h"
+#include "fe_utils/exit_codes.h"
 #include "fe_utils/option_utils.h"
 #include "fe_utils/string_utils.h"
 #include "getopt_long.h"
@@ -4944,9 +4946,6 @@ initGenerateDataClientSide(PGconn *con)
 		if (PQputline(con, sql.data))
 			pg_fatal("PQputline failed");
 
-		if (CancelRequested)
-			break;
-
 		/*
 		 * If we want to stick with the original logging, print a message each
 		 * 100k inserted rows.
@@ -5139,6 +5138,16 @@ checkInitSteps(const char *initialize_steps)
 	}
 }
 
+/*
+ * This function is called within a signal handler. Only use signal-safe
+ * functions. See signal-safety(7) for more information.
+ */
+static void
+post_cancel_callback(void)
+{
+	_exit(EXIT_USER);
+}
+
 /*
  * Invoke each initialization step in the given string
  */
@@ -5156,7 +5165,7 @@ runInitSteps(const char *initialize_steps)
 	if ((con = doConnect()) == NULL)
 		pg_fatal("could not create connection for initialization");
 
-	setup_cancel_handler(NULL);
+	setup_cancel_handler(NULL, post_cancel_callback);
 	SetCancelConn(con);
 
 	for (step = initialize_steps; *step != '\0'; step++)
diff --git a/src/bin/scripts/clusterdb.c b/src/bin/scripts/clusterdb.c
index e3585b3272..f7a4785ec5 100644
--- a/src/bin/scripts/clusterdb.c
+++ b/src/bin/scripts/clusterdb.c
@@ -140,7 +140,7 @@ main(int argc, char *argv[])
 	cparams.prompt_password = prompt_password;
 	cparams.override_dbname = NULL;
 
-	setup_cancel_handler(NULL);
+	setup_cancel_handler(NULL, NULL);
 
 	if (alldb)
 	{
diff --git a/src/bin/scripts/reindexdb.c b/src/bin/scripts/reindexdb.c
index 3e8f6aca40..9a6af80f9e 100644
--- a/src/bin/scripts/reindexdb.c
+++ b/src/bin/scripts/reindexdb.c
@@ -201,7 +201,7 @@ main(int argc, char *argv[])
 	cparams.prompt_password = prompt_password;
 	cparams.override_dbname = NULL;
 
-	setup_cancel_handler(NULL);
+	setup_cancel_handler(NULL, NULL);
 
 	if (alldb)
 	{
diff --git a/src/bin/scripts/vacuumdb.c b/src/bin/scripts/vacuumdb.c
index 687af9c1f3..4a49f83489 100644
--- a/src/bin/scripts/vacuumdb.c
+++ b/src/bin/scripts/vacuumdb.c
@@ -364,7 +364,7 @@ main(int argc, char *argv[])
 	cparams.prompt_password = prompt_password;
 	cparams.override_dbname = NULL;
 
-	setup_cancel_handler(NULL);
+	setup_cancel_handler(NULL, NULL);
 
 	/* Avoid opening extra connections. */
 	if (tbl_count && (concurrentCons > tbl_count))
diff --git a/src/fe_utils/cancel.c b/src/fe_utils/cancel.c
index 10c0cd4554..dcefde9315 100644
--- a/src/fe_utils/cancel.c
+++ b/src/fe_utils/cancel.c
@@ -63,10 +63,14 @@ static CRITICAL_SECTION cancelConnLock;
 #endif
 
 /*
- * Additional callback for cancellations.
+ * Additional callback for cancellations which is called prior to PQcancel.
  */
-static void (*cancel_callback) (void) = NULL;
+static cancel_callback *pre_cancel_callback = NULL;
 
+/*
+ * Additional callback for cancellations which is called after PQcancel.
+ */
+static cancel_callback *post_cancel_callback = NULL;
 
 /*
  * SetCancelConn
@@ -157,8 +161,8 @@ handle_sigint(SIGNAL_ARGS)
 
 	CancelRequested = true;
 
-	if (cancel_callback != NULL)
-		cancel_callback();
+	if (pre_cancel_callback != NULL)
+		pre_cancel_callback();
 
 	/* Send QueryCancel if we are processing a database query */
 	if (cancelConn != NULL)
@@ -174,6 +178,9 @@ handle_sigint(SIGNAL_ARGS)
 		}
 	}
 
+	if (post_cancel_callback != NULL)
+		post_cancel_callback();
+
 	errno = save_errno;			/* just in case the write changed it */
 }
 
@@ -183,9 +190,11 @@ handle_sigint(SIGNAL_ARGS)
  * Register query cancellation callback for SIGINT.
  */
 void
-setup_cancel_handler(void (*query_cancel_callback) (void))
+setup_cancel_handler(cancel_callback *query_pre_cancel_callback,
+					 cancel_callback *query_post_cancel_callback)
 {
-	cancel_callback = query_cancel_callback;
+	pre_cancel_callback = query_pre_cancel_callback;
+	post_cancel_callback = query_post_cancel_callback;
 	cancel_sent_msg = _("Cancel request sent\n");
 	cancel_not_sent_msg = _("Could not send cancel request: ");
 
@@ -204,8 +213,8 @@ consoleHandler(DWORD dwCtrlType)
 	{
 		CancelRequested = true;
 
-		if (cancel_callback != NULL)
-			cancel_callback();
+		if (pre_cancel_callback != NULL)
+			pre_cancel_callback();
 
 		/* Send QueryCancel if we are processing a database query */
 		EnterCriticalSection(&cancelConnLock);
@@ -224,6 +233,9 @@ consoleHandler(DWORD dwCtrlType)
 
 		LeaveCriticalSection(&cancelConnLock);
 
+		if (post_cancel_callback != NULL)
+			post_cancel_callback();
+
 		return TRUE;
 	}
 	else
@@ -232,9 +244,11 @@ consoleHandler(DWORD dwCtrlType)
 }
 
 void
-setup_cancel_handler(void (*callback) (void))
+setup_cancel_handler(cancel_callback *query_pre_cancel_callback,
+					 cancel_callback *query_post_cancel_callback)
 {
-	cancel_callback = callback;
+	pre_cancel_callback = query_pre_cancel_callback;
+	post_cancel_callback = query_post_cancel_callback;
 	cancel_sent_msg = _("Cancel request sent\n");
 	cancel_not_sent_msg = _("Could not send cancel request: ");
 
diff --git a/src/include/fe_utils/cancel.h b/src/include/fe_utils/cancel.h
index 6e8b9ddfc4..0597186bed 100644
--- a/src/include/fe_utils/cancel.h
+++ b/src/include/fe_utils/cancel.h
@@ -18,6 +18,8 @@
 
 #include "libpq-fe.h"
 
+typedef void cancel_callback(void);
+
 extern PGDLLIMPORT volatile sig_atomic_t CancelRequested;
 
 extern void SetCancelConn(PGconn *conn);
@@ -27,6 +29,7 @@ extern void ResetCancelConn(void);
  * A callback can be optionally set up to be called at cancellation
  * time.
  */
-extern void setup_cancel_handler(void (*query_cancel_callback) (void));
+extern void setup_cancel_handler(cancel_callback *query_pre_cancel_callback,
+								 cancel_callback *query_post_cancel_callback);
 
 #endif							/* CANCEL_H */
-- 
-- 
Tristan Partin
Neon (https://neon.tech)

#3Michael Paquier
michael@paquier.xyz
In reply to: Tristan Partin (#1)
Re: Make pgbench exit on SIGINT more reliably

On Mon, May 22, 2023 at 10:02:02AM -0500, Tristan Partin wrote:

The way that pgbench handled SIGINT changed in
1d468b9ad81b9139b4a0b16b416c3597925af4b0. Unfortunately this had a
couple of unintended consequences, at least from what I can tell[1].

- CTRL-C no longer stops the program unless the right point in pgbench
execution is hit
- pgbench no longer exits with a non-zero exit code

An easy reproduction of these problems is to run with a large scale
factor like: pgbench -i -s 500000. Then try to CTRL-C the program.

This comes from the code path where the data is generated client-side,
and where the current CancelRequested may not be that responsive,
isn't it?
--
Michael

#4Tristan Partin
tristan@neon.tech
In reply to: Michael Paquier (#3)
Re: Make pgbench exit on SIGINT more reliably

On Tue May 23, 2023 at 7:31 PM CDT, Michael Paquier wrote:

On Mon, May 22, 2023 at 10:02:02AM -0500, Tristan Partin wrote:

The way that pgbench handled SIGINT changed in
1d468b9ad81b9139b4a0b16b416c3597925af4b0. Unfortunately this had a
couple of unintended consequences, at least from what I can tell[1].

- CTRL-C no longer stops the program unless the right point in pgbench
execution is hit
- pgbench no longer exits with a non-zero exit code

An easy reproduction of these problems is to run with a large scale
factor like: pgbench -i -s 500000. Then try to CTRL-C the program.

This comes from the code path where the data is generated client-side,
and where the current CancelRequested may not be that responsive,
isn't it?

It comes from the fact that CancelRequested is only checked within the
generation of the pgbench_accounts table, but with a large enough scale
factor or enough latency, say cross-continent communication, generating
the other tables can be time consuming too. But, yes you are more likely
to run into this issue when generating client-side data.

--
Tristan Partin
Neon (https://neon.tech)

#5Tristan Partin
tristan@neon.tech
In reply to: Michael Paquier (#3)
Re: Make pgbench exit on SIGINT more reliably

On Tue May 23, 2023 at 7:31 PM CDT, Michael Paquier wrote:

On Mon, May 22, 2023 at 10:02:02AM -0500, Tristan Partin wrote:

The way that pgbench handled SIGINT changed in
1d468b9ad81b9139b4a0b16b416c3597925af4b0. Unfortunately this had a
couple of unintended consequences, at least from what I can tell[1].

- CTRL-C no longer stops the program unless the right point in pgbench
execution is hit
- pgbench no longer exits with a non-zero exit code

An easy reproduction of these problems is to run with a large scale
factor like: pgbench -i -s 500000. Then try to CTRL-C the program.

This comes from the code path where the data is generated client-side,
and where the current CancelRequested may not be that responsive,
isn't it?

Yes, that is exactly it. There is only a single check for
CancelRequested in the client-side data generation at the moment.

--
Tristan Partin
Neon (https://neon.tech)

#6Tristan Partin
tristan@neon.tech
In reply to: Tristan Partin (#5)
Re: Make pgbench exit on SIGINT more reliably

Did not even remember sending an original reply. Disregard.

--
Tristan Partin
Neon (https://neon.tech)

#7Yugo NAGATA
nagata@sraoss.co.jp
In reply to: Tristan Partin (#4)
Re: Make pgbench exit on SIGINT more reliably

On Wed, 24 May 2023 08:58:46 -0500
"Tristan Partin" <tristan@neon.tech> wrote:

On Tue May 23, 2023 at 7:31 PM CDT, Michael Paquier wrote:

On Mon, May 22, 2023 at 10:02:02AM -0500, Tristan Partin wrote:

The way that pgbench handled SIGINT changed in
1d468b9ad81b9139b4a0b16b416c3597925af4b0. Unfortunately this had a
couple of unintended consequences, at least from what I can tell[1].

- CTRL-C no longer stops the program unless the right point in pgbench
execution is hit
- pgbench no longer exits with a non-zero exit code

An easy reproduction of these problems is to run with a large scale
factor like: pgbench -i -s 500000. Then try to CTRL-C the program.

This comes from the code path where the data is generated client-side,
and where the current CancelRequested may not be that responsive,
isn't it?

It comes from the fact that CancelRequested is only checked within the
generation of the pgbench_accounts table, but with a large enough scale
factor or enough latency, say cross-continent communication, generating
the other tables can be time consuming too. But, yes you are more likely
to run into this issue when generating client-side data.

If I understand correctly, your patch allows to exit pgbench immediately
during a client-side data generation even while the tables other than
pgbench_accounts are processed. It can be useful when the scale factor
is large. However, the same purpose seems to be achieved by you other patch [1]/messages/by-id/CSTU5P82ONZ1.19XFUGHMXHBRY@c3po.
Is the latter your latest proposal that replaces the patch in this thread?ad?

[1]: /messages/by-id/CSTU5P82ONZ1.19XFUGHMXHBRY@c3po

Also, your proposal includes to change the exit code when pgbench is cancelled by
SIGINT. I think it is nice because this will make it easy to understand and clarify
the result of the initialization.

Currently, pgbench returns 0 when the initialization is cancelled while populating
pgbench_branches or pgbench_tellers, but the resultant pgbench_accounts has only
one row, which is an inconsistent result. Returning the specific value for SIGINT
cancel can let user know it explicitly. In addition, it seems better if an error
message to inform would be output.

For the above purpose, the patch moved exit codes of psql to fe_utils to be shared.
However, I am not sure this is good way. Each frontend-tool may have it own exit code,
for example. "2" means "bad connection" in psql [2]https://www.postgresql.org/docs/current/app-psql.html#id-1.9.4.20.7, but "errors during the run" in
pgbench [3]https://www.postgresql.org/docs/current/pgbench.html#id=id-1.9.4.11.7. So, I think it is better to define them separately.

[2]: https://www.postgresql.org/docs/current/app-psql.html#id-1.9.4.20.7
[3]: https://www.postgresql.org/docs/current/pgbench.html#id=id-1.9.4.11.7

Regards,
Yugo Nagata

--
Tristan Partin
Neon (https://neon.tech)

--
Yugo NAGATA <nagata@sraoss.co.jp>

#8Tristan Partin
tristan@neon.tech
In reply to: Yugo NAGATA (#7)
Re: Make pgbench exit on SIGINT more reliably

On Mon Jun 19, 2023 at 6:39 AM PDT, Yugo NAGATA wrote:

On Wed, 24 May 2023 08:58:46 -0500
"Tristan Partin" <tristan@neon.tech> wrote:

On Tue May 23, 2023 at 7:31 PM CDT, Michael Paquier wrote:

On Mon, May 22, 2023 at 10:02:02AM -0500, Tristan Partin wrote:

The way that pgbench handled SIGINT changed in
1d468b9ad81b9139b4a0b16b416c3597925af4b0. Unfortunately this had a
couple of unintended consequences, at least from what I can tell[1].

- CTRL-C no longer stops the program unless the right point in pgbench
execution is hit
- pgbench no longer exits with a non-zero exit code

An easy reproduction of these problems is to run with a large scale
factor like: pgbench -i -s 500000. Then try to CTRL-C the program.

This comes from the code path where the data is generated client-side,
and where the current CancelRequested may not be that responsive,
isn't it?

It comes from the fact that CancelRequested is only checked within the
generation of the pgbench_accounts table, but with a large enough scale
factor or enough latency, say cross-continent communication, generating
the other tables can be time consuming too. But, yes you are more likely
to run into this issue when generating client-side data.

If I understand correctly, your patch allows to exit pgbench immediately
during a client-side data generation even while the tables other than
pgbench_accounts are processed. It can be useful when the scale factor
is large. However, the same purpose seems to be achieved by you other patch [1].
Is the latter your latest proposal that replaces the patch in this thread?ad?

[1] /messages/by-id/CSTU5P82ONZ1.19XFUGHMXHBRY@c3po

The other patch does not replace this one. They are entirely separate.

Also, your proposal includes to change the exit code when pgbench is cancelled by
SIGINT. I think it is nice because this will make it easy to understand and clarify
the result of the initialization.

Currently, pgbench returns 0 when the initialization is cancelled while populating
pgbench_branches or pgbench_tellers, but the resultant pgbench_accounts has only
one row, which is an inconsistent result. Returning the specific value for SIGINT
cancel can let user know it explicitly. In addition, it seems better if an error
message to inform would be output.

For the above purpose, the patch moved exit codes of psql to fe_utils to be shared.
However, I am not sure this is good way. Each frontend-tool may have it own exit code,
for example. "2" means "bad connection" in psql [2], but "errors during the run" in
pgbench [3]. So, I think it is better to define them separately.

[2] https://www.postgresql.org/docs/current/app-psql.html#id-1.9.4.20.7
[3] https://www.postgresql.org/docs/current/pgbench.html#id=id-1.9.4.11.7

I can change it. I just figured that sharing exit code definitions
would've been preferrable. I will post a v3 some time soon which removes
that patch.

Thanks for your review :).

--
Tristan Partin
Neon (https://neon.tech)

#9Yugo NAGATA
nagata@sraoss.co.jp
In reply to: Tristan Partin (#8)
Re: Make pgbench exit on SIGINT more reliably

On Mon, 19 Jun 2023 16:49:05 -0700
"Tristan Partin" <tristan@neon.tech> wrote:

On Mon Jun 19, 2023 at 6:39 AM PDT, Yugo NAGATA wrote:

On Wed, 24 May 2023 08:58:46 -0500
"Tristan Partin" <tristan@neon.tech> wrote:

On Tue May 23, 2023 at 7:31 PM CDT, Michael Paquier wrote:

On Mon, May 22, 2023 at 10:02:02AM -0500, Tristan Partin wrote:

The way that pgbench handled SIGINT changed in
1d468b9ad81b9139b4a0b16b416c3597925af4b0. Unfortunately this had a
couple of unintended consequences, at least from what I can tell[1].

- CTRL-C no longer stops the program unless the right point in pgbench
execution is hit
- pgbench no longer exits with a non-zero exit code

An easy reproduction of these problems is to run with a large scale
factor like: pgbench -i -s 500000. Then try to CTRL-C the program.

This comes from the code path where the data is generated client-side,
and where the current CancelRequested may not be that responsive,
isn't it?

It comes from the fact that CancelRequested is only checked within the
generation of the pgbench_accounts table, but with a large enough scale
factor or enough latency, say cross-continent communication, generating
the other tables can be time consuming too. But, yes you are more likely
to run into this issue when generating client-side data.

If I understand correctly, your patch allows to exit pgbench immediately
during a client-side data generation even while the tables other than
pgbench_accounts are processed. It can be useful when the scale factor
is large. However, the same purpose seems to be achieved by you other patch [1].
Is the latter your latest proposal that replaces the patch in this thread?ad?

[1] /messages/by-id/CSTU5P82ONZ1.19XFUGHMXHBRY@c3po

The other patch does not replace this one. They are entirely separate.

After applying the other patch, CancelRequested would be checked enough frequently
even during initialization of pgbench_branches and pgbench_tellers, and it will
allow the initialization step to be cancelled immediately even if the scale factor
is high. So, is it unnecessary to modify setup_cancel_handler anyway?

I think it would be still nice to introduce a new exit code for query cancel separately.

Also, your proposal includes to change the exit code when pgbench is cancelled by
SIGINT. I think it is nice because this will make it easy to understand and clarify
the result of the initialization.

Currently, pgbench returns 0 when the initialization is cancelled while populating
pgbench_branches or pgbench_tellers, but the resultant pgbench_accounts has only
one row, which is an inconsistent result. Returning the specific value for SIGINT
cancel can let user know it explicitly. In addition, it seems better if an error
message to inform would be output.

For the above purpose, the patch moved exit codes of psql to fe_utils to be shared.
However, I am not sure this is good way. Each frontend-tool may have it own exit code,
for example. "2" means "bad connection" in psql [2], but "errors during the run" in
pgbench [3]. So, I think it is better to define them separately.

[2] https://www.postgresql.org/docs/current/app-psql.html#id-1.9.4.20.7
[3] https://www.postgresql.org/docs/current/pgbench.html#id=id-1.9.4.11.7

I can change it. I just figured that sharing exit code definitions
would've been preferrable. I will post a v3 some time soon which removes
that patch.

Great!

Thanks for your review :).

--
Tristan Partin
Neon (https://neon.tech)

--
Yugo NAGATA <nagata@sraoss.co.jp>

#10Michael Paquier
michael@paquier.xyz
In reply to: Yugo NAGATA (#9)
Re: Make pgbench exit on SIGINT more reliably

On Thu, Jun 22, 2023 at 02:58:14PM +0900, Yugo NAGATA wrote:

On Mon, 19 Jun 2023 16:49:05 -0700
"Tristan Partin" <tristan@neon.tech> wrote:

On Mon Jun 19, 2023 at 6:39 AM PDT, Yugo NAGATA wrote:

[1] /messages/by-id/CSTU5P82ONZ1.19XFUGHMXHBRY@c3po

The other patch does not replace this one. They are entirely separate.

After applying the other patch, CancelRequested would be checked enough frequently
even during initialization of pgbench_branches and pgbench_tellers, and it will
allow the initialization step to be cancelled immediately even if the scale factor
is high. So, is it unnecessary to modify setup_cancel_handler anyway?

I think it would be still nice to introduce a new exit code for query cancel separately.

I have the same impression as Nagata-san while going again through
the proposal of this thread. COPY is more responsive to
interruptions, and is always going to have a much better performance
than individual or multi-value INSERTs for the bulk-loading of data,
so we could just move on with what's proposed on the other thread and
solve the problem of this thread on top of improving the loading
performance.

What are the benefits we gain with the proposal of this thread once we
switch to COPY as method for the client-side data generation? If
the argument is to be able switch to a different error code on
cancellations for pgbench, that sounds a bit thin to me versus the
argument of keeping the cancellation callback path as simple as
possible.
--
Michael

#11Tristan Partin
tristan@neon.tech
In reply to: Michael Paquier (#10)
Re: Make pgbench exit on SIGINT more reliably

On Thu Jun 22, 2023 at 6:19 PM CDT, Michael Paquier wrote:

On Thu, Jun 22, 2023 at 02:58:14PM +0900, Yugo NAGATA wrote:

On Mon, 19 Jun 2023 16:49:05 -0700
"Tristan Partin" <tristan@neon.tech> wrote:

On Mon Jun 19, 2023 at 6:39 AM PDT, Yugo NAGATA wrote:

[1] /messages/by-id/CSTU5P82ONZ1.19XFUGHMXHBRY@c3po

The other patch does not replace this one. They are entirely separate.

After applying the other patch, CancelRequested would be checked enough frequently
even during initialization of pgbench_branches and pgbench_tellers, and it will
allow the initialization step to be cancelled immediately even if the scale factor
is high. So, is it unnecessary to modify setup_cancel_handler anyway?

I think it would be still nice to introduce a new exit code for query cancel separately.

I have the same impression as Nagata-san while going again through
the proposal of this thread. COPY is more responsive to
interruptions, and is always going to have a much better performance
than individual or multi-value INSERTs for the bulk-loading of data,
so we could just move on with what's proposed on the other thread and
solve the problem of this thread on top of improving the loading
performance.

What are the benefits we gain with the proposal of this thread once we
switch to COPY as method for the client-side data generation? If
the argument is to be able switch to a different error code on
cancellations for pgbench, that sounds a bit thin to me versus the
argument of keeping the cancellation callback path as simple as
possible.

I would say there probably isn't much benefit if you think the polling
for CancelRequested is fine. Looking at the other patch, I think it
might be simple to add an exit code for SIGINT. But I think it might be
best to do it after that patch is merged. I added the other patch to the
commitfest for July. Holding off on this one.

--
Tristan Partin
Neon (https://neon.tech)

#12Michael Paquier
michael@paquier.xyz
In reply to: Tristan Partin (#11)
Re: Make pgbench exit on SIGINT more reliably

On Tue, Jun 27, 2023 at 09:42:05AM -0500, Tristan Partin wrote:

I would say there probably isn't much benefit if you think the polling
for CancelRequested is fine. Looking at the other patch, I think it
might be simple to add an exit code for SIGINT. But I think it might be
best to do it after that patch is merged. I added the other patch to the
commitfest for July. Holding off on this one.

Okay, I am going to jump on the patch to switch to COPY, then.
--
Michael

#13Tristan Partin
tristan@neon.tech
In reply to: Michael Paquier (#12)
Re: Make pgbench exit on SIGINT more reliably

On Mon Jul 10, 2023 at 10:29 PM CDT, Michael Paquier wrote:

On Tue, Jun 27, 2023 at 09:42:05AM -0500, Tristan Partin wrote:

I would say there probably isn't much benefit if you think the polling
for CancelRequested is fine. Looking at the other patch, I think it
might be simple to add an exit code for SIGINT. But I think it might be
best to do it after that patch is merged. I added the other patch to the
commitfest for July. Holding off on this one.

Okay, I am going to jump on the patch to switch to COPY, then.

After looking at the other patch some more, I don't think there is a
good way to reliably exit from SIGINT. The only time pgbench ever polls
for CancelRequested is in initPopulateTable(). So if you hit CTRL+C at
any time during the execution of the other initalization steps, you're
out of luck. The default initialization steps are dtgvp, so after g we
have vacuuming and primary keys. From what I can tell pgbench will
continue to run these steps even after it has received a SIGINT.

This behavior seems unintended and odd to me. Though, maybe I am missing
something.

--
Tristan Partin
Neon (https://neon.tech)