From b194426726f39404a0e630ef480c69ee6f0fe7f0 Mon Sep 17 00:00:00 2001
From: "kuroda.hayato%40jp.fujitsu.com" <kuroda.hayato@jp.fujitsu.com>
Date: Wed, 21 Sep 2022 06:30:56 +0000
Subject: [PATCH v17 2/4] postgres_fdw: Implement health check feature

This patch adds a new GUC parameter postgres_fdw.health_check_interval.
This defines a time interval between checking remote servers.

In the checking function we use a socket event WL_SOCKET_CLOSED.
---
 contrib/postgres_fdw/connection.c   | 177 +++++++++++++++++++++++++++-
 contrib/postgres_fdw/option.c       |  63 ++++++++++
 contrib/postgres_fdw/postgres_fdw.h |   3 +
 3 files changed, 240 insertions(+), 3 deletions(-)

diff --git a/contrib/postgres_fdw/connection.c b/contrib/postgres_fdw/connection.c
index 939d114f02..99ee663614 100644
--- a/contrib/postgres_fdw/connection.c
+++ b/contrib/postgres_fdw/connection.c
@@ -23,12 +23,14 @@
 #include "postgres_fdw.h"
 #include "storage/fd.h"
 #include "storage/latch.h"
+#include "tcop/tcopprot.h"
 #include "utils/builtins.h"
 #include "utils/datetime.h"
 #include "utils/hsearch.h"
 #include "utils/inval.h"
 #include "utils/memutils.h"
 #include "utils/syscache.h"
+#include "utils/timeout.h"
 
 /*
  * Connection cache hash table entry
@@ -63,6 +65,7 @@ typedef struct ConnCacheEntry
 	bool		keep_connections;	/* setting value of keep_connections
 									 * server option */
 	Oid			serverid;		/* foreign server OID used to get server name */
+	char	   *servername;		/* foreign server name used to report ERROR */
 	uint32		server_hashvalue;	/* hash value of foreign server OID */
 	uint32		mapping_hashvalue;	/* hash value of user mapping OID */
 	PgFdwConnState state;		/* extra per-connection state */
@@ -80,6 +83,9 @@ static unsigned int prep_stmt_number = 0;
 /* tracks whether any work is needed in callback functions */
 static bool xact_got_connection = false;
 
+/* Timeout identifier for health check  */
+TimeoutId pgfdw_health_check_timeout = MAX_TIMEOUTS;
+
 /*
  * SQL functions
  */
@@ -117,6 +123,11 @@ static void pgfdw_finish_pre_subcommit_cleanup(List *pending_entries,
 static bool UserMappingPasswordRequired(UserMapping *user);
 static bool disconnect_cached_connections(Oid serverid);
 
+/* Functions for checking remote servers */
+static void pgfdw_connection_check(void *arg);
+static bool pgfdw_connection_check_internal(PGconn *conn);
+static void pgfdw_checking_remote_servers_timeout_handler(void);
+
 /*
  * Get a PGconn which can be used to execute queries on the remote PostgreSQL
  * server with the user's authorization.  A new connection is established
@@ -160,6 +171,13 @@ GetConnection(UserMapping *user, bool will_prep_stmt, PgFdwConnState **state)
 									  pgfdw_inval_callback, (Datum) 0);
 		CacheRegisterSyscacheCallback(USERMAPPINGOID,
 									  pgfdw_inval_callback, (Datum) 0);
+		CacheRegisterSyscacheCallback(FOREIGNSERVERNAME,
+									  pgfdw_inval_callback, (Datum) 0);
+
+		/* Register a timeout and a callback for checking remote servers */
+		pgfdw_health_check_timeout = RegisterTimeout(USER_TIMEOUT,
+													 pgfdw_checking_remote_servers_timeout_handler);
+		RegisterCheckingRemoteServersCallback(pgfdw_connection_check, NULL);
 	}
 
 	/* Set flag that we did GetConnection during the current transaction */
