From a7d47f6eba34077c0b9da7551b5be7c360131a6a Mon Sep 17 00:00:00 2001
From: Hayato Kuroda <kuroda.hayato@fujitsu.com>
Date: Fri, 27 Jan 2023 03:17:18 +0000
Subject: [PATCH v37 1/3] Add PQconnCheck and PQconnCheckable to libpq

PQconnCheck() function allows to check the status of the connection by
polling the socket. This function is currently available only on systems
that support the non-standard POLLRDHUP extension to the poll system call,
including Linux.

PQconnCheckable() checks whether the above function is available or not.
---
 doc/src/sgml/libpq.sgml          | 38 ++++++++++++++
 src/interfaces/libpq/exports.txt |  2 +
 src/interfaces/libpq/fe-misc.c   | 86 ++++++++++++++++++++++++++------
 src/interfaces/libpq/libpq-fe.h  |  6 +++
 4 files changed, 117 insertions(+), 15 deletions(-)

diff --git a/doc/src/sgml/libpq.sgml b/doc/src/sgml/libpq.sgml
index 3ccd8ff942..b1a4471403 100644
--- a/doc/src/sgml/libpq.sgml
+++ b/doc/src/sgml/libpq.sgml
@@ -2679,6 +2679,44 @@ void *PQgetssl(const PGconn *conn);
      </listitem>
     </varlistentry>
 
+    <varlistentry id="libpq-PQconnCheck">
+     <term><function>PQconnCheck</function><indexterm><primary>PQconnCheck</primary></indexterm></term>
+     <listitem>
+      <para>
+       Returns the health of the connection.
+
+<synopsis>
+int PQconnCheck(PGconn *conn);
+</synopsis>
+      </para>
+
+      <para>
+       This function checks the health of the connection. Unlike <xref linkend="libpq-PQstatus"/>,
+       this check is performed by polling the corresponding socket. 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. <xref linkend="libpq-PQconnCheck"/>
+       returns greater than zero if the remote peer seems to be closed, returns
+       <literal>0</literal> if the connection is not closed, and returns
+       <literal>-1</literal> if an error has occurred.
+      </para>
+     </listitem>
+    </varlistentry>
+
+    <varlistentry id="libpq-PQconnCheckable">
+     <term><function>PQconnCheckable</function><indexterm><primary>PQconnCheckable</primary></indexterm></term>
+     <listitem>
+      <para>
+       Returns true (1) or false (0) to indicate if the <xref linkend="libpq-PQconnCheck"/>
+       function is supported on this platform.
+
+<synopsis>
+int PQconnCheckable(void);
+</synopsis>
+      </para>
+     </listitem>
+    </varlistentry>
+
    </variablelist>
   </para>
 
diff --git a/src/interfaces/libpq/exports.txt b/src/interfaces/libpq/exports.txt
index e8bcc88370..e7f0d435bd 100644
--- a/src/interfaces/libpq/exports.txt
+++ b/src/interfaces/libpq/exports.txt
@@ -186,3 +186,5 @@ PQpipelineStatus          183
 PQsetTraceFlags           184
 PQmblenBounded            185
 PQsendFlushRequest        186
+PQconnCheck               187
+PQconnCheckable           188
diff --git a/src/interfaces/libpq/fe-misc.c b/src/interfaces/libpq/fe-misc.c
index 3653a1a8a6..ccb9a0ae31 100644
--- a/src/interfaces/libpq/fe-misc.c
+++ b/src/interfaces/libpq/fe-misc.c
@@ -53,9 +53,10 @@
 
 static int	pqPutMsgBytes(const void *buf, size_t len, PGconn *conn);
 static int	pqSendSome(PGconn *conn, int len);
-static int	pqSocketCheck(PGconn *conn, int forRead, int forWrite,
-						  time_t end_time);
-static int	pqSocketPoll(int sock, int forRead, int forWrite, time_t end_time);
+static int	pqSocketCheck(PGconn *conn, int forRead,
+						  int forWrite, int forConnCheck, time_t end_time);
+static int	pqSocketPoll(int sock, int forRead,
+						 int forWrite, int forConnCheck, time_t end_time);
 
 /*
  * PQlibVersion: return the libpq version number
@@ -993,7 +994,7 @@ pqWaitTimed(int forRead, int forWrite, PGconn *conn, time_t finish_time)
 {
 	int			result;
 
-	result = pqSocketCheck(conn, forRead, forWrite, finish_time);
+	result = pqSocketCheck(conn, forRead, forWrite, 0, finish_time);
 
 	if (result < 0)
 		return -1;				/* errorMessage is already set */
