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-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      386494 : 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      386494 :     oldCancelConn = cancelConn;
      87             : 
      88             :     /* be sure handle_sigint doesn't use pointer while freeing */
      89      386494 :     cancelConn = NULL;
      90             : 
      91      386494 :     if (oldCancelConn != NULL)
      92        1422 :         PQfreeCancel(oldCancelConn);
      93             : 
      94      386494 :     cancelConn = PQgetCancel(conn);
      95             : 
      96             : #ifdef WIN32
      97             :     LeaveCriticalSection(&cancelConnLock);
      98             : #endif
      99      386494 : }
     100             : 
     101             : /*
     102             :  * ResetCancelConn
     103             :  *
     104             :  * Free the current cancel connection, if any, and set to NULL.
     105             :  */
     106             : void
     107      386472 : ResetCancelConn(void)
     108             : {
     109             :     PGcancel   *oldCancelConn;
     110             : 
     111             : #ifdef WIN32
     112             :     EnterCriticalSection(&cancelConnLock);
     113             : #endif
     114             : 
     115      386472 :     oldCancelConn = cancelConn;
     116             : 
     117             :     /* be sure handle_sigint doesn't use pointer while freeing */
     118      386472 :     cancelConn = NULL;
     119             : 
     120      386472 :     if (oldCancelConn != NULL)
     121      385006 :         PQfreeCancel(oldCancelConn);
     122             : 
     123             : #ifdef WIN32
     124             :     LeaveCriticalSection(&cancelConnLock);
     125             : #endif
     126      386472 : }
     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             :     char        errbuf[256];
     156             : 
     157           2 :     CancelRequested = true;
     158             : 
     159           2 :     if (cancel_callback != NULL)
     160           2 :         cancel_callback();
     161             : 
     162             :     /* Send QueryCancel if we are processing a database query */
     163           2 :     if (cancelConn != NULL)
     164             :     {
     165           2 :         if (PQcancel(cancelConn, errbuf, sizeof(errbuf)))
     166             :         {
     167           2 :             write_stderr(cancel_sent_msg);
     168             :         }
     169             :         else
     170             :         {
     171           0 :             write_stderr(cancel_not_sent_msg);
     172           0 :             write_stderr(errbuf);
     173             :         }
     174             :     }
     175           2 : }
     176             : 
     177             : /*
     178             :  * setup_cancel_handler
     179             :  *
     180             :  * Register query cancellation callback for SIGINT.
     181             :  */
     182             : void
     183       18608 : setup_cancel_handler(void (*query_cancel_callback) (void))
     184             : {
     185       18608 :     cancel_callback = query_cancel_callback;
     186       18608 :     cancel_sent_msg = _("Cancel request sent\n");
     187       18608 :     cancel_not_sent_msg = _("Could not send cancel request: ");
     188             : 
     189       18608 :     pqsignal(SIGINT, handle_sigint);
     190       18608 : }
     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 */