From 4aa771e49095730cd1c5751c0517d147a2db6eeb Mon Sep 17 00:00:00 2001 From: Tristan Partin 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 #include #include +#include #include #include /* 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)