From a8f7be269dcfe6cd8c4771840faa290e2c20efdb Mon Sep 17 00:00:00 2001
From: "kuroda.hayato%40jp.fujitsu.com" <kuroda.hayato@jp.fujitsu.com>
Date: Wed, 21 Sep 2022 06:02:45 +0000
Subject: [PATCH v17 1/4] Add an infrastracture for checking remote servers

This patch adds a mechanism for registering callback functions.
They should be used for checking health of remote servers.

These functions will be called when flags CheckingRemoteServersTimeoutPending and InterruptPending are set to true.
CheckingRemoteServersTimeoutPending is expected that it is set from signal handlers, which is registered by FDWs.
Inside the function a signal SIGINT should be raised and a message should be set to QueryCancelMessage
if one of remote servers is disconnected. To fill the message new API TrySetQueryCancelMessage can be used.

When a query is canceled and a string is set to QueryCancelMessage,
the server will output the given message to the log instead of the normal message.

Note that TrySetQueryCancelMessage does not allow to override the QueryCancelMessage.
If the function is called but the message is already filled, returns immediately
---
 src/backend/foreign/foreign.c    | 60 ++++++++++++++++++++++++++++++++
 src/backend/tcop/postgres.c      | 50 ++++++++++++++++++++++++++
 src/backend/utils/init/globals.c |  1 +
 src/include/foreign/foreign.h    | 19 ++++++++++
 src/include/miscadmin.h          |  2 ++
 src/include/tcop/tcopprot.h      |  3 ++
 src/tools/pgindent/typedefs.list |  1 +
 7 files changed, 136 insertions(+)

diff --git a/src/backend/foreign/foreign.c b/src/backend/foreign/foreign.c
index 353e20a0cf..58353abc40 100644
--- a/src/backend/foreign/foreign.c
+++ b/src/backend/foreign/foreign.c
@@ -28,7 +28,9 @@
 #include "utils/rel.h"
 #include "utils/syscache.h"
 #include "utils/varlena.h"
+#include "utils/timeout.h"
 
+static CheckingRemoteServersCallbackItem *remote_check_callbacks = NULL;
 
 /*
  * GetForeignDataWrapper -	look up the foreign-data wrapper by OID.
@@ -810,3 +812,61 @@ GetExistingLocalJoinPath(RelOptInfo *joinrel)
 	}
 	return NULL;
 }
+
+
+/*
+ * Register callbacks for checking remote servers.
+ *
+ * This function is intended for use by FDW extensions.
+ */
+void
+RegisterCheckingRemoteServersCallback(CheckingRemoteServersCallback callback,
+									  void *arg)
+{
+	CheckingRemoteServersCallbackItem *item;
+
+	item = (CheckingRemoteServersCallbackItem *)
+			MemoryContextAlloc(TopMemoryContext,
+			sizeof(CheckingRemoteServersCallbackItem));
+	item->callback = callback;
+	item->arg = arg;
+	item->next = remote_check_callbacks;
+	remote_check_callbacks = item;
+}
+
+/*
+ * Deregister callbacks for checking remote servers.
+ */
+void
+UnregisterCheckingRemoteServersCallback(CheckingRemoteServersCallback callback,
+										void *arg)
+{
+	CheckingRemoteServersCallbackItem *item;
+	CheckingRemoteServersCallbackItem *prev;
+
+	prev = NULL;
+	for (item = remote_check_callbacks; item; prev = item, item = item->next)
+	{
+			if (item->callback == callback && item->arg == arg)
+			{
+					if (prev)
+							prev->next = item->next;
+					else
+							remote_check_callbacks = item->next;
+					pfree(item);
+					break;
+			}
+	}
+}
+
+/*
+ * Call callbacks for checking remote servers.
+ */
+void
+CallCheckingRemoteServersCallbacks(void)
+{
+	CheckingRemoteServersCallbackItem *item;
+
+	for (item = remote_check_callbacks; item; item = item->next)
+		item->callback(item->arg);
+}
diff --git a/src/backend/tcop/postgres.c b/src/backend/tcop/postgres.c
index a9a1851c94..f0f00d316e 100644
--- a/src/backend/tcop/postgres.c
+++ b/src/backend/tcop/postgres.c
@@ -35,6 +35,7 @@
 #include "commands/async.h"
 #include "commands/prepare.h"
 #include "common/pg_prng.h"
+#include "foreign/foreign.h"
 #include "jit/jit.h"
 #include "libpq/libpq.h"
 #include "libpq/pqformat.h"
@@ -166,6 +167,9 @@ static ProcSignalReason RecoveryConflictReason;
 static MemoryContext row_description_context = NULL;
 static StringInfoData row_description_buf;
 
+/* Message string for canceling qurery caused by extensions */
+static char		*QueryCancelMessage = NULL;
+
 /* ----------------------------------------------------------------
  *		decls for routines only used in this file
  * ----------------------------------------------------------------
@@ -3226,6 +3230,13 @@ ProcessInterrupts(void)
 				 errmsg("connection to client lost")));
 	}
 
+	if (CheckingRemoteServersTimeoutPending)
+	{
+		CheckingRemoteServersTimeoutPending = false;
+
+		CallCheckingRemoteServersCallbacks();
+	}
+
 	/*
 	 * If a recovery conflict happens while we are waiting for input from the
 	 * client, the client is presumably just sitting idle in a transaction,
@@ -3330,8 +3341,21 @@ ProcessInterrupts(void)
 			LockErrorCleanup();
 			ereport(ERROR,
 					(errcode(ERRCODE_QUERY_CANCELED),
+					 QueryCancelMessage ?
+					 errmsg("%s", QueryCancelMessage) :
 					 errmsg("canceling statement due to user request")));
 		}
+
+		/*
+		 * If a cancel request from FDW is ignored, we expect that it raise SIGINT and
+		 * fill QueryCancelMessage again. But some FDWs may skip its health check if
+		 * we already have a cancel message. To avoid that we clean up it.
+		 */
+		if (QueryCancelMessage != NULL)
+		{
+			pfree(QueryCancelMessage);
+			QueryCancelMessage = NULL;
+		}
 	}
 
 	if (IdleInTransactionSessionTimeoutPending)
