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