From f36effa4eb0045682249c2113afa354b3733a309 Mon Sep 17 00:00:00 2001
From: Hayato Kuroda <kuroda.hayato@fujitsu.com>
Date: Fri, 27 Jan 2023 03:17:30 +0000
Subject: [PATCH v28 2/4] postgres_fdw: add
 postgres_fdw_verify_connection_states

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 checked connection is still valid, or the checking is
not supported on this platform. False is returned if the connection seems
to be closed.
---
 contrib/postgres_fdw/Makefile                 |   2 +-
 contrib/postgres_fdw/connection.c             | 124 ++++++++++++++++++
 contrib/postgres_fdw/meson.build              |   1 +
 .../postgres_fdw/postgres_fdw--1.1--1.2.sql   |  19 +++
 contrib/postgres_fdw/postgres_fdw.control     |   2 +-
 doc/src/sgml/postgres-fdw.sgml                |  69 ++++++++++
 6 files changed, 215 insertions(+), 2 deletions(-)
 create mode 100644 contrib/postgres_fdw/postgres_fdw--1.1--1.2.sql

diff --git a/contrib/postgres_fdw/Makefile b/contrib/postgres_fdw/Makefile
index c1b0cad453..6d23768389 100644
--- a/contrib/postgres_fdw/Makefile
+++ b/contrib/postgres_fdw/Makefile
@@ -14,7 +14,7 @@ PG_CPPFLAGS = -I$(libpq_srcdir)
 SHLIB_LINK_INTERNAL = $(libpq)
 
 EXTENSION = postgres_fdw
-DATA = postgres_fdw--1.0.sql postgres_fdw--1.0--1.1.sql
+DATA = postgres_fdw--1.0.sql postgres_fdw--1.0--1.1.sql postgres_fdw--1.1--1.2.sql
 
 REGRESS = postgres_fdw
 
diff --git a/contrib/postgres_fdw/connection.c b/contrib/postgres_fdw/connection.c
index 7760380f00..1d7e350b33 100644
--- a/contrib/postgres_fdw/connection.c
+++ b/contrib/postgres_fdw/connection.c
@@ -87,6 +87,9 @@ 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_connection_states);
+PG_FUNCTION_INFO_V1(postgres_fdw_verify_connection_states_all);
+PG_FUNCTION_INFO_V1(postgres_fdw_can_verify_connection_states);
 
 /* prototypes of private functions */
 static void make_new_connection(ConnCacheEntry *entry, UserMapping *user);
@@ -117,6 +120,7 @@ static void pgfdw_finish_pre_subcommit_cleanup(List *pending_entries,
 											   int curlevel);
 static bool UserMappingPasswordRequired(UserMapping *user);
 static bool disconnect_cached_connections(Oid serverid);