@@ -175,10 +193,11 @@ GetConnection(UserMapping *user, bool will_prep_stmt, PgFdwConnState **state)
 	if (!found)
 	{
 		/*
-		 * We need only clear "conn" here; remaining fields will be filled
-		 * later when "conn" is set.
+		 * We need clear "conn" and "servername" here; remaining fields will be filled
+		 * later when they are set.
 		 */
 		entry->conn = NULL;
+		entry->servername = NULL;
 	}
 
 	/* Reject further use of connections which failed abort cleanup. */
@@ -203,6 +222,19 @@ GetConnection(UserMapping *user, bool will_prep_stmt, PgFdwConnState **state)
 	if (entry->conn == NULL)
 		make_new_connection(entry, user);
 
+	/*
+	 * If cache entry doesn't have a name, duplicate it to the TopMemoryContext
+	 */
+	if (entry->servername == NULL)
+	{
+		MemoryContext old;
+		ForeignServer *server = GetForeignServer(user->serverid);
+
+		old = MemoryContextSwitchTo(TopMemoryContext);
+		entry->servername = pstrdup(server->servername);
+		MemoryContextSwitchTo(old);
+	}
+
 	/*
 	 * We check the health of the cached connection here when using it.  In
 	 * cases where we're out of all transactions, if a broken connection is
@@ -283,6 +315,12 @@ GetConnection(UserMapping *user, bool will_prep_stmt, PgFdwConnState **state)
 	if (state)
 		*state = &entry->state;
 
+	/* Start health-check timer if needed */
+	if (pgfdw_health_check_interval > 0 &&
+		!get_timeout_active(pgfdw_health_check_timeout))
+		enable_timeout_after(pgfdw_health_check_timeout,
+							 pgfdw_health_check_interval);
+
 	return entry->conn;
 }
 
@@ -532,6 +570,12 @@ disconnect_pg_server(ConnCacheEntry *entry)
 		entry->conn = NULL;
 		ReleaseExternalFD();
 	}
+
+	if (entry->servername != NULL)
+	{
+		pfree(entry->servername);
+		entry->servername = NULL;
+	}
 }
 
 /*
@@ -1040,6 +1084,9 @@ pgfdw_xact_callback(XactEvent event, void *arg)
 	 */
 	xact_got_connection = false;
 
+	/* Stop timer because checking is no more needed. */
+	disable_timeout(pgfdw_health_check_timeout, false);
+
 	/* Also reset cursor numbering for next transaction */
 	cursor_number = 0;
 }
@@ -1148,7 +1195,7 @@ pgfdw_inval_callback(Datum arg, int cacheid, uint32 hashvalue)
 	HASH_SEQ_STATUS scan;
 	ConnCacheEntry *entry;
 
-	Assert(cacheid == FOREIGNSERVEROID || cacheid == USERMAPPINGOID);
+	Assert(cacheid == FOREIGNSERVEROID || cacheid == USERMAPPINGOID || cacheid == FOREIGNSERVERNAME);
 
 	/* ConnectionHash must exist already, if we're registered */
 	hash_seq_init(&scan, ConnectionHash);
@@ -1179,6 +1226,21 @@ pgfdw_inval_callback(Datum arg, int cacheid, uint32 hashvalue)
 			else
 				entry->invalidated = true;
 		}
+
+		/* If ALTER SERVER RENAME is executed, the entry must be also modified. */
+		if (cacheid == FOREIGNSERVERNAME &&
+			entry->server_hashvalue == hashvalue)
+		{
+			ForeignServer *server = GetForeignServer(entry->serverid);
+			MemoryContext old;
+
+			if (entry->servername != NULL)
+				pfree(entry->servername);
+
+			old = MemoryContextSwitchTo(TopMemoryContext);
+			entry->servername = pstrdup(server->servername);
+			MemoryContextSwitchTo(old);
+		}
 	}
 }
 
@@ -1862,3 +1924,112 @@ disconnect_cached_connections(Oid serverid)
 
 	return result;
 }
