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-2024, 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  * Predetermined localized error strings --- needed to avoid trying
  47  * to call gettext() from a signal handler.
  48  */
  49 static const char *cancel_sent_msg = NULL;
  50 static const char *cancel_not_sent_msg = NULL;
  51
  52 /*
  53  * CancelRequested is set when we receive SIGINT (or local equivalent).
  54  * There is no provision in this module for resetting it; but applications
  55  * might choose to clear it after successfully recovering from a cancel.
  56  * Note that there is no guarantee that we successfully sent a Cancel request,
  57  * or that the request will have any effect if we did send it.
  58  */
  59 volatile sig_atomic_t CancelRequested = false;
  60
  61 #ifdef WIN32
  62 static CRITICAL_SECTION cancelConnLock;
  63 #endif
  64
  65 /*
  66  * Additional callback for cancellations.
  67  */
  68 static void (*cancel_callback) (void) = NULL;
  69
  70
  71 /*
  72  * SetCancelConn
  73  *
  74  * Set cancelConn to point to the current database connection.
  75  */
  76 void
  77 SetCancelConn(PGconn *conn)
  78 {
  79         PGcancel   *oldCancelConn;
  80
  81 #ifdef WIN32
  82         EnterCriticalSection(&cancelConnLock);
  83 #endif
  84
  85         /* Free the old one if we have one */
  86         oldCancelConn = cancelConn;
  87
  88         /* be sure handle_sigint doesn't use pointer while freeing */
  89         cancelConn = NULL;
  90
  91         if (oldCancelConn != NULL)
  92                 PQfreeCancel(oldCancelConn);
  93
  94         cancelConn = PQgetCancel(conn);
  95
  96 #ifdef WIN32
  97         LeaveCriticalSection(&cancelConnLock);
  98 #endif
  99 }
 100
 101 /*
 102  * ResetCancelConn
 103  *
 104  * Free the current cancel connection, if any, and set to NULL.
 105  */
 106 void
 107 ResetCancelConn(void)
 108 {
 109         PGcancel   *oldCancelConn;
 110
 111 #ifdef WIN32
 112         EnterCriticalSection(&cancelConnLock);
 113 #endif
 114
 115         oldCancelConn = cancelConn;
 116
 117         /* be sure handle_sigint doesn't use pointer while freeing */
 118         cancelConn = NULL;
 119
 120         if (oldCancelConn != NULL)
 121                 PQfreeCancel(oldCancelConn);
 122
 123 #ifdef WIN32
 124         LeaveCriticalSection(&cancelConnLock);
 125 #endif
 126 }
 127
 128
 129 /*
 130  * Code to support query cancellation
 131  *
 132  * Note that sending the cancel directly from the signal handler is safe
 133  * because PQcancel() is written to make it so.  We use write() to report
 134  * to stderr because it's better to use simple facilities in a signal
 135  * handler.
 136  *
 137  * On Windows, the signal canceling happens on a separate thread, because
 138  * that's how SetConsoleCtrlHandler works.  The PQcancel function is safe
 139  * for this (unlike PQrequestCancel).  However, a CRITICAL_SECTION is required
 140  * to protect the PGcancel structure against being changed while the signal
 141  * thread is using it.
 142  */
 143
 144 #ifndef WIN32
 145
 146 /*
 147  * handle_sigint
 148  *
 149  * Handle interrupt signals by canceling the current command, if cancelConn
 150  * is set.
 151  */
 152 static void
 153 handle_sigint(SIGNAL_ARGS)
 154 {
 155         char            errbuf[256];
 156
 157         CancelRequested = true;
 158
 159         if (cancel_callback != NULL)
 160                 cancel_callback();
 161
 162         /* Send QueryCancel if we are processing a database query */
 163         if (cancelConn != NULL)
 164         {
 165                 if (PQcancel(cancelConn, errbuf, sizeof(errbuf)))
 166                 {
 167                         write_stderr(cancel_sent_msg);
 168                 }
 169                 else
 170                 {
 171                         write_stderr(cancel_not_sent_msg);
 172                         write_stderr(errbuf);
 173                 }
 174         }
 175 }
 176
 177 /*
 178  * setup_cancel_handler
 179  *
 180  * Register query cancellation callback for SIGINT.
 181  */
 182 void
 183 setup_cancel_handler(void (*query_cancel_callback) (void))
 184 {
 185         cancel_callback = query_cancel_callback;
 186         cancel_sent_msg = _("Cancel request sent\n");
 187         cancel_not_sent_msg = _("Could not send cancel request: ");
 188
 189         pqsignal(SIGINT, handle_sigint);
 190 }
 191
 192 #else                                                   /* WIN32 */
 193
 194 static BOOL WINAPI
 195 consoleHandler(DWORD dwCtrlType)
 196 {
 197         char            errbuf[256];
 198
 199         if (dwCtrlType == CTRL_C_EVENT ||
 200                 dwCtrlType == CTRL_BREAK_EVENT)
 201         {
 202                 CancelRequested = true;
 203
 204                 if (cancel_callback != NULL)
 205                         cancel_callback();
 206
 207                 /* Send QueryCancel if we are processing a database query */
 208                 EnterCriticalSection(&cancelConnLock);
 209                 if (cancelConn != NULL)
 210                 {
 211                         if (PQcancel(cancelConn, errbuf, sizeof(errbuf)))
 212                         {
 213                                 write_stderr(cancel_sent_msg);
 214                         }
 215                         else
 216                         {
 217                                 write_stderr(cancel_not_sent_msg);
 218                                 write_stderr(errbuf);
 219                         }
 220                 }
 221
 222                 LeaveCriticalSection(&cancelConnLock);
 223
 224                 return TRUE;
 225         }
 226         else
 227                 /* Return FALSE for any signals not being handled */
 228                 return FALSE;
 229 }
 230
 231 void
 232 setup_cancel_handler(void (*callback) (void))
 233 {
 234         cancel_callback = callback;
 235         cancel_sent_msg = _("Cancel request sent\n");
 236         cancel_not_sent_msg = _("Could not send cancel request: ");
 237
 238         InitializeCriticalSection(&cancelConnLock);
 239
 240         SetConsoleCtrlHandler(consoleHandler, TRUE);
 241 }
 242
 243 #endif                                                  /* WIN32 */