From 2da884f1dab6278a1a3fd71faf02f659da71934f 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 v34 2/3] 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 existing connection is not closed by the remote peer.
---
 contrib/postgres_fdw/Makefile                 |   2 +-
 contrib/postgres_fdw/connection.c             | 162 ++++++++++++++++++
 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                |  74 ++++++++
 6 files changed, 258 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 12b54f15cd..99aff6304e 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, bool *checked);
 
 /*
  * Get a PGconn which can be used to execute queries on the remote PostgreSQL
@@ -1832,3 +1836,161 @@ 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. This return true
+ * if disconnections cannot be found, otherwise return false.
+ *
+ * checked will be set to true if PQconnCheck() is called at least once.
+ */
+static bool
+verify_cached_connections(Oid serverid, bool *checked)
+{
+	HASH_SEQ_STATUS scan;
+	ConnCacheEntry *entry;
+	bool		all = !OidIsValid(serverid);
+	bool		result = true;
+	StringInfoData str;
+
+	*checked = false;
+
+	Assert(ConnectionHash);
+
+	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);
+			}
+
+			/* Set a flag to notify the caller */
+			*checked = true;
+		}
+	}
+
+	/* Raise a warning if disconnections are found */
+	if (!result)
+	{
+		Assert(str.len);
+		ereport(WARNING,
+				(errcode(ERRCODE_CONNECTION_FAILURE),
+				 errmsg("%s", str.data),
+				 errdetail("Connection 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. This returns true
+ * if existing connection is not closed by the remote peer. false is returned
+ * if the local session seems to be disconnected from other servers. NULL is
+ * returned if a connection to the specified foreign server has not been
+ * established yet, or this function is not available on this platform.
+ */
+Datum
+postgres_fdw_verify_connection_states(PG_FUNCTION_ARGS)
+{
+	ForeignServer *server;
+	char	   *servername;
+	bool		result,
+				checked = false;
+
+	/* quick exit if the checking does not work well on this platfrom */
+	if (!PQconnCheckable())
+		PG_RETURN_NULL();
+
+	/* quick exit if connection cache has not been initialized yet */
+	if (!ConnectionHash)
+		PG_RETURN_NULL();
+
+	servername = text_to_cstring(PG_GETARG_TEXT_PP(0));
+	server = GetForeignServerByName(servername, false);
+
+	result = verify_cached_connections(server->serverid, &checked);
+
+	/* Return the result if checking function was called, otherwise NULL */
+	if (checked)
+		PG_RETURN_BOOL(result);
+	else
+		PG_RETURN_NULL();
+}
+
+/*
+ * 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)
+{
+	bool		result,
+				checked = false;
+
+	/* quick exit if the checking does not work well on this platfrom */
+	if (!PQconnCheckable())
+		PG_RETURN_NULL();
+
+	/* quick exit if connection cache has not been initialized yet */
+	if (!ConnectionHash)
+		PG_RETURN_NULL();
+
+	result = verify_cached_connections(InvalidOid, &checked);
+
+	/* Return the result if checking function was called, otherwise NULL */
+	if (checked)
+		PG_RETURN_BOOL(result);
+	else
+		PG_RETURN_NULL();
+}
+
+/*
+ * Check whether functions for verifying cached connections work well or not
+ */
+Datum
+postgres_fdw_can_verify_connection_states(PG_FUNCTION_ARGS)
+{
+	PG_RETURN_BOOL(PQconnCheckable());
+}
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..85e21f8904 100644
--- a/doc/src/sgml/postgres-fdw.sgml
+++ b/doc/src/sgml/postgres-fdw.sgml
@@ -826,6 +826,80 @@ 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 existing connection is not closed by
+      the remote peer. <literal>false</literal> is returned if the local
+      session seems to be disconnected from other servers. <literal>NULL</literal>
+      is returned if a connection to the specified foreign server has not been
+      established yet, or this function is not available on this platform.
+      If no foreign server with the given name is found, an error is reported.
+      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 connections are not closed by the
+      remote peer. <literal>false</literal> is returned if the local session
+      seems to be disconnected from at least one remote server. <literal>NULL</literal>
+      is returned if the local session does not have connection caches, or this
+      function is not available on this platform. 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 <function>postgres_fdw_verify_connection_states</function>
+      and <function>postgres_fdw_verify_connection_states</function> 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