+
+/*
+ * Signal handler for triggering health check
+ */
+static void
+pgfdw_checking_remote_servers_timeout_handler(void)
+{
+	CheckingRemoteServersTimeoutPending = true;
+	InterruptPending = true;
+	SetLatch(MyLatch);
+}
+
+/*
+ * Function for checking remote servers.
+ *
+ * This function searches the hash table from the beginning
+ * and performs a health-check on each entry.
+ *
+ * Raise SIGINT if someone might be down, otherwise do nothing.
+ */
+static void
+pgfdw_connection_check(void *arg)
+{
+	HASH_SEQ_STATUS scan;
+	ConnCacheEntry *entry;
+	bool			killed = false;
+
+	Assert(ConnectionHash);
+
+	/*
+	 * checking will be done by waiting WL_SOCKET_CLOSED event,
+	 * so exit immediately if it cannot be used in this system.
+	 */
+	if (!WaitEventSetCanReportClosed())
+		return;
+
+	/* raise SIGINT if QueryCancelMessage has been already set. */
+	if (HasQueryCancelMessage())
+	{
+		kill(MyProcPid, SIGINT);
+		return;
+	}
+
+	if (!killed)
+	{
+		hash_seq_init(&scan, ConnectionHash);
+		while ((entry = (ConnCacheEntry *) hash_seq_search(&scan)))
+		{
+			if (entry->conn == NULL || entry->xact_depth == 0)
+				continue;
+			if (!pgfdw_connection_check_internal(entry->conn))
+			{
+				/*
+				 * Foreign server might be down, so raise SIGINT.
+				 * Note that error message is passed to QueryCancelMessage
+				 * for reporting error in ProcessInterrupts().
+				 */
+				char msg[31 + MAXDATELEN];
+				snprintf(msg, sizeof(msg), "Foreign Server \"%s\" might be down.", entry->servername);
+				TrySetQueryCancelMessage(msg);
+
+				hash_seq_term(&scan);
+
+				kill(MyProcPid, SIGINT);
+				break;
+			}
+		}
+	}
+
+	/* re-schedule timer if needed. */
+	if (pgfdw_health_check_interval > 0)
+		enable_timeout_after(pgfdw_health_check_timeout,
+							 pgfdw_health_check_interval);
+
+	return;
+}
+
+/*
+ * Helper function for pgfdw_connection_check
+ */
+static bool
+pgfdw_connection_check_internal(PGconn *conn)
+{
+	WaitEventSet *eventset;
+	WaitEvent events;
+
+	Assert(WaitEventSetCanReportClosed());
+
+	/* Raise an ERROR if invalid socket has come */
+	if (conn == NULL ||
+		PQsocket(conn) == PGINVALID_SOCKET)
+		elog(ERROR, "Invalid connection has been arrived");
+
+	eventset = CreateWaitEventSet(TopMemoryContext, 1);
+
+	(void) AddWaitEventToSet(eventset, WL_SOCKET_CLOSED, PQsocket(conn), NULL, NULL);
+
+	WaitEventSetWait(eventset, 0, &events, 1, 0);
+
+	/* return false if the socket seems to be closed */
+	if (events.events & WL_SOCKET_CLOSED)
+	{
+		FreeWaitEventSet(eventset);
+		return false;
+	}
+
+	FreeWaitEventSet(eventset);
+	return true;
+}
diff --git a/contrib/postgres_fdw/option.c b/contrib/postgres_fdw/option.c
index fa80ee2a55..9cd5a32e77 100644
--- a/contrib/postgres_fdw/option.c
+++ b/contrib/postgres_fdw/option.c
@@ -20,6 +20,7 @@
 #include "commands/extension.h"
 #include "libpq/libpq-be.h"
 #include "postgres_fdw.h"
+#include "storage/latch.h"
 #include "utils/builtins.h"
 #include "utils/guc.h"
 #include "utils/varlena.h"
