From 0353eef23f1603a7ae27333bcca5379a5a7a301e 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 v19 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 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/connection.c             | 100 ++++++++++++++++++
 .../postgres_fdw/postgres_fdw--1.0--1.1.sql   |  10 ++
 doc/src/sgml/postgres-fdw.sgml                |  49 +++++++++
 3 files changed, 159 insertions(+)

diff --git a/contrib/postgres_fdw/connection.c b/contrib/postgres_fdw/connection.c
index f0c45b00db..69959a0f12 100644
--- a/contrib/postgres_fdw/connection.c
+++ b/contrib/postgres_fdw/connection.c
@@ -86,6 +86,8 @@ 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);
 
 /* prototypes of private functions */
 static void make_new_connection(ConnCacheEntry *entry, UserMapping *user);
@@ -116,6 +118,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
@@ -1862,3 +1865,100 @@ 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;
+
+	/* quick exit if connection cache has been not 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))
+			{
+				/*
+				 * Foreign server might be down, so throw ereport(WARNING).
+				 */
+				ForeignServer *server;
+
+				server = GetForeignServer(entry->serverid);
+				ereport(WARNING,
+						(errcode(ERRCODE_CONNECTION_FAILURE),
+						 errmsg("could not connect to server \"%s\"",
+								server->servername),
+						 errdetail("Socket close is detected."),
+						 errhint("Plsease check the health of the server.")));
+
+				result = false;
+			}
+		}
+	}
+
+	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));
+}
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..f1e2ee442f 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,13 @@ CREATE FUNCTION postgres_fdw_disconnect_all ()
 RETURNS bool
 AS 'MODULE_PATHNAME'
 LANGUAGE C STRICT PARALLEL RESTRICTED;
+
+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;
diff --git a/doc/src/sgml/postgres-fdw.sgml b/doc/src/sgml/postgres-fdw.sgml
index 527f4deaaa..d535973237 100644
--- a/doc/src/sgml/postgres-fdw.sgml
+++ b/doc/src/sgml/postgres-fdw.sgml
@@ -790,6 +790,55 @@ 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 that are 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 a 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_foreign_servers(false);
+ 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 remote connections that are 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 a lastly 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_foreign_servers(false);
+ postgres_fdw_verify_connection_states_all
+-------------------------------------------
+ t
+</screen>
+     </para>
+    </listitem>
+   </varlistentry>
    </variablelist>
 
 </sect2>
-- 
2.27.0

