From 90cc57c6a55f279df1a427f665c272bbe0d9d1a8 Mon Sep 17 00:00:00 2001
From: "kuroda.hayato%40jp.fujitsu.com" <kuroda.hayato@jp.fujitsu.com>
Date: Tue, 1 Nov 2022 09:13:42 +0000
Subject: [PATCH v18 2/3] postgres_fdw: add postgres_fdw_verify_foreign_servers

This function can verify the status of connections that are establieshed by
postgres_fdw. This check wil be done by PQConncheck(), which means this is
available only on systems that support the non-standard POLLRDHUP extension
to the poll system call, including Linux.
This returns true if all checked connections are still valid. False is returned
if the checking is not supported on this platform.
---
 contrib/postgres_fdw/connection.c             | 48 +++++++++++++++++++
 .../postgres_fdw/postgres_fdw--1.0--1.1.sql   |  5 ++
 doc/src/sgml/postgres-fdw.sgml                | 26 ++++++++++
 3 files changed, 79 insertions(+)

diff --git a/contrib/postgres_fdw/connection.c b/contrib/postgres_fdw/connection.c
index f0c45b00db..584b914814 100644
--- a/contrib/postgres_fdw/connection.c
+++ b/contrib/postgres_fdw/connection.c
@@ -86,6 +86,7 @@ static bool xact_got_connection = false;
 PG_FUNCTION_INFO_V1(postgres_fdw_get_connections);
 PG_FUNCTION_INFO_V1(postgres_fdw_disconnect);
 PG_FUNCTION_INFO_V1(postgres_fdw_disconnect_all);
+PG_FUNCTION_INFO_V1(postgres_fdw_verify_foreign_servers);
 
 /* prototypes of private functions */
 static void make_new_connection(ConnCacheEntry *entry, UserMapping *user);
@@ -1862,3 +1863,50 @@ disconnect_cached_connections(Oid serverid)
 
 	return result;
 }
+
+/*
+ * Check health of remote servers.
+ *
+ * This function searches the hash table from the beginning and performs a
+ * health-check on each entry.
+ *
+ * Return true if all servers seem good, return false if we could not check on
+ * this platform, and raise an ERROR if a disconnection is found.
+ *
+ */
+Datum
+postgres_fdw_verify_foreign_servers(PG_FUNCTION_ARGS)
+{
+	HASH_SEQ_STATUS scan;
+	ConnCacheEntry *entry;
+	bool check_all = PG_GETARG_BOOL(0);
+
+	/* quick exit if connection cache has been not initialized yet. */
+	if (ConnectionHash == NULL)
+		PG_RETURN_BOOL(true);
+
+	hash_seq_init(&scan, ConnectionHash);
+	while ((entry = (ConnCacheEntry *) hash_seq_search(&scan)))
+	{
+		if (entry->conn == NULL ||
+			(entry->xact_depth == 0 && !check_all))
+			continue;
+		if (PQConncheck(entry->conn))
+		{
+			/*
+			 * Foreign server might be down, so throw ereport(ERROR).
+			 */
+			ForeignServer *server;
+
+			server = GetForeignServer(entry->serverid);
+			ereport(ERROR,
+					(errcode(ERRCODE_CONNECTION_FAILURE),
+					 errmsg("could not connect to server \"%s\"",
+							server->servername),
+					 errdetail("Socket close is detected."),
+					 errhint("Please check the health of the server.")));
+		}
+	}
+
+	PG_RETURN_BOOL(true);
+}
diff --git a/contrib/postgres_fdw/postgres_fdw--1.0--1.1.sql b/contrib/postgres_fdw/postgres_fdw--1.0--1.1.sql
index ed4ca378d4..d0a82410c9 100644
--- a/contrib/postgres_fdw/postgres_fdw--1.0--1.1.sql
+++ b/contrib/postgres_fdw/postgres_fdw--1.0--1.1.sql
@@ -18,3 +18,8 @@ CREATE FUNCTION postgres_fdw_disconnect_all ()
 RETURNS bool
 AS 'MODULE_PATHNAME'
 LANGUAGE C STRICT PARALLEL RESTRICTED;
+
+CREATE FUNCTION postgres_fdw_verify_foreign_servers (boolean)
+RETURNS bool
+AS 'MODULE_PATHNAME'
+LANGUAGE C STRICT PARALLEL RESTRICTED;
diff --git a/doc/src/sgml/postgres-fdw.sgml b/doc/src/sgml/postgres-fdw.sgml
index 527f4deaaa..790a5256ac 100644
--- a/doc/src/sgml/postgres-fdw.sgml
+++ b/doc/src/sgml/postgres-fdw.sgml
@@ -790,6 +790,32 @@ postgres=# SELECT postgres_fdw_disconnect_all();
      </para>
     </listitem>
    </varlistentry>
+
+   <varlistentry>
+    <term><function>postgres_fdw_verify_foreign_servers(check_all boolean) returns boolean</function></term>
+    <listitem>
+     <para>
+     This function checks the status of remote open connections that are
+     established by <filename>postgres_fdw</filename>. This check is performed
+     by polling the socket and allows long-running transactions to be aborted
+     sooner if the kernel reports that the connection is closed. This function is
+     currently available only on systems that support the non-standard <symbol>POLLRDHUP</symbol>
+     extension to the <symbol>poll</symbol> system call, including Linux. This
+     returns <literal>true</literal> if all checked connections are still valid.
+     <literal>false</literal> is returned if the checking is not supported on
+     this platform. If the local session seems to be disconnected from other
+     servers, an error is reported. When check_all is set to true, all
+     established connections are checked. Otherwise connections that have been
+     used in the transaction are checked. Example usage of the function:
+<screen>
+postgres=# SELECT postgres_fdw_verify_foreign_servers(false);
+ postgres_fdw_verify_foreign_servers
+-------------------------------------
+ t
+</screen>
+     </para>
+    </listitem>
+   </varlistentry>
    </variablelist>
 
 </sect2>
-- 
2.27.0