@@ -1014,7 +1015,7 @@ pqWaitTimed(int forRead, int forWrite, PGconn *conn, time_t finish_time)
 int
 pqReadReady(PGconn *conn)
 {
-	return pqSocketCheck(conn, 1, 0, (time_t) 0);
+	return pqSocketCheck(conn, 1, 0, 0, (time_t) 0);
 }
 
 /*
@@ -1024,19 +1025,52 @@ pqReadReady(PGconn *conn)
 int
 pqWriteReady(PGconn *conn)
 {
-	return pqSocketCheck(conn, 0, 1, (time_t) 0);
+	return pqSocketCheck(conn, 0, 1, 0, (time_t) 0);
+}
+
+/*
+ * Check whether the socket peer closed the connection or not.
+ *
+ * Returns >0 if remote peer seems to be closed, 0 if it is valid,
+ * -1 if the input connection is bad or an error occurred.
+ */
+int
+PQconnCheck(PGconn *conn)
+{
+	return pqSocketCheck(conn, 0, 0, 1, (time_t) 0);
+}
+
+/*
+ * Check whether PQconnCheck() can work on this platform.
+ *
+ * Returns true (1) if this can use PQconnCheck(), otherwise false (0).
+ */
+int
+PQconnCheckable(void)
+{
+#if (defined(HAVE_POLL) && defined(POLLRDHUP))
+	return true;
+#else
+	return false;
+#endif
 }
 
 /*
  * Checks a socket, using poll or select, for data to be read, written,
- * or both.  Returns >0 if one or more conditions are met, 0 if it timed
- * out, -1 if an error occurred.
+ * or both. Moreover, this function can check the health of the connetion on
+ * some limited platforms if forConnCehck is specified.
+ *
+ * Returns >0 if one or more conditions are met, 0 if it timed out, -1 if an
+ * error occurred. Note that if 0 is returned when forConnCheck is requested,
+ * it means that the socket has not matched POLLRDHUP event and the connection
+ * is not closed by the socket peer.
  *
  * If SSL is in use, the SSL buffer is checked prior to checking the socket
  * for read data directly.
  */
 static int
-pqSocketCheck(PGconn *conn, int forRead, int forWrite, time_t end_time)
+pqSocketCheck(PGconn *conn, int forRead,
+			  int forWrite, int forConnCheck, time_t end_time)
 {
 	int			result;
 
@@ -1059,7 +1093,8 @@ pqSocketCheck(PGconn *conn, int forRead, int forWrite, time_t end_time)
 
 	/* We will retry as long as we get EINTR */
 	do
-		result = pqSocketPoll(conn->sock, forRead, forWrite, end_time);
+		result = pqSocketPoll(conn->sock, forRead,
+							  forWrite, forConnCheck, end_time);
 	while (result < 0 && SOCK_ERRNO == EINTR);
 
 	if (result < 0)
@@ -1076,22 +1111,33 @@ pqSocketCheck(PGconn *conn, int forRead, int forWrite, time_t end_time)
 
 /*
  * Check a file descriptor for read and/or write data, possibly waiting.
- * If neither forRead nor forWrite are set, immediately return a timeout
- * condition (without waiting).  Return >0 if condition is met, 0
- * if a timeout occurred, -1 if an error or interrupt occurred.
+ * Moreover, this function can check the health of connection on some limited
+ * platform if forConnCehck is specified.
+ *
+ * If neither forRead, forWrite nor forConnCheck are set, immediately return a
+ * timeout condition (without waiting). Return >0 if condition is met, 0 if a
+ * timeout occurred, -1 if an error or interrupt occurred. Note that if 0 is
+ * returned when forConnCheck is requested, it means that the socket has not
+ * matched POLLRDHUP event and the connection is not closed by the socket peer.
  *
  * Timeout is infinite if end_time is -1.  Timeout is immediate (no blocking)
  * if end_time is 0 (or indeed, any time before now).
  */
 static int
-pqSocketPoll(int sock, int forRead, int forWrite, time_t end_time)
+pqSocketPoll(int sock, int forRead,
+			 int forWrite, int forConnCheck, time_t end_time)
 {
 	/* We use poll(2) if available, otherwise select(2) */
 #ifdef HAVE_POLL
 	struct pollfd input_fd;
 	int			timeout_ms;
 
-	if (!forRead && !forWrite)
+	/*
+	 * Quick exit if no effective options are given. Note that connection
+	 * check can be available on some limted platforms.
+	 */
+	if (!forRead && !forWrite &&
+		!(forConnCheck && PQconnCheckable()))
 		return 0;
 
 	input_fd.fd = sock;
@@ -1102,6 +1148,10 @@ pqSocketPoll(int sock, int forRead, int forWrite, time_t end_time)
 		input_fd.events |= POLLIN;
 	if (forWrite)
 		input_fd.events |= POLLOUT;
+#if defined(POLLRDHUP)
+	if (forConnCheck)
+		input_fd.events |= POLLRDHUP;
+#endif
 
 	/* Compute appropriate timeout interval */
 	if (end_time == ((time_t) -1))
@@ -1136,6 +1186,12 @@ pqSocketPoll(int sock, int forRead, int forWrite, time_t end_time)
 
 	if (forWrite)
 		FD_SET(sock, &output_mask);
+
+	/*
+	 * forConncheck is not checked because the feature is not supported on
+	 * this platform.
+	 */
+
 	FD_SET(sock, &except_mask);
 
 	/* Compute appropriate timeout interval */
diff --git a/src/interfaces/libpq/libpq-fe.h b/src/interfaces/libpq/libpq-fe.h
index f3d9220496..3759a87225 100644
--- a/src/interfaces/libpq/libpq-fe.h
+++ b/src/interfaces/libpq/libpq-fe.h
@@ -648,6 +648,12 @@ extern int	PQdsplen(const char *s, int encoding);
 /* Get encoding id from environment variable PGCLIENTENCODING */
 extern int	PQenv2encoding(void);
 
+/* Check whether the postgres server is still alive or not */
+extern int	PQconnCheck(PGconn *conn);
+
+/* Check whether PQconnCheck() can work on this platform */
+extern int	PQconnCheckable(void);
+
 /* === in fe-auth.c === */
 
 extern char *PQencryptPassword(const char *passwd, const char *user);
-- 
2.27.0

