Refactor query cancellation code into src/fe_utils/

Originally, this code was duplicated in src/bin/psql/ and
src/bin/scripts/, but it can be useful for other frontend applications,
like pgbench.  This refactoring offers the possibility to setup a custom
callback which would get called in the signal handler for SIGINT or when
the interruption console events happen on Windows.

Author: Fabien Coelho, with contributions from Michael Paquier
Reviewed-by: Álvaro Herrera, Ibrar Ahmed
Discussion: https://postgr.es/m/alpine.DEB.2.21.1910311939430.27369@lancre
This commit is contained in:
Michael Paquier 2019-12-02 11:18:56 +09:00
parent c01ac6dcba
commit a4fd3aa719
14 changed files with 298 additions and 330 deletions

View File

@ -29,6 +29,7 @@
#include "copy.h"
#include "crosstabview.h"
#include "describe.h"
#include "fe_utils/cancel.h"
#include "fe_utils/print.h"
#include "fe_utils/string_utils.h"
#include "help.h"

View File

@ -23,6 +23,7 @@
#include "common/logging.h"
#include "copy.h"
#include "crosstabview.h"
#include "fe_utils/cancel.h"
#include "fe_utils/mbprint.h"
#include "fe_utils/string_utils.h"
#include "portability/instr_time.h"
@ -228,58 +229,28 @@ NoticeProcessor(void *arg, const char *message)
* Code to support query cancellation
*
* Before we start a query, we enable the SIGINT signal catcher to send a
* cancel request to the backend. Note that sending the cancel directly from
* the signal handler is safe because PQcancel() is written to make it
* so. We use write() to report to stderr because it's better to use simple
* facilities in a signal handler.
*
* On win32, the signal canceling happens on a separate thread, because
* that's how SetConsoleCtrlHandler works. The PQcancel function is safe
* for this (unlike PQrequestCancel). However, a CRITICAL_SECTION is required
* to protect the PGcancel structure against being changed while the signal
* thread is using it.
* cancel request to the backend.
*
* SIGINT is supposed to abort all long-running psql operations, not only
* database queries. In most places, this is accomplished by checking
* cancel_pressed during long-running loops. However, that won't work when
* CancelRequested during long-running loops. However, that won't work when
* blocked on user input (in readline() or fgets()). In those places, we
* set sigint_interrupt_enabled true while blocked, instructing the signal
* catcher to longjmp through sigint_interrupt_jmp. We assume readline and
* fgets are coded to handle possible interruption. (XXX currently this does
* not work on win32, so control-C is less useful there)
* fgets are coded to handle possible interruption.
*
* On Windows, currently this does not work, so control-C is less useful
* there, and the callback is just a no-op.
*/
volatile bool sigint_interrupt_enabled = false;
sigjmp_buf sigint_interrupt_jmp;
static PGcancel *volatile cancelConn = NULL;
#ifdef WIN32
static CRITICAL_SECTION cancelConnLock;
#endif
/*
* Write a simple string to stderr --- must be safe in a signal handler.
* We ignore the write() result since there's not much we could do about it.
* Certain compilers make that harder than it ought to be.
*/
#define write_stderr(str) \
do { \
const char *str_ = (str); \
int rc_; \
rc_ = write(fileno(stderr), str_, strlen(str_)); \
(void) rc_; \
} while (0)
#ifndef WIN32
static void
handle_sigint(SIGNAL_ARGS)
psql_cancel_callback(void)
{
int save_errno = errno;
char errbuf[256];
/* if we are waiting for input, longjmp out of it */
if (sigint_interrupt_enabled)
{
@ -288,74 +259,24 @@ handle_sigint(SIGNAL_ARGS)
}
/* else, set cancel flag to stop any long-running loops */
cancel_pressed = true;
/* and send QueryCancel if we are processing a database query */
if (cancelConn != NULL)
{
if (PQcancel(cancelConn, errbuf, sizeof(errbuf)))
write_stderr("Cancel request sent\n");
else
{
write_stderr("Could not send cancel request: ");
write_stderr(errbuf);
}
}
errno = save_errno; /* just in case the write changed it */
CancelRequested = true;
}
#else
static void
psql_cancel_callback(void)
{
/* nothing to do here */
}
#endif /* !WIN32 */
void
setup_cancel_handler(void)
psql_setup_cancel_handler(void)
{
pqsignal(SIGINT, handle_sigint);
setup_cancel_handler(psql_cancel_callback);
}
#else /* WIN32 */
static BOOL WINAPI
consoleHandler(DWORD dwCtrlType)
{
char errbuf[256];
if (dwCtrlType == CTRL_C_EVENT ||
dwCtrlType == CTRL_BREAK_EVENT)
{
/*
* Can't longjmp here, because we are in wrong thread :-(
*/
/* set cancel flag to stop any long-running loops */
cancel_pressed = true;
/* and send QueryCancel if we are processing a database query */
EnterCriticalSection(&cancelConnLock);
if (cancelConn != NULL)
{
if (PQcancel(cancelConn, errbuf, sizeof(errbuf)))
write_stderr("Cancel request sent\n");
else
{
write_stderr("Could not send cancel request: ");
write_stderr(errbuf);
}
}
LeaveCriticalSection(&cancelConnLock);
return TRUE;
}
else
/* Return FALSE for any signals not being handled */
return FALSE;
}
void
setup_cancel_handler(void)
{
InitializeCriticalSection(&cancelConnLock);
SetConsoleCtrlHandler(consoleHandler, TRUE);
}
#endif /* WIN32 */
/* ConnectionUp
@ -428,62 +349,6 @@ CheckConnection(void)
/*
* SetCancelConn
*
* Set cancelConn to point to the current database connection.
*/
void
SetCancelConn(void)
{
PGcancel *oldCancelConn;
#ifdef WIN32
EnterCriticalSection(&cancelConnLock);
#endif
/* Free the old one if we have one */
oldCancelConn = cancelConn;
/* be sure handle_sigint doesn't use pointer while freeing */
cancelConn = NULL;
if (oldCancelConn != NULL)
PQfreeCancel(oldCancelConn);
cancelConn = PQgetCancel(pset.db);
#ifdef WIN32
LeaveCriticalSection(&cancelConnLock);
#endif
}
/*
* ResetCancelConn
*
* Free the current cancel connection, if any, and set to NULL.
*/
void
ResetCancelConn(void)
{
PGcancel *oldCancelConn;
#ifdef WIN32
EnterCriticalSection(&cancelConnLock);
#endif
oldCancelConn = cancelConn;
/* be sure handle_sigint doesn't use pointer while freeing */
cancelConn = NULL;
if (oldCancelConn != NULL)
PQfreeCancel(oldCancelConn);
#ifdef WIN32
LeaveCriticalSection(&cancelConnLock);
#endif
}
/*
* AcceptResult
@ -707,7 +572,7 @@ PSQLexec(const char *query)
return NULL;
}
SetCancelConn();
SetCancelConn(pset.db);
res = PQexec(pset.db, query);
@ -746,7 +611,7 @@ PSQLexecWatch(const char *query, const printQueryOpt *opt)
return 0;
}
SetCancelConn();
SetCancelConn(pset.db);
if (pset.timing)
INSTR_TIME_SET_CURRENT(before);
@ -773,7 +638,7 @@ PSQLexecWatch(const char *query, const printQueryOpt *opt)
* consumed. The user's intention, though, is to cancel the entire watch
* process, so detect a sent cancellation request and exit in this case.
*/
if (cancel_pressed)
if (CancelRequested)
{
PQclear(res);
return 0;
@ -973,8 +838,8 @@ ExecQueryTuples(const PGresult *result)
{
const char *query = PQgetvalue(result, r, c);
/* Abandon execution if cancel_pressed */
if (cancel_pressed)
/* Abandon execution if CancelRequested */
if (CancelRequested)
goto loop_exit;
/*
@ -1091,7 +956,7 @@ ProcessResult(PGresult **results)
FILE *copystream;
PGresult *copy_result;
SetCancelConn();
SetCancelConn(pset.db);
if (result_status == PGRES_COPY_OUT)
{
bool need_close = false;
@ -1342,7 +1207,7 @@ SendQuery(const char *query)
if (fgets(buf, sizeof(buf), stdin) != NULL)
if (buf[0] == 'x')
goto sendquery_cleanup;
if (cancel_pressed)
if (CancelRequested)
goto sendquery_cleanup;
}
else if (pset.echo == PSQL_ECHO_QUERIES)
@ -1360,7 +1225,7 @@ SendQuery(const char *query)
fflush(pset.logfile);
}
SetCancelConn();
SetCancelConn(pset.db);
transaction_status = PQtransactionStatus(pset.db);
@ -1886,7 +1751,7 @@ ExecQueryUsingCursor(const char *query, double *elapsed_msec)
* writing things to the stream, we presume $PAGER has disappeared and
* stop bothering to pull down more data.
*/
if (ntuples < fetch_count || cancel_pressed || flush_error ||
if (ntuples < fetch_count || CancelRequested || flush_error ||
ferror(fout))
break;
}

View File

@ -26,10 +26,7 @@ extern volatile bool sigint_interrupt_enabled;
extern sigjmp_buf sigint_interrupt_jmp;
extern void setup_cancel_handler(void);
extern void SetCancelConn(void);
extern void ResetCancelConn(void);
extern void psql_setup_cancel_handler(void);
extern PGresult *PSQLexec(const char *query);
extern int PSQLexecWatch(const char *query, const printQueryOpt *opt);

View File

@ -9,6 +9,7 @@
#include "common.h"
#include "common/logging.h"
#include "fe_utils/cancel.h"
#include "large_obj.h"
#include "settings.h"
@ -146,7 +147,7 @@ do_lo_export(const char *loid_arg, const char *filename_arg)
if (!start_lo_xact("\\lo_export", &own_transaction))
return false;
SetCancelConn();
SetCancelConn(NULL);
status = lo_export(pset.db, atooid(loid_arg), filename_arg);
ResetCancelConn();
@ -182,7 +183,7 @@ do_lo_import(const char *filename_arg, const char *comment_arg)
if (!start_lo_xact("\\lo_import", &own_transaction))
return false;
SetCancelConn();
SetCancelConn(NULL);
loid = lo_import(pset.db, filename_arg);
ResetCancelConn();
@ -244,7 +245,7 @@ do_lo_unlink(const char *loid_arg)
if (!start_lo_xact("\\lo_unlink", &own_transaction))
return false;
SetCancelConn();
SetCancelConn(NULL);
status = lo_unlink(pset.db, loid);
ResetCancelConn();

View File

@ -301,7 +301,7 @@ main(int argc, char *argv[])
exit(EXIT_BADCONN);
}
setup_cancel_handler();
psql_setup_cancel_handler();
PQsetNoticeProcessor(pset.db, NoticeProcessor, NULL);

View File

@ -133,7 +133,7 @@ main(int argc, char *argv[])
exit(1);
}
setup_cancel_handler();
setup_cancel_handler(NULL);
if (alldb)
{

View File

@ -24,14 +24,6 @@
#define ERRCODE_UNDEFINED_TABLE "42P01"
static PGcancel *volatile cancelConn = NULL;
bool CancelRequested = false;
#ifdef WIN32
static CRITICAL_SECTION cancelConnLock;
#endif
/*
* Provide strictly harmonized handling of --help and --version
* options.
@ -465,142 +457,3 @@ yesno_prompt(const char *question)
_(PG_YESLETTER), _(PG_NOLETTER));
}
}
/*
* SetCancelConn
*
* Set cancelConn to point to the current database connection.
*/
void
SetCancelConn(PGconn *conn)
{
PGcancel *oldCancelConn;
#ifdef WIN32
EnterCriticalSection(&cancelConnLock);
#endif
/* Free the old one if we have one */
oldCancelConn = cancelConn;
/* be sure handle_sigint doesn't use pointer while freeing */
cancelConn = NULL;
if (oldCancelConn != NULL)
PQfreeCancel(oldCancelConn);
cancelConn = PQgetCancel(conn);
#ifdef WIN32
LeaveCriticalSection(&cancelConnLock);
#endif
}
/*
* ResetCancelConn
*
* Free the current cancel connection, if any, and set to NULL.
*/
void
ResetCancelConn(void)
{
PGcancel *oldCancelConn;
#ifdef WIN32
EnterCriticalSection(&cancelConnLock);
#endif
oldCancelConn = cancelConn;
/* be sure handle_sigint doesn't use pointer while freeing */
cancelConn = NULL;
if (oldCancelConn != NULL)
PQfreeCancel(oldCancelConn);
#ifdef WIN32
LeaveCriticalSection(&cancelConnLock);
#endif
}
#ifndef WIN32
/*
* Handle interrupt signals by canceling the current command, if a cancelConn
* is set.
*/
static void
handle_sigint(SIGNAL_ARGS)
{
int save_errno = errno;
char errbuf[256];
/* Send QueryCancel if we are processing a database query */
if (cancelConn != NULL)
{
if (PQcancel(cancelConn, errbuf, sizeof(errbuf)))
{
CancelRequested = true;
fprintf(stderr, _("Cancel request sent\n"));
}
else
fprintf(stderr, _("Could not send cancel request: %s"), errbuf);
}
else
CancelRequested = true;
errno = save_errno; /* just in case the write changed it */
}
void
setup_cancel_handler(void)
{
pqsignal(SIGINT, handle_sigint);
}
#else /* WIN32 */
/*
* Console control handler for Win32. Note that the control handler will
* execute on a *different thread* than the main one, so we need to do
* proper locking around those structures.
*/
static BOOL WINAPI
consoleHandler(DWORD dwCtrlType)
{
char errbuf[256];
if (dwCtrlType == CTRL_C_EVENT ||
dwCtrlType == CTRL_BREAK_EVENT)
{
/* Send QueryCancel if we are processing a database query */
EnterCriticalSection(&cancelConnLock);
if (cancelConn != NULL)
{
if (PQcancel(cancelConn, errbuf, sizeof(errbuf)))
{
fprintf(stderr, _("Cancel request sent\n"));
CancelRequested = true;
}
else
fprintf(stderr, _("Could not send cancel request: %s"), errbuf);
}
else
CancelRequested = true;
LeaveCriticalSection(&cancelConnLock);
return TRUE;
}
else
/* Return FALSE for any signals not being handled */
return FALSE;
}
void
setup_cancel_handler(void)
{
InitializeCriticalSection(&cancelConnLock);
SetConsoleCtrlHandler(consoleHandler, TRUE);
}
#endif /* WIN32 */

View File

@ -10,6 +10,7 @@
#define COMMON_H
#include "common/username.h"
#include "fe_utils/cancel.h"
#include "getopt_long.h" /* pgrminclude ignore */
#include "libpq-fe.h"
#include "pqexpbuffer.h" /* pgrminclude ignore */
@ -60,10 +61,4 @@ extern void appendQualifiedRelation(PQExpBuffer buf, const char *name,
extern bool yesno_prompt(const char *question);
extern void setup_cancel_handler(void);
extern void SetCancelConn(PGconn *conn);
extern void ResetCancelConn(void);
#endif /* COMMON_H */

View File

@ -187,7 +187,7 @@ main(int argc, char *argv[])
exit(1);
}
setup_cancel_handler();
setup_cancel_handler(NULL);
if (alldb)
{

View File

@ -257,7 +257,7 @@ main(int argc, char *argv[])
/* allow 'and_analyze' with 'analyze_only' */
}
setup_cancel_handler();
setup_cancel_handler(NULL);
/* Avoid opening extra connections. */
if (tbl_count && (concurrentCons > tbl_count))

View File

@ -20,6 +20,7 @@ include $(top_builddir)/src/Makefile.global
override CPPFLAGS := -DFRONTEND -I$(libpq_srcdir) $(CPPFLAGS)
OBJS = \
cancel.o \
conditional.o \
mbprint.o \
print.o \

225
src/fe_utils/cancel.c Normal file
View File

@ -0,0 +1,225 @@
/*------------------------------------------------------------------------
*
* Query cancellation support for frontend code
*
* Assorted utility functions to control query cancellation with signal
* handler for SIGINT.
*
*
* Portions Copyright (c) 1996-2019, PostgreSQL Global Development Group
* Portions Copyright (c) 1994, Regents of the University of California
*
* src/fe-utils/cancel.c
*
*------------------------------------------------------------------------
*/
#include "postgres_fe.h"
#include <signal.h>
#include <unistd.h>
#include "fe_utils/cancel.h"
#include "fe_utils/connect.h"
#include "fe_utils/string_utils.h"
/*
* Write a simple string to stderr --- must be safe in a signal handler.
* We ignore the write() result since there's not much we could do about it.
* Certain compilers make that harder than it ought to be.
*/
#define write_stderr(str) \
do { \
const char *str_ = (str); \
int rc_; \
rc_ = write(fileno(stderr), str_, strlen(str_)); \
(void) rc_; \
} while (0)
static PGcancel *volatile cancelConn = NULL;
bool CancelRequested = false;
#ifdef WIN32
static CRITICAL_SECTION cancelConnLock;
#endif
/*
* Additional callback for cancellations.
*/
static void (*cancel_callback) (void) = NULL;
/*
* SetCancelConn
*
* Set cancelConn to point to the current database connection.
*/
void
SetCancelConn(PGconn *conn)
{
PGcancel *oldCancelConn;
#ifdef WIN32
EnterCriticalSection(&cancelConnLock);
#endif
/* Free the old one if we have one */
oldCancelConn = cancelConn;
/* be sure handle_sigint doesn't use pointer while freeing */
cancelConn = NULL;
if (oldCancelConn != NULL)
PQfreeCancel(oldCancelConn);
cancelConn = PQgetCancel(conn);
#ifdef WIN32
LeaveCriticalSection(&cancelConnLock);
#endif
}
/*
* ResetCancelConn
*
* Free the current cancel connection, if any, and set to NULL.
*/
void
ResetCancelConn(void)
{
PGcancel *oldCancelConn;
#ifdef WIN32
EnterCriticalSection(&cancelConnLock);
#endif
oldCancelConn = cancelConn;
/* be sure handle_sigint doesn't use pointer while freeing */
cancelConn = NULL;
if (oldCancelConn != NULL)
PQfreeCancel(oldCancelConn);
#ifdef WIN32
LeaveCriticalSection(&cancelConnLock);
#endif
}
/*
* Code to support query cancellation
*
* Note that sending the cancel directly from the signal handler is safe
* because PQcancel() is written to make it so. We use write() to report
* to stderr because it's better to use simple facilities in a signal
* handler.
*
* On Windows, the signal canceling happens on a separate thread, because
* that's how SetConsoleCtrlHandler works. The PQcancel function is safe
* for this (unlike PQrequestCancel). However, a CRITICAL_SECTION is required
* to protect the PGcancel structure against being changed while the signal
* thread is using it.
*/
#ifndef WIN32
/*
* handle_sigint
*
* Handle interrupt signals by canceling the current command, if cancelConn
* is set.
*/
static void
handle_sigint(SIGNAL_ARGS)
{
int save_errno = errno;
char errbuf[256];
if (cancel_callback != NULL)
cancel_callback();
/* Send QueryCancel if we are processing a database query */
if (cancelConn != NULL)
{
if (PQcancel(cancelConn, errbuf, sizeof(errbuf)))
{
CancelRequested = true;
write_stderr(_("Cancel request sent\n"));
}
else
{
write_stderr(_("Could not send cancel request: "));
write_stderr(errbuf);
}
}
else
CancelRequested = true;
errno = save_errno; /* just in case the write changed it */
}
/*
* setup_cancel_handler
*
* Register query cancellation callback for SIGINT.
*/
void
setup_cancel_handler(void (*callback) (void))
{
cancel_callback = callback;
pqsignal(SIGINT, handle_sigint);
}
#else /* WIN32 */
static BOOL WINAPI
consoleHandler(DWORD dwCtrlType)
{
char errbuf[256];
if (dwCtrlType == CTRL_C_EVENT ||
dwCtrlType == CTRL_BREAK_EVENT)
{
if (cancel_callback != NULL)
cancel_callback();
/* Send QueryCancel if we are processing a database query */
EnterCriticalSection(&cancelConnLock);
if (cancelConn != NULL)
{
if (PQcancel(cancelConn, errbuf, sizeof(errbuf)))
{
write_stderr(_("Cancel request sent\n"));
CancelRequested = true;
}
else
{
write_stderr(_("Could not send cancel request: %s"));
write_stderr(errbuf);
}
}
else
CancelRequested = true;
LeaveCriticalSection(&cancelConnLock);
return TRUE;
}
else
/* Return FALSE for any signals not being handled */
return FALSE;
}
void
setup_cancel_handler(void (*callback) (void))
{
cancel_callback = callback;
InitializeCriticalSection(&cancelConnLock);
SetConsoleCtrlHandler(consoleHandler, TRUE);
}
#endif /* WIN32 */

View File

@ -0,0 +1,30 @@
/*-------------------------------------------------------------------------
*
* Query cancellation support for frontend code
*
*
* Portions Copyright (c) 1996-2019, PostgreSQL Global Development Group
* Portions Copyright (c) 1994, Regents of the University of California
*
* src/include/fe_utils/cancel.h
*
*-------------------------------------------------------------------------
*/
#ifndef CANCEL_H
#define CANCEL_H
#include "libpq-fe.h"
extern bool CancelRequested;
extern void SetCancelConn(PGconn *conn);
extern void ResetCancelConn(void);
/*
* A callback can be optionally set up to be called at cancellation
* time.
*/
extern void setup_cancel_handler(void (*cancel_callback) (void));
#endif /* CANCEL_H */

View File

@ -142,7 +142,7 @@ sub mkvcbuild
our @pgcommonbkndfiles = @pgcommonallfiles;
our @pgfeutilsfiles = qw(
conditional.c mbprint.c print.c psqlscan.l psqlscan.c
cancel.c conditional.c mbprint.c print.c psqlscan.l psqlscan.c
simple_list.c string_utils.c recovery_gen.c);
$libpgport = $solution->AddProject('libpgport', 'lib', 'misc');