mirror of https://github.com/postgres/postgres
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@lancrepull/44/head
Michael Paquier
6 years ago
parent
c01ac6dcba
commit
a4fd3aa719
@ -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 */ |
||||
@ -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 */ |
||||
Loading…
Reference in new issue