Line data    Source code

       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-2023, 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      337668 : 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      337668 :     oldCancelConn = cancelConn;
      87             : 
      88             :     /* be sure handle_sigint doesn't use pointer while freeing */
      89      337668 :     cancelConn = NULL;
      90             : 
      91      337668 :     if (oldCancelConn != NULL)
      92        1278 :         PQfreeCancel(oldCancelConn);
      93             : 
      94      337668 :     cancelConn = PQgetCancel(conn);
      95             : 
      96             : #ifdef WIN32
      97             :     LeaveCriticalSection(&cancelConnLock);
      98             : #endif
      99      337668 : }
     100             : 
     101             : /*
     102             :  * ResetCancelConn
     103             :  *
     104             :  * Free the current cancel connection, if any, and set to NULL.
     105             :  */
     106             : void
     107      337652 : ResetCancelConn(void)
     108             : {
     109             :     PGcancel   *oldCancelConn;
     110             : 
     111             : #ifdef WIN32
     112             :     EnterCriticalSection(&cancelConnLock);
     113             : #endif
     114             : 
     115      337652 :     oldCancelConn = cancelConn;
     116             : 
     117             :     /* be sure handle_sigint doesn't use pointer while freeing */
     118      337652 :     cancelConn = NULL;
     119             : 
     120      337652 :     if (oldCancelConn != NULL)
     121      336330 :         PQfreeCancel(oldCancelConn);
     122             : 
     123             : #ifdef WIN32
     124             :     LeaveCriticalSection(&cancelConnLock);
     125             : #endif
     126      337652 : }
     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           2 : handle_sigint(SIGNAL_ARGS)
     154             : {
     155           2 :     int         save_errno = errno;
     156             :     char        errbuf[256];
     157             : 
     158           2 :     CancelRequested = true;
     159             : 
     160           2 :     if (cancel_callback != NULL)
     161           2 :         cancel_callback();
     162             : 
     163             :     /* Send QueryCancel if we are processing a database query */
     164           2 :     if (cancelConn != NULL)
     165             :     {
     166           2 :         if (PQcancel(cancelConn, errbuf, sizeof(errbuf)))
     167             :         {
     168           2 :             write_stderr(cancel_sent_msg);
     169             :         }
     170             :         else
     171             :         {
     172           0 :             write_stderr(cancel_not_sent_msg);
     173           0 :             write_stderr(errbuf);
     174             :         }
     175             :     }
     176             : 
     177           2 :     errno = save_errno;         /* just in case the write changed it */
     178           2 : }
     179             : 
     180             : /*
     181             :  * setup_cancel_handler
     182             :  *
     183             :  * Register query cancellation callback for SIGINT.
     184             :  */
     185             : void
     186       13248 : setup_cancel_handler(void (*query_cancel_callback) (void))
     187             : {
     188       13248 :     cancel_callback = query_cancel_callback;
     189       13248 :     cancel_sent_msg = _("Cancel request sent\n");
     190       13248 :     cancel_not_sent_msg = _("Could not send cancel request: ");
     191             : 
     192       13248 :     pqsignal(SIGINT, handle_sigint);
     193       13248 : }
     194             : 
     195             : #else                           /* WIN32 */
     196             : 
     197             : static BOOL WINAPI
     198             : consoleHandler(DWORD dwCtrlType)
     199             : {
     200             :     char        errbuf[256];
     201             : 
     202             :     if (dwCtrlType == CTRL_C_EVENT ||
     203             :         dwCtrlType == CTRL_BREAK_EVENT)
     204             :     {
     205             :         CancelRequested = true;
     206             : 
     207             :         if (cancel_callback != NULL)
     208             :             cancel_callback();
     209             : 
     210             :         /* Send QueryCancel if we are processing a database query */
     211             :         EnterCriticalSection(&cancelConnLock);
     212             :         if (cancelConn != NULL)
     213             :         {
     214             :             if (PQcancel(cancelConn, errbuf, sizeof(errbuf)))
     215             :             {
     216             :                 write_stderr(cancel_sent_msg);
     217             :             }
     218             :             else
     219             :             {
     220             :                 write_stderr(cancel_not_sent_msg);
     221             :                 write_stderr(errbuf);
     222             :             }
     223             :         }
     224             : 
     225             :         LeaveCriticalSection(&cancelConnLock);
     226             : 
     227             :         return TRUE;
     228             :     }
     229             :     else
     230             :         /* Return FALSE for any signals not being handled */
     231             :         return FALSE;
     232             : }
     233             : 
     234             : void
     235             : setup_cancel_handler(void (*callback) (void))
     236             : {
     237             :     cancel_callback = callback;
     238             :     cancel_sent_msg = _("Cancel request sent\n");
     239             :     cancel_not_sent_msg = _("Could not send cancel request: ");
     240             : 
     241             :     InitializeCriticalSection(&cancelConnLock);
     242             : 
     243             :     SetConsoleCtrlHandler(consoleHandler, TRUE);
     244             : }
     245             : 
     246             : #endif                          /* WIN32 */