@@ -50,6 +51,7 @@ static PQconninfoOption *libpq_options;
  * GUC parameters
  */
 char	   *pgfdw_application_name = NULL;
+int			pgfdw_health_check_interval;
 
 /*
  * Helper functions
@@ -58,6 +60,13 @@ static void InitPgFdwOptions(void);
 static bool is_valid_option(const char *keyword, Oid context);
 static bool is_libpq_option(const char *keyword);
 
+/*
+ * GUC hooks
+ */
+static bool check_pgfdw_health_check_interval(int *newval, void **extra,
+											  GucSource source);
+static void assign_pgfdw_health_check_interval(int newval, void *extra);
+
 #include "miscadmin.h"
 
 /*
@@ -518,6 +527,47 @@ process_pgfdw_appname(const char *appname)
 	return buf.data;
 }
 
+/*
+ * Check hook for pgfdw_health_check_interval
+ */
+static bool
+check_pgfdw_health_check_interval(int *newval, void **extra, GucSource source)
+{
+	if (!WaitEventSetCanReportClosed() && *newval != 0)
+	{
+		GUC_check_errdetail("postgres_fdw.health_check_interval must be set to 0 on this platform");
+		return false;
+	}
+	return true;
+}
+
+/*
+ * Assign hook for pgfdw_health_check_interval
+ */
+static void
+assign_pgfdw_health_check_interval(int newval, void *extra)
+{
+	/* Quick return if timeout is not registered yet. */
+	if (pgfdw_health_check_timeout == MAX_TIMEOUTS)
+		return;
+
+	if (get_timeout_active(pgfdw_health_check_timeout))
+	{
+		if (newval == 0)
+			disable_timeout(pgfdw_health_check_timeout, false);
+
+		/*
+		 * we don't have to do anything because
+		 * new value will be used in pgfdw_connection_check().
+		 */
+		return;
+	}
+
+	/* Start timeout if wants to */
+	if (newval > 0)
+		enable_timeout_after(pgfdw_health_check_timeout, newval);
+}
+
 /*
  * Module load callback
  */
@@ -543,5 +593,18 @@ _PG_init(void)
 							   NULL,
 							   NULL);
 
+	DefineCustomIntVariable("postgres_fdw.health_check_interval",
+							"Sets the time interval between checks of remote servers.",
+							NULL,
+							&pgfdw_health_check_interval,
+							0,
+							0,
+							INT_MAX,
+							PGC_USERSET,
+							GUC_UNIT_MS,
+							check_pgfdw_health_check_interval,
+							assign_pgfdw_health_check_interval,
+							NULL);
+
 	MarkGUCPrefixReserved("postgres_fdw");
 }
diff --git a/contrib/postgres_fdw/postgres_fdw.h b/contrib/postgres_fdw/postgres_fdw.h
index a11d45bedf..4670e8d391 100644
--- a/contrib/postgres_fdw/postgres_fdw.h
+++ b/contrib/postgres_fdw/postgres_fdw.h
@@ -18,6 +18,7 @@
 #include "libpq-fe.h"
 #include "nodes/execnodes.h"
 #include "nodes/pathnodes.h"
+#include "utils/timeout.h"
 #include "utils/relcache.h"
 
 /*
@@ -151,6 +152,7 @@ extern PGresult *pgfdw_exec_query(PGconn *conn, const char *query,
 								  PgFdwConnState *state);
 extern void pgfdw_report_error(int elevel, PGresult *res, PGconn *conn,
 							   bool clear, const char *sql);
+extern TimeoutId pgfdw_health_check_timeout;
 
 /* in option.c */
 extern int	ExtractConnectionOptions(List *defelems,
@@ -160,6 +162,7 @@ extern List *ExtractExtensionList(const char *extensionsString,
 								  bool warnOnMissing);
 extern char *process_pgfdw_appname(const char *appname);
 extern char *pgfdw_application_name;
+extern int	pgfdw_health_check_interval;
 
 /* in deparse.c */
 extern void classifyConditions(PlannerInfo *root,
-- 
2.27.0