+static bool verify_cached_connections(Oid serverid);
 
 /*
  * Get a PGconn which can be used to execute queries on the remote PostgreSQL
@@ -1832,3 +1836,123 @@ disconnect_cached_connections(Oid serverid)
 
 	return result;
 }
+
+/*
+ * Workhorse to verify cached connections.
+ *
+ * This function scans all the connection cache entries and verifies the
+ * connections whose foreign server OID matches with the specified one. If
+ * InvalidOid is specified, it verifies all the cached connections.
+ *
+ * This function emits warnings if a disconnection is found, and this returns
+ * false only when the lastly verified server seems to be disconnected.
+ *
+ * Note that the verification can be used on some limited platforms. If this
+ * server does not support it, this function alwayse returns true.
+ */
+static bool
+verify_cached_connections(Oid serverid)
+{
+	HASH_SEQ_STATUS scan;
+	ConnCacheEntry *entry;
+	bool		all = !OidIsValid(serverid);
+	bool		result = true;
+	StringInfoData str;
+
+	/* quick exit if connection cache has not been initialized yet. */
+	if (!ConnectionHash)
+		return true;
+
+	hash_seq_init(&scan, ConnectionHash);
+	while ((entry = (ConnCacheEntry *) hash_seq_search(&scan)))
+	{
+		/* Ignore cache entry if no open connection right now. */
+		if (!entry->conn)
+			continue;
+
+		/* Skip if the entry is invalidated. */
+		if (entry->invalidated)
+			continue;
+
+		if (all || entry->serverid == serverid)
+		{
+			if (PQconnCheck(entry->conn))
+			{
+				/* A foreign server might be down, so construct a message. */
+				ForeignServer *server = GetForeignServer(entry->serverid);
+
+				if (result)
+				{
+					/*
+					 * Initialize and add a prefix if this is the first
+					 * disconnection we found.
+					 */
+					initStringInfo(&str);
+					appendStringInfo(&str, "could not connect to server ");
+
+					result = false;
+				}
+				else
+					appendStringInfo(&str, ", ");
+
+				appendStringInfo(&str, "\"%s\"", server->servername);
+			}
+		}
+	}
+
+	/* Raise a warning if disconnections are found. */
+	if (!result)
+	{
+		Assert(str.len);
+		ereport(WARNING,
+				(errcode(ERRCODE_CONNECTION_FAILURE),
+				 errmsg("%s", str.data),
+				 errdetail("Socket close is detected."),
+				 errhint("Plsease check the health of server.")));
+		pfree(str.data);
+	}
+
+	return result;
+}
+
+/*
+ * Verify the specified cached connections.
+ *
+ * This function verifies the connections that are established by postgres_fdw
+ * from the local session to the foreign server with the given name.
+ *
+ * This function emits a warning if a disconnection is found, and this returns
+ * false only when the verified server seems to be disconnected.
+ *
+ * Note that the verification can be used on some limited platforms. If this
+ * server does not support it, this function alwayse returns true.
+ */
+Datum
+postgres_fdw_verify_connection_states(PG_FUNCTION_ARGS)
+{
+	ForeignServer *server;
+	char	   *servername;
+
+	servername = text_to_cstring(PG_GETARG_TEXT_PP(0));
+	server = GetForeignServerByName(servername, false);
+
+	PG_RETURN_BOOL(verify_cached_connections(server->serverid));
+}
+
+/*
+ * Verify all the cached connections.
+ *
+ * This function verifies all the connections that are established by postgres_fdw
+ * from the local session to the foreign servers.
+ */
+Datum
+postgres_fdw_verify_connection_states_all(PG_FUNCTION_ARGS)
+{
+	PG_RETURN_BOOL(verify_cached_connections(InvalidOid));
+}
+
+Datum
+postgres_fdw_can_verify_connection_states(PG_FUNCTION_ARGS)
+{
+	PG_RETURN_BOOL(PQcanConnCheck());
+}
diff --git a/contrib/postgres_fdw/meson.build b/contrib/postgres_fdw/meson.build
index 2b451f165e..29118d47bb 100644
--- a/contrib/postgres_fdw/meson.build
+++ b/contrib/postgres_fdw/meson.build
@@ -26,6 +26,7 @@ install_data(
   'postgres_fdw.control',
   'postgres_fdw--1.0.sql',
   'postgres_fdw--1.0--1.1.sql',
+  'postgres_fdw--1.1--1.2.sql',
   kwargs: contrib_data_args,
 )
 
