From 0b3a6479c4ccadf97ab48925bfd42fe3088887e8 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 v34 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   | 79 ++++++++++++++++++++++++++------
 src/interfaces/libpq/libpq-fe.h  |  4 ++
 4 files changed, 108 insertions(+), 15 deletions(-)

diff --git a/doc/src/sgml/libpq.sgml b/doc/src/sgml/libpq.sgml
index 3ccd8ff942..b6461cf82f 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 socket is valid, and returns <literal>-1</literal>
+       if the connection has already been closed or 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..fba5359dcb 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 and forConnCheck is requested, it
+ * means that the socket has not matched POLLRDHUP event and the socket has
+ * still survived.
  *
  * 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,15 +1111,21 @@ 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 and forConnCheck is requested, it means that the socket has not
+ * matched POLLRDHUP event and the socket has still survived.
  *
  * 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
@@ -1092,7 +1133,11 @@ pqSocketPoll(int sock, int forRead, int forWrite, time_t end_time)
 	int			timeout_ms;
 
 	if (!forRead && !forWrite)
-		return 0;
+	{
+		/* Connection check can be available on some limted platforms */
+		if (!(forConnCheck && PQconnCheckable()))
+			return 0;
+	}
 
 	input_fd.fd = sock;
 	input_fd.events = POLLERR;
@@ -1102,6 +1147,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))
diff --git a/src/interfaces/libpq/libpq-fe.h b/src/interfaces/libpq/libpq-fe.h
index f3d9220496..e1bd0cd7b7 100644
--- a/src/interfaces/libpq/libpq-fe.h
+++ b/src/interfaces/libpq/libpq-fe.h
@@ -648,6 +648,10 @@ 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);
+extern int PQconnCheckable(void);
+
 /* === in fe-auth.c === */
 
 extern char *PQencryptPassword(const char *passwd, const char *user);
-- 
2.27.0

