From f363393138f78d087d8e319d0eccdf04adc2bc00 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 v31 1/4] Add PQconnCheck and PQcanConnCheck to libpq

PQconnCheck() function allows to check the status of socket 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.

PQcanConnCheck() checks whether above function is available or not.
---
 doc/src/sgml/libpq.sgml          | 41 +++++++++++++++++++++
 src/interfaces/libpq/exports.txt |  2 ++
 src/interfaces/libpq/fe-misc.c   | 61 +++++++++++++++++++++++++++-----
 src/interfaces/libpq/libpq-fe.h  |  4 +++
 4 files changed, 99 insertions(+), 9 deletions(-)

diff --git a/doc/src/sgml/libpq.sgml b/doc/src/sgml/libpq.sgml
index 0e7ae70c70..f49b6e7452 100644
--- a/doc/src/sgml/libpq.sgml
+++ b/doc/src/sgml/libpq.sgml
@@ -2679,6 +2679,47 @@ 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 socket.
+
+<synopsis>
+int PQconnCheck(PGconn *conn);
+</synopsis>
+      </para>
+
+      <para>
+       Unlike <xref linkend="libpq-PQstatus"/>, this function checks socket
+       health. This check is performed by polling the 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 been already invalid or
+       an error is error occurred.
+      </para>
+     </listitem>
+    </varlistentry>
+
+    <varlistentry id="libpq-PQcanConnCheck">
+     <term><function>PQcanConnCheck</function><indexterm><primary>PQcanConnCheck</primary></indexterm></term>
+     <listitem>
+      <para>
+       Returns whether <xref linkend="libpq-PQconnCheck"/> is available on
+       this platform. <xref linkend="libpq-PQcanConnCheck"/> returns
+       <literal>1</literal> if the function is supported, otherwise returns
+       <literal>0</literal>.
+
+<synopsis>
+int PQcanConnCheck(void);
+</synopsis>
+      </para>
+     </listitem>
+    </varlistentry>
+
    </variablelist>
   </para>
 
diff --git a/src/interfaces/libpq/exports.txt b/src/interfaces/libpq/exports.txt
index e8bcc88370..5c908bfe6e 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
+PQcanConnCheck            188
diff --git a/src/interfaces/libpq/fe-misc.c b/src/interfaces/libpq/fe-misc.c
index 3653a1a8a6..3bcd760cd7 100644
--- a/src/interfaces/libpq/fe-misc.c
+++ b/src/interfaces/libpq/fe-misc.c
@@ -53,8 +53,8 @@
 
 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	pqSocketIsReadableOrWritableOrValid(PGconn *conn, int forRead,
+												int forWrite, time_t end_time);
 static int	pqSocketPoll(int sock, int forRead, int forWrite, time_t end_time);
 
 /*
@@ -993,7 +993,7 @@ pqWaitTimed(int forRead, int forWrite, PGconn *conn, time_t finish_time)
 {
 	int			result;
 
-	result = pqSocketCheck(conn, forRead, forWrite, finish_time);
+	result = pqSocketIsReadableOrWritableOrValid(conn, forRead, forWrite, finish_time);
 
 	if (result < 0)
 		return -1;				/* errorMessage is already set */
@@ -1014,7 +1014,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 pqSocketIsReadableOrWritableOrValid(conn, 1, 0, (time_t) 0);
 }
 
 /*
@@ -1024,7 +1024,7 @@ pqReadReady(PGconn *conn)
 int
 pqWriteReady(PGconn *conn)
 {
-	return pqSocketCheck(conn, 0, 1, (time_t) 0);
+	return pqSocketIsReadableOrWritableOrValid(conn, 0, 1, (time_t) 0);
 }
 
 /*
@@ -1034,9 +1034,13 @@ pqWriteReady(PGconn *conn)
  *
  * If SSL is in use, the SSL buffer is checked prior to checking the socket
  * for read data directly.
+ *
+ * Moreover, when neither forRead nor forWrite is requested and timeout is
+ * disabled, try to check the health of socket.
  */
 static int
-pqSocketCheck(PGconn *conn, int forRead, int forWrite, time_t end_time)
+pqSocketIsReadableOrWritableOrValid(PGconn *conn, int forRead, int forWrite,
+									time_t end_time)
 {
 	int			result;
 
@@ -1082,20 +1086,33 @@ pqSocketCheck(PGconn *conn, int forRead, int forWrite, time_t end_time)
  *
  * Timeout is infinite if end_time is -1.  Timeout is immediate (no blocking)
  * if end_time is 0 (or indeed, any time before now).
+ *
+ * Moreover, when neither forRead nor forWrite is requested and timeout is
+ * disabled, try to check the health of socket.
  */
 static int
 pqSocketPoll(int sock, int forRead, int forWrite, time_t end_time)
 {
 	/* We use poll(2) if available, otherwise select(2) */
 #ifdef HAVE_POLL
-	struct pollfd input_fd;
+	struct pollfd input_fd = {0};
 	int			timeout_ms;
 
 	if (!forRead && !forWrite)
-		return 0;
+	{
+		/* Try to check the health if requested */
+		if (end_time == 0)
+#if defined(POLLRDHUP)
+			input_fd.events = POLLRDHUP | POLLHUP | POLLNVAL;
+#else
+			return 0;
+#endif /* defined(POLLRDHUP) */
+		else
+			return 0;
+	}
 
 	input_fd.fd = sock;
-	input_fd.events = POLLERR;
+	input_fd.events |= POLLERR;
 	input_fd.revents = 0;
 
 	if (forRead)
@@ -1218,6 +1235,32 @@ PQenv2encoding(void)
 	return encoding;
 }
 
+/*
+ * Check whether the socket peer closed 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 pqSocketIsReadableOrWritableOrValid(conn, 0, 0, (time_t) 0);
+}
+
+/*
+ * Check whether PQconnCheck() can work well on this platform.
+ *
+ * Returns 1 if this can use PQconnCheck(), otherwise 0.
+ */
+int
+PQcanConnCheck(void)
+{
+#if (defined(HAVE_POLL) && defined(POLLRDHUP))
+	return true;
+#else
+	return false;
+#endif
+}
 
 #ifdef ENABLE_NLS
 
diff --git a/src/interfaces/libpq/libpq-fe.h b/src/interfaces/libpq/libpq-fe.h
index f3d9220496..04a0395efb 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 PQcanConnCheck(void);
+
 /* === in fe-auth.c === */
 
 extern char *PQencryptPassword(const char *passwd, const char *user);
-- 
2.27.0

