From 2f17e7ce900c17dd06b03aa891890b345917d303 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 v28 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          | 40 ++++++++++++++++++
 src/interfaces/libpq/exports.txt |  2 +
 src/interfaces/libpq/fe-misc.c   | 69 ++++++++++++++++++++++++++++++++
 src/interfaces/libpq/libpq-fe.h  |  4 ++
 4 files changed, 115 insertions(+)

diff --git a/doc/src/sgml/libpq.sgml b/doc/src/sgml/libpq.sgml
index 0e7ae70c70..5e07a252ce 100644
--- a/doc/src/sgml/libpq.sgml
+++ b/doc/src/sgml/libpq.sgml
@@ -2679,6 +2679,46 @@ 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 status 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 <literal>1</literal>
+       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..6e807e7c6a 100644
--- a/src/interfaces/libpq/fe-misc.c
+++ b/src/interfaces/libpq/fe-misc.c
@@ -1218,6 +1218,75 @@ PQenv2encoding(void)
 	return encoding;
 }
 
+/*
+ * Helper function for PQconnCheck().
+ *
+ * Return >0 if opposite side seems to be disconnected, 0 the socket is valid,
+ * and -1 if an error occurred.
+ */
+static int
+pqconnCheck_internal(int sock)
+{
+#if (defined(HAVE_POLL) && defined(POLLRDHUP))
+	struct pollfd input_fd;
+	int errflags = POLLHUP | POLLERR | POLLNVAL;
+	int result;
+
+	/* Prepare pollfd entry */
+	input_fd.fd = sock;
+	input_fd.events = POLLRDHUP | errflags;
+	input_fd.revents = 0;
+
+	/*
+	 * Check the status of socket. Note that we will retry as long as we get
+	 * EINTR.
+	 */
+	do
+		result = poll(&input_fd, 1, 0);
+	while (result < 0 && errno == EINTR);
+
+	if (result < 0)
+			return -1;
+
+	return input_fd.revents;
+#else
+	/* Do not support socket checking on this platform, return 0 */
+	return 0;
+#endif
+}
+
+/*
+ * 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)
+{
+	/* quick exit if invalid connection has come */
+	if (conn == NULL ||
+		conn->sock == PGINVALID_SOCKET ||
+		conn->status != CONNECTION_OK)
+		return -1;
+
+	return pqconnCheck_internal(conn->sock);
+}
+
+/*
+ * 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