src/fe_utils/cancel.c

   1 /*------------------------------------------------------------------------
   2  *
   3  * Query cancellation support for frontend code
   4  *
   5  * Assorted utility functions to control query cancellation with signal
   6  * handler for SIGINT.
   7  *
   8  *
   9  * Portions Copyright (c) 1996-2021, PostgreSQL Global Development Group
  10  * Portions Copyright (c) 1994, Regents of the University of California
  11  *
  12  * src/fe-utils/cancel.c
  13  *
  14  *------------------------------------------------------------------------
  15  */
  16
  17 #include "postgres_fe.h"
  18
  19 #include <unistd.h>
  20
  21 #include "common/connect.h"
  22 #include "fe_utils/cancel.h"
  23 #include "fe_utils/string_utils.h"
  24
  25
  26 /*
  27  * Write a simple string to stderr --- must be safe in a signal handler.
  28  * We ignore the write() result since there's not much we could do about it.
  29  * Certain compilers make that harder than it ought to be.
  30  */
  31 #define write_stderr(str) \
  32         do { \
  33                 const char *str_ = (str); \
  34                 int             rc_; \
  35                 rc_ = write(fileno(stderr), str_, strlen(str_)); \
  36                 (void) rc_; \
  37         } while (0)
  38
  39 /*
  40  * Contains all the information needed to cancel a query issued from
  41  * a database connection to the backend.
  42  */
  43 static PGcancel *volatile cancelConn = NULL;
  44
  45 /*
  46  * CancelRequested is set when we receive SIGINT (or local equivalent).
  47  * There is no provision in this module for resetting it; but applications
  48  * might choose to clear it after successfully recovering from a cancel.
  49  * Note that there is no guarantee that we successfully sent a Cancel request,
  50  * or that the request will have any effect if we did send it.
  51  */
  52 volatile sig_atomic_t CancelRequested = false;
  53
  54 #ifdef WIN32
  55 static CRITICAL_SECTION cancelConnLock;
  56 #endif
  57
  58 /*
  59  * Additional callback for cancellations.
  60  */
  61 static void (*cancel_callback) (void) = NULL;
  62
  63
  64 /*
  65  * SetCancelConn
  66  *
  67  * Set cancelConn to point to the current database connection.
  68  */
  69 void
  70 SetCancelConn(PGconn *conn)
  71 {
  72         PGcancel   *oldCancelConn;
  73
  74 #ifdef WIN32
  75         EnterCriticalSection(&cancelConnLock);
  76 #endif
  77
  78         /* Free the old one if we have one */
  79         oldCancelConn = cancelConn;
  80
  81         /* be sure handle_sigint doesn't use pointer while freeing */
  82         cancelConn = NULL;
  83
  84         if (oldCancelConn != NULL)
  85                 PQfreeCancel(oldCancelConn);
  86
  87         cancelConn = PQgetCancel(conn);
  88
  89 #ifdef WIN32
  90         LeaveCriticalSection(&cancelConnLock);
  91 #endif
  92 }
  93
  94 /*
  95  * ResetCancelConn
  96  *
  97  * Free the current cancel connection, if any, and set to NULL.
  98  */
  99 void
 100 ResetCancelConn(void)
 101 {
 102         PGcancel   *oldCancelConn;
 103
 104 #ifdef WIN32
 105         EnterCriticalSection(&cancelConnLock);
 106 #endif
 107
 108         oldCancelConn = cancelConn;
 109
 110         /* be sure handle_sigint doesn't use pointer while freeing */
 111         cancelConn = NULL;
 112
 113         if (oldCancelConn != NULL)
 114                 PQfreeCancel(oldCancelConn);
 115
 116 #ifdef WIN32
 117         LeaveCriticalSection(&cancelConnLock);
 118 #endif
 119 }
 120
 121
 122 /*
 123  * Code to support query cancellation
 124  *
 125  * Note that sending the cancel directly from the signal handler is safe
 126  * because PQcancel() is written to make it so.  We use write() to report
 127  * to stderr because it's better to use simple facilities in a signal
 128  * handler.
 129  *
 130  * On Windows, the signal canceling happens on a separate thread, because
 131  * that's how SetConsoleCtrlHandler works.  The PQcancel function is safe
 132  * for this (unlike PQrequestCancel).  However, a CRITICAL_SECTION is required
 133  * to protect the PGcancel structure against being changed while the signal
 134  * thread is using it.
 135  */
 136
 137 #ifndef WIN32
 138
 139 /*
 140  * handle_sigint
 141  *
 142  * Handle interrupt signals by canceling the current command, if cancelConn
 143  * is set.
 144  */
 145 static void
 146 handle_sigint(SIGNAL_ARGS)
 147 {
 148         int                     save_errno = errno;
 149         char            errbuf[256];
 150
 151         CancelRequested = true;
 152
 153         if (cancel_callback != NULL)
 154                 cancel_callback();
 155
 156         /* Send QueryCancel if we are processing a database query */
 157         if (cancelConn != NULL)
 158         {
 159                 if (PQcancel(cancelConn, errbuf, sizeof(errbuf)))
 160                 {
 161                         write_stderr(_("Cancel request sent\n"));
 162                 }
 163                 else
 164                 {
 165                         write_stderr(_("Could not send cancel request: "));
 166                         write_stderr(errbuf);
 167                 }
 168         }
 169
 170         errno = save_errno;                     /* just in case the write changed it */
 171 }
 172
 173 /*
 174  * setup_cancel_handler
 175  *
 176  * Register query cancellation callback for SIGINT.
 177  */
 178 void
 179 setup_cancel_handler(void (*callback) (void))
 180 {
 181         cancel_callback = callback;
 182         pqsignal(SIGINT, handle_sigint);
 183 }
 184
 185 #else                                                   /* WIN32 */
 186
 187 static BOOL WINAPI
 188 consoleHandler(DWORD dwCtrlType)
 189 {
 190         char            errbuf[256];
 191
 192         if (dwCtrlType == CTRL_C_EVENT ||
 193                 dwCtrlType == CTRL_BREAK_EVENT)
 194         {
 195                 CancelRequested = true;
 196
 197                 if (cancel_callback != NULL)
 198                         cancel_callback();
 199
 200                 /* Send QueryCancel if we are processing a database query */
 201                 EnterCriticalSection(&cancelConnLock);
 202                 if (cancelConn != NULL)
 203                 {
 204                         if (PQcancel(cancelConn, errbuf, sizeof(errbuf)))
 205                         {
 206                                 write_stderr(_("Cancel request sent\n"));
 207                         }
 208                         else
 209                         {
 210                                 write_stderr(_("Could not send cancel request: "));
 211                                 write_stderr(errbuf);
 212                         }
 213                 }
 214
 215                 LeaveCriticalSection(&cancelConnLock);
 216
 217                 return TRUE;
 218         }
 219         else
 220                 /* Return FALSE for any signals not being handled */
 221                 return FALSE;
 222 }
 223
 224 void
 225 setup_cancel_handler(void (*callback) (void))
 226 {
 227         cancel_callback = callback;
 228
 229         InitializeCriticalSection(&cancelConnLock);
 230
 231         SetConsoleCtrlHandler(consoleHandler, TRUE);
 232 }
 233
 234 #endif                                                  /* WIN32 */