postgres_fdw: Add connection status check to postgres_fdw_get_connections().

This commit extends the postgres_fdw_get_connections() function
to check if connections are closed. This is useful for detecting closed
postgres_fdw connections that could prevent successful transaction
commits. Users can roll back transactions immediately upon detecting
closed connections, avoiding unnecessary processing of failed
transactions.

This feature is available only on systems supporting the non-standard
POLLRDHUP extension to the poll system call, including Linux.

Author: Hayato Kuroda
Reviewed-by: Shinya Kato, Zhihong Yu, Kyotaro Horiguchi, Andres Freund
Reviewed-by: Onder Kalaci, Takamichi Osumi, Vignesh C, Tom Lane, Ted Yu
Reviewed-by: Katsuragi Yuta, Peter Smith, Shubham Khanna, Fujii Masao
Discussion: https://postgr.es/m/TYAPR01MB58662809E678253B90E82CE5F5889@TYAPR01MB5866.jpnprd01.prod.outlook.com

1 files changed, 78 insertions, 6 deletions

diff --git a/contrib/postgres_fdw/connection.c b/contrib/postgres_fdw/connection.c
index 3e902a723d4..78cea232d50 100644
--- a/contrib/postgres_fdw/connection.c
+++ b/contrib/postgres_fdw/connection.c
@@ -12,6 +12,10 @@
  */
 #include "postgres.h"
 
+#if HAVE_POLL_H
+#include <poll.h>
+#endif
+
 #include "access/htup_details.h"
 #include "access/xact.h"
 #include "catalog/pg_user_mapping.h"
@@ -171,6 +175,8 @@ static bool UserMappingPasswordRequired(UserMapping *user);
 static bool disconnect_cached_connections(Oid serverid);
 static void postgres_fdw_get_connections_internal(FunctionCallInfo fcinfo,
 												  enum pgfdwVersion api_version);
+static int	pgfdw_conn_check(PGconn *conn);
+static bool pgfdw_conn_checkable(void);
 
 /*
  * Get a PGconn which can be used to execute queries on the remote PostgreSQL
@@ -1991,14 +1997,14 @@ pgfdw_finish_abort_cleanup(List *pending_entries, List *cancel_requested,
 
 /* Number of output arguments (columns) for various API versions */
 #define POSTGRES_FDW_GET_CONNECTIONS_COLS_V1_1	2
-#define POSTGRES_FDW_GET_CONNECTIONS_COLS_V1_2	3
-#define POSTGRES_FDW_GET_CONNECTIONS_COLS	3	/* maximum of above */
+#define POSTGRES_FDW_GET_CONNECTIONS_COLS_V1_2	4
+#define POSTGRES_FDW_GET_CONNECTIONS_COLS	4	/* maximum of above */
 
 /*
  * Internal function used by postgres_fdw_get_connections variants.
  *
- * For API version 1.1, this function returns a set of records with
- * the following values:
+ * For API version 1.1, this function takes no input parameter and
+ * returns a set of records with the following values:
  *
  * - server_name - server name of active connection. In case the foreign server
  *   is dropped but still the connection is active, then the server name will
@@ -2006,10 +2012,12 @@ pgfdw_finish_abort_cleanup(List *pending_entries, List *cancel_requested,
  * - valid - true/false representing whether the connection is valid or not.
  *   Note that connections can become invalid in pgfdw_inval_callback.
  *
- * For API version 1.2 and later, this function returns the following
- * additional value along with the two values from version 1.1:
+ * For API version 1.2 and later, this function takes an input parameter
+ * to check a connection status and returns the following
+ * additional values along with the two values from version 1.1:
  *
  * - used_in_xact - true if the connection is used in the current transaction.
+ * - closed: true if the connection is closed.
  *
  * No records are returned when there are no cached connections at all.
  */
@@ -2101,8 +2109,19 @@ postgres_fdw_get_connections_internal(FunctionCallInfo fcinfo,
 
 		if (api_version >= PGFDW_V1_2)
 		{
+			bool		check_conn = PG_GETARG_BOOL(0);
+
 			/* Is this connection used in the current transaction? */
 			values[2] = BoolGetDatum(entry->xact_depth > 0);
+
+			/*
+			 * If a connection status check is requested and supported, return
+			 * whether the connection is closed. Otherwise, return NULL.
+			 */
+			if (check_conn && pgfdw_conn_checkable())
+				values[3] = BoolGetDatum(pgfdw_conn_check(entry->conn) != 0);
+			else
+				nulls[3] = true;
 		}
 
 		tuplestore_putvalues(rsinfo->setResult, rsinfo->setDesc, values, nulls);
@@ -2258,3 +2277,56 @@ disconnect_cached_connections(Oid serverid)
 
 	return result;
 }
+
+/*
+ * Check if the remote server closed the connection.
+ *
+ * Returns 1 if the connection is closed, -1 if an error occurred,
+ * and 0 if it's not closed or if the connection check is unavailable
+ * on this platform.
+ */
+static int
+pgfdw_conn_check(PGconn *conn)
+{
+	int			sock = PQsocket(conn);
+
+	if (PQstatus(conn) != CONNECTION_OK || sock == -1)
+		return -1;
+
+#if (defined(HAVE_POLL) && defined(POLLRDHUP))
+	{
+		struct pollfd input_fd;
+		int			result;
+
+		input_fd.fd = sock;
+		input_fd.events = POLLRDHUP;
+		input_fd.revents = 0;
+
+		do
+			result = poll(&input_fd, 1, 0);
+		while (result < 0 && errno == EINTR);
+
+		if (result < 0)
+			return -1;
+
+		return (input_fd.revents & POLLRDHUP) ? 1 : 0;
+	}
+#else
+	return 0;
+#endif
+}
+
+/*
+ * Check if connection status checking is available on this platform.
+ *
+ * Returns true if available, false otherwise.
+ */
+static bool
+pgfdw_conn_checkable(void)
+{
+#if (defined(HAVE_POLL) && defined(POLLRDHUP))
+	return true;
+#else
+	return false;
+#endif
+}