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)
From 3cff6aec8f154eb8a47524efc15f16a6b7e95f37 Mon Sep 17 00:00:00 2001
From: Tristan Partin <tris...@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)
From b90da559dc27d671095ebafa9c805f1562029508 Mon Sep 17 00:00:00 2001
From: Tristan Partin <tris...@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)
