add IsHoldOffCancellationReceived utility function (#3290)

src/backend/distributed/connection/connection_management.c

@@ -25,6 +25,7 @@
 #include "distributed/hash_helpers.h"
 #include "distributed/placement_connection.h"
 #include "distributed/run_from_same_connection.h"
+#include "distributed/cancel_utils.h"
 #include "distributed/remote_commands.h"
 #include "distributed/version_compat.h"
 #include "mb/pg_wchar.h"
@@ -738,7 +739,7 @@ FinishConnectionListEstablishment(List *multiConnectionList)
 
 				CHECK_FOR_INTERRUPTS();
 
-				if (InterruptHoldoffCount > 0 && (QueryCancelPending || ProcDiePending))
+				if (IsHoldOffCancellationReceived())
 				{
 					/*
 					 * because we can't break from 2 loops easily we need to not forget to

src/backend/distributed/connection/remote_commands.c

@@ -17,6 +17,7 @@
 #include "distributed/errormessage.h"
 #include "distributed/log_utils.h"
 #include "distributed/remote_commands.h"
+#include "distributed/cancel_utils.h"
 #include "lib/stringinfo.h"
 #include "miscadmin.h"
 #include "storage/latch.h"
@@ -757,7 +758,7 @@ FinishConnectionIO(MultiConnection *connection, bool raiseInterrupts)
 			 * interrupts held, return instead, and mark the transaction as
 			 * failed.
 			 */
-			if (InterruptHoldoffCount > 0 && (QueryCancelPending || ProcDiePending))
+			if (IsHoldOffCancellationReceived())
 			{
 				connection->remoteTransaction.transactionFailed = true;
 				break;
@@ -863,8 +864,7 @@ WaitForAllConnections(List *connectionList, bool raiseInterrupts)
 						CHECK_FOR_INTERRUPTS();
 					}
 
-					if (InterruptHoldoffCount > 0 && (QueryCancelPending ||
+					if (IsHoldOffCancellationReceived())
-													  ProcDiePending))
 					{
 						/*
 						 * Break out of event loop immediately in case of cancellation.

src/backend/distributed/executor/adaptive_executor.c

@@ -143,6 +143,7 @@
 #include "distributed/placement_access.h"
 #include "distributed/placement_connection.h"
 #include "distributed/relation_access_tracking.h"
+#include "distributed/cancel_utils.h"
 #include "distributed/remote_commands.h"
 #include "distributed/resource_lock.h"
 #include "distributed/subplan_execution.h"
@@ -1810,7 +1811,7 @@ SequentialRunDistributedExecution(DistributedExecution *execution)
 
 		CHECK_FOR_INTERRUPTS();
 
-		if (InterruptHoldoffCount > 0 && (QueryCancelPending || ProcDiePending))
+		if (IsHoldOffCancellationReceived())
 		{
 			break;
 		}
@@ -1918,8 +1919,7 @@ RunDistributedExecution(DistributedExecution *execution)
 						CHECK_FOR_INTERRUPTS();
 					}
 
-					if (InterruptHoldoffCount > 0 && (QueryCancelPending ||
+					if (IsHoldOffCancellationReceived())
-													  ProcDiePending))
 					{
 						/*
 						 * Break out of event loop immediately in case of cancellation.

src/backend/distributed/utils/cancel_utils.c

@@ -0,0 +1,25 @@
+/*
+ * cancel_utils.c
+ *
+ * Utilities related to query cancellation
+ *
+ * Copyright (c) Citus Data, Inc.
+ */
+
+
+#include "postgres.h"
+#include "miscadmin.h"
+#include "distributed/cancel_utils.h"
+
+
+/*
+ * IsHoldOffCancellationReceived returns true if a cancel signal
+ * was sent and HOLD_INTERRUPTS was called prior to this. The motivation
+ * here is that since our queries can take a long time, in some places
+ * we do not want to wait even if HOLD_INTERRUPTS was called.
+ */
+bool
+IsHoldOffCancellationReceived()
+{
+	return InterruptHoldoffCount > 0 && (QueryCancelPending || ProcDiePending);
+}

src/include/distributed/cancel_utils.h

@@ -0,0 +1,17 @@
+/*-------------------------------------------------------------------------
+ *
+ * cencel_utils.h
+ *	  Utilities related to query cancellation
+ *
+ * Copyright (c) Citus Data, Inc.
+ *
+ *-------------------------------------------------------------------------
+ */
+
+#ifndef CANCEL_UTILS_H
+#define CANCEL_UTILS_H
+
+
+extern bool IsHoldOffCancellationReceived(void);
+
+#endif /* CANCEL_UTILS_H */