@@ -4266,6 +4290,9 @@ PostgresMain(const char *dbname, const char *username)
 		/* Report the error to the client and/or server log */
 		EmitErrorReport();
 
+		/* Make sure QueryCancelMessage is reset. */
+		QueryCancelMessage = NULL;
+
 		/*
 		 * Make sure debug_query_string gets reset before we possibly clobber
 		 * the storage it points at.
@@ -5024,3 +5051,26 @@ disable_statement_timeout(void)
 	if (get_timeout_active(STATEMENT_TIMEOUT))
 		disable_timeout(STATEMENT_TIMEOUT, false);
 }
+
+bool
+TrySetQueryCancelMessage(char *message)
+{
+	if (!HasQueryCancelMessage())
+	{
+		MemoryContext old;
+
+		old = MemoryContextSwitchTo(CurTransactionContext);
+		QueryCancelMessage = pstrdup(message);
+		MemoryContextSwitchTo(old);
+
+		return true;
+	}
+
+	return false;
+}
+
+bool
+HasQueryCancelMessage(void)
+{
+	return QueryCancelMessage != NULL;
+}
diff --git a/src/backend/utils/init/globals.c b/src/backend/utils/init/globals.c
index 1a5d29ac9b..bb94adfea8 100644
--- a/src/backend/utils/init/globals.c
+++ b/src/backend/utils/init/globals.c
@@ -37,6 +37,7 @@ volatile sig_atomic_t IdleSessionTimeoutPending = false;
 volatile sig_atomic_t ProcSignalBarrierPending = false;
 volatile sig_atomic_t LogMemoryContextPending = false;
 volatile sig_atomic_t IdleStatsUpdateTimeoutPending = false;
+volatile sig_atomic_t CheckingRemoteServersTimeoutPending = false;
 volatile uint32 InterruptHoldoffCount = 0;
 volatile uint32 QueryCancelHoldoffCount = 0;
 volatile uint32 CritSectionCount = 0;
diff --git a/src/include/foreign/foreign.h b/src/include/foreign/foreign.h
index ac82125530..22f8143bdd 100644
--- a/src/include/foreign/foreign.h
+++ b/src/include/foreign/foreign.h
@@ -82,4 +82,23 @@ extern List *GetForeignColumnOptions(Oid relid, AttrNumber attnum);
 extern Oid	get_foreign_data_wrapper_oid(const char *fdwname, bool missing_ok);
 extern Oid	get_foreign_server_oid(const char *servername, bool missing_ok);
 
+
+/* Functions and variables for fdw checking */
+typedef void (*CheckingRemoteServersCallback) (void *arg);
+
+typedef struct CheckingRemoteServersCallbackItem CheckingRemoteServersCallbackItem;
+
+struct CheckingRemoteServersCallbackItem
+{
+	CheckingRemoteServersCallbackItem	*next;
+	CheckingRemoteServersCallback		callback;
+	void								*arg;
+};
+
+extern void RegisterCheckingRemoteServersCallback(CheckingRemoteServersCallback callback,
+												  void *arg);
+extern void UnregisterCheckingRemoteServersCallback(CheckingRemoteServersCallback callback,
+													void *arg);
+extern void CallCheckingRemoteServersCallbacks(void);
+
 #endif							/* FOREIGN_H */
diff --git a/src/include/miscadmin.h b/src/include/miscadmin.h
index e7ebea4ff4..75697e2be0 100644
--- a/src/include/miscadmin.h
+++ b/src/include/miscadmin.h
@@ -99,6 +99,8 @@ extern PGDLLIMPORT volatile sig_atomic_t IdleStatsUpdateTimeoutPending;
 extern PGDLLIMPORT volatile sig_atomic_t CheckClientConnectionPending;
 extern PGDLLIMPORT volatile sig_atomic_t ClientConnectionLost;
 
+extern PGDLLIMPORT volatile sig_atomic_t CheckingRemoteServersTimeoutPending;
+
 /* these are marked volatile because they are examined by signal handlers: */
 extern PGDLLIMPORT volatile uint32 InterruptHoldoffCount;
 extern PGDLLIMPORT volatile uint32 QueryCancelHoldoffCount;
diff --git a/src/include/tcop/tcopprot.h b/src/include/tcop/tcopprot.h
index 5d34978f32..5d1216f170 100644
--- a/src/include/tcop/tcopprot.h
+++ b/src/include/tcop/tcopprot.h
@@ -91,4 +91,7 @@ extern bool set_plan_disabling_options(const char *arg,
 									   GucContext context, GucSource source);
 extern const char *get_stats_option_name(const char *arg);
 
+extern bool TrySetQueryCancelMessage(char *message);
+extern bool HasQueryCancelMessage(void);
+
 #endif							/* TCOPPROT_H */
diff --git a/src/tools/pgindent/typedefs.list b/src/tools/pgindent/typedefs.list
index d9b839c979..3a1f9ac271 100644
--- a/src/tools/pgindent/typedefs.list
+++ b/src/tools/pgindent/typedefs.list
@@ -363,6 +363,7 @@ CatalogId
 CatalogIdMapEntry
 CatalogIndexState
 ChangeVarNodes_context
+CheckingRemoteServersCallbackItem
 CheckPoint
 CheckPointStmt
 CheckpointStatsData
-- 
2.27.0