diff --git a/contrib/postgres_fdw/postgres_fdw--1.1--1.2.sql b/contrib/postgres_fdw/postgres_fdw--1.1--1.2.sql
new file mode 100644
index 0000000000..a8556b4c9f
--- /dev/null
+++ b/contrib/postgres_fdw/postgres_fdw--1.1--1.2.sql
@@ -0,0 +1,19 @@
+/* contrib/postgres_fdw/postgres_fdw--1.1--1.2.sql */
+
+-- complain if script is sourced in psql, rather than via ALTER EXTENSION
+\echo Use "ALTER EXTENSION postgres_fdw UPDATE TO '1.2'" to load this file. \quit
+
+CREATE FUNCTION postgres_fdw_verify_connection_states (text)
+RETURNS bool
+AS 'MODULE_PATHNAME'
+LANGUAGE C STRICT PARALLEL RESTRICTED;
+
+CREATE FUNCTION postgres_fdw_verify_connection_states_all ()
+RETURNS bool
+AS 'MODULE_PATHNAME'
+LANGUAGE C STRICT PARALLEL RESTRICTED;
+
+CREATE FUNCTION postgres_fdw_can_verify_connection_states ()
+RETURNS bool
+AS 'MODULE_PATHNAME'
+LANGUAGE C STRICT PARALLEL SAFE;
diff --git a/contrib/postgres_fdw/postgres_fdw.control b/contrib/postgres_fdw/postgres_fdw.control
index d489382064..a4b800be4f 100644
--- a/contrib/postgres_fdw/postgres_fdw.control
+++ b/contrib/postgres_fdw/postgres_fdw.control
@@ -1,5 +1,5 @@
 # postgres_fdw extension
 comment = 'foreign-data wrapper for remote PostgreSQL servers'
-default_version = '1.1'
+default_version = '1.2'
 module_pathname = '$libdir/postgres_fdw'
 relocatable = true
diff --git a/doc/src/sgml/postgres-fdw.sgml b/doc/src/sgml/postgres-fdw.sgml
index 644f51835b..05265a1ef8 100644
--- a/doc/src/sgml/postgres-fdw.sgml
+++ b/doc/src/sgml/postgres-fdw.sgml
@@ -826,6 +826,75 @@ postgres=# SELECT postgres_fdw_disconnect_all();
      </para>
     </listitem>
    </varlistentry>
+
+   <varlistentry>
+    <term><function>postgres_fdw_verify_connection_states(server_name text) returns boolean</function></term>
+    <listitem>
+     <para>
+      This function checks the status of remote connections established by
+      <filename>postgres_fdw</filename> from the local session to the foreign
+      server with the given name. 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 checked connections are still valid,
+      or the checking is not supported on this platform. <literal>false</literal>
+      is returned if the local session seems to be disconnected from other
+      servers. Example usage of the function:
+<screen>
+postgres=# SELECT postgres_fdw_verify_connection_states('loopback1');
+ postgres_fdw_verify_connection_states
+---------------------------------------
+ t
+</screen>
+     </para>
+    </listitem>
+   </varlistentry>
+   <varlistentry>
+    <term><function>postgres_fdw_verify_connection_states_all() returns boolean</function></term>
+    <listitem>
+     <para>
+      This function checks the status of all the remote connections established
+      by <filename>postgres_fdw</filename> from the local session to the
+      foreign servers. 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 the remote connections are still
+      valid, or the checking is not supported on this platform.
+      <literal>false</literal> is returned if the local session seems to be
+      disconnected from at least one remote server. Example usage of the
+      function:
+<screen>
+postgres=# SELECT postgres_fdw_verify_connection_states_all();
+ postgres_fdw_verify_connection_states_all
+-------------------------------------------
+ t
+</screen>
+     </para>
+    </listitem>
+   </varlistentry>
+
+   <varlistentry>
+    <term><function>postgres_fdw_can_verify_connection_states() returns boolean</function></term>
+    <listitem>
+     <para>
+      This function checks whether functions the health of remote connectio
+      work well or not. This returns <literal>true</literal> if it can be used,
+      otherwise returns <literal>false</literal>. Example usage of the function:
+
+<screen>
+postgres=# SELECT postgres_fdw_can_verify_connection_states();
+ postgres_fdw_can_verify_connection_states
+-------------------------------------------
+ t
+</screen>
+     </para>
+    </listitem>
+   </varlistentry>
+
    </variablelist>
 
 </sect2>
-- 
2.27.